Merge 441b1aa into 004c5d7

olavloite · web-flow · commit 95b40b5fc541 · 2021-02-01T09:23:45.000Z
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
@@ -25,6 +25,7 @@
 import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.cloud.spanner.SessionClient.SessionId;
 import com.google.cloud.spanner.SpannerOptions.CloseableExecutorProvider;
+import com.google.cloud.spanner.spi.v1.GapicSpannerRpc;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
 import com.google.cloud.spanner.spi.v1.SpannerRpc.Paginated;
 import com.google.common.annotations.VisibleForTesting;
@@ -269,7 +270,11 @@ void close(long timeout, TimeUnit unit) {
       sessionClients.clear();
       asyncExecutorProvider.close();
       try {
-        gapicRpc.shutdown();
+        if (timeout == Long.MAX_VALUE || !(gapicRpc instanceof GapicSpannerRpc)) {
+          gapicRpc.shutdown();
+        } else {
+          ((GapicSpannerRpc) gapicRpc).shutdownNow();
+        }
       } catch (RuntimeException e) {
         logger.log(Level.WARNING, "Failed to close channels", e);
       }
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
@@ -1718,6 +1718,22 @@ public void shutdown() {
     }
   }
 
+  public void shutdownNow() {
+    this.rpcIsClosed = true;
+    this.spannerStub.close();
+    this.partitionedDmlStub.close();
+    this.instanceAdminStub.close();
+    this.databaseAdminStub.close();
+    this.spannerWatchdog.shutdown();
+    this.executorProvider.shutdown();
+
+    this.spannerStub.shutdownNow();
+    this.partitionedDmlStub.shutdownNow();
+    this.instanceAdminStub.shutdownNow();
+    this.databaseAdminStub.shutdownNow();
+    this.spannerWatchdog.shutdownNow();
+  }
+
   @Override
   public boolean isClosed() {
     return rpcIsClosed;
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BackendExhaustedTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BackendExhaustedTest.java
@@ -136,7 +136,6 @@ public void setUp() throws Exception {
                 SessionPoolOptions.newBuilder()
                     .setMinSessions(executor.getCorePoolSize())
                     .setMaxSessions(executor.getCorePoolSize() * 3)
-                    .setWriteSessionsFraction(0.0f)
                     .build())
             .build();
     executorFactory.release(executor);
@@ -159,7 +158,7 @@ public void tearDown() {
     // This test case force-closes the Spanner instance as it would otherwise wait
     // forever on the BatchCreateSessions requests that are 'stuck'.
     try {
-      ((SpannerImpl) spanner).close(100L, TimeUnit.MILLISECONDS);
+      ((SpannerImpl) spanner).close(10L, TimeUnit.MILLISECONDS);
     } catch (SpannerException e) {
       // ignore any errors during close as they are expected.
     }
