test: force-close to speed up timeout tests (#840)

* test: force-close to speed up timeout tests

* fix: mark connection as closed before waiting

* fix: set closed before shutting down

* cleanup: add test cases + remove unused code

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

@@ -259,11 +259,17 @@ public void close() {
             // Ignore as we are closing the connection.
           }
         }
-        statementExecutor.shutdownNow();
-        spannerPool.removeConnection(options, this);
+        // Try to wait for the current statement to finish (if any) before we actually close the
+        // connection.
+        this.closed = true;
+        statementExecutor.shutdown();
         leakedException = null;
+        spannerPool.removeConnection(options, this);
+        statementExecutor.awaitTermination(10L, TimeUnit.SECONDS);
+      } catch (InterruptedException e) {
+        // ignore and continue to close the connection.
       } finally {
-        this.closed = true;
+        statementExecutor.shutdownNow();
       }
     }
   }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java

@@ -25,6 +25,7 @@
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.SpannerOptions;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicates;
@@ -62,6 +63,15 @@ public class SpannerPool {
   private static final String CONNECTION_API_CLIENT_LIB_TOKEN = "sp-jdbc";
   private static final Logger logger = Logger.getLogger(SpannerPool.class.getName());
 
+  private static final Function<Spanner, Void> DEFAULT_CLOSE_FUNCTION =
+      new Function<Spanner, Void>() {
+        @Override
+        public Void apply(Spanner spanner) {
+          spanner.close();
+          return null;
+        }
+      };
+
   /**
    * Closes the default {@link SpannerPool} and all {@link Spanner} instances that have been opened
    * by connections and that are still open. Call this method at the end of your application to
@@ -395,6 +405,12 @@ void checkAndCloseSpanners() {
 
   @VisibleForTesting
   void checkAndCloseSpanners(CheckAndCloseSpannersMode mode) {
+    checkAndCloseSpanners(mode, DEFAULT_CLOSE_FUNCTION);
+  }
+
+  @VisibleForTesting
+  void checkAndCloseSpanners(
+      CheckAndCloseSpannersMode mode, Function<Spanner, Void> closeSpannerFunction) {
     List<SpannerPoolKey> keysStillInUse = new ArrayList<>();
     synchronized (this) {
       for (Entry<SpannerPoolKey, Spanner> entry : spanners.entrySet()) {
@@ -416,7 +432,7 @@ void checkAndCloseSpanners(CheckAndCloseSpannersMode mode) {
           // Force close all Spanner instances by passing in a value that will always be less than
           // the
           // difference between the current time and the close time of a connection.
-          closeUnusedSpanners(Long.MIN_VALUE);
+          closeUnusedSpanners(Long.MIN_VALUE, closeSpannerFunction);
         } else {
           logLeakedConnections(keysStillInUse);
           throw SpannerExceptionFactory.newSpannerException(
@@ -456,6 +472,11 @@ private void logLeakedConnections(List<SpannerPoolKey> keysStillInUse) {
    */
   @VisibleForTesting
   void closeUnusedSpanners(long closeSpannerAfterMillisecondsUnused) {
+    closeUnusedSpanners(closeSpannerAfterMillisecondsUnused, DEFAULT_CLOSE_FUNCTION);
+  }
+
+  void closeUnusedSpanners(
+      long closeSpannerAfterMillisecondsUnused, Function<Spanner, Void> closeSpannerFunction) {
     List<SpannerPoolKey> keysToBeRemoved = new ArrayList<>();
     synchronized (this) {
       for (Entry<SpannerPoolKey, Long> entry : lastConnectionClosedAt.entrySet()) {
@@ -469,7 +490,9 @@ void closeUnusedSpanners(long closeSpannerAfterMillisecondsUnused) {
           Spanner spanner = spanners.get(entry.getKey());
           if (spanner != null) {
             try {
-              spanner.close();
+              closeSpannerFunction.apply(spanner);
+            } catch (Throwable t) {
+              // Ignore any errors and continue with the next one in the pool.
             } finally {
               // Even if the close operation failed, we should remove the spanner object as it is no
               // longer valid.
@@ -484,11 +507,4 @@ void closeUnusedSpanners(long closeSpannerAfterMillisecondsUnused) {
       }
     }
   }
-
-  @VisibleForTesting
-  int getCurrentSpannerCount() {
-    synchronized (this) {
-      return spanners.size();
-    }
-  }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementExecutor.java

@@ -164,6 +164,14 @@ private static ListeningExecutorService createExecutorService() {
     this.interceptors = Collections.unmodifiableList(interceptors);
   }
 
+  void shutdown() {
+    executor.shutdown();
+  }
+
+  void awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
+    executor.awaitTermination(timeout, unit);
+  }
+
   /**
    * Shutdown this executor now and do not wait for any statement that is being executed to finish.
    */

google-cloud-spanner/src/test/java/com/google/cloud/spanner/ForceCloseSpannerFunction.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner;
+
+import com.google.common.base.Function;
+import java.util.concurrent.TimeUnit;
+
+/** Class for tests that need to be able to force-close a {@link Spanner} instance. */
+public class ForceCloseSpannerFunction implements Function<Spanner, Void> {
+  private final long timeout;
+  private final TimeUnit unit;
+
+  public ForceCloseSpannerFunction(long timeout, TimeUnit unit) {
+    this.timeout = timeout;
+    this.unit = unit;
+  }
+
+  public Void apply(Spanner spanner) {
+    if (spanner instanceof SpannerImpl) {
+      ((SpannerImpl) spanner).close(timeout, unit);
+    } else {
+      spanner.close();
+    }
+    return null;
+  }
+}

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java

@@ -16,13 +16,16 @@
 
 package com.google.cloud.spanner.connection;
 
+import com.google.cloud.spanner.ForceCloseSpannerFunction;
 import com.google.cloud.spanner.MockSpannerServiceImpl;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.RandomResultSetGenerator;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.admin.database.v1.MockDatabaseAdminImpl;
 import com.google.cloud.spanner.admin.instance.v1.MockInstanceAdminImpl;
 import com.google.cloud.spanner.connection.ITAbstractSpannerTest.AbortInterceptor;
 import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection;
+import com.google.cloud.spanner.connection.SpannerPool.CheckAndCloseSpannersMode;
 import com.google.common.util.concurrent.AbstractFuture;
 import com.google.longrunning.GetOperationRequest;
 import com.google.longrunning.Operation;
@@ -49,6 +52,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 import org.junit.After;
 import org.junit.AfterClass;
@@ -113,6 +117,7 @@ public abstract class AbstractMockServerTest {
   private boolean futureParentHandlers;
   private boolean exceptionRunnableParentHandlers;
   private boolean nettyServerParentHandlers;
+  private boolean clientStreamParentHandlers;
 
   @BeforeClass
   public static void startStaticServer() throws IOException {
@@ -152,9 +157,7 @@ public void getOperation(
 
   @AfterClass
   public static void stopServer() throws Exception {
-    SpannerPool.closeSpannerPool();
     server.shutdown();
-    server.awaitTermination();
   }
 
   @Before
@@ -169,22 +172,30 @@ public void setupResults() {
     nettyServerParentHandlers =
         Logger.getLogger("io.grpc.netty.shaded.io.grpc.netty.NettyServerHandler")
             .getUseParentHandlers();
+    clientStreamParentHandlers =
+        Logger.getLogger("io.grpc.netty.shaded.io.grpc.netty.NettyServerHandler")
+            .getUseParentHandlers();
     Logger.getLogger(AbstractFuture.class.getName()).setUseParentHandlers(false);
     Logger.getLogger(LogExceptionRunnable.class.getName()).setUseParentHandlers(false);
     Logger.getLogger("io.grpc.netty.shaded.io.grpc.netty.NettyServerHandler")
         .setUseParentHandlers(false);
+    Logger.getLogger("io.grpc.internal.AbstractClientStream").setUseParentHandlers(false);
   }
 
   @After
   public void closeSpannerPool() {
     try {
-      SpannerPool.closeSpannerPool();
+      SpannerPool.INSTANCE.checkAndCloseSpanners(
+          CheckAndCloseSpannersMode.ERROR,
+          new ForceCloseSpannerFunction(100L, TimeUnit.MILLISECONDS));
     } finally {
       Logger.getLogger(AbstractFuture.class.getName()).setUseParentHandlers(futureParentHandlers);
       Logger.getLogger(LogExceptionRunnable.class.getName())
           .setUseParentHandlers(exceptionRunnableParentHandlers);
       Logger.getLogger("io.grpc.netty.shaded.io.grpc.netty.NettyServerHandler")
           .setUseParentHandlers(nettyServerParentHandlers);
+      Logger.getLogger("io.grpc.internal.AbstractClientStream")
+          .setUseParentHandlers(clientStreamParentHandlers);
     }
   }
 

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java

@@ -17,12 +17,15 @@
 package com.google.cloud.spanner.connection;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.auth.Credentials;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SessionPoolOptions;
@@ -430,8 +433,9 @@ public void testSpannerPoolKeyEquality() {
     ConnectionOptions options1 =
         ConnectionOptions.newBuilder()
             .setUri(
-                "cloudspanner:/projects/p/instances/i/databases/d?minSessions=200;maxSessions=400")
-            .setCredentials(NoCredentials.getInstance())
+                "cloudspanner://localhost:9010/projects/p1/instances/i/databases/d"
+                    + "?minSessions=200;maxSessions=400;numChannels=8;usePlainText=true;userAgent=test-agent")
+            .setCredentials(mock(Credentials.class))
             .build();
     // options2 equals the default session pool options, and is therefore equal to ConnectionOptions
     // without any session pool configuration.
@@ -451,8 +455,9 @@ public void testSpannerPoolKeyEquality() {
     SpannerPoolKey key2 = SpannerPoolKey.of(options2);
     SpannerPoolKey key3 = SpannerPoolKey.of(options3);
 
-    assertThat(key1).isNotEqualTo(key2);
-    assertThat(key2).isEqualTo(key3);
-    assertThat(key1).isNotEqualTo(key3);
+    assertFalse(key1.equals(key2));
+    assertTrue(key2.equals(key3));
+    assertFalse(key1.equals(key3));
+    assertFalse(key1.equals(new Object()));
   }
 }