diff --git a/google-cloud-clients/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java b/google-cloud-clients/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java index 709e01a25064..22407f902c92 100644 --- a/google-cloud-clients/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java +++ b/google-cloud-clients/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/google-cloud-clients/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITStorageTest.java b/google-cloud-clients/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITStorageTest.java index 0bf20e58dcf2..aba71eb7ddb9 100644 --- a/google-cloud-clients/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITStorageTest.java +++ b/google-cloud-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) {