Storage: Close response body after request (#4116)

* 4107: Disconnect response after request * 4107: Added integration test testWriteChannelWithConnectionPool * 4107: Renamed blobName. * Fix lint
googleapis · Dec 4, 2018 · 442b41d · 442b41d
1 parent ccc1532
commit 442b41d
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 1 deletion.
diff --git a/...ts/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java b/...ts/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java
@@ -715,14 +715,19 @@ public void write(
       int code;
       String message;
       IOException exception = null;
+      HttpResponse response = null;
       try {
-        HttpResponse response = httpRequest.execute();
+        response = httpRequest.execute();
         code = response.getStatusCode();
         message = response.getStatusMessage();
       } catch (HttpResponseException ex) {
         exception = ex;
         code = ex.getStatusCode();
         message = ex.getStatusMessage();
+      } finally {
+        if (response != null) {
+          response.disconnect();
+        }
       }
       if (!last && code != 308 || last && !(code == 200 || code == 201)) {
         if (exception != null) {

diff --git a/...clients/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITStorageTest.java b/...clients/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITStorageTest.java
@@ -28,15 +28,20 @@
 import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.apache.ApacheHttpTransport;
 import com.google.api.client.util.DateTime;
 import com.google.api.gax.paging.Page;
 import com.google.auth.ServiceAccountSigner;
+import com.google.auth.http.HttpTransportFactory;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.Identity;
 import com.google.cloud.Policy;
 import com.google.cloud.ReadChannel;
 import com.google.cloud.RestorableState;
+import com.google.cloud.TransportOptions;
 import com.google.cloud.WriteChannel;
+import com.google.cloud.http.HttpTransportOptions;
 import com.google.cloud.kms.v1.CreateCryptoKeyRequest;
 import com.google.cloud.kms.v1.CreateKeyRingRequest;
 import com.google.cloud.kms.v1.CryptoKey;
@@ -111,6 +116,8 @@
 import java.util.logging.Logger;
 import java.util.zip.GZIPInputStream;
 import javax.crypto.spec.SecretKeySpec;
+import org.apache.http.impl.client.HttpClients;
+import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -184,6 +191,15 @@ public static void afterClass() throws ExecutionException, InterruptedException
     }
   }
 
+  private static class CustomHttpTransportFactory implements HttpTransportFactory {
+    @Override
+    public HttpTransport create() {
+      PoolingHttpClientConnectionManager manager = new PoolingHttpClientConnectionManager();
+      manager.setMaxTotal(1);
+      return new ApacheHttpTransport(HttpClients.createMinimal(manager));
+    }
+  }
+
   private static void prepareKmsKeys() throws IOException {
     String projectId = remoteStorageHelper.getOptions().getProjectId();
     GoogleCredentials credentials = GoogleCredentials.getApplicationDefault();
@@ -1744,6 +1760,29 @@ public void testWriteChannelExistingBlob() throws IOException {
     assertTrue(storage.delete(BUCKET, blobName));
   }
 
+  @Test(timeout = 5000)
+  public void testWriteChannelWithConnectionPool() throws IOException {
+    TransportOptions transportOptions =
+        HttpTransportOptions.newBuilder()
+            .setHttpTransportFactory(new CustomHttpTransportFactory())
+            .build();
+    Storage storageWithPool =
+        StorageOptions.newBuilder().setTransportOptions(transportOptions).build().getService();
+    String blobName = "test-custom-pool-management";
+    BlobInfo blob = BlobInfo.newBuilder(BUCKET, blobName).build();
+    byte[] stringBytes;
+    try (WriteChannel writer = storageWithPool.writer(blob)) {
+      stringBytes = BLOB_STRING_CONTENT.getBytes(UTF_8);
+      writer.write(ByteBuffer.wrap(BLOB_BYTE_CONTENT));
+      writer.write(ByteBuffer.wrap(stringBytes));
+    }
+    try (WriteChannel writer = storageWithPool.writer(blob)) {
+      stringBytes = BLOB_STRING_CONTENT.getBytes(UTF_8);
+      writer.write(ByteBuffer.wrap(BLOB_BYTE_CONTENT));
+      writer.write(ByteBuffer.wrap(stringBytes));
+    }
+  }
+
   @Test
   public void testGetSignedUrl() throws IOException {
     if (storage.getOptions().getCredentials() != null) {