fix: Flush pool as late as possible during schema cache reloading

retryingSchemaCacheLoad flushes the pool upon every retry before it starts reloading the schema. This is too early as schema reloading might take some time during which new connections might be acquired. The consequence is that:
* upon successful schema cache reload we might have some connections created with the old schema cache
* we close connections upon each retry and under load we will keep closing and re-opening connections until schema cache load succeeds

This change is to make sure we flush the pool only after successful schema cache querying but before loading (so that connections acquired during loading wait for it and do not interfere with timing the loading process).

Author: mkleczek

src/PostgREST/AppState.hs

@@ -368,8 +368,6 @@ retryingSchemaCacheLoad appState@AppState{stateObserver=observer, stateMainThrea
       observer $ ConnectionRetryObs delay
       putNextListenerDelay appState delay
 
-    flushPool appState
-
     (,) <$> qPgVersion <*> (qInDbConfig *> qSchemaCache)
   )
   where
@@ -417,8 +415,12 @@ retryingSchemaCacheLoad appState@AppState{stateObserver=observer, stateMainThrea
           -- IMPORTANT: While the pending schema cache state starts from running the above querySchemaCache, only at this stage we block API requests due to the usage of an
           -- IORef on putSchemaCache. This is why SCacheStatus is put at SCPending here to signal the Admin server (using isPending) that we're on a recovery state.
           putSCacheStatus appState SCPending
-          putSchemaCache appState $ Just sCache
           observer $ SchemaCacheQueriedObs resultTime
+          putSchemaCache appState $ Just sCache
+          -- Flush the pool after loading the schema cache to reset any stale session cache entries
+          -- We do it after successfully querying the schema cache
+          -- and after marking sCacheStatus as pending,
+          flushPool appState
           (t, _) <- timeItT $ observer $ SchemaCacheSummaryObs $ showSummary sCache
           observer $ SchemaCacheLoadedObs t
           putSCacheStatus appState SCLoaded

test/io/postgrest.py

@@ -187,6 +187,7 @@ def wait_until_exit(postgrest):
 def wait_until_status_code(url, max_seconds, status_code):
     "Wait for the given HTTP endpoint to return a status code"
     session = requests_unixsocket.Session()
+    response = None
 
     for _ in range(max_seconds * 10):
         try:

test/io/test_io.py

@@ -18,6 +18,7 @@
     sleep_until_postgrest_full_reload,
     sleep_until_postgrest_scache_reload,
     wait_until_exit,
+    wait_until_status_code,
 )
 
 
@@ -105,6 +106,88 @@ def sleep():
         t.join()
 
 
+def test_pool_flushes_metric(defaultenv):
+    "Should increase pool flushes metric when the pool is flushed"
+
+    with run(env=defaultenv, port=freeport()) as postgrest:
+        response = postgrest.admin.get("/metrics")
+        assert response.status_code == 200
+        before = float(
+            re.search(r"pgrst_db_pool_flushes_total ([0-9.e+-]+)", response.text).group(
+                1
+            )
+        )
+
+        postgrest.process.send_signal(signal.SIGUSR1)
+        sleep_until_postgrest_scache_reload()
+
+        response = postgrest.admin.get("/metrics")
+        assert response.status_code == 200
+        after = float(
+            re.search(r"pgrst_db_pool_flushes_total ([0-9.e+-]+)", response.text).group(
+                1
+            )
+        )
+        assert after == before + 1
+
+
+def test_pool_flushes_metric_with_schema_cache_retries(defaultenv, metapostgrest):
+    "Should flush the pool exactly once even when schema cache reload retries"
+
+    role = "timeout_authenticator"
+    app_name = "pool-flush-retry"
+    env = {
+        **defaultenv,
+        "PGUSER": role,
+        "PGAPPNAME": app_name,
+        "PGRST_INTERNAL_SCHEMA_CACHE_QUERY_SLEEP": "50",
+    }
+
+    with run(env=env, port=freeport()) as postgrest:
+        response = postgrest.admin.get("/metrics")
+        assert response.status_code == 200
+        before = float(
+            re.search(r"pgrst_db_pool_flushes_total ([0-9.e+-]+)", response.text).group(
+                1
+            )
+        )
+
+        try:
+            # Force schema cache reload failures to trigger retries.
+            set_statement_timeout(metapostgrest, role, 20)
+            postgrest.process.send_signal(signal.SIGUSR1)
+
+            postgrest.wait_until_scache_starts_loading(max_seconds=2)
+
+            # Give retry loop time to run and verify it doesn't flush the pool.
+            time.sleep(1)
+
+            response = postgrest.admin.get("/metrics")
+            assert response.status_code == 200
+            during_retries = float(
+                re.search(
+                    r"pgrst_db_pool_flushes_total ([0-9.e+-]+)", response.text
+                ).group(1)
+            )
+            assert during_retries == before
+
+            reset_statement_timeout(metapostgrest, role)
+            # Ensure next retry establishes fresh sessions with the reset timeout.
+            metapostgrest.session.get(f"/rpc/terminate_pgrst?appname={app_name}")
+            wait_until_status_code(postgrest.admin.baseurl + "/ready", 12, 200)
+        finally:
+            reset_statement_timeout(metapostgrest, role)
+
+        response = postgrest.admin.get("/metrics")
+        assert response.status_code == 200
+        after = float(
+            re.search(r"pgrst_db_pool_flushes_total ([0-9.e+-]+)", response.text).group(
+                1
+            )
+        )
+        assert after == before + 1
+
+
 def test_random_port_bound(defaultenv):
     "PostgREST should bind to a random port when PGRST_SERVER_PORT is 0."
 
@@ -661,55 +744,53 @@ 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=9)
+
+        def match_log(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"
+                )
 
         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(
+                [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(
+                [
+                    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(
+                [
+                    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(
+                [
+                    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 len(output) == 9
             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 +1427,16 @@ 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=8)
 
         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]
+            assert " 500 " in output[7]
+            assert "canceling statement due to statement timeout" in output[6]
         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)
 
@@ -1460,6 +1541,7 @@ def test_admin_metrics(defaultenv):
         assert "pgrst_db_pool_waiting" in response.text
         assert "pgrst_db_pool_available" in response.text
         assert "pgrst_db_pool_timeouts_total" in response.text
+        assert "pgrst_db_pool_flushes_total" in response.text
 
 
 def test_schema_cache_startup_load_with_in_db_config(defaultenv, metapostgrest):
@@ -1557,10 +1639,10 @@ def test_log_pool_req_observation(level, defaultenv):
         postgrest.session.get("/authors_only", headers=headers)
 
         if level == "debug":
-            output = postgrest.read_stdout(nlines=5)
+            output = postgrest.read_stdout(nlines=7)
+            assert len(output) == 7
             assert pool_req in output[1]
-            assert pool_req_fullfill in output[4]
-            assert len(output) == 5
+            assert pool_req_fullfill in output[6]
         elif level == "info":
             output = postgrest.read_stdout(nlines=4)
             assert len(output) == 1

test/spec/Feature/MetricsSpec.hs

@@ -16,19 +16,14 @@ import           Test.Hspec            (SpecWith, describe,
                                         expectationFailure, it)
 import           Test.Hspec.Wai        (getState)
 
-untilM :: Int -> (a -> Bool) -> IO a -> IO (Maybe a)
+untilM :: Int -> (a -> Maybe b) -> IO a -> IO (Maybe b)
 untilM timeout cond act = rightToMaybe <$> race (threadDelay timeout) (fix $
-  \loop -> do
-    value <- act
-    if cond value then
-      pure value
-    else
-      loop)
+  \loop -> act >>= maybe loop pure . cond)
 
 waitForSchemaReload :: Int -> IO Observation -> IO (Maybe Observation)
 waitForSchemaReload timeout = untilM timeout $ \case
-  SchemaCacheLoadedObs _ -> True
-  _ -> False
+  o@(SchemaCacheLoadedObs _) -> pure o
+  _ -> empty
 
 spec :: SpecWith ((MetricsState, AppState.AppState, IO Observation), Application)
 spec = describe "Server started with metrics enabled" $