HADOOP-18410. S3AInputStream.unbuffer() not releasing http connections

hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/S3AFileSystem.java

@@ -1606,8 +1606,11 @@ public <T> CompletableFuture<T> submit(final CallableRaisingIOE<T> operation) {
       CompletableFuture<T> result = new CompletableFuture<>();
       unboundedThreadPool.submit(() ->
           LambdaUtils.eval(result, () -> {
+            LOG.debug("Starting submitted operation in {}", auditSpan.getSpanId());
             try (AuditSpan span = auditSpan.activate()) {
               return operation.apply();
+            } finally {
+              LOG.debug("Completed submitted operation in {}", auditSpan.getSpanId());
             }
           }));
       return result;

hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/S3AInputStream.java

@@ -604,7 +604,7 @@ public synchronized void close() throws IOException {
       try {
         stopVectoredIOOperations.set(true);
         // close or abort the stream; blocking
-        awaitFuture(closeStream("close() operation", false, true));
+        closeStream("close() operation", false, true);
         // end the client+audit span.
         client.close();
         // this is actually a no-op
@@ -666,16 +666,21 @@ private CompletableFuture<Boolean> closeStream(
     CompletableFuture<Boolean> operation;
 
     if (blocking || shouldAbort || remaining <= asyncDrainThreshold) {
-      // don't bother with async io.
+      // don't bother with async IO if the caller plans to wait for
+      // the result, there's an abort (which is fast), or
+      // there is not much data to read.
       operation = CompletableFuture.completedFuture(
-          drain(shouldAbort, reason, remaining, object, wrappedStream));
+          drain(uri, streamStatistics, shouldAbort, reason, remaining, object, wrappedStream));
 
     } else {
       LOG.debug("initiating asynchronous drain of {} bytes", remaining);
       // schedule an async drain/abort with references to the fields so they
       // can be reused
+      S3Object s3Object = object;
+      S3ObjectInputStream stream = wrappedStream;
       operation = client.submit(
-          () -> drain(false, reason, remaining, object, wrappedStream));
+          () -> drain(uri, streamStatistics, false, reason, remaining, s3Object,
+              stream));
     }
 
     // either the stream is closed in the blocking call or the async call is
@@ -689,14 +694,17 @@ private CompletableFuture<Boolean> closeStream(
    * drain the stream. This method is intended to be
    * used directly or asynchronously, and measures the
    * duration of the operation in the stream statistics.
+   * @param uri URI for messages
+   * @param streamStatistics stats to update
    * @param shouldAbort force an abort; used if explicitly requested.
    * @param reason reason for stream being closed; used in messages
    * @param remaining remaining bytes
    * @param requestObject http request object; needed to avoid GC issues.
    * @param inner stream to close.
    * @return was the stream aborted?
    */
-  private boolean drain(
+  private static boolean drain(final String uri,
+      final S3AInputStreamStatistics streamStatistics,
       final boolean shouldAbort,
       final String reason,
       final long remaining,
@@ -707,6 +715,8 @@ private boolean drain(
       return invokeTrackingDuration(
           streamStatistics.initiateInnerStreamClose(shouldAbort),
           () -> drainOrAbortHttpStream(
+              uri,
+              streamStatistics,
               shouldAbort,
               reason,
               remaining,
@@ -730,42 +740,67 @@ private boolean drain(
    * A reference to the stream is passed in so that the instance
    * {@link #wrappedStream} field can be reused as soon as this
    * method is submitted;
+   * @param uri URI for messages
+   * @param streamStatistics stats to update
    * @param shouldAbort force an abort; used if explicitly requested.
    * @param reason reason for stream being closed; used in messages
    * @param remaining remaining bytes
    * @param requestObject http request object; needed to avoid GC issues.
    * @param inner stream to close.
    * @return was the stream aborted?
    */
-  private boolean drainOrAbortHttpStream(
+  private static boolean drainOrAbortHttpStream(
+      final String uri,
+      final S3AInputStreamStatistics streamStatistics,
       boolean shouldAbort,
       final String reason,
-      final long remaining,
-      final S3Object requestObject,
-      final S3ObjectInputStream inner) {
+      long remaining,
+      S3Object requestObject,
+      S3ObjectInputStream inner) {
     // force a use of the request object so IDEs don't warn of
     // lack of use.
     requireNonNull(requestObject);
+    LOG.debug("drain or abort reason {} remaining={} abort={}",
+        reason, remaining, shouldAbort);
 
     if (!shouldAbort) {
       try {
         // clean close. This will read to the end of the stream,
         // so, while cleaner, can be pathological on a multi-GB object
-
-        // explicitly drain the stream
         long drained = 0;
-        byte[] buffer = new byte[DRAIN_BUFFER_SIZE];
-        while (true) {
-          final int count = inner.read(buffer);
-          if (count < 0) {
-            // no more data is left
-            break;
+        if (remaining > 0) {
+          // explicitly drain the stream
+          LOG.debug("draining {} bytes", remaining);
+          drained = 0;
+          int size  = DRAIN_BUFFER_SIZE;
+          if (remaining < size) {
+            size = (int) remaining;
+          }
+          byte[] buffer = new byte[size];
+          // read the data; bail out early if
+          // the connection breaks
+          while (remaining > 0) {
+            final int count = inner.read(buffer);
+            if (count < 0) {
+              // no more data is left
+              break;
+            }
+            drained += count;
+            remaining -= count;
           }
-          drained += count;
+          LOG.debug("Drained stream of {} bytes", drained);
+        }
+
+        if (remaining != 0) {
+          // fewer bytes than expected came back; not treating as a
+          // reason to escalate to an abort().
+          // just log.
+          LOG.debug("drained fewer bytes than expected; {} remaining",
+              remaining);
         }
-        LOG.debug("Drained stream of {} bytes", drained);
 
         // now close it
+        LOG.debug("Closing stream");
         inner.close();
         // this MUST come after the close, so that if the IO operations fail
         // and an abort is triggered, the initial attempt's statistics
@@ -779,7 +814,7 @@ private boolean drainOrAbortHttpStream(
       }
     }
     if (shouldAbort) {
-      // Abort, rather than just close, the underlying stream.  Otherwise, the
+      // Abort, rather than just close, the underlying stream. Otherwise, the
       // remaining object payload is read from S3 while closing the stream.
       LOG.debug("Aborting stream {}", uri);
       try {
@@ -1345,6 +1380,10 @@ public synchronized void unbuffer() {
       closeStream("unbuffer()", false, false);
     } finally {
       streamStatistics.unbuffered();
+      if (inputPolicy.isAdaptive()) {
+        LOG.debug("Switching to Random IO seek policy after unbuffer() invoked");
+        setInputPolicy(S3AInputPolicy.Random);
+      }
     }
   }
 

hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/S3AUtils.java

@@ -378,6 +378,7 @@ private static InterruptedIOException translateInterruptedException(
     } else {
       String name = innerCause.getClass().getName();
       if (name.endsWith(".ConnectTimeoutException")
+          || name.endsWith(".ConnectionPoolTimeoutException")
           || name.endsWith("$ConnectTimeoutException")) {
         // TCP connection http timeout from the shaded or unshaded filenames
         // com.amazonaws.thirdparty.apache.http.conn.ConnectTimeoutException

hadoop-tools/hadoop-aws/src/main/java/org/apache/hadoop/fs/s3a/impl/InternalConstants.java

@@ -97,6 +97,7 @@ private InternalConstants() {
 
   static {
     Set<String> keys = Stream.of(
+        Constants.ASYNC_DRAIN_THRESHOLD,
         Constants.INPUT_FADVISE,
         Constants.READAHEAD_RANGE)
         .collect(Collectors.toSet());

hadoop-tools/hadoop-aws/src/test/java/org/apache/hadoop/fs/s3a/performance/ITestUnbufferDraining.java

@@ -0,0 +1,188 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.fs.s3a.performance;
+
+import java.io.IOException;
+
+import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataInputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.contract.ContractTestUtils;
+import org.apache.hadoop.fs.s3a.S3AFileSystem;
+import org.apache.hadoop.io.IOUtils;
+
+import static org.apache.hadoop.fs.Options.OpenFileOptions.FS_OPTION_OPENFILE_READ_POLICY;
+import static org.apache.hadoop.fs.Options.OpenFileOptions.FS_OPTION_OPENFILE_READ_POLICY_WHOLE_FILE;
+import static org.apache.hadoop.fs.contract.ContractTestUtils.dataset;
+import static org.apache.hadoop.fs.s3a.Constants.ASYNC_DRAIN_THRESHOLD;
+import static org.apache.hadoop.fs.s3a.Constants.ESTABLISH_TIMEOUT;
+import static org.apache.hadoop.fs.s3a.Constants.INPUT_FADVISE;
+import static org.apache.hadoop.fs.s3a.Constants.MAXIMUM_CONNECTIONS;
+import static org.apache.hadoop.fs.s3a.Constants.MAX_ERROR_RETRIES;
+import static org.apache.hadoop.fs.s3a.Constants.PREFETCH_ENABLED_KEY;
+import static org.apache.hadoop.fs.s3a.Constants.READAHEAD_RANGE;
+import static org.apache.hadoop.fs.s3a.Constants.REQUEST_TIMEOUT;
+import static org.apache.hadoop.fs.s3a.Constants.RETRY_LIMIT;
+import static org.apache.hadoop.fs.s3a.Constants.SOCKET_TIMEOUT;
+import static org.apache.hadoop.fs.s3a.S3ATestUtils.removeBaseAndBucketOverrides;
+
+/**
+ * Test stream unbuffer performance/behavior with stream draining
+ * and aborting.
+ */
+public class ITestUnbufferDraining extends AbstractS3ACostTest {
+
+  private static final Logger LOG =
+      LoggerFactory.getLogger(ITestUnbufferDraining.class);
+
+  public static final int READAHEAD = 1000;
+
+  public static final int FILE_SIZE = 50_000;
+
+  public static final int ATTEMPTS = 10;
+
+  private FileSystem brittleFS;
+
+  /**
+   * Create with markers kept, always.
+   */
+  public ITestUnbufferDraining() {
+    super(false);
+  }
+
+  @Override
+  public Configuration createConfiguration() {
+    Configuration conf = super.createConfiguration();
+    removeBaseAndBucketOverrides(conf,
+        ASYNC_DRAIN_THRESHOLD,
+        ESTABLISH_TIMEOUT,
+        INPUT_FADVISE,
+        MAX_ERROR_RETRIES,
+        MAXIMUM_CONNECTIONS,
+        PREFETCH_ENABLED_KEY,
+        READAHEAD_RANGE,
+        REQUEST_TIMEOUT,
+        RETRY_LIMIT,
+        SOCKET_TIMEOUT);
+
+    return conf;
+  }
+
+  @Override
+  public void setup() throws Exception {
+    super.setup();
+
+    // now create a new FS with minimal http capacity and recovery
+    // a separate one is used to avoid test teardown suffering
+    // from the lack of http connections and short timeouts.
+    Configuration conf = getConfiguration();
+    // kick off async drain for any data
+    conf.setInt(ASYNC_DRAIN_THRESHOLD, 1);
+    conf.setInt(MAXIMUM_CONNECTIONS, 2);
+    conf.setInt(MAX_ERROR_RETRIES, 1);
+    conf.setInt(ESTABLISH_TIMEOUT, 1000);
+    conf.setInt(READAHEAD_RANGE, READAHEAD);
+    conf.setInt(RETRY_LIMIT, 1);
+
+    brittleFS = FileSystem.newInstance(getFileSystem().getUri(), conf);
+  }
+
+  @Override
+  public void teardown() throws Exception {
+    super.teardown();
+    IOUtils.cleanupWithLogger(LOG, brittleFS);
+  }
+
+  public FileSystem getBrittleFS() {
+    return brittleFS;
+  }
+
+  /**
+   * Test stream close performance/behavior with stream draining
+   * and unbuffer.
+   */
+  @Test
+  public void testUnbufferDraining() throws Throwable {
+
+    describe("unbuffer draining");
+    FileStatus st = createTestFile();
+
+    int offset = FILE_SIZE - READAHEAD + 1;
+    try (FSDataInputStream in = getBrittleFS().openFile(st.getPath())
+        .withFileStatus(st)
+        .must(ASYNC_DRAIN_THRESHOLD, 1)
+        .build().get()) {
+      describe("Initiating unbuffer with async drain\n");
+      for (int i = 0; i < ATTEMPTS; i++) {
+        describe("Starting read/unbuffer #%d", i);
+        in.seek(offset);
+        in.read();
+        in.unbuffer();
+      }
+    }
+  }
+
+  /**
+   * Test stream close performance/behavior with stream draining
+   * and unbuffer.
+   */
+  @Test
+  public void testUnbufferAborting() throws Throwable {
+
+    describe("unbuffer draining");
+    FileStatus st = createTestFile();
+
+
+    // open the file at the beginning with a whole file read policy,
+    // so even with s3a switching to random on unbuffer,
+    // this always does a full GET
+    try (FSDataInputStream in = getBrittleFS().openFile(st.getPath())
+        .withFileStatus(st)
+        .must(ASYNC_DRAIN_THRESHOLD, 1)
+        .must(FS_OPTION_OPENFILE_READ_POLICY,
+            FS_OPTION_OPENFILE_READ_POLICY_WHOLE_FILE)
+        .build().get()) {
+
+      describe("Initiating unbuffer with async drain\n");
+      for (int i = 0; i < ATTEMPTS; i++) {
+        describe("Starting read/unbuffer #%d", i);
+        in.read();
+        in.unbuffer();
+      }
+    }
+  }
+
+  private FileStatus createTestFile() throws IOException {
+    byte[] data = dataset(FILE_SIZE, '0', 10);
+    S3AFileSystem fs = getFileSystem();
+
+    Path path = methodPath();
+    ContractTestUtils.createFile(fs, path, true, data);
+    FileStatus st = fs.getFileStatus(path);
+    return st;
+  }
+
+
+}