Skip to content

Commit c72cfeb

Browse files
mkleczekmkleczek
committed
experimental: dynamicly adjust lock count
1 parent a96c490 commit c72cfeb

3 files changed

Lines changed: 98 additions & 36 deletions

File tree

‎postgrest.cabal‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,7 @@ library
158158
, stm-hamt >= 1.2 && < 2
159159
, focus >= 1.0 && < 2
160160
, some >= 1.0.4.1 && < 2
161+
, uuid >= 1.3 && < 2
161162
-- -fno-spec-constr may help keep compile time memory use in check,
162163
-- see https://gitlab.haskell.org/ghc/ghc/issues/16017#note_219304
163164
-- -optP-Wno-nonportable-include-path

‎src/PostgREST/AppState.hs‎

Lines changed: 89 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
1-
{-# LANGUAGE LambdaCase #-}
2-
{-# LANGUAGE NamedFieldPuns #-}
3-
{-# LANGUAGE QuasiQuotes #-}
4-
{-# LANGUAGE RecordWildCards #-}
1+
{-# LANGUAGE LambdaCase #-}
2+
{-# LANGUAGE NamedFieldPuns #-}
3+
{-# LANGUAGE QuasiQuotes #-}
4+
{-# LANGUAGE RecordWildCards #-}
5+
{-# LANGUAGE TypeApplications #-}
56

67
module PostgREST.AppState
78
( AppState
@@ -74,15 +75,21 @@ import PostgREST.SchemaCache (SchemaCache (..),
7475
import PostgREST.SchemaCache.Identifiers (quoteQi)
7576
import PostgREST.Unix (createAndBindDomainSocket)
7677

77-
import Data.Functor.Contravariant ((>$<))
78-
import Data.Streaming.Network (bindPortTCP,
79-
bindRandomPortTCP)
80-
import Data.String (IsString (..))
81-
import qualified Hasql.Decoders as HD
82-
import qualified Hasql.Encoders as HE
83-
import qualified Hasql.Statement as SQL
84-
import NeatInterpolation (trimming)
78+
import Control.Arrow ((&&&))
79+
import Data.Bitraversable (bisequenceA)
80+
import Data.Streaming.Network (bindPortTCP,
81+
bindRandomPortTCP)
82+
import Data.String (IsString (..))
83+
import Data.Tuple.Extra (both)
84+
import Data.UUID hiding (fromString)
85+
import qualified Focus
86+
import qualified Hasql.Decoders as HD
87+
import qualified Hasql.Encoders as HE
88+
import Hasql.Pool.Observation
89+
import qualified Hasql.Statement as SQL
90+
import NeatInterpolation (trimming)
8591
import Protolude
92+
import qualified StmHamt.SizedHamt as SH
8693

8794

8895
data AppState = AppState
@@ -118,6 +125,7 @@ data AppState = AppState
118125
, stateJwtCache :: JwtCache.JwtCacheState
119126
, stateLogger :: Logger.LoggerState
120127
, stateMetrics :: Metrics.MetricsState
128+
, stateConnTrack :: ConnTrack
121129
}
122130

123131
-- | Schema cache status
@@ -132,7 +140,7 @@ init :: AppConfig -> IO AppState
132140
init conf@AppConfig{configLogLevel, configDbPoolSize} = do
133141
loggerState <- Logger.init
134142
metricsState <- Metrics.init configDbPoolSize
135-
let observer = liftA2 (>>) (Logger.observationLogger loggerState configLogLevel) (Metrics.observationMetrics metricsState)
143+
let observer = Logger.observationLogger loggerState configLogLevel <> Metrics.observationMetrics metricsState
136144

137145
observer $ AppStartObs prettyVersion
138146

@@ -144,6 +152,7 @@ init conf@AppConfig{configLogLevel, configDbPoolSize} = do
144152
initWithPool :: AppSockets -> SQL.Pool -> AppConfig -> Logger.LoggerState -> Metrics.MetricsState -> ObservationHandler -> IO AppState
145153
initWithPool (sock, adminSock) pool conf loggerState metricsState observer = do
146154

155+
connTrack <- ConnTrack <$> SH.newIO <*> SH.newIO
147156
appState <- AppState pool
148157
<$> newIORef minimumPgVersion -- assume we're in a supported version when starting, this will be corrected on a later step
149158
<*> newIORef Nothing
@@ -157,10 +166,11 @@ initWithPool (sock, adminSock) pool conf loggerState metricsState observer = do
157166
<*> newIORef 1
158167
<*> pure sock
159168
<*> pure adminSock
160-
<*> pure observer
169+
<*> pure (observer <> trackConnections connTrack)
161170
<*> JwtCache.init conf observer
162171
<*> pure loggerState
163172
<*> pure metricsState
173+
<*> pure connTrack
164174

165175
deb <-
166176
let decisecond = 100000 in
@@ -410,11 +420,22 @@ retryingSchemaCacheLoad appState@AppState{stateObserver=observer, stateMainThrea
410420
qSchemaCache :: IO (Maybe SchemaCache)
411421
qSchemaCache = do
412422
conf@AppConfig{..} <- getConfig appState
413-
-- Allow 10 concurrent schema cache loads, guarded by advisory locks.
414-
-- This is to prevent thundering herd problem on startup or when many PostgREST
415-
-- instances receive "reload schema" notifications at the same time
416-
let withTxLock = SQL.statement (50168275, 10) $
417-
SQL.Statement get_lock_sql get_lock_params HD.noResult configDbPreparedStatements
423+
-- Throttle concurrent schema cache loads, guarded by advisory locks.
424+
-- This is to prevent thundering herd problem on startup or when many PostgREST
425+
-- instances receive "reload schema" notifications at the same time
426+
-- schema reloading session + listener session
427+
-- See get_lock_sql for details of the algorithm.
428+
-- Here we calculate the number of open connections passed to the query.
429+
(connected, inUse) <- sessionCounts appState
430+
-- Determine whether schema cache loading will create a new session
431+
let
432+
scLoadingSessions = case (connected <= inUse, inUse >= configDbPoolSize) of
433+
(_, True) -> 0 -- all connections in use and pool full - schema cache loading won't create session
434+
(True, _) -> 1 -- all connections in use but pool not full - schema cache loading will create session
435+
(_, _) -> 0 -- not all connections in use - schema cache loading won't create session
436+
withTxLock = SQL.statement
437+
(fromIntegral $ connected + scLoadingSessions)
438+
(SQL.Statement get_lock_sql get_lock_params HD.noResult configDbPreparedStatements)
418439

419440
(resultTime, result) <-
420441
let transaction = if configDbPreparedStatements then SQL.transaction else SQL.unpreparedTransaction in
@@ -437,23 +458,42 @@ retryingSchemaCacheLoad appState@AppState{stateObserver=observer, stateMainThrea
437458
putSCacheStatus appState SCLoaded
438459
return $ Just sCache
439460
where
440-
-- recursive query that tries acquiring locks in order
441-
-- and waits for randomly selected lock if no attempt succeeded
442-
-- parameters are lock number and number of locks to try
461+
-- Recursive query that tries acquiring locks in order
462+
-- and waits for randomly selected lock if no attempt succeeded.
463+
-- It has a single parameter: this node open connection count.
464+
-- It is used to estimate the number of nodes
465+
-- by counting the number of active sessions for current session_user
466+
-- and dividing it by this node open connections.
467+
-- Assuming load is uniform among cluster nodes, all should have
468+
-- statistically the same number of open connections.
469+
-- Once the number of nodes is known we calculate the number
470+
-- of locks as ceil(log(2, number_of_nodes))
443471
get_lock_sql = encodeUtf8 [trimming|
444472
WITH RECURSIVE attempts AS (
445-
SELECT 1 AS lock_number, pg_try_advisory_xact_lock($$1, 1) AS success WHERE $$2 > 0
473+
SELECT 1 AS lock_number, pg_try_advisory_xact_lock(lock_id, 1) AS success FROM parameters
446474
UNION ALL
447-
SELECT next_lock_number AS lock_number, pg_try_advisory_xact_lock($$1, next_lock_number) AS success FROM (
448-
SELECT lock_number + 1 AS next_lock_number FROM attempts
449-
WHERE NOT success AND lock_number < $$2
450-
ORDER BY lock_number DESC
451-
LIMIT 1
452-
) AS previous_attempt
475+
SELECT next_lock_number AS lock_number, pg_try_advisory_xact_lock(lock_id, next_lock_number) AS success
476+
FROM
477+
parameters CROSS JOIN LATERAL (
478+
SELECT lock_number + 1 AS next_lock_number FROM attempts
479+
WHERE NOT success AND lock_number < locks_count
480+
ORDER BY lock_number DESC
481+
LIMIT 1
482+
) AS previous_attempt
483+
),
484+
counts AS (
485+
SELECT round(log(2, round(count(*)::double precision/$$1)::numeric))::int AS locks_count
486+
FROM
487+
pg_stat_activity WHERE usename = SESSION_USER
488+
),
489+
parameters AS (
490+
SELECT locks_count, 50168275 AS lock_id FROM counts WHERE locks_count > 0
453491
)
454-
SELECT pg_advisory_xact_lock($$1, floor(random() * $$2)::int + 1) WHERE NOT EXISTS (SELECT 1 FROM attempts WHERE success) |]
492+
SELECT pg_advisory_xact_lock(lock_id, floor(random() * locks_count)::int + 1)
493+
FROM
494+
parameters WHERE NOT EXISTS (SELECT 1 FROM attempts WHERE success) |]
455495

456-
get_lock_params = (fst >$< HE.param (HE.nonNullable HE.int4)) <> (snd >$< HE.param (HE.nonNullable HE.int4))
496+
get_lock_params = HE.param (HE.nonNullable HE.int4)
457497

458498
shouldRetry :: RetryStatus -> (Maybe PgVersion, Maybe SchemaCache) -> IO Bool
459499
shouldRetry _ (pgVer, sCache) = do
@@ -511,3 +551,21 @@ readInDbConfig startingUp appState@AppState{stateObserver=observer} = do
511551
pass
512552
else
513553
observer ConfigSucceededObs
554+
555+
556+
data ConnTrack = ConnTrack { connTrackConnected :: SH.SizedHamt UUID, connTrackInUse :: SH.SizedHamt UUID }
557+
558+
trackConnections :: ConnTrack -> ObservationHandler
559+
trackConnections ConnTrack{..} (HasqlPoolObs (ConnectionObservation uuid status)) = case status of
560+
ReadyForUseConnectionStatus -> atomically $
561+
SH.insert identity uuid connTrackConnected *>
562+
SH.focus Focus.delete identity uuid connTrackInUse
563+
TerminatedConnectionStatus _ -> atomically $
564+
SH.focus Focus.delete identity uuid connTrackConnected *>
565+
SH.focus Focus.delete identity uuid connTrackInUse
566+
InUseConnectionStatus -> atomically $ SH.insert identity uuid connTrackInUse
567+
_ -> mempty
568+
trackConnections _ _ = mempty
569+
570+
sessionCounts :: AppState -> IO (Int, Int)
571+
sessionCounts = atomically . bisequenceA . both SH.size . (connTrackConnected &&& connTrackInUse) . stateConnTrack

‎test/io/test_io.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1060,9 +1060,11 @@ def test_schema_cache_concurrent_notifications(slow_schema_cache_env):
10601060
assert response.status_code == 200
10611061

10621062

1063-
@pytest.mark.parametrize("instance_count", [10, 11])
1063+
@pytest.mark.parametrize(
1064+
"instance_count, expected_concurrency", [(2, 2), (4, 3), (6, 4), (8, 4), (16, 5)]
1065+
)
10641066
def test_schema_cache_reload_throttled_with_advisory_locks(
1065-
instance_count, slow_schema_cache_env
1067+
instance_count, expected_concurrency, slow_schema_cache_env
10661068
):
10671069
"schema cache reloads should be throttled across instances if instance count > 10"
10681070

@@ -1112,7 +1114,7 @@ def read_available_output_lines(postgrest):
11121114
reload_durations_ms = []
11131115
for postgrest in instances:
11141116
output_lines = []
1115-
for _ in range(5):
1117+
for _ in range(instance_count * 2):
11161118
output_lines.extend(read_available_output_lines(postgrest))
11171119
if any(query_log_pattern.search(line) for line in output_lines):
11181120
break
@@ -1131,14 +1133,15 @@ def read_available_output_lines(postgrest):
11311133

11321134
# 10 instances should be fast, remaining instances should be slow
11331135
assert (
1134-
len(
1136+
instance_count
1137+
- len(
11351138
[
11361139
duration
11371140
for duration in reload_durations_ms
11381141
if duration > lock_wait_threshold_ms
11391142
]
11401143
)
1141-
== instance_count - 10
1144+
== expected_concurrency
11421145
)
11431146

11441147

0 commit comments

Comments
 (0)