Repository: googleapis/java-storage
Pull request #3054 (2025-04-23 19:01:40)

chore: ensure grpc client thread pool is shutdown


Patch:
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java
index 606d0c6f..8a838604 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java
@@ -203,7 +203,8 @@ final class GrpcStorageImpl extends BaseService<StorageOptions>
 
   @Override
   public void close() throws Exception {
-    try (StorageClient s = storageClient) {
+    try (StorageClient s = storageClient;
+        StorageDataClient ignore = storageDataClient) {
       s.shutdownNow();
       java.time.Duration terminationAwaitDuration = getOptions().getTerminationAwaitDuration();
       s.awaitTermination(terminationAwaitDuration.toMillis(), TimeUnit.MILLISECONDS);
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java
index 93deda5e..7d116e16 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java
@@ -911,6 +911,7 @@ public final class GrpcStorageOptions extends StorageOptions
             StorageDataClient dataClient =
                 StorageDataClient.create(
                     executor,
+                    grpcStorageOptions.terminationAwaitDuration,
                     new ZeroCopyBidiStreamingCallable<>(
                         stub.bidiReadObjectCallable(), stub.bidiReadObjectResponseMarshaller),
                     retryContextProvider,
@@ -932,6 +933,7 @@ public final class GrpcStorageOptions extends StorageOptions
             StorageDataClient dataClient =
                 StorageDataClient.create(
                     executor,
+                    grpcStorageOptions.terminationAwaitDuration,
                     new ZeroCopyBidiStreamingCallable<>(
                         client.bidiReadObjectCallable(),
                         ResponseContentLifecycleManager.noopBidiReadObjectResponse()),
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageDataClient.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageDataClient.java
index 2f84c048..19fa8b5c 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageDataClient.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageDataClient.java
@@ -29,12 +29,15 @@ import com.google.cloud.storage.RetryContext.RetryContextProvider;
 import com.google.storage.v2.BidiReadObjectRequest;
 import com.google.storage.v2.BidiReadObjectResponse;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 
 @InternalApi
-final class StorageDataClient implements IOAutoCloseable {
+final class StorageDataClient implements AutoCloseable {
 
   private final ScheduledExecutorService executor;
+  private final Duration terminationAwaitDuration;
   private final ZeroCopyBidiStreamingCallable<BidiReadObjectRequest, BidiReadObjectResponse>
       bidiReadObject;
   private final RetryContextProvider retryContextProvider;
@@ -42,10 +45,12 @@ final class StorageDataClient implements IOAutoCloseable {
 
   private StorageDataClient(
       ScheduledExecutorService executor,
+      Duration terminationAwaitDuration,
       ZeroCopyBidiStreamingCallable<BidiReadObjectRequest, BidiReadObjectResponse> bidiReadObject,
       RetryContextProvider retryContextProvider,
       IOAutoCloseable onClose) {
     this.executor = executor;
+    this.terminationAwaitDuration = terminationAwaitDuration;
     this.bidiReadObject = bidiReadObject;
     this.retryContextProvider = retryContextProvider;
     this.onClose = onClose;
@@ -112,11 +117,14 @@ final class StorageDataClient implements IOAutoCloseable {
     return objectReadSessionFuture;
   }
 
+  @SuppressWarnings("ResultOfMethodCallIgnored")
   @Override
-  public void close() throws IOException {
-    //noinspection EmptyTryBlock
+  public void close() throws Exception {
     try (IOAutoCloseable ignore = onClose) {
-      // intentional
+      // today, we own the executor service. If StorageDataClient is ever standalone, this code will
+      // need to be re-evaluated. Especially if a customer is able to provide the executor.
+      executor.shutdownNow();
+      executor.awaitTermination(terminationAwaitDuration.toNanos(), TimeUnit.NANOSECONDS);
     }
   }
 
@@ -127,10 +135,12 @@ final class StorageDataClient implements IOAutoCloseable {
 
   static StorageDataClient create(
       ScheduledExecutorService executor,
+      Duration terminationAwaitDuration,
       ZeroCopyBidiStreamingCallable<BidiReadObjectRequest, BidiReadObjectResponse> read,
       RetryContextProvider retryContextProvider,
       IOAutoCloseable onClose) {
-    return new StorageDataClient(executor, read, retryContextProvider, onClose);
+    return new StorageDataClient(
+        executor, terminationAwaitDuration, read, retryContextProvider, onClose);
   }
 
   @FunctionalInterface


Test patch:
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageDataClientTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageDataClientTest.java
index c3011200..c3ec75fd 100644
--- a/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageDataClientTest.java
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageDataClientTest.java
@@ -16,20 +16,31 @@
 
 package com.google.cloud.storage;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.storage.v2.BidiReadObjectRequest;
 import com.google.storage.v2.ReadRange;
-import java.io.IOException;
+import java.time.Duration;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
 
+@RunWith(MockitoJUnitRunner.class)
 public final class StorageDataClientTest {
 
+  @Mock public ScheduledExecutorService exec;
+
   @Test
-  public void readSession_requestWithRangeRead_noAllowed() throws IOException {
+  public void readSession_requestWithRangeRead_noAllowed() throws Exception {
     try (StorageDataClient dc =
-        StorageDataClient.create(null, null, null, IOAutoCloseable.noOp())) {
+        StorageDataClient.create(exec, Duration.ofSeconds(2), null, null, IOAutoCloseable.noOp())) {
       assertThrows(
           IllegalArgumentException.class,
           () -> {
@@ -41,4 +52,15 @@ public final class StorageDataClientTest {
           });
     }
   }
+
+  @Test
+  public void executorServiceProvidedShouldBeClosed() throws Exception {
+    assertThat(exec).isNotNull();
+    StorageDataClient sdc =
+        StorageDataClient.create(exec, Duration.ofSeconds(2), null, null, IOAutoCloseable.noOp());
+
+    sdc.close();
+    verify(exec, times(1)).shutdownNow();
+    verify(exec, times(1)).awaitTermination(TimeUnit.SECONDS.toNanos(2), TimeUnit.NANOSECONDS);
+  }
 }