feat: introduce pool acquisition timeout (fixes PostgREST#2348)

The configuration option db-pool-acquisition-timeout specifies the time in seconds to wait for the pool to free up a connection slot. Otherwise, a 504 error is returned. By default, there is no timeout.
robx · Aug 30, 2022 · 24924df · 24924df
1 parent e2aa227
commit 24924df
Show file tree

Hide file tree

Showing 16 changed files with 84 additions and 35 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 ### Fixed
 
  - #2401, #2444, Fix SIGUSR1 to fully flush connections pool, remove `db-pool-timeout`. - @robx
+ - #2348, Add `db-pool-acquisition-timeout` configuration option, time in seconds to wait to acquire a connection. - @robx
 
 ### Deprecated
 

diff --git a/src/PostgREST/AppState.hs b/src/PostgREST/AppState.hs
@@ -97,7 +97,7 @@ destroy = destroyPool
 
 initPool :: AppConfig -> IO SQL.Pool
 initPool AppConfig{..} =
-  SQL.acquire configDbPoolSize Nothing $ toUtf8 configDbUri
+  SQL.acquire configDbPoolSize ((1000000 *) <$> configDbPoolAcquisitionTimeout) $ toUtf8 configDbUri
 
 -- | Run an action with a database connection.
 usePool :: AppState -> SQL.Session a -> IO (Either SQL.UsageError a)

diff --git a/src/PostgREST/CLI.hs b/src/PostgREST/CLI.hs
@@ -148,6 +148,9 @@ exampleConfigFile =
       |## Number of open connections in the pool
       |db-pool = 10
       |
+      |## Time in seconds to wait to acquire a slot from the connection pool
+      |# db-pool-acquisition-timeout = 10
+      |
       |## Stored proc to exec immediately after auth
       |# db-pre-request = "stored_proc_name"
       |

diff --git a/src/PostgREST/Config.hs b/src/PostgREST/Config.hs
@@ -62,39 +62,40 @@ import Protolude hiding (Proxy, toList)
 
 
 data AppConfig = AppConfig
-  { configAppSettings           :: [(Text, Text)]
-  , configDbAnonRole            :: Maybe Text
-  , configDbChannel             :: Text
-  , configDbChannelEnabled      :: Bool
-  , configDbExtraSearchPath     :: [Text]
-  , configDbMaxRows             :: Maybe Integer
-  , configDbPlanEnabled         :: Bool
-  , configDbPoolSize            :: Int
-  , configDbPreRequest          :: Maybe QualifiedIdentifier
-  , configDbPreparedStatements  :: Bool
-  , configDbRootSpec            :: Maybe QualifiedIdentifier
-  , configDbSchemas             :: NonEmpty Text
-  , configDbConfig              :: Bool
-  , configDbTxAllowOverride     :: Bool
-  , configDbTxRollbackAll       :: Bool
-  , configDbUri                 :: Text
-  , configDbUseLegacyGucs       :: Bool
-  , configFilePath              :: Maybe FilePath
-  , configJWKS                  :: Maybe JWKSet
-  , configJwtAudience           :: Maybe StringOrURI
-  , configJwtRoleClaimKey       :: JSPath
-  , configJwtSecret             :: Maybe BS.ByteString
-  , configJwtSecretIsBase64     :: Bool
-  , configLogLevel              :: LogLevel
-  , configOpenApiMode           :: OpenAPIMode
-  , configOpenApiSecurityActive :: Bool
-  , configOpenApiServerProxyUri :: Maybe Text
-  , configRawMediaTypes         :: [MediaType]
-  , configServerHost            :: Text
-  , configServerPort            :: Int
-  , configServerUnixSocket      :: Maybe FilePath
-  , configServerUnixSocketMode  :: FileMode
-  , configAdminServerPort       :: Maybe Int
+  { configAppSettings              :: [(Text, Text)]
+  , configDbAnonRole               :: Maybe Text
+  , configDbChannel                :: Text
+  , configDbChannelEnabled         :: Bool
+  , configDbExtraSearchPath        :: [Text]
+  , configDbMaxRows                :: Maybe Integer
+  , configDbPlanEnabled            :: Bool
+  , configDbPoolSize               :: Int
+  , configDbPoolAcquisitionTimeout :: Maybe Int
+  , configDbPreRequest             :: Maybe QualifiedIdentifier
+  , configDbPreparedStatements     :: Bool
+  , configDbRootSpec               :: Maybe QualifiedIdentifier
+  , configDbSchemas                :: NonEmpty Text
+  , configDbConfig                 :: Bool
+  , configDbTxAllowOverride        :: Bool
+  , configDbTxRollbackAll          :: Bool
+  , configDbUri                    :: Text
+  , configDbUseLegacyGucs          :: Bool
+  , configFilePath                 :: Maybe FilePath
+  , configJWKS                     :: Maybe JWKSet
+  , configJwtAudience              :: Maybe StringOrURI
+  , configJwtRoleClaimKey          :: JSPath
+  , configJwtSecret                :: Maybe BS.ByteString
+  , configJwtSecretIsBase64        :: Bool
+  , configLogLevel                 :: LogLevel
+  , configOpenApiMode              :: OpenAPIMode
+  , configOpenApiSecurityActive    :: Bool
+  , configOpenApiServerProxyUri    :: Maybe Text
+  , configRawMediaTypes            :: [MediaType]
+  , configServerHost               :: Text
+  , configServerPort               :: Int
+  , configServerUnixSocket         :: Maybe FilePath
+  , configServerUnixSocketMode     :: FileMode
+  , configAdminServerPort          :: Maybe Int
   }
 
 data LogLevel = LogCrit | LogError | LogWarn | LogInfo
@@ -129,6 +130,7 @@ toText conf =
       ,("db-max-rows",                   maybe "\"\"" show . configDbMaxRows)
       ,("db-plan-enabled",               T.toLower . show . configDbPlanEnabled)
       ,("db-pool",                       show . configDbPoolSize)
+      ,("db-pool-acquisition-timeout",   maybe "\"\"" show . configDbPoolAcquisitionTimeout)
       ,("db-pre-request",            q . maybe mempty dumpQi . configDbPreRequest)
       ,("db-prepared-statements",        T.toLower . show . configDbPreparedStatements)
       ,("db-root-spec",              q . maybe mempty dumpQi . configDbRootSpec)
@@ -217,6 +219,7 @@ parser optPath env dbSettings =
                      (optInt "max-rows")
     <*> (fromMaybe False <$> optBool "db-plan-enabled")
     <*> (fromMaybe 10 <$> optInt "db-pool")
+    <*> optInt "db-pool-acquisition-timeout"
     <*> (fmap toQi <$> optWithAlias (optString "db-pre-request")
                                     (optString "pre-request"))
     <*> (fromMaybe True <$> optBool "db-prepared-statements")
@@ -352,7 +355,7 @@ parser optPath env dbSettings =
           let dbSettingName = T.pack $ dashToUnderscore <$> toS key in
           if dbSettingName `notElem` [
             "server_host", "server_port", "server_unix_socket", "server_unix_socket_mode", "admin_server_port", "log_level",
-            "db_uri", "db_channel_enabled", "db_channel", "db_pool", "db_config"]
+            "db_uri", "db_channel_enabled", "db_channel", "db_pool", "db_pool_acquisition_timeout", "db_config"]
           then lookup dbSettingName dbSettings
           else Nothing
 

diff --git a/test/io/configs/expected/aliases.config b/test/io/configs/expected/aliases.config
@@ -5,6 +5,7 @@ db-extra-search-path = "public"
 db-max-rows = 1000
 db-plan-enabled = false
 db-pool = 10
+db-pool-acquisition-timeout = ""
 db-pre-request = "check_alias"
 db-prepared-statements = true
 db-root-spec = "open_alias"

diff --git a/test/io/configs/expected/boolean-numeric.config b/test/io/configs/expected/boolean-numeric.config
@@ -5,6 +5,7 @@ db-extra-search-path = "public"
 db-max-rows = ""
 db-plan-enabled = false
 db-pool = 10
+db-pool-acquisition-timeout = ""
 db-pre-request = ""
 db-prepared-statements = false
 db-root-spec = ""

diff --git a/test/io/configs/expected/boolean-string.config b/test/io/configs/expected/boolean-string.config
@@ -5,6 +5,7 @@ db-extra-search-path = "public"
 db-max-rows = ""
 db-plan-enabled = false
 db-pool = 10
+db-pool-acquisition-timeout = ""
 db-pre-request = ""
 db-prepared-statements = false
 db-root-spec = ""

diff --git a/test/io/configs/expected/defaults.config b/test/io/configs/expected/defaults.config
@@ -5,6 +5,7 @@ db-extra-search-path = "public"
 db-max-rows = ""
 db-plan-enabled = false
 db-pool = 10
+db-pool-acquisition-timeout = ""
 db-pre-request = ""
 db-prepared-statements = true
 db-root-spec = ""

diff --git a/test/io/configs/expected/no-defaults-with-db-other-authenticator.config b/test/io/configs/expected/no-defaults-with-db-other-authenticator.config
@@ -5,6 +5,7 @@ db-extra-search-path = "public,extensions,other"
 db-max-rows = 100
 db-plan-enabled = true
 db-pool = 1
+db-pool-acquisition-timeout = 10
 db-pre-request = "test.other_custom_headers"
 db-prepared-statements = false
 db-root-spec = "other_root"

diff --git a/test/io/configs/expected/no-defaults-with-db.config b/test/io/configs/expected/no-defaults-with-db.config
@@ -5,6 +5,7 @@ db-extra-search-path = "public,extensions,private"
 db-max-rows = 1000
 db-plan-enabled = true
 db-pool = 1
+db-pool-acquisition-timeout = 10
 db-pre-request = "test.custom_headers"
 db-prepared-statements = false
 db-root-spec = "root"

diff --git a/test/io/configs/expected/no-defaults.config b/test/io/configs/expected/no-defaults.config
@@ -5,6 +5,7 @@ db-extra-search-path = "public,test"
 db-max-rows = 1000
 db-plan-enabled = true
 db-pool = 1
+db-pool-acquisition-timeout = 10
 db-pre-request = "please_run_fast"
 db-prepared-statements = false
 db-root-spec = "openapi_v3"

diff --git a/test/io/configs/expected/types.config b/test/io/configs/expected/types.config
@@ -5,6 +5,7 @@ db-extra-search-path = "public"
 db-max-rows = ""
 db-plan-enabled = false
 db-pool = 10
+db-pool-acquisition-timeout = ""
 db-pre-request = ""
 db-prepared-statements = true
 db-root-spec = ""

diff --git a/test/io/configs/no-defaults-env.yaml b/test/io/configs/no-defaults-env.yaml
@@ -7,6 +7,7 @@ PGRST_DB_EXTRA_SEARCH_PATH: public, test
 PGRST_DB_MAX_ROWS: 1000
 PGRST_DB_PLAN_ENABLED: true
 PGRST_DB_POOL: 1
+PGRST_DB_POOL_ACQUISITION_TIMEOUT: 10
 PGRST_DB_PREPARED_STATEMENTS: false
 PGRST_DB_PRE_REQUEST: please_run_fast
 PGRST_DB_ROOT_SPEC: openapi_v3

diff --git a/test/io/configs/no-defaults.config b/test/io/configs/no-defaults.config
@@ -5,6 +5,7 @@ db-extra-search-path = "public, test"
 db-max-rows = 1000
 db-plan-enabled = true
 db-pool = 1
+db-pool-acquisition-timeout = 10
 db-pre-request = "please_run_fast"
 db-prepared-statements = false
 db-root-spec = "openapi_v3"

diff --git a/test/io/test_io.py b/test/io/test_io.py
@@ -993,6 +993,37 @@ def sleep(i=i):
         assert delta > 1 and delta < 1.5
 
 
+def test_pool_acquisition_timeout(defaultenv, metapostgrest):
+    "Verify that PGRST_DB_POOL_ACQUISITON_TIMEOUT times out when the pool is empty"
+
+    env = {
+        **defaultenv,
+        "PGRST_DB_POOL": "1",
+        "PGRST_DB_POOL_ACQUISITION_TIMEOUT": "1",  # 1 second
+    }
+
+    with run(env=env) as postgrest:
+
+        # start a slow request that holds the only pool connection
+        def hold_connection():
+            response = postgrest.session.get("/rpc/sleep?seconds=1.5")
+            assert response.status_code == 204
+
+        hold = Thread(target=hold_connection)
+        hold.start()
+
+        # give the thread time to start
+        time.sleep(0.1)
+
+        response = postgrest.session.get("/authors_only")
+        assert response.status_code == 504
+        data = response.json()
+        assert data["message"] == "Timed out acquiring connection from connection pool."
+
+        # collect our thread
+        hold.join()
+
+
 def test_change_statement_timeout_held_connection(defaultenv, metapostgrest):
     "Statement timeout changes take effect immediately, even with a request outliving the reconfiguration"
 

diff --git a/test/spec/SpecHelper.hs b/test/spec/SpecHelper.hs
@@ -79,6 +79,7 @@ baseCfg = let secret = Just $ encodeUtf8 "reallyreallyreallyreallyverysafe" in
   , configDbMaxRows             = Nothing
   , configDbPlanEnabled         = False
   , configDbPoolSize            = 10
+  , configDbPoolAcquisitionTimeout = Nothing
   , configDbPreRequest          = Just $ QualifiedIdentifier "test" "switch_role"
   , configDbPreparedStatements  = True
   , configDbRootSpec            = Nothing