1- {-# LANGUAGE LambdaCase #-}
2- {-# LANGUAGE NamedFieldPuns #-}
3- {-# LANGUAGE RecordWildCards #-}
4- {-# LANGUAGE RecursiveDo #-}
1+ {-# LANGUAGE LambdaCase #-}
2+ {-# LANGUAGE NamedFieldPuns #-}
3+ {-# LANGUAGE QuasiQuotes #-}
4+ {-# LANGUAGE RecordWildCards #-}
5+ {-# LANGUAGE RecursiveDo #-}
6+ {-# LANGUAGE TypeApplications #-}
57
68module PostgREST.AppState
79 ( AppState
@@ -33,7 +35,8 @@ import qualified Data.ByteString.Char8 as BS
3335import Data.Either.Combinators (whenLeft )
3436import qualified Hasql.Pool as SQL
3537import qualified Hasql.Pool.Config as SQL
36- import qualified Hasql.Session as SQL
38+ import qualified Hasql.Session as SQL hiding (statement )
39+ import qualified Hasql.Transaction as SQL hiding (sql )
3740import qualified Hasql.Transaction.Sessions as SQL
3841import qualified Network.HTTP.Types.Status as HTTP
3942import qualified PostgREST.Auth.JwtCache as JwtCache
@@ -63,11 +66,17 @@ import PostgREST.Config.Database (queryDbSettings,
6366import PostgREST.Config.PgVersion (PgVersion (.. ),
6467 minimumPgVersion )
6568import PostgREST.Debounce (makeDebouncer )
69+ import PostgREST.Metrics (MetricsState (connTrack ))
6670import PostgREST.SchemaCache (SchemaCache (.. ),
6771 querySchemaCache ,
6872 showSummary )
6973import PostgREST.SchemaCache.Identifiers (quoteQi )
7074
75+ import qualified Hasql.Decoders as HD
76+ import qualified Hasql.Encoders as HE
77+ import qualified Hasql.Statement as SQL
78+ import NeatInterpolation (trimming )
79+
7180import Protolude
7281
7382data AppState = AppState
@@ -296,7 +305,7 @@ getObserver = stateObserver
296305-- + Because connections cache the pg catalog(see #2620)
297306-- + For rapid recovery. Otherwise, the pool idle or lifetime timeout would have to be reached for new healthy connections to be acquired.
298307retryingSchemaCacheLoad :: AppState -> IO ()
299- retryingSchemaCacheLoad appState@ AppState {stateObserver= observer, stateMainThreadId= mainThreadId} =
308+ retryingSchemaCacheLoad appState@ AppState {stateObserver= observer, stateMainThreadId= mainThreadId, stateMetrics } =
300309 void $ retrying retryPolicy shouldRetry (\ RetryStatus {rsIterNumber, rsPreviousDelay} -> do
301310 when (rsIterNumber > 0 ) $ do
302311 let delay = fromMaybe 0 rsPreviousDelay `div` oneSecondInUs
@@ -335,9 +344,23 @@ retryingSchemaCacheLoad appState@AppState{stateObserver=observer, stateMainThrea
335344 qSchemaCache :: IO (Maybe SchemaCache )
336345 qSchemaCache = do
337346 conf@ AppConfig {.. } <- getConfig appState
347+ -- Throttle concurrent schema cache loads, guarded by advisory locks.
348+ -- This is to prevent thundering herd problem on startup or when many PostgREST
349+ -- instances receive "reload schema" notifications at the same time
350+ -- See get_lock_sql for details of the algorithm.
351+ -- Here we calculate the number of open connections passed to the query.
352+ Metrics. ConnStats connected inUse <- Metrics. connectionCounts $ connTrack stateMetrics
353+ -- Determine whether schema cache loading will create a new session
354+ let
355+ -- if all connections in use but pool not full - schema cache loading will create session
356+ scLoadingSessions = if connected <= inUse && inUse < configDbPoolSize then 1 else 0
357+ withTxLock = SQL. statement
358+ (fromIntegral $ connected + scLoadingSessions)
359+ (SQL. Statement get_lock_sql get_lock_params HD. noResult configDbPreparedStatements)
360+
338361 (resultTime, result) <-
339362 let transaction = if configDbPreparedStatements then SQL. transaction else SQL. unpreparedTransaction in
340- timeItT $ usePool appState (transaction SQL. ReadCommitted SQL. Read $ querySchemaCache conf)
363+ timeItT $ usePool appState (transaction SQL. ReadCommitted SQL. Read $ withTxLock *> querySchemaCache conf)
341364 case result of
342365 Left e -> do
343366 markSchemaCachePending appState
@@ -359,6 +382,43 @@ retryingSchemaCacheLoad appState@AppState{stateObserver=observer, stateMainThrea
359382 observer $ SchemaCacheLoadedObs loadTime summary
360383 markSchemaCacheLoaded appState
361384 return $ Just sCache
385+ where
386+ -- Recursive query that tries acquiring locks in order
387+ -- and waits for randomly selected lock if no attempt succeeded.
388+ -- It has a single parameter: this node open connection count.
389+ -- It is used to estimate the number of nodes
390+ -- by counting the number of active sessions for current session_user
391+ -- and dividing it by this node open connections.
392+ -- Assuming load is uniform among cluster nodes, all should have
393+ -- statistically the same number of open connections.
394+ -- Once the number of nodes is known we calculate the number
395+ -- of locks as ceil(log(2, number_of_nodes))
396+ get_lock_sql = encodeUtf8 [trimming |
397+ WITH RECURSIVE attempts AS (
398+ SELECT 1 AS lock_number, pg_try_advisory_xact_lock(lock_id, 1) AS success FROM parameters
399+ UNION ALL
400+ SELECT next_lock_number AS lock_number, pg_try_advisory_xact_lock(lock_id, next_lock_number) AS success
401+ FROM
402+ parameters CROSS JOIN LATERAL (
403+ SELECT lock_number + 1 AS next_lock_number FROM attempts
404+ WHERE NOT success AND lock_number < locks_count
405+ ORDER BY lock_number DESC
406+ LIMIT 1
407+ ) AS previous_attempt
408+ ),
409+ counts AS (
410+ SELECT round(log(2, round(count(*)::double precision/$$1)::numeric))::int AS locks_count
411+ FROM
412+ pg_stat_activity WHERE usename = SESSION_USER
413+ ),
414+ parameters AS (
415+ SELECT locks_count, 50168275 AS lock_id FROM counts WHERE locks_count > 0
416+ )
417+ SELECT pg_advisory_xact_lock(lock_id, floor(random() * locks_count)::int + 1)
418+ FROM
419+ parameters WHERE NOT EXISTS (SELECT 1 FROM attempts WHERE success) |]
420+
421+ get_lock_params = HE. param (HE. nonNullable HE. int4)
362422
363423 shouldRetry :: RetryStatus -> (Maybe PgVersion , Maybe SchemaCache ) -> IO Bool
364424 shouldRetry _ (pgVer, sCache) = do
0 commit comments