fix: clean up connection handler if connection fails (#3405)

Remove ConnectionHandler instances from the CONNECTION_HANDLERS map when the connection is
closed, regardless how the connection would be closed. Before this change, connection handlers
that failed at startup, for example connection attempts from simple TCP probers, would not
be cleaned up, and would cause a memory leak if repated TCP probes are executed.

Updates #3402

Author: olavloite

src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java

@@ -14,6 +14,8 @@
 
 package com.google.cloud.spanner.pgadapter;
 
+import static com.google.cloud.spanner.pgadapter.ProxyServer.CONNECTION_HANDLERS;
+
 import com.google.api.core.InternalApi;
 import com.google.auth.Credentials;
 import com.google.cloud.spanner.Database;
@@ -90,7 +92,6 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.UUID;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
@@ -123,8 +124,6 @@ public class ConnectionHandler implements Runnable {
           .concurrencyLevel(1)
           .build();
   private final Map<String, IntermediatePortalStatement> portalsMap = new HashMap<>();
-  private static final Map<Integer, ConnectionHandler> CONNECTION_HANDLERS =
-      new ConcurrentHashMap<>();
   private volatile ConnectionStatus status = ConnectionStatus.UNAUTHENTICATED;
   private Thread thread;
   private final int connectionId;
@@ -165,7 +164,6 @@ public class ConnectionHandler implements Runnable {
     this.socket = socket;
     this.secret = new SecureRandom().nextInt();
     this.connectionId = incrementingConnectionId.incrementAndGet();
-    CONNECTION_HANDLERS.put(this.connectionId, this);
     this.spannerConnection = spannerConnection;
   }
 
@@ -384,6 +382,7 @@ public void run() {
                 String.format(
                     "Connection handler with ID %s starting for client %s with thread %s",
                     getName(), socket.getInetAddress().getHostAddress(), thread.toString())));
+    CONNECTION_HANDLERS.put(this.connectionId, this);
     if (runConnection(false) == RunConnectionState.RESTART_WITH_SSL) {
       logger.log(
           Level.INFO,
@@ -522,6 +521,7 @@ && getServer().getOptions().getSslMode().isSslEnabled()
                       String.format(
                           "Exception while closing connection handler with ID %s", getName())));
         } finally {
+          CONNECTION_HANDLERS.remove(this.connectionId);
           this.server.deregister(this);
           logger.log(
               Level.INFO,

src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java

@@ -28,6 +28,7 @@
 import com.google.cloud.spanner.pgadapter.utils.Metrics;
 import com.google.cloud.spanner.pgadapter.wireprotocol.MessageReader;
 import com.google.cloud.spanner.pgadapter.wireprotocol.WireMessage;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Tracer;
@@ -43,6 +44,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
@@ -65,6 +67,10 @@
 public class ProxyServer extends AbstractApiService {
 
   private static final Logger logger = Logger.getLogger(ProxyServer.class.getName());
+
+  @VisibleForTesting
+  static final Map<Integer, ConnectionHandler> CONNECTION_HANDLERS = new ConcurrentHashMap<>();
+
   private final OptionsMetadata options;
   private final OpenTelemetry openTelemetry;
   private final Metrics metrics;

src/test/java/com/google/cloud/spanner/pgadapter/ConnectionHandlerTest.java

@@ -18,6 +18,7 @@
 import static com.google.cloud.spanner.pgadapter.ConnectionHandler.buildConnectionURL;
 import static com.google.cloud.spanner.pgadapter.ConnectionHandler.listDatabasesOrInstances;
 import static com.google.cloud.spanner.pgadapter.EmulatedPsqlMockServerTest.newStatusResourceNotFoundException;
+import static com.google.cloud.spanner.pgadapter.ProxyServer.CONNECTION_HANDLERS;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -329,30 +330,35 @@ public void testCancelActiveStatement() {
     ConnectionHandler connectionHandlerToCancel =
         new ConnectionHandler(server, socket, spannerConnection);
     connectionHandlerToCancel.setThread(mock(Thread.class));
+    CONNECTION_HANDLERS.put(connectionHandlerToCancel.getConnectionId(), connectionHandlerToCancel);
 
-    // Cancelling yourself is not allowed.
-    assertFalse(
-        connectionHandler.cancelActiveStatement(
-            connectionHandler.getConnectionId(), connectionHandler.getSecret()));
-    // Cancelling a random non-existing connection should not work.
-    assertFalse(connectionHandler.cancelActiveStatement(100, 100));
-    // Cancelling another connecting using the wrong secret is not allowed.
-    assertFalse(
-        connectionHandler.cancelActiveStatement(
-            connectionHandlerToCancel.getConnectionId(),
-            connectionHandlerToCancel.getSecret() - 1));
-
-    assertTrue(
-        connectionHandler.cancelActiveStatement(
-            connectionHandlerToCancel.getConnectionId(), connectionHandlerToCancel.getSecret()));
-
-    // The method should just return false if an error occurs.
-    doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.INTERNAL, "test error"))
-        .when(spannerConnection)
-        .cancel();
-    assertFalse(
-        connectionHandler.cancelActiveStatement(
-            connectionHandlerToCancel.getConnectionId(), connectionHandlerToCancel.getSecret()));
+    try {
+      // Cancelling yourself is not allowed.
+      assertFalse(
+          connectionHandler.cancelActiveStatement(
+              connectionHandler.getConnectionId(), connectionHandler.getSecret()));
+      // Cancelling a random non-existing connection should not work.
+      assertFalse(connectionHandler.cancelActiveStatement(100, 100));
+      // Cancelling another connecting using the wrong secret is not allowed.
+      assertFalse(
+          connectionHandler.cancelActiveStatement(
+              connectionHandlerToCancel.getConnectionId(),
+              connectionHandlerToCancel.getSecret() - 1));
+
+      assertTrue(
+          connectionHandler.cancelActiveStatement(
+              connectionHandlerToCancel.getConnectionId(), connectionHandlerToCancel.getSecret()));
+
+      // The method should just return false if an error occurs.
+      doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.INTERNAL, "test error"))
+          .when(spannerConnection)
+          .cancel();
+      assertFalse(
+          connectionHandler.cancelActiveStatement(
+              connectionHandlerToCancel.getConnectionId(), connectionHandlerToCancel.getSecret()));
+    } finally {
+      CONNECTION_HANDLERS.remove(connectionHandlerToCancel.getConnectionId());
+    }
   }
 
   @Test

src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java

@@ -14,6 +14,7 @@
 
 package com.google.cloud.spanner.pgadapter;
 
+import static com.google.cloud.spanner.pgadapter.ProxyServer.CONNECTION_HANDLERS;
 import static com.google.cloud.spanner.pgadapter.statements.BackendConnection.TRANSACTION_ABORTED_ERROR;
 import static com.google.cloud.spanner.pgadapter.statements.PgCatalog.PgNamespace.PG_NAMESPACE_CTE;
 import static org.junit.Assert.assertArrayEquals;
@@ -649,6 +650,12 @@ public void testSelectHelloWorld() throws SQLException {
         }
       }
     }
+    Stopwatch stopwatch = Stopwatch.createStarted();
+    // The cleanup of the connection is async, so it can take a few milliseconds to reach the state
+    // that we expect.
+    //noinspection StatementWithEmptyBody
+    while (stopwatch.elapsed(TimeUnit.SECONDS) < 3 && !CONNECTION_HANDLERS.isEmpty()) {}
+    assertEquals(0, CONNECTION_HANDLERS.size());
   }
 
   @Test

src/test/java/com/google/cloud/spanner/pgadapter/ProxyServerTest.java

@@ -14,13 +14,17 @@
 
 package com.google.cloud.spanner.pgadapter;
 
+import static com.google.cloud.spanner.pgadapter.ProxyServer.CONNECTION_HANDLERS;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.TestOptionsMetadataBuilder;
+import com.google.common.base.Stopwatch;
 import java.net.Socket;
 import java.time.Duration;
+import java.util.concurrent.TimeUnit;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -30,16 +34,24 @@ public class ProxyServerTest {
 
   @Test
   public void testProbeConnection() throws Exception {
+    assertEquals(0, CONNECTION_HANDLERS.size());
     // This test verifies that doing a simple TCP aliveness probe to check whether PGAdapter is
     // running can be done without any errors.
     ProxyServer server = new ProxyServer(OptionsMetadata.newBuilder().setPort(0).build());
     server.startServer();
     server.awaitRunning();
 
-    //noinspection EmptyTryBlock
     try (Socket ignore = new Socket("localhost", server.getLocalPort())) {
       // Do nothing, just verify that we can connect without any errors.
+      Stopwatch stopwatch = Stopwatch.createStarted();
+      //noinspection StatementWithEmptyBody
+      while (stopwatch.elapsed(TimeUnit.SECONDS) < 3 && CONNECTION_HANDLERS.isEmpty()) {}
+      assertEquals(1, CONNECTION_HANDLERS.size());
     }
+    Stopwatch stopwatch = Stopwatch.createStarted();
+    //noinspection StatementWithEmptyBody
+    while (stopwatch.elapsed(TimeUnit.SECONDS) < 3 && !CONNECTION_HANDLERS.isEmpty()) {}
+    assertEquals(0, CONNECTION_HANDLERS.size());
     server.stopServer();
     server.awaitTerminated();
   }