add(observation): db pool flushes observation

mkleczek · mkleczek · commit cefc5d33cd58 · 2026-02-24T18:03:49.000+01:00
Emit a dedicated PoolFlushed observation when the DB pool is released during schema cache reload.
diff --git a/postgrest.cabal b/postgrest.cabal
@@ -258,6 +258,7 @@ test-suite spec
                       Feature.Query.UpsertSpec
                       Feature.RollbackSpec
                       Feature.RpcPreRequestGucsSpec
+                      Feature.SchemaCacheSpec
                       SpecHelper
   build-depends:      base              >= 4.9 && < 4.20
                     , aeson             >= 2.0.3 && < 2.3
diff --git a/src/PostgREST/AppState.hs b/src/PostgREST/AppState.hs
@@ -18,6 +18,7 @@ module PostgREST.AppState
   , init
   , initSockets
   , initWithPool
+  , putConfig -- For testing, TODO refactoring
   , putNextListenerDelay
   , putSchemaCache
   , putPgVersion
@@ -270,10 +271,14 @@ usePool AppState{stateObserver=observer, stateMainThreadId=mainThreadId, ..} ses
 
 -- | Flush the connection pool so that any future use of the pool will
 -- use connections freshly established after this call.
+-- | Emits PoolFlushed observation
 flushPool :: AppState -> IO ()
-flushPool AppState{..} = SQL.release statePool
+flushPool AppState{..} = do
+  SQL.release statePool
+  stateObserver PoolFlushed
 
 -- | Destroy the pool on shutdown.
+-- | Differs from flushPool in not emiting PoolFlushed observation.
 destroyPool :: AppState -> IO ()
 destroyPool AppState{..} = SQL.release statePool
 
diff --git a/src/PostgREST/Logger.hs b/src/PostgREST/Logger.hs
@@ -102,6 +102,9 @@ observationLogger loggerState logLevel obs = case obs of
   o@PoolRequestFullfilled ->
     when (logLevel >= LogDebug) $ do
       logWithZTime loggerState $ observationMessage o
+  o@PoolFlushed ->
+    when (logLevel >= LogDebug) $ do
+      logWithZTime loggerState $ observationMessage o
   o@JwtCacheEviction ->
     when (logLevel >= LogDebug) $ do
       logWithZTime loggerState $ observationMessage o
diff --git a/src/PostgREST/Observation.hs b/src/PostgREST/Observation.hs
@@ -63,6 +63,7 @@ data Observation
   | HasqlPoolObs SQL.Observation
   | PoolRequest
   | PoolRequestFullfilled
+  | PoolFlushed
   | JwtCacheLookup Bool
   | JwtCacheEviction
   | WarpErrorObs Text
@@ -157,6 +158,8 @@ observationMessage = \case
     "Trying to borrow a connection from pool"
   PoolRequestFullfilled ->
     "Borrowed a connection from the pool"
+  PoolFlushed ->
+    "Database connection pool flushed"
   JwtCacheLookup _ ->
     "Looked up a JWT in JWT cache"
   JwtCacheEviction ->
diff --git a/test/io/test_io.py b/test/io/test_io.py
@@ -21,6 +21,19 @@
 )
 
 
+def match_log(output, matchers):
+    ito = iter(output)
+    itm = iter(matchers)
+    nextMatcher = next(itm, None)
+    while nextMatcher is not None and (line := next(ito, None)) is not None:
+        if re.match(nextMatcher, line) is not None:
+            nextMatcher = next(itm, None)
+    if nextMatcher is not None:
+        raise AssertionError(
+            f"Expected log line matching {nextMatcher} not found in output"
+        )
+
+
 def test_connect_with_dburi(dburi, defaultenv):
     "Connecting with db-uri instead of LIPQ* environment variables should work."
     defaultenv_without_libpq = {
@@ -661,54 +674,44 @@ def test_log_level(level, defaultenv):
         response = postgrest.session.get("/")
         assert response.status_code == 200
 
-        output = sorted(postgrest.read_stdout(nlines=7))
+        output = postgrest.read_stdout(nlines=7)
 
         if level == "crit":
             assert len(output) == 0
         elif level == "error":
-            assert re.match(
-                r'- - - \[.+\] "GET / HTTP/1.1" 500 \d+ "" "python-requests/.+"',
-                output[0],
+            match_log(
+                output,
+                [r'- - - \[.+\] "GET / HTTP/1.1" 500 \d+ "" "python-requests/.+"'],
             )
             assert len(output) == 1
         elif level == "warn":
-            assert re.match(
-                r'- - - \[.+\] "GET / HTTP/1.1" 500 \d+ "" "python-requests/.+"',
-                output[0],
-            )
-            assert re.match(
-                r'- - postgrest_test_anonymous \[.+\] "GET /unknown HTTP/1.1" 404 \d+ "" "python-requests/.+"',
-                output[1],
+            match_log(
+                output,
+                [
+                    r'- - - \[.+\] "GET / HTTP/1.1" 500 \d+ "" "python-requests/.+"',
+                    r'- - postgrest_test_anonymous \[.+\] "GET /unknown HTTP/1.1" 404 \d+ "" "python-requests/.+"',
+                ],
             )
             assert len(output) == 2
         elif level == "info":
-            assert re.match(
-                r'- - - \[.+\] "GET / HTTP/1.1" 500 \d+ "" "python-requests/.+"',
-                output[0],
-            )
-            assert re.match(
-                r'- - postgrest_test_anonymous \[.+\] "GET / HTTP/1.1" 200 \d+ "" "python-requests/.+"',
-                output[1],
-            )
-            assert re.match(
-                r'- - postgrest_test_anonymous \[.+\] "GET /unknown HTTP/1.1" 404 \d+ "" "python-requests/.+"',
-                output[2],
+            match_log(
+                output,
+                [
+                    r'- - - \[.+\] "GET / HTTP/1.1" 500 \d+ "" "python-requests/.+"',
+                    r'- - postgrest_test_anonymous \[.+\] "GET /unknown HTTP/1.1" 404 \d+ "" "python-requests/.+"',
+                    r'- - postgrest_test_anonymous \[.+\] "GET / HTTP/1.1" 200 \d+ "" "python-requests/.+"',
+                ],
             )
             assert len(output) == 3
         elif level == "debug":
-            assert re.match(
-                r'- - - \[.+\] "GET / HTTP/1.1" 500 \d+ "" "python-requests/.+"',
-                output[0],
+            match_log(
+                output,
+                [
+                    r'- - - \[.+\] "GET / HTTP/1.1" 500 \d+ "" "python-requests/.+"',
+                    r'- - postgrest_test_anonymous \[.+\] "GET /unknown HTTP/1.1" 404 \d+ "" "python-requests/.+"',
+                    r'- - postgrest_test_anonymous \[.+\] "GET / HTTP/1.1" 200 \d+ "" "python-requests/.+"',
+                ],
             )
-            assert re.match(
-                r'- - postgrest_test_anonymous \[.+\] "GET / HTTP/1.1" 200 \d+ "" "python-requests/.+"',
-                output[1],
-            )
-            assert re.match(
-                r'- - postgrest_test_anonymous \[.+\] "GET /unknown HTTP/1.1" 404 \d+ "" "python-requests/.+"',
-                output[2],
-            )
-
             assert len(output) == 7
             assert any("Connection" and "is available" in line for line in output)
             assert any("Connection" and "is used" in line for line in output)
@@ -1346,16 +1349,21 @@ def test_db_error_logging_to_stderr(level, defaultenv, metapostgrest):
         assert response.status_code == 500
 
         # ensure the message appears on the logs
-        output = sorted(postgrest.read_stdout(nlines=6))
+        output = postgrest.read_stdout(nlines=6)
 
         if level == "crit":
             assert len(output) == 0
         elif level == "debug":
-            assert " 500 " in output[0]
-            assert "canceling statement due to statement timeout" in output[5]
+            match_log(
+                output,
+                [
+                    r".*canceling statement due to statement timeout.*",
+                    r".*500.*",
+                ],
+            )
         else:
-            assert " 500 " in output[0]
-            assert "canceling statement due to statement timeout" in output[1]
+            assert " 500 " in output[1]
+            assert "canceling statement due to statement timeout" in output[0]
 
     reset_statement_timeout(metapostgrest, role)
 
@@ -1549,18 +1557,17 @@ def test_log_pool_req_observation(level, defaultenv):
 
     headers = jwtauthheader({"role": "postgrest_test_author"}, SECRET)
 
-    pool_req = "Trying to borrow a connection from pool"
-    pool_req_fullfill = "Borrowed a connection from the pool"
+    pool_req = r".*Trying to borrow a connection from pool.*"
+    pool_req_fullfill = r".*Borrowed a connection from the pool.*"
 
     with run(env=env) as postgrest:
 
         postgrest.session.get("/authors_only", headers=headers)
 
         if level == "debug":
-            output = postgrest.read_stdout(nlines=5)
-            assert pool_req in output[1]
-            assert pool_req_fullfill in output[4]
-            assert len(output) == 5
+            output = postgrest.read_stdout(nlines=7)
+            assert len(output) == 6
+            match_log(output, [pool_req, pool_req_fullfill])
         elif level == "info":
             output = postgrest.read_stdout(nlines=4)
             assert len(output) == 1
diff --git a/test/spec/Feature/SchemaCacheSpec.hs b/test/spec/Feature/SchemaCacheSpec.hs
@@ -0,0 +1,90 @@
+{-# LANGUAGE DataKinds           #-}
+{-# LANGUAGE DeriveAnyClass      #-}
+{-# LANGUAGE MonadComprehensions #-}
+module Feature.SchemaCacheSpec
+
+where
+
+import Network.Wai (Application)
+
+import qualified PostgREST.AppState      as AppState
+import           PostgREST.Config        (configDbSchemas)
+import           PostgREST.Observation
+import           Protolude
+import           System.Timeout          (timeout)
+import           Test.Hspec              (SpecWith, describe, it)
+import           Test.Hspec.Expectations
+import           Test.Hspec.Wai          (getState)
+
+data TimeoutException = TimeoutException deriving (Show, Exception)
+
+accumulateUntilTimeout :: Int -> (s -> a -> s) -> s -> IO a -> IO s
+accumulateUntilTimeout t f start act = do
+  tid <- myThreadId
+  -- mask to make sure TimeoutException is not thrown before starting the loop
+  mask $ \unmask -> do
+    -- start timeout thread unmasking exceptions
+    ttid <- forkIOWithUnmask ($ (threadDelay t *> throwTo tid TimeoutException))
+    -- unmask effect
+    unmask $ fix (\loop accum -> (act >>= loop . f accum) `onTimeout` pure accum) start
+      -- make sure we catch timeout if happens bifore entering the loop
+      `onTimeout` pure start
+      -- make sure timer thread is killed on other exceptions
+      -- so that it won't throw TimeoutException later
+      `onException` killThread ttid
+  where
+    onTimeout m a = m `catch` \TimeoutException -> a
+
+spec :: SpecWith ((AppState.AppState, Chan Observation), Application)
+spec = describe "Server started with metrics enabled" $ do
+
+  it "Should emit PoolFlushed, SchemaCacheQueriedObs and SchemaCacheLoadedObs when schema cache is reloaded" $ do
+    (appState, waitFor) <- prepareState
+
+    liftIO $ do
+      AppState.schemaCacheLoader appState
+
+      waitFor (1 * sec) "PoolFlushed" $ \x -> [ o | o@PoolFlushed <- pure x ]
+      waitFor (1 * sec) "SchemaCacheQueriedObs" $ \x -> [ o | o@SchemaCacheQueriedObs{} <- pure x ]
+      waitFor (1 * sec) "SchemaCacheLoadedObs" $ \x -> [ o | o@SchemaCacheLoadedObs{} <- pure x ]
+
+
+  it "Should flush pool multiple times when schema reloading retries" $ do
+    (appState, waitFor) <- prepareState
+
+    liftIO $ do
+        AppState.getConfig appState >>= \cfg -> do
+          AppState.putConfig appState $ cfg { configDbSchemas = pure "bad_schema" }
+          AppState.schemaCacheLoader appState
+
+          waitFor (1 * sec) "PoolFlushed 1" $ \x -> [ o | o@PoolFlushed <- pure x ]
+          waitFor (1 * sec) "SchemaCacheErrorObs" $ \x -> [ o | o@SchemaCacheErrorObs{} <- pure x ]
+
+          -- Restore configuration
+          AppState.putConfig appState cfg
+
+        -- Wait for 2 seconds so that retry can happen
+        waitFor (2 * sec) "PoolFlushed 2" $ \x -> [ o | o@PoolFlushed <- pure x ]
+        waitFor (1 * sec) "SchemaCacheQueriedObs" $ \x -> [ o | o@SchemaCacheQueriedObs{} <- pure x ]
+        waitFor (1 * sec) "SchemaCacheLoadedObs" $ \x -> [ o | o@SchemaCacheLoadedObs{} <- pure x ]
+
+  where
+    sec = 1000000
+    -- duplicate the provided channel and construct wairFor function binding both channels
+    prepareState = getState >>= traverse (liftA2 (<$>) waitFor (liftIO . dupChan))
+      where
+        -- execute effectful computation until result meets provided condition
+        untilM cond m = fix $ \loop -> m >>= \a -> ifM (cond a) (pure a) loop
+        -- accumulate effecful computation results into a list for specified time
+        takeUntilTimeout t = fmap reverse . accumulateUntilTimeout t (flip (:)) []
+        -- read messages from copy chan and once condition is met drain original to the same point
+        -- upon timeout report error and messages remaining in the original chan
+        -- that way we report messages since last successful read
+        waitFor original copy t msg f =
+          timeout t (readUntil copy *> readUntil original) >>= maybe failTimeout mempty
+          where
+            failTimeout = takeUntilTimeout 100000 (readChan original) >>=
+              expectationFailure .
+                ("Timeout waiting for " <> msg <> ". Remaining observations:\n" ++) .
+                foldMap ((++ "\n") . show . observationMessage)
+            readUntil = void . untilM (pure . isJust . f) . readChan
diff --git a/test/spec/Main.hs b/test/spec/Main.hs
@@ -68,6 +68,7 @@ import qualified Feature.Query.UpdateSpec
 import qualified Feature.Query.UpsertSpec
 import qualified Feature.RollbackSpec
 import qualified Feature.RpcPreRequestGucsSpec
+import qualified Feature.SchemaCacheSpec
 
 
 main :: IO ()
@@ -95,6 +96,13 @@ main = do
       AppState.putSchemaCache appState (Just sCache)
       return (st, postgrest (configLogLevel config) appState (pure ()))
 
+    initObservationsApp sCache config = do
+      obsChan <- newChan
+      appState <- AppState.initWithPool sockets pool config loggerState metricsState (Metrics.observationMetrics metricsState <> writeChan obsChan)
+      AppState.putPgVersion appState actualPgVersion
+      AppState.putSchemaCache appState (Just sCache)
+      return ((appState, obsChan), postgrest (configLogLevel config) appState (pure ()))
+
     -- For tests that run with the same schema cache
     app = initApp baseSchemaCache ()
 
@@ -123,6 +131,7 @@ main = do
       obsApp               = app testObservabilityCfg
       serverTiming         = app testCfgServerTiming
       aggregatesEnabled    = app testCfgAggregatesEnabled
+      observationsApp      = initObservationsApp baseSchemaCache testCfg
 
       extraSearchPathApp   = appDbs testCfgExtraSearchPath
       unicodeApp           = appDbs testUnicodeCfg
@@ -278,6 +287,9 @@ main = do
     before (initApp baseSchemaCache metricsState testCfgJwtCache) $
       describe "Feature.Auth.JwtCacheSpec" Feature.Auth.JwtCacheSpec.spec
 
+    before observationsApp $
+      describe "Feature.SchemaCacheSpec" Feature.SchemaCacheSpec.spec
+
   where
     loadSCache pool conf =
       either (panic.show) id <$> P.use pool (HT.transaction HT.ReadCommitted HT.Read $ querySchemaCache conf)
