test: reduce wait time for BackendExhaustedTest (#835)

olavloite · web-flow · commit 61ea2b940cad · 2021-02-01T10:06:03.000Z
The BackendExhaustedTest deliberately causes requests to get stuck on the mock server. This causes the Spanner instance to refuse to shutdown nicely and causes it to wait for 10 seconds. This can be circumvented by using shutdownNow() to force the closure of all gRPC transport channels. This reduces the overall build time with approx 10 seconds.
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.
     }
