Cleanly close the read channel in the GoogleHadoopFSInputStream so that read to EOF does not happen on calling close() unnecessarily

gcsio/src/main/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageReadChannel.java

@@ -147,6 +147,16 @@ public class GoogleCloudStorageReadChannel implements SeekableByteChannel {
   // 3. Test that footer prefetch always disabled for gzipped files.
   private byte[] footerContent;
 
+  /**
+   * The {@link HttpResponse} for the media stream currently returned from {@link #openStream}, if
+   * any. Held so that {@link #closeContentChannel} can call {@link HttpResponse#disconnect()}
+   * before closing the {@link #contentChannel}. Otherwise Apache HTTP's {@code
+   * ContentLengthInputStream} drains the entire remaining entity on {@code InputStream#close()},
+   * which can take minutes on a large object when the caller closes after only reading a small
+   * prefix (e.g. Hadoop distcp aborting a copy).
+   */
+  @Nullable private HttpResponse openMediaResponse;
+
   @VisibleForTesting protected boolean metadataInitialized = false;
 
   /**
@@ -362,6 +372,7 @@ public int read(ByteBuffer buffer) throws IOException {
           if (contentChannelEnd != size && currentPosition == contentChannelEnd) {
             closeContentChannel();
           } else {
+            openMediaResponse = null; // response body fully consumed; no need to abort on close()
             break;
           }
         }
@@ -525,6 +536,7 @@ public boolean isOpen() {
    * already responsible for performing local cleanup at the time the exception was raised.
    */
   protected void closeContentChannel() {
+    disconnectOpenMediaResponse();
     if (contentChannel != null) {
       logger.atFiner().log("Closing internal contentChannel for '%s'", resourceId);
       try {
@@ -541,6 +553,36 @@ protected void closeContentChannel() {
     }
   }
 
+  /**
+   * Aborts the underlying HTTP request for the current media body, if any. Must run before closing
+   * the {@link #contentChannel} so that closing the Apache {@code ContentLengthInputStream} does
+   * not synchronously read the rest of the response body.
+   *
+   * <p>Failures from {@link HttpResponse#disconnect()} are expected during normal teardown (for
+   * example when the socket is already closed) and are only logged here; they are not reported via
+   * {@link GoogleCloudStorageEventBus#postOnException()}.
+   */
+  void disconnectHttpResponse(@Nullable HttpResponse response) {
+    if (response == null) {
+      return;
+    }
+    try {
+      response.disconnect();
+    } catch (Exception e) {
+      logger.atFine().withCause(e).log(
+          "Got an exception on HttpResponse.disconnect() for '%s'; ignoring it.", resourceId);
+    }
+  }
+
+  private void disconnectOpenMediaResponse() {
+    if (openMediaResponse == null) {
+      return;
+    }
+    HttpResponse response = openMediaResponse;
+    openMediaResponse = null;
+    disconnectHttpResponse(response);
+  }
+
   private void resetContentChannel() {
     checkState(contentChannel == null, "contentChannel should be null for '%s'", resourceId);
     contentChannelPosition = -1;
@@ -1020,6 +1062,7 @@ protected InputStream openStream(long bytesToRead) throws IOException {
           metadataInitialized, "metadata should be initialized already for '%s'", resourceId);
       if (size == 0) {
         resetContentChannel();
+        disconnectHttpResponse(response);
         return new ByteArrayInputStream(new byte[0]);
       }
       if (gzipEncoded) {
@@ -1030,6 +1073,7 @@ protected InputStream openStream(long bytesToRead) throws IOException {
           contentChannelEnd = size;
         } else {
           resetContentChannel();
+          disconnectHttpResponse(response);
           return openStream(bytesToRead);
         }
       }
@@ -1118,11 +1162,14 @@ protected InputStream openStream(long bytesToRead) throws IOException {
           currentPosition,
           resourceId);
 
-      return new GcsReadDurationTrackerStream(
-          contentStream,
-          UriPaths.fromResourceId(resourceId, /* allowEmptyObjectName= */ false),
-          response.getHeaders(),
-          readOptions.getLatencyLoggingThreshold());
+      GcsReadDurationTrackerStream tracked =
+          new GcsReadDurationTrackerStream(
+              contentStream,
+              UriPaths.fromResourceId(resourceId, /* allowEmptyObjectName= */ false),
+              response.getHeaders(),
+              readOptions.getLatencyLoggingThreshold());
+      openMediaResponse = response;
+      return tracked;
     } catch (IOException e) {
       GoogleCloudStorageEventBus.postOnException();
       try {

gcsio/src/test/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageReadChannelTest.java

@@ -37,6 +37,7 @@
 import static org.junit.Assert.assertThrows;
 
 import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpResponse;
 import com.google.api.client.json.gson.GsonFactory;
 import com.google.api.client.testing.http.MockHttpTransport;
 import com.google.api.client.testing.http.MockLowLevelHttpResponse;
@@ -63,6 +64,8 @@
 import java.util.List;
 import java.util.Random;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import javax.annotation.Nullable;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -1060,6 +1063,118 @@ public void readBeyondChannelLength() throws Exception {
     assertThat(readChannel.position()).isEqualTo(20);
   }
 
+  @Test
+  public void close_afterPartialMediaRead_invokesHttpDisconnectBeforeClosingChannel()
+      throws IOException {
+    byte[] testData = new byte[100];
+    Arrays.fill(testData, (byte) 0x55);
+    MockHttpTransport transport = mockTransport(dataRangeResponse(testData, 0, testData.length));
+
+    Storage storage = new Storage(transport, GsonFactory.getDefaultInstance(), r -> {});
+
+    GoogleCloudStorageReadOptions options =
+        newLazyReadOptionsBuilder().setFadvise(Fadvise.SEQUENTIAL).build();
+
+    HttpDisconnectCountingReadChannel readChannel =
+        new HttpDisconnectCountingReadChannel(
+            storage,
+            new StorageResourceId(BUCKET_NAME, OBJECT_NAME),
+            ApiErrorExtractor.INSTANCE,
+            new ClientRequestHelper<>(),
+            options);
+
+    readChannel.position(0);
+    assertThat(readChannel.read(ByteBuffer.allocate(1))).isEqualTo(1);
+
+    readChannel.close();
+
+    assertThat(readChannel.getDisconnectHttpResponseCallCount()).isEqualTo(1);
+  }
+
+  @Test
+  public void close_withoutMediaRead_doesNotInvokeHttpDisconnect() throws IOException {
+    StorageObject object = newStorageObject(BUCKET_NAME, OBJECT_NAME);
+    MockHttpTransport transport = mockTransport(jsonDataResponse(object));
+
+    Storage storage = new Storage(transport, GsonFactory.getDefaultInstance(), r -> {});
+
+    HttpDisconnectCountingReadChannel readChannel =
+        new HttpDisconnectCountingReadChannel(
+            storage,
+            new StorageResourceId(BUCKET_NAME, OBJECT_NAME),
+            ApiErrorExtractor.INSTANCE,
+            new ClientRequestHelper<>(),
+            newLazyReadOptionsBuilder().setFastFailOnNotFoundEnabled(true).build());
+
+    assertThat(readChannel.size()).isEqualTo(object.getSize().longValue());
+    readChannel.close();
+
+    assertThat(readChannel.getDisconnectHttpResponseCallCount()).isEqualTo(0);
+  }
+
+  @Test
+  public void closeContentChannel_whenReplacingStream_invokesHttpDisconnectEachTime()
+      throws IOException {
+    byte[] testData = new byte[100];
+    Arrays.fill(testData, (byte) 0x33);
+    MockHttpTransport transport =
+        mockTransport(
+            dataRangeResponse(testData, 0, testData.length),
+            dataRangeResponse(Arrays.copyOfRange(testData, 50, 100), 50, testData.length));
+
+    Storage storage = new Storage(transport, GsonFactory.getDefaultInstance(), r -> {});
+
+    GoogleCloudStorageReadOptions options =
+        newLazyReadOptionsBuilder().setFadvise(Fadvise.SEQUENTIAL).setInplaceSeekLimit(0).build();
+
+    HttpDisconnectCountingReadChannel readChannel =
+        new HttpDisconnectCountingReadChannel(
+            storage,
+            new StorageResourceId(BUCKET_NAME, OBJECT_NAME),
+            ApiErrorExtractor.INSTANCE,
+            new ClientRequestHelper<>(),
+            options);
+
+    readChannel.position(0);
+    assertThat(readChannel.read(ByteBuffer.allocate(1))).isEqualTo(1);
+    readChannel.position(50);
+    assertThat(readChannel.read(ByteBuffer.allocate(1))).isEqualTo(1);
+
+    readChannel.close();
+
+    assertThat(readChannel.getDisconnectHttpResponseCallCount()).isEqualTo(2);
+  }
+
+  /**
+   * Subclass for asserting {@link GoogleCloudStorageReadChannel#disconnectHttpResponse} is invoked
+   * when tearing down a partially consumed media response (package-private override).
+   */
+  private static final class HttpDisconnectCountingReadChannel
+      extends GoogleCloudStorageReadChannel {
+
+    private final AtomicInteger disconnectHttpResponseCallCount = new AtomicInteger();
+
+    HttpDisconnectCountingReadChannel(
+        Storage storage,
+        StorageResourceId resourceId,
+        ApiErrorExtractor errorExtractor,
+        ClientRequestHelper<StorageObject> requestHelper,
+        GoogleCloudStorageReadOptions readOptions)
+        throws IOException {
+      super(storage, resourceId, errorExtractor, requestHelper, readOptions);
+    }
+
+    int getDisconnectHttpResponseCallCount() {
+      return disconnectHttpResponseCallCount.get();
+    }
+
+    @Override
+    void disconnectHttpResponse(@Nullable HttpResponse response) {
+      disconnectHttpResponseCallCount.incrementAndGet();
+      super.disconnectHttpResponse(response);
+    }
+  }
+
   private static GoogleCloudStorageReadOptions.Builder newLazyReadOptionsBuilder() {
     return GoogleCloudStorageReadOptions.builder().setFastFailOnNotFoundEnabled(false);
   }