Repository: GoogleCloudPlatform/pgadapter
Pull request #2272 (2024-09-05 19:02:33)

deps: update dependency sqlalchemy to v1.4.54
This PR contains the following updates:

| Package | Change | Age | Adoption | Passing | Confidence |
|---|---|---|---|---|---|
| [sqlalchemy](https://www.sqlalchemy.org) ([changelog](https://docs.sqlalchemy.org/en/latest/changelog/)) | `==1.4.53` -> `==1.4.54` | [![age](https://developer.mend.io/api/mc/badges/age/pypi/sqlalchemy/1.4.54?slim=true)](https://docs.renovatebot.com/merge-confidence/) | [![adoption](https://developer.mend.io/api/mc/badges/adoption/pypi/sqlalchemy/1.4.54?slim=true)](https://docs.renovatebot.com/merge-confidence/) | [![passing](https://developer.mend.io/api/mc/badges/compatibility/pypi/sqlalchemy/1.4.53/1.4.54?slim=true)](https://docs.renovatebot.com/merge-confidence/) | [![confidence](https://developer.mend.io/api/mc/badges/confidence/pypi/sqlalchemy/1.4.53/1.4.54?slim=true)](https://docs.renovatebot.com/merge-confidence/) |

---

### Configuration

📅 **Schedule**: Branch creation - At any time (no schedule defined), Automerge - At any time (no schedule defined).

🚦 **Automerge**: Disabled by config. Please merge this manually once you are satisfied.

♻ **Rebasing**: Whenever PR becomes conflicted, or you tick the rebase/retry checkbox.

🔕 **Ignore**: Close this PR and you won't be reminded about this update again.

---

 - [ ] <!-- rebase-check -->If you want to rebase/retry this PR, check this box

---

This PR was generated by [Mend Renovate](https://mend.io/renovate/). View the [repository job log](https://developer.mend.io/github/GoogleCloudPlatform/pgadapter).
<!--renovate-debug:eyJjcmVhdGVkSW5WZXIiOiIzOC41OS4yIiwidXBkYXRlZEluVmVyIjoiMzguNTkuMiIsInRhcmdldEJyYW5jaCI6InBvc3RncmVzcWwtZGlhbGVjdCIsImxhYmVscyI6W119-->



Patch:
diff --git a/pom.xml b/pom.xml
index 9589ab0d..53093413 100644
--- a/pom.xml
+++ b/pom.xml
@@ -214,6 +214,12 @@
       <artifactId>bcpkix-jdk18on</artifactId>
       <version>1.78.1</version>
     </dependency>
+    <dependency>
+      <groupId>org.codehaus.mojo</groupId>
+      <artifactId>animal-sniffer-annotations</artifactId>
+      <version>1.24</version>
+      <optional>true</optional>
+    </dependency>
 
     <!-- Test dependencies -->
     <dependency>
diff --git a/samples/python/sqlalchemy-sample/requirements.txt b/samples/python/sqlalchemy-sample/requirements.txt
index 01788fd8..bedc5d5a 100644
--- a/samples/python/sqlalchemy-sample/requirements.txt
+++ b/samples/python/sqlalchemy-sample/requirements.txt
@@ -1,3 +1,3 @@
 psycopg2-binary~=2.9.3
 pytz~=2024.1
-sqlalchemy==1.4.53
+sqlalchemy==1.4.54
diff --git a/samples/sidecar-proxy/Dockerfile b/samples/sidecar-proxy/Dockerfile
index c7fe4fe7..f7e2b85d 100644
--- a/samples/sidecar-proxy/Dockerfile
+++ b/samples/sidecar-proxy/Dockerfile
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-FROM golang:1.23.0 as builder
+FROM golang:1.23.1 as builder
 WORKDIR /app
 RUN go mod init example-app
 RUN go get github.com/jackc/pgx/v4
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
index a4caf07e..918109a6 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ConnectionHandler.java
@@ -457,13 +457,14 @@ public class ConnectionHandler implements Runnable {
                   () ->
                       String.format(
                           "Exception while closing connection handler with ID %s", getName())));
+        } finally {
+          this.server.deregister(this);
+          logger.log(
+              Level.INFO,
+              Logging.format(
+                  "RunConnection",
+                  () -> String.format("Connection handler with ID %s closed", getName())));
         }
-        this.server.deregister(this);
-        logger.log(
-            Level.INFO,
-            Logging.format(
-                "RunConnection",
-                () -> String.format("Connection handler with ID %s closed", getName())));
       }
     }
     return result;
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java b/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
index 9dbf2f1c..da4760aa 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ProxyServer.java
@@ -44,9 +44,12 @@ import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.newsclub.net.unix.AFUNIXServerSocket;
@@ -63,7 +66,7 @@ public class ProxyServer extends AbstractApiService {
   private final OpenTelemetry openTelemetry;
   private final Metrics metrics;
   private final Properties properties;
-  private final List<ConnectionHandler> handlers = Collections.synchronizedList(new LinkedList<>());
+  private final List<ConnectionHandler> handlers = new LinkedList<>();
 
   /**
    * Latch that is closed when the TCP server has started. We need this to know the exact port that
@@ -88,6 +91,34 @@ public class ProxyServer extends AbstractApiService {
 
   private final ThreadFactory threadFactory;
 
+  /**
+   * Shutdown mode determines how the server handles any open connections when the server is being
+   * terminated. These modes correspond with the <a
+   * href="https://www.postgresql.org/docs/current/server-shutdown.html">shutdown modes of
+   * PostgreSQL</a>.
+   */
+  public enum ShutdownMode {
+    /**
+     * Stop accepting incoming connections and wait for all existing connections to finish before
+     * shutting down.
+     */
+    SMART,
+
+    /** Actively terminate all existing connections and shut down the server. */
+    FAST,
+
+    /**
+     * Actively terminate all existing connections and shut down the server. There is currently no
+     * difference between FAST and IMMEDIATE, but this could change in the future.
+     */
+    IMMEDIATE,
+  }
+
+  private final AtomicReference<ShutdownMode> shutdownMode = new AtomicReference<>();
+
+  private final AtomicReference<CountDownLatch> allHandlersTerminatedLatch =
+      new AtomicReference<>();
+
   /**
    * Instantiates the ProxyServer from CLI-gathered metadata.
    *
@@ -213,6 +244,7 @@ public class ProxyServer extends AbstractApiService {
 
   @Override
   protected void doStop() {
+    logger.log(Level.INFO, "Stopping server using shutdown mode {0}", this.shutdownMode.get());
     for (ServerSocket serverSocket : this.serverSockets) {
       try {
         logger.log(
@@ -227,22 +259,72 @@ public class ProxyServer extends AbstractApiService {
             () -> String.format("Closing server socket %s failed: %s", serverSocket, exception));
       }
     }
-    for (ConnectionHandler handler : getConnectionHandlers()) {
-      handler.terminate();
+    if (this.shutdownMode.get() == ShutdownMode.SMART) {
+      try {
+        waitForAllConnectionsToTerminate();
+      } catch (InterruptedException interruptedException) {
+        // The waiting thread will be interrupted if a new shutdown request with mode FAST or
+        // IMMEDIATE comes in. In that case, we actively terminate all connection handlers.
+        terminateAllConnectionHandlers();
+      }
+    } else {
+      terminateAllConnectionHandlers();
     }
     try {
-      SpannerPool.closeSpannerPool();
+      // Do a shoot-and-forget close of the underlying Spanner pool.
+      ExecutorService service = Executors.newSingleThreadExecutor(threadFactory);
+      service.submit(SpannerPool::closeSpannerPool);
+      service.shutdown();
+      if (shutdownMode.get() != ShutdownMode.IMMEDIATE) {
+        if (!service.awaitTermination(1L, TimeUnit.SECONDS)) {
+          service.shutdownNow();
+          logger.log(Level.INFO, "SpannerPool was not closed after waiting for 1 second");
+        }
+      }
     } catch (Throwable ignore) {
     }
+    if (openTelemetry instanceof Closeable) {
+      try {
+        ((Closeable) openTelemetry).close();
+      } catch (IOException ioException) {
+        logger.log(Level.WARNING, "Failed to close OpenTelemetry", ioException);
+      }
+    }
     notifyStopped();
   }
 
-  /** Safely stops the server (iff started), closing specific socket and cleaning up. */
+  private void terminateAllConnectionHandlers() {
+    logger.log(Level.INFO, "Terminating {0} connections", getNumberOfConnections());
+    for (ConnectionHandler handler : getConnectionHandlers()) {
+      handler.terminate();
+    }
+  }
+
+  private void waitForAllConnectionsToTerminate() throws InterruptedException {
+    logger.log(Level.INFO, "Waiting for {0} connections to terminate", getNumberOfConnections());
+    createConnectionHandlersTerminatedLatch();
+    this.allHandlersTerminatedLatch.get().await();
+  }
+
+  /**
+   * Stops the server using the FAST shutdown mode, terminating all existing connections and
+   * cleaning up.
+   */
   public void stopServer() {
+    stopServer(ShutdownMode.FAST);
+  }
+
+  /** Stops the server using the specified shutdown mode. */
+  public void stopServer(ShutdownMode shutdownMode) {
+    setShutdownMode(shutdownMode);
     stopAsync();
     awaitTerminated();
   }
 
+  void setShutdownMode(ShutdownMode shutdownMode) {
+    this.shutdownMode.set(shutdownMode);
+  }
+
   /**
    * Thread logic: opens the TCP listening socket and instantiates the connection handler.
    *
@@ -313,18 +395,15 @@ public class ProxyServer extends AbstractApiService {
         createConnectionHandler(serverSocket.accept());
       }
     } catch (SocketException e) {
-      // This is a normal exception, as this will occur when Server#stopServer() is called.
+      // This is a normal and expected exception when the server is shutting down.
       logger.log(
-          Level.INFO,
+          shutdownMode.get() == null ? Level.WARNING : Level.FINEST,
           () ->
               String.format(
                   "Socket exception on socket %s: %s. This is normal when the server is stopped.",
                   serverSocket, e));
     } finally {
       logger.log(Level.INFO, () -> String.format("Socket %s stopped", serverSocket));
-      if (openTelemetry instanceof Closeable) {
-        ((Closeable) openTelemetry).close();
-      }
       stoppedLatch.countDown();
     }
   }
@@ -350,7 +429,15 @@ public class ProxyServer extends AbstractApiService {
 
   /** Returns an immutable copy of the current connection handlers at this server. */
   ImmutableList<ConnectionHandler> getConnectionHandlers() {
-    return ImmutableList.copyOf(this.handlers);
+    synchronized (this.handlers) {
+      return ImmutableList.copyOf(this.handlers);
+    }
+  }
+
+  private void createConnectionHandlersTerminatedLatch() {
+    synchronized (this.handlers) {
+      this.allHandlersTerminatedLatch.set(new CountDownLatch(this.handlers.size()));
+    }
   }
 
   /**
@@ -359,7 +446,9 @@ public class ProxyServer extends AbstractApiService {
    * @param handler The handler currently in use.
    */
   private void register(ConnectionHandler handler) {
-    this.handlers.add(handler);
+    synchronized (this.handlers) {
+      this.handlers.add(handler);
+    }
   }
 
   /**
@@ -368,7 +457,12 @@ public class ProxyServer extends AbstractApiService {
    * @param handler The handler to revoke.
    */
   void deregister(ConnectionHandler handler) {
-    this.handlers.remove(handler);
+    synchronized (this.handlers) {
+      this.handlers.remove(handler);
+      if (this.allHandlersTerminatedLatch.get() != null) {
+        this.allHandlersTerminatedLatch.get().countDown();
+      }
+    }
   }
 
   public OptionsMetadata getOptions() {
@@ -396,7 +490,9 @@ public class ProxyServer extends AbstractApiService {
 
   /** @return the current number of connections. */
   public int getNumberOfConnections() {
-    return this.handlers.size();
+    synchronized (this.handlers) {
+      return this.handlers.size();
+    }
   }
 
   /** @return the local TCP port that this server is using. */
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/Server.java b/src/main/java/com/google/cloud/spanner/pgadapter/Server.java
index edf41253..958a9a3b 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/Server.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/Server.java
@@ -21,6 +21,7 @@ import com.google.cloud.opentelemetry.metric.GoogleCloudMetricExporter;
 import com.google.cloud.opentelemetry.metric.MetricConfiguration;
 import com.google.cloud.opentelemetry.trace.TraceConfiguration;
 import com.google.cloud.opentelemetry.trace.TraceExporter;
+import com.google.cloud.spanner.pgadapter.ProxyServer.ShutdownMode;
 import com.google.cloud.spanner.pgadapter.logging.DefaultLogConfiguration;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.common.collect.ImmutableMap;
@@ -40,15 +41,28 @@ import java.io.File;
 import java.io.FileReader;
 import java.io.IOException;
 import java.io.PrintStream;
+import java.lang.invoke.CallSite;
+import java.lang.invoke.LambdaMetafactory;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
 import java.time.Duration;
 import java.util.Locale;
 import java.util.Objects;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import org.apache.maven.model.Model;
 import org.apache.maven.model.io.xpp3.MavenXpp3Reader;
+import org.codehaus.mojo.animal_sniffer.IgnoreJRERequirement;
 
 /** Effectively this is the main class */
 public class Server {
+  private static final Logger logger = Logger.getLogger(Server.class.getName());
+
+  private static ShutdownHandler shutdownHandler;
 
   /**
    * Main method for running a Spanner PostgreSQL Adapter {@link Server} as a stand-alone
@@ -59,13 +73,142 @@ public class Server {
       DefaultLogConfiguration.configureLogging(args);
       OptionsMetadata optionsMetadata = extractMetadata(args, System.out);
       OpenTelemetry openTelemetry = setupOpenTelemetry(optionsMetadata);
-      ProxyServer server = new ProxyServer(optionsMetadata, openTelemetry);
-      server.startServer();
+      ProxyServer proxyServer = new ProxyServer(optionsMetadata, openTelemetry);
+      proxyServer.startServer();
+
+      // Create a shutdown handler and register signal handlers for the signals that should
+      // terminate the server.
+      Server.shutdownHandler = ShutdownHandler.createForServer(proxyServer);
+      registerSignalHandlers();
     } catch (Exception e) {
       printError(e, System.err, System.out);
     }
   }
 
+  /**
+   * Registers signal handlers for TERM, INT, and QUIT. This method uses reflection and fails
+   * gracefully if signal handling is not available on this JVM.
+   */
+  static void registerSignalHandlers() {
+    Class<?> signalClass = getSignalClass();
+    Class<?> signalHandlerClass = getSignalHandlerClass();
+    if (signalClass == null || signalHandlerClass == null) {
+      return;
+    }
+    try {
+      Method handleMethod =
+          signalClass.getDeclaredMethod("handle", signalClass, signalHandlerClass);
+      Constructor<?> signalConstructor = signalClass.getConstructor(String.class);
+
+      Object termSignal = signalConstructor.newInstance("TERM");
+      Object intSignal = signalConstructor.newInstance("INT");
+      Object quitSignal = signalConstructor.newInstance("QUIT");
+
+      Object termHandler = createSignalHandler(signalClass, signalHandlerClass, "handleTerm");
+      Object intHandler = createSignalHandler(signalClass, signalHandlerClass, "handleInt");
+      Object quitHandler = createSignalHandler(signalClass, signalHandlerClass, "handleQuit");
+
+      // Register the actual signal handlers. This will fail if the JVM has already registered a
+      // signal handler for the given signal. Normally, TERM and INT can be registered. QUIT is
+      // normally registered by the JVM.
+      try {
+        handleMethod.invoke(null, termSignal, termHandler);
+      } catch (Throwable throwable) {
+        logger.log(Level.WARNING, "Failed to register signal handler for TERM", throwable);
+      }
+      try {
+        handleMethod.invoke(null, intSignal, intHandler);
+      } catch (Throwable throwable) {
+        logger.log(Level.WARNING, "Failed to register signal handler for INT", throwable);
+      }
+      try {
+        handleMethod.invoke(null, quitSignal, quitHandler);
+      } catch (Throwable throwable) {
+        // Log this at FINE level, as QUIT is normally already registered by the JVM. This means
+        // that registering QUIT will fail on most JVMs. This means that QUIT signals will still
+        // lead to the server stopping quickly, just that it is handled directly by the JVM instead
+        // of PGAdapter.
+        logger.log(Level.FINE, "Failed to register signal handler for QUIT", throwable);
+      }
+    } catch (Throwable exception) {
+      logger.log(Level.WARNING, "Failed to register signal handlers", exception);
+    }
+  }
+
+  private static Class<?> getSignalClass() {
+    try {
+      return Class.forName("sun.misc.Signal");
+    } catch (ClassNotFoundException exception) {
+      logger.log(
+          Level.INFO,
+          "Cannot register shutdown signal handlers as sun.misc.Signal is not available on this JVM");
+      return null;
+    }
+  }
+
+  private static Class<?> getSignalHandlerClass() {
+    try {
+      return Class.forName("sun.misc.SignalHandler");
+    } catch (ClassNotFoundException exception) {
+      logger.log(
+          Level.INFO,
+          "Cannot register shutdown signal handlers as sun.misc.SignalHandler is not available on this JVM");
+      return null;
+    }
+  }
+
+  /**
+   * This method dynamically creates the signal handler lambdas that call the handleTerm, handleInt,
+   * and handleQuit methods. We need to do this through reflection, because we cannot use the
+   * sun.misc.SignalHandler interface directly to create the lambda expressions.
+   */
+  @IgnoreJRERequirement
+  private static Object createSignalHandler(
+      Class<?> signalClass, Class<?> signalHandlerClass, String handleMethodName) throws Throwable {
+    MethodHandles.Lookup caller = MethodHandles.lookup();
+    MethodType signalHandlerHandleMethodType = MethodType.methodType(void.class, signalClass);
+    MethodType invokedType = MethodType.methodType(signalHandlerClass);
+
+    CallSite handleTermSite =
+        LambdaMetafactory.metafactory(
+            caller,
+            "handle",
+            invokedType,
+            signalHandlerHandleMethodType,
+            caller.findStatic(
+                Server.class, handleMethodName, MethodType.methodType(void.class, Object.class)),
+            signalHandlerHandleMethodType);
+    MethodHandle factory = handleTermSite.getTarget();
+    return factory.invoke();
+  }
+
+  /** This method is called by the signal handler that is registered for TERM. */
+  static void handleTerm(Object ignoredSignal) {
+    if (Server.shutdownHandler == null) {
+      return;
+    }
+    logger.log(Level.INFO, "Server received TERM");
+    Server.shutdownHandler.shutdown(ShutdownMode.SMART);
+  }
+
+  /** This method is called by the signal handler that is registered for INT. */
+  static void handleInt(Object ignoredSignal) {
+    if (Server.shutdownHandler == null) {
+      return;
+    }
+    logger.log(Level.INFO, "Server received INT");
+    Server.shutdownHandler.shutdown(ShutdownMode.FAST);
+  }
+
+  /** This method is called by the signal handler that is registered for QUIT. */
+  static void handleQuit(Object ignoredSignal) {
+    if (Server.shutdownHandler == null) {
+      return;
+    }
+    logger.log(Level.INFO, "Server received QUIT");
+    Server.shutdownHandler.shutdown(ShutdownMode.IMMEDIATE);
+  }
+
   /** Creates an {@link OpenTelemetry} object from the given options. */
   static OpenTelemetry setupOpenTelemetry(OptionsMetadata optionsMetadata) {
     if (!optionsMetadata.isEnableOpenTelemetry()
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/ShutdownHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/ShutdownHandler.java
new file mode 100644
index 00000000..23c67d8b
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/ShutdownHandler.java
@@ -0,0 +1,61 @@
+// Copyright 2024 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.pgadapter;
+
+import com.google.cloud.spanner.pgadapter.ProxyServer.ShutdownMode;
+import com.google.common.base.Preconditions;
+import java.util.concurrent.atomic.AtomicReference;
+import javax.annotation.Nonnull;
+
+/**
+ * Handles the orderly shutdown of a {@link ProxyServer}. A {@link ShutdownHandler} can only be used
+ * once.
+ */
+public class ShutdownHandler {
+  private final ProxyServer proxyServer;
+  private final Thread shutdownThread;
+  private final AtomicReference<ShutdownMode> shutdownMode = new AtomicReference<>();
+
+  static ShutdownHandler createForServer(@Nonnull ProxyServer proxyServer) {
+    return new ShutdownHandler(proxyServer);
+  }
+
+  /** Constructor for a {@link ShutdownHandler} for a given {@link ProxyServer}. */
+  private ShutdownHandler(@Nonnull ProxyServer proxyServer) {
+    this.proxyServer = Preconditions.checkNotNull(proxyServer);
+    this.shutdownThread =
+        new Thread("server-shutdown-handler") {
+          @Override
+          public void run() {
+            ShutdownHandler.this.proxyServer.stopServer(ShutdownHandler.this.shutdownMode.get());
+          }
+        };
+  }
+
+  /** Shuts down the proxy server using the given {@link ShutdownMode}. */
+  public synchronized void shutdown(@Nonnull ShutdownMode shutdownMode) {
+    if (this.shutdownMode.get() != null) {
+      if (this.shutdownMode.get() == ShutdownMode.SMART
+          && (shutdownMode == ShutdownMode.FAST || shutdownMode == ShutdownMode.IMMEDIATE)) {
+        // Upgrade shutdown mode and interrupt the shutdown thread.
+        this.proxyServer.setShutdownMode(shutdownMode);
+        this.shutdownThread.interrupt();
+      }
+      return;
+    }
+    this.shutdownMode.set(Preconditions.checkNotNull(shutdownMode));
+    this.shutdownThread.start();
+  }
+}


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java
new file mode 100644
index 00000000..e3b58109
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java
@@ -0,0 +1,229 @@
+// Copyright 2024 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.pgadapter;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.core.ApiService.State;
+import com.google.cloud.NoCredentials;
+import com.google.cloud.spanner.pgadapter.ProxyServer.ShutdownMode;
+import com.google.cloud.spanner.pgadapter.metadata.TestOptionsMetadataBuilder;
+import com.google.common.base.Stopwatch;
+import io.opentelemetry.api.OpenTelemetry;
+import java.io.IOException;
+import java.net.Socket;
+import java.net.SocketException;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.concurrent.TimeUnit;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ShutdownModeMockServerTest extends AbstractMockServerTest {
+  private ProxyServer proxyServer;
+
+  @Before
+  public void createProxyServer() {
+    TestOptionsMetadataBuilder builder = new TestOptionsMetadataBuilder();
+    builder.setProject("p").setInstance("i").setDatabase("d");
+    builder
+        .enableDebugMode()
+        .setUsePlainText()
+        .setEndpoint(String.format("localhost:%d", spannerServer.getPort()))
+        .setCredentials(NoCredentials.getInstance());
+    proxyServer = new ProxyServer(builder.build(), OpenTelemetry.noop());
+    proxyServer.startServer();
+  }
+
+  @After
+  public void stopProxyServer() {
+    proxyServer.stopServer();
+  }
+
+  private String createUrl() {
+    return String.format("jdbc:postgresql://localhost:%d/d", proxyServer.getLocalPort());
+  }
+
+  @Test
+  public void testSmartShutdown() throws SQLException {
+    String sql = "SELECT 1";
+
+    ShutdownHandler shutdownHandler = ShutdownHandler.createForServer(proxyServer);
+    // Verify that the server is running.
+    assertEquals(State.RUNNING, proxyServer.state());
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      // Verify that the connection is valid.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+      // Initiate a smart shutdown.
+      shutdownHandler.shutdown(ShutdownMode.SMART);
+      // Wait until the server port is no longer open.
+      Stopwatch stopwatch = Stopwatch.createStarted();
+      while (isPortListening(proxyServer.getLocalPort())
+          && stopwatch.elapsed(TimeUnit.MILLISECONDS) < 2000) {
+        Thread.yield();
+      }
+      assertEquals(State.STOPPING, proxyServer.state());
+      // Verify that we can still use the existing connection.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+      // Verify that we cannot open a new connection.
+      SQLException exception =
+          assertThrows(SQLException.class, () -> DriverManager.getConnection(createUrl()));
+      assertEquals(
+          String.format(
+              "Connection to localhost:%d refused. Check that the hostname and port are correct and that the postmaster is accepting TCP/IP connections.",
+              proxyServer.getLocalPort()),
+          exception.getMessage());
+      // Verify that the original connection is still valid.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+      // Verify that the server is still stopping, but not yet terminated.
+      assertEquals(State.STOPPING, proxyServer.state());
+    }
+    Stopwatch stopwatch = Stopwatch.createStarted();
+    while (proxyServer.state() == State.STOPPING
+        && stopwatch.elapsed(TimeUnit.MILLISECONDS) < 2000) {
+      Thread.yield();
+    }
+    // Verify that closing the last connection terminates the server.
+    assertEquals(State.TERMINATED, proxyServer.state());
+  }
+
+  @Test
+  public void testFastShutdown() throws SQLException {
+    String sql = "SELECT 1";
+
+    ShutdownHandler shutdownHandler = ShutdownHandler.createForServer(proxyServer);
+    // Verify that the server is running.
+    assertEquals(State.RUNNING, proxyServer.state());
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      // Verify that the connection is valid.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+      // Initiate a fast shutdown.
+      shutdownHandler.shutdown(ShutdownMode.FAST);
+      // Verify that the server stops.
+      Stopwatch stopwatch = Stopwatch.createStarted();
+      while (proxyServer.state() != State.TERMINATED
+          && stopwatch.elapsed(TimeUnit.MILLISECONDS) < 2000) {
+        Thread.yield();
+      }
+      assertEquals(State.TERMINATED, proxyServer.state());
+
+      // Verify that existing connection has been invalidated.
+      SQLException exception =
+          assertThrows(SQLException.class, () -> connection.createStatement().executeQuery(sql));
+      assertEquals("An I/O error occurred while sending to the backend.", exception.getMessage());
+
+      // Verify that we cannot open a new connection.
+      exception = assertThrows(SQLException.class, () -> DriverManager.getConnection(createUrl()));
+      assertEquals(
+          String.format(
+              "Connection to localhost:%d refused. Check that the hostname and port are correct and that the postmaster is accepting TCP/IP connections.",
+              proxyServer.getLocalPort()),
+          exception.getMessage());
+    }
+  }
+
+  @Test
+  public void testSmartShutdownFollowedByFastShutdown() throws SQLException {
+    String sql = "SELECT 1";
+
+    ShutdownHandler shutdownHandler = ShutdownHandler.createForServer(proxyServer);
+    // Verify that the server is running.
+    assertEquals(State.RUNNING, proxyServer.state());
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      // Verify that the connection is valid.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+      // Initiate a smart shutdown.
+      shutdownHandler.shutdown(ShutdownMode.SMART);
+      // Wait until the server port is no longer open.
+      Stopwatch stopwatch = Stopwatch.createStarted();
+      while (isPortListening(proxyServer.getLocalPort())
+          && stopwatch.elapsed(TimeUnit.MILLISECONDS) < 2000) {
+        Thread.yield();
+      }
+      // Verify that we can still use the existing connection.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+      // Verify that we cannot open a new connection.
+      SQLException exception =
+          assertThrows(SQLException.class, () -> DriverManager.getConnection(createUrl()));
+      assertEquals(
+          String.format(
+              "Connection to localhost:%d refused. Check that the hostname and port are correct and that the postmaster is accepting TCP/IP connections.",
+              proxyServer.getLocalPort()),
+          exception.getMessage());
+      // Verify that the server is stopping, but not yet terminated.
+      assertEquals(State.STOPPING, proxyServer.state());
+
+      // Now force the server to stop by initiating a fast shutdown.
+      shutdownHandler.shutdown(ShutdownMode.FAST);
+      stopwatch = Stopwatch.createStarted();
+      while (proxyServer.state() != State.TERMINATED
+          && stopwatch.elapsed(TimeUnit.MILLISECONDS) < 2000) {
+        Thread.yield();
+      }
+      assertEquals(State.TERMINATED, proxyServer.state());
+
+      // Verify that existing connection has been invalidated.
+      exception =
+          assertThrows(SQLException.class, () -> connection.createStatement().executeQuery(sql));
+      assertEquals("An I/O error occurred while sending to the backend.", exception.getMessage());
+    }
+  }
+
+  boolean isPortListening(int port) {
+    try (Socket ignored = new Socket("localhost", port)) {
+      return true;
+    } catch (SocketException e) {
+      return false;
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+}