add: add option to delay shutdown initiation after receiving a SIGTERM

CHANGELOG.md

@@ -4,6 +4,8 @@ All notable changes to this project will be documented in this file. From versio
 
 ## Unreleased
 
+- Add option to delay shutdown initiation after receiving a `SIGTERM` by @lforst in #4580
+
 ### Added
 
 - Log error when `db-schemas` config contains schema `pg_catalog` or `information_schema` by @taimoorzaeem in #4359

src/PostgREST/App.hs

@@ -67,7 +67,7 @@ run appState = do
   conf@AppConfig{..} <- AppState.getConfig appState
 
   AppState.schemaCacheLoader appState -- Loads the initial SchemaCache
-  Unix.installSignalHandlers (AppState.getMainThreadId appState) (AppState.schemaCacheLoader appState) (AppState.readInDbConfig False appState)
+  Unix.installSignalHandlers configServerShutdownWaitPeriod (AppState.getMainThreadId appState) (AppState.schemaCacheLoader appState) (AppState.readInDbConfig False appState)
 
   Listener.runListener appState
 

src/PostgREST/Config.hs

@@ -112,6 +112,7 @@ data AppConfig = AppConfig
   , configServerPort               :: Int
   , configServerTraceHeader        :: Maybe (CI.CI BS.ByteString)
   , configServerTimingEnabled      :: Bool
+  , configServerShutdownWaitPeriod :: Int
   , configServerUnixSocket         :: Maybe FilePath
   , configServerUnixSocketMode     :: FileMode
   , configAdminServerHost          :: Text
@@ -186,6 +187,7 @@ toText conf =
       ,("server-port",                   show . configServerPort)
       ,("server-trace-header",       q . T.decodeUtf8 . maybe mempty CI.original . configServerTraceHeader)
       ,("server-timing-enabled",         T.toLower . show . configServerTimingEnabled)
+      ,("server-shutdown-wait-period",   show . configServerShutdownWaitPeriod)
       ,("server-unix-socket",        q . maybe mempty T.pack . configServerUnixSocket)
       ,("server-unix-socket-mode",   q . T.pack . showSocketMode)
       ,("admin-server-host",         q . configAdminServerHost)
@@ -299,6 +301,7 @@ parser optPath env dbSettings roleSettings roleIsolationLvl =
     <*> parseServerPort "server-port"
     <*> (fmap (CI.mk . encodeUtf8) <$> optString "server-trace-header")
     <*> (fromMaybe False <$> optBool "server-timing-enabled")
+    <*> (fromMaybe 0 <$> optInt "server-shutdown-wait-period")
     <*> (fmap T.unpack <$> optString "server-unix-socket")
     <*> parseSocketFileMode "server-unix-socket-mode"
     <*> (defaultServerHost <$> optWithAlias (optString "admin-server-host")
@@ -742,6 +745,10 @@ exampleConfigFile = S.unlines
   , "## Allow getting the request-response timing information through the `Server-Timing` header"
   , "server-timing-enabled = false"
   , ""
+  , "## Time in seconds to delay the shutdown after receiving SIGTERM"
+  , "## Useful for waiting on the load balancer to deregister the service"
+  , "# server-shutdown-wait-period = 0"
+  , ""
   , "## Unix socket location"
   , "## if specified it takes precedence over server-port"
   , "# server-unix-socket = \"/tmp/pgrst.sock\""

src/PostgREST/Unix.hs

@@ -17,20 +17,31 @@ import           Protolude
 import           System.Directory (removeFile)
 import           System.IO.Error  (isDoesNotExistError)
 
--- | Set signal handlers, only for systems with signals
-installSignalHandlers :: ThreadId -> IO () -> IO () -> IO ()
+-- | Set signal handlers
+--
+-- SIGINT:  Immediately throws UserInterrupt to the main thread.
+-- SIGTERM: Delays shutdown by 'shutdownWaitPeriod' seconds before throwing UserInterrupt.
+--          This allows load balancers (e.g., in AWS ECS) time to deregister the service
+--          before the process terminates, reducing the risk of rejected connections.
+-- SIGUSR1: Reloads the schema cache.
+-- SIGUSR2: Reloads the configuration from the database.
+installSignalHandlers :: Int -> ThreadId -> IO () -> IO () -> IO ()
 #ifndef mingw32_HOST_OS
-installSignalHandlers tid usr1 usr2 = do
-  let interrupt = throwTo tid UserInterrupt
-  install Signals.sigINT interrupt
-  install Signals.sigTERM interrupt
+installSignalHandlers shutdownWaitPeriod tid usr1 usr2 = do
+  let interruptImmediately = throwTo tid UserInterrupt
+      interruptWithDelay = do
+        when (shutdownWaitPeriod > 0) $
+          threadDelay (shutdownWaitPeriod * 1000000)
+        throwTo tid UserInterrupt
+  install Signals.sigINT interruptImmediately
+  install Signals.sigTERM interruptWithDelay
   install Signals.sigUSR1 usr1
   install Signals.sigUSR2 usr2
   where
     install signal handler =
       void $ Signals.installHandler signal (Signals.Catch handler) Nothing
 #else
-installSignalHandlers _ _ _ = pass
+installSignalHandlers _ _ _ _ = pass
 #endif
 
 -- | Create a unix domain socket and bind it to the given path.

test/io/configs/expected/aliases.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/boolean-numeric.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/boolean-string.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/defaults.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/jwt-role-claim-key1.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/jwt-role-claim-key2.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/jwt-role-claim-key3.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/jwt-role-claim-key4.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/jwt-role-claim-key5.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/no-defaults-with-db-other-authenticator.config

@@ -34,6 +34,7 @@ server-host = "0.0.0.0"
 server-port = 80
 server-trace-header = "traceparent"
 server-timing-enabled = true
+server-shutdown-wait-period = 5
 server-unix-socket = "/tmp/pgrst_io_test.sock"
 server-unix-socket-mode = "777"
 admin-server-host = "127.0.0.1"

test/io/configs/expected/no-defaults-with-db.config

@@ -34,6 +34,7 @@ server-host = "0.0.0.0"
 server-port = 80
 server-trace-header = "CF-Ray"
 server-timing-enabled = false
+server-shutdown-wait-period = 5
 server-unix-socket = "/tmp/pgrst_io_test.sock"
 server-unix-socket-mode = "777"
 admin-server-host = "127.0.0.1"

test/io/configs/expected/no-defaults.config

@@ -34,6 +34,7 @@ server-host = "0.0.0.0"
 server-port = 80
 server-trace-header = "X-Request-Id"
 server-timing-enabled = true
+server-shutdown-wait-period = 5
 server-unix-socket = "/tmp/pgrst_io_test.sock"
 server-unix-socket-mode = "777"
 admin-server-host = "127.0.0.1"

test/io/configs/expected/types.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/expected/utf-8.config

@@ -34,6 +34,7 @@ server-host = "!4"
 server-port = 3000
 server-trace-header = ""
 server-timing-enabled = false
+server-shutdown-wait-period = 0
 server-unix-socket = ""
 server-unix-socket-mode = "660"
 admin-server-host = "!4"

test/io/configs/no-defaults-env.yaml

@@ -37,6 +37,7 @@ PGRST_SERVER_HOST: 0.0.0.0
 PGRST_SERVER_PORT: 80
 PGRST_SERVER_TRACE_HEADER: X-Request-Id
 PGRST_SERVER_TIMING_ENABLED: true
+PGRST_SERVER_SHUTDOWN_WAIT_PERIOD: 5
 PGRST_SERVER_UNIX_SOCKET: /tmp/pgrst_io_test.sock
 PGRST_SERVER_UNIX_SOCKET_MODE: 777
 PGRST_ADMIN_SERVER_HOST: 127.0.0.1

test/io/configs/no-defaults.config

@@ -34,6 +34,7 @@ server-host = "0.0.0.0"
 server-port = 80
 server-trace-header = "X-Request-Id"
 server-timing-enabled = true
+server-shutdown-wait-period = 5
 server-unix-socket = "/tmp/pgrst_io_test.sock"
 server-unix-socket-mode = "777"
 admin-server-port = 3001

test/io/postgrest.py

@@ -176,10 +176,10 @@ def freeport(used_ports=None):
                 return port
 
 
-def wait_until_exit(postgrest):
+def wait_until_exit(postgrest, timeout=1):
     "Wait for PostgREST to exit, or times out"
     try:
-        return postgrest.process.wait(timeout=1)
+        return postgrest.process.wait(timeout=timeout)
     except subprocess.TimeoutExpired:
         raise PostgrestTimedOut()
 

test/io/test_io.py

@@ -1247,6 +1247,40 @@ def test_fail_with_automatic_recovery_disabled_and_terminated_using_query(defaul
         assert exitCode == 1
 
 
+def test_shutdown_wait_period_delays_sigterm(defaultenv):
+    "SIGTERM should wait server-shutdown-wait-period seconds before exiting"
+
+    env = {
+        **defaultenv,
+        "PGRST_SERVER_SHUTDOWN_WAIT_PERIOD": "1",
+    }
+
+    with run(env=env, wait_max_seconds=3) as postgrest:
+        start_time = time.time()
+        postgrest.process.send_signal(signal.SIGTERM)
+        wait_until_exit(postgrest, timeout=3)
+        elapsed = time.time() - start_time
+
+        assert elapsed >= 1.0, f"Should delay at least 1 second, but only waited {elapsed}s"
+
+
+def test_sigint_exits_immediately_with_shutdown_wait_period(defaultenv):
+    "SIGINT should exit immediately even with shutdown wait period configured"
+
+    env = {
+        **defaultenv,
+        "PGRST_SERVER_SHUTDOWN_WAIT_PERIOD": "5",
+    }
+
+    with run(env=env) as postgrest:
+        start_time = time.time()
+        postgrest.process.send_signal(signal.SIGINT)
+        wait_until_exit(postgrest, timeout=2)
+        elapsed = time.time() - start_time
+
+        assert elapsed < 1.0, f"SIGINT should exit immediately, but waited {elapsed}s"
+
+
 def test_preflight_request_with_cors_allowed_origin_config(defaultenv):
     "OPTIONS preflight request should return Access-Control-Allow-Origin equal to origin"
 

test/spec/SpecHelper.hs

@@ -161,6 +161,7 @@ baseCfg = let secret = encodeUtf8 "reallyreallyreallyreallyverysafe" in
   , configInternalSCLoadSleep       = Nothing
   , configInternalSCRelLoadSleep    = Nothing
   , configServerTimingEnabled       = True
+  , configServerShutdownWaitPeriod  = 0
   }
 
 testCfg :: AppConfig