experimental: dynamicly adjust lock count

mkleczek · mkleczek · commit 8b8ef5cc24cb · 2026-02-20T17:56:43.000+01:00
diff --git a/postgrest.cabal b/postgrest.cabal
@@ -158,6 +158,7 @@ library
                     , stm-hamt                  >= 1.2 && < 2
                     , focus                     >= 1.0 && < 2
                     , some                      >= 1.0.4.1 && < 2
+                    , uuid                      >= 1.3 && < 2
                       -- -fno-spec-constr may help keep compile time memory use in check,
                       --   see https://gitlab.haskell.org/ghc/ghc/issues/16017#note_219304
                       -- -optP-Wno-nonportable-include-path
diff --git a/src/PostgREST/AppState.hs b/src/PostgREST/AppState.hs
@@ -1,7 +1,8 @@
-{-# LANGUAGE LambdaCase      #-}
-{-# LANGUAGE NamedFieldPuns  #-}
-{-# LANGUAGE QuasiQuotes     #-}
-{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE LambdaCase       #-}
+{-# LANGUAGE NamedFieldPuns   #-}
+{-# LANGUAGE QuasiQuotes      #-}
+{-# LANGUAGE RecordWildCards  #-}
+{-# LANGUAGE TypeApplications #-}
 
 module PostgREST.AppState
   ( AppState
@@ -74,15 +75,21 @@ import PostgREST.SchemaCache             (SchemaCache (..),
 import PostgREST.SchemaCache.Identifiers (quoteQi)
 import PostgREST.Unix                    (createAndBindDomainSocket)
 
-import           Data.Functor.Contravariant ((>$<))
-import           Data.Streaming.Network     (bindPortTCP,
-                                             bindRandomPortTCP)
-import           Data.String                (IsString (..))
-import qualified Hasql.Decoders             as HD
-import qualified Hasql.Encoders             as HE
-import qualified Hasql.Statement            as SQL
-import           NeatInterpolation          (trimming)
+import           Control.Arrow          ((&&&))
+import           Data.Bitraversable     (bisequenceA)
+import           Data.Streaming.Network (bindPortTCP,
+                                         bindRandomPortTCP)
+import           Data.String            (IsString (..))
+import           Data.Tuple.Extra       (both)
+import           Data.UUID              hiding (fromString)
+import qualified Focus
+import qualified Hasql.Decoders         as HD
+import qualified Hasql.Encoders         as HE
+import           Hasql.Pool.Observation
+import qualified Hasql.Statement        as SQL
+import           NeatInterpolation      (trimming)
 import           Protolude
+import qualified StmHamt.SizedHamt      as SH
 
 
 data AppState = AppState
@@ -118,6 +125,7 @@ data AppState = AppState
   , stateJwtCache          :: JwtCache.JwtCacheState
   , stateLogger            :: Logger.LoggerState
   , stateMetrics           :: Metrics.MetricsState
+  , stateConnTrack         :: ConnTrack
   }
 
 -- | Schema cache status
@@ -132,7 +140,7 @@ init :: AppConfig -> IO AppState
 init conf@AppConfig{configLogLevel, configDbPoolSize} = do
   loggerState  <- Logger.init
   metricsState <- Metrics.init configDbPoolSize
-  let observer = liftA2 (>>) (Logger.observationLogger loggerState configLogLevel) (Metrics.observationMetrics metricsState)
+  let observer = Logger.observationLogger loggerState configLogLevel <> Metrics.observationMetrics metricsState
 
   observer $ AppStartObs prettyVersion
 
@@ -144,6 +152,7 @@ init conf@AppConfig{configLogLevel, configDbPoolSize} = do
 initWithPool :: AppSockets -> SQL.Pool -> AppConfig -> Logger.LoggerState -> Metrics.MetricsState -> ObservationHandler -> IO AppState
 initWithPool (sock, adminSock) pool conf loggerState metricsState observer = do
 
+  connTrack <- ConnTrack <$> SH.newIO <*> SH.newIO
   appState <- AppState pool
     <$> newIORef minimumPgVersion -- assume we're in a supported version when starting, this will be corrected on a later step
     <*> newIORef Nothing
@@ -157,10 +166,11 @@ initWithPool (sock, adminSock) pool conf loggerState metricsState observer = do
     <*> newIORef 1
     <*> pure sock
     <*> pure adminSock
-    <*> pure observer
+    <*> pure (observer <> trackConnections connTrack)
     <*> JwtCache.init conf observer
     <*> pure loggerState
     <*> pure metricsState
+    <*> pure connTrack
 
   deb <-
     let decisecond = 100000 in
@@ -410,11 +420,21 @@ retryingSchemaCacheLoad appState@AppState{stateObserver=observer, stateMainThrea
     qSchemaCache :: IO (Maybe SchemaCache)
     qSchemaCache = do
       conf@AppConfig{..} <- getConfig appState
-        -- Allow 10 concurrent schema cache loads, guarded by advisory locks.
-        -- This is to prevent thundering herd problem on startup or when many PostgREST
-        -- instances receive "reload schema" notifications at the same time
-      let withTxLock = SQL.statement (50168275, 10) $
-            SQL.Statement get_lock_sql get_lock_params HD.noResult configDbPreparedStatements
+      -- Throttle concurrent schema cache loads, guarded by advisory locks.
+      -- This is to prevent thundering herd problem on startup or when many PostgREST
+      -- instances receive "reload schema" notifications at the same time
+      -- schema reloading session + listener session
+      -- See get_lock_sql for details of the algorithm.
+      -- Here we calculate the number of open connections passed to the query.
+      (connected, inUse) <- sessionCounts appState
+      -- Determine whether schema cache loading will create a new session
+      let
+        scLoadingSessions = case (connected <= inUse, inUse >= configDbPoolSize) of
+          (True, False) -> 1 -- all connections in use but pool not full - schema cache loading will create session
+          _             -> 0
+        withTxLock = SQL.statement
+          (fromIntegral $ connected + scLoadingSessions)
+          (SQL.Statement get_lock_sql get_lock_params HD.noResult configDbPreparedStatements)
 
       (resultTime, result) <-
         let transaction = if configDbPreparedStatements then SQL.transaction else SQL.unpreparedTransaction in
@@ -437,23 +457,42 @@ retryingSchemaCacheLoad appState@AppState{stateObserver=observer, stateMainThrea
           putSCacheStatus appState SCLoaded
           return $ Just sCache
       where
-        -- recursive query that tries acquiring locks in order
-        -- and waits for randomly selected lock if no attempt succeeded
-        -- parameters are lock number and number of locks to try
+        -- Recursive query that tries acquiring locks in order
+        -- and waits for randomly selected lock if no attempt succeeded.
+        -- It has a single parameter: this node open connection count.
+        -- It is used to estimate the number of nodes
+        -- by counting the number of active sessions for current session_user
+        -- and dividing it by this node open connections.
+        -- Assuming load is uniform among cluster nodes, all should have
+        -- statistically the same number of open connections.
+        -- Once the number of nodes is known we calculate the number
+        -- of locks as ceil(log(2, number_of_nodes))
         get_lock_sql = encodeUtf8 [trimming|
           WITH RECURSIVE attempts AS (
-            SELECT 1 AS lock_number, pg_try_advisory_xact_lock($$1, 1) AS success WHERE $$2 > 0
+            SELECT 1 AS lock_number, pg_try_advisory_xact_lock(lock_id, 1) AS success FROM parameters
             UNION ALL
-            SELECT next_lock_number AS lock_number, pg_try_advisory_xact_lock($$1, next_lock_number) AS success FROM (
-                SELECT lock_number + 1 AS next_lock_number FROM attempts
-                WHERE NOT success AND lock_number < $$2
-                ORDER BY lock_number DESC
-                LIMIT 1
-            ) AS previous_attempt
+            SELECT next_lock_number AS lock_number, pg_try_advisory_xact_lock(lock_id, next_lock_number) AS success
+            FROM
+              parameters CROSS JOIN LATERAL (
+                  SELECT lock_number + 1 AS next_lock_number FROM attempts
+                  WHERE NOT success AND lock_number < locks_count
+                  ORDER BY lock_number DESC
+                  LIMIT 1
+                ) AS previous_attempt
+          ),
+          counts AS (
+            SELECT round(log(2, round(count(*)::double precision/$$1)::numeric))::int AS locks_count
+            FROM
+              pg_stat_activity WHERE usename = SESSION_USER
+          ),
+          parameters AS (
+            SELECT locks_count, 50168275 AS lock_id FROM counts WHERE locks_count > 0
           )
-          SELECT pg_advisory_xact_lock($$1, floor(random() * $$2)::int + 1) WHERE NOT EXISTS (SELECT 1 FROM attempts WHERE success) |]
+          SELECT pg_advisory_xact_lock(lock_id, floor(random() * locks_count)::int + 1)
+          FROM
+            parameters WHERE NOT EXISTS (SELECT 1 FROM attempts WHERE success) |]
 
-        get_lock_params = (fst >$< HE.param (HE.nonNullable HE.int4)) <> (snd >$< HE.param (HE.nonNullable HE.int4))
+        get_lock_params = HE.param (HE.nonNullable HE.int4)
 
     shouldRetry :: RetryStatus -> (Maybe PgVersion, Maybe SchemaCache) -> IO Bool
     shouldRetry _ (pgVer, sCache) = do
@@ -511,3 +550,21 @@ readInDbConfig startingUp appState@AppState{stateObserver=observer} = do
         pass
       else
         observer ConfigSucceededObs
+
+
+data ConnTrack = ConnTrack { connTrackConnected :: SH.SizedHamt UUID, connTrackInUse :: SH.SizedHamt UUID }
+
+trackConnections :: ConnTrack -> ObservationHandler
+trackConnections ConnTrack{..} (HasqlPoolObs (ConnectionObservation uuid status)) = case status of
+  ReadyForUseConnectionStatus -> atomically $
+    SH.insert identity uuid connTrackConnected *>
+    SH.focus Focus.delete identity uuid connTrackInUse
+  TerminatedConnectionStatus _ -> atomically $
+    SH.focus Focus.delete identity uuid connTrackConnected *>
+    SH.focus Focus.delete identity uuid connTrackInUse
+  InUseConnectionStatus -> atomically $ SH.insert identity uuid connTrackInUse
+  _ -> mempty
+trackConnections _ _ = mempty
+
+sessionCounts :: AppState -> IO (Int, Int)
+sessionCounts = atomically . bisequenceA . both SH.size . (connTrackConnected &&& connTrackInUse) . stateConnTrack
diff --git a/test/io/test_io.py b/test/io/test_io.py
@@ -1060,9 +1060,11 @@ def test_schema_cache_concurrent_notifications(slow_schema_cache_env):
         assert response.status_code == 200
 
 
-@pytest.mark.parametrize("instance_count", [10, 11])
+@pytest.mark.parametrize(
+    "instance_count, expected_concurrency", [(2, 2), (4, 3), (6, 4), (8, 4), (16, 5)]
+)
 def test_schema_cache_reload_throttled_with_advisory_locks(
-    instance_count, slow_schema_cache_env
+    instance_count, expected_concurrency, slow_schema_cache_env
 ):
     "schema cache reloads should be throttled across instances if instance count > 10"
 
@@ -1112,7 +1114,7 @@ def read_available_output_lines(postgrest):
         reload_durations_ms = []
         for postgrest in instances:
             output_lines = []
-            for _ in range(5):
+            for _ in range(instance_count * 2):
                 output_lines.extend(read_available_output_lines(postgrest))
                 if any(query_log_pattern.search(line) for line in output_lines):
                     break
@@ -1131,14 +1133,15 @@ def read_available_output_lines(postgrest):
 
         # 10 instances should be fast, remaining instances should be slow
         assert (
-            len(
+            instance_count
+            - len(
                 [
                     duration
                     for duration in reload_durations_ms
                     if duration > lock_wait_threshold_ms
                 ]
             )
-            == instance_count - 10
+            == expected_concurrency
         )
 
 
