fix: Start listening after schema cache load

This change ensures PostgREST starts listening on a server socket only after it loaded the schema cache and is ready to handle requests. It is no longer going to return 503 errors during startup until the schema cache is loaded.

Author: mkleczek

CHANGELOG.md

@@ -21,6 +21,7 @@ All notable changes to this project will be documented in this file. From versio
 - Shutdown should wait for in flight requests by @mkleczek in #4702
 - Remove automatic transaction retries on `40001 (serialization_failure)` errors to prevent replication lag by @laurenceisla in #3673
 - Fix unexpected results when embedding and filtering the same table more than once by @laurenceisla in #4075
+- Do not listen on main socket until schema cache is loaded by @mkleczek in #4880
 
 
 ### Changed

src/PostgREST/Admin.hs

@@ -22,20 +22,20 @@ import qualified PostgREST.AppState as AppState
 import qualified Network.Socket as NS
 import           Protolude
 
-runAdmin :: AppState -> Maybe NS.Socket -> NS.Socket -> Warp.Settings -> IO ()
-runAdmin appState maybeAdminSocket socketREST settings = do
+runAdmin :: AppState -> Maybe NS.Socket -> IO (Maybe NS.Socket) -> Warp.Settings -> IO ()
+runAdmin appState maybeAdminSocket getSocketREST settings = do
   whenJust maybeAdminSocket $ \adminSocket -> do
     address <- resolveSocketToAddress adminSocket
     observer $ AdminStartObs address
     void . forkIO $ Warp.runSettingsSocket settings adminSocket adminApp
   where
-    adminApp = admin appState socketREST
+    adminApp = admin appState getSocketREST
     observer = AppState.getObserver appState
 
 -- | PostgREST admin application
-admin :: AppState.AppState -> NS.Socket -> Wai.Application
-admin appState socketREST req respond  = do
-  isMainAppReachable  <- isRight <$> reachMainApp socketREST
+admin :: AppState.AppState -> IO (Maybe NS.Socket) -> Wai.Application
+admin appState getSocketREST req respond = do
+  isMainAppReachable <- getSocketREST >>= maybe (pure False) (fmap isRight . reachMainApp)
   isLoaded <- AppState.isLoaded appState
   isPending <- AppState.isPending appState
 
@@ -44,8 +44,8 @@ admin appState socketREST req respond  = do
       respond $ Wai.responseLBS (if isMainAppReachable then HTTP.status200 else HTTP.status500) [] mempty
     ["ready"] ->
       let
-        status | not isMainAppReachable = HTTP.status500
-               | isPending              = HTTP.status503
+        status | isPending              = HTTP.status503
+               | not isMainAppReachable = HTTP.status500
                | isLoaded               = HTTP.status200
                | otherwise              = HTTP.status500
       in

src/PostgREST/App.hs

@@ -25,6 +25,8 @@ import System.IO.Error  (ioeGetErrorType)
 import Control.Monad.Except     (liftEither)
 import Control.Monad.Extra      (whenJust)
 import Data.Either.Combinators  (mapLeft, whenLeft)
+import Data.IORef               (atomicWriteIORef, newIORef,
+                                 readIORef)
 import Data.String              (IsString (..))
 import Network.Wai.Handler.Warp (defaultSettings, setHost,
                                  setOnException, setPort,
@@ -63,11 +65,10 @@ import PostgREST.Version              (docsVersion, prettyVersion)
 
 import qualified Data.ByteString.Char8     as BS
 import qualified Data.List                 as L
-import           Data.Streaming.Network    (bindPortTCP,
-                                            bindRandomPortTCP)
+import           Data.Streaming.Network    (bindPortTCP)
 import qualified Data.Text                 as T
 import qualified Network.HTTP.Types        as HTTP
-import qualified Network.HTTP.Types.Header as HTTP (hVary)
+import qualified Network.HTTP.Types.Header as HTTP
 import qualified Network.Socket            as NS
 import           PostgREST.Unix            (createAndBindDomainSocket)
 import           Protolude                 hiding (Handler)
@@ -76,22 +77,30 @@ run :: AppState -> IO ()
 run appState = do
   conf <- AppState.getConfig appState
 
-  AppState.schemaCacheLoader appState -- Loads the initial SchemaCache
-  (mainSocket, adminSocket) <- initSockets conf
+  mainSocketRef <- newIORef Nothing
+  adminSocket <- initAdminServerSocket conf
+
   let closeSockets = do
         whenJust adminSocket NS.close
-        NS.close mainSocket
+        readIORef mainSocketRef >>= foldMap NS.close
   Unix.installSignalHandlers observer closeSockets (AppState.schemaCacheLoader appState) (AppState.readInDbConfig False appState)
 
+  Admin.runAdmin appState adminSocket (readIORef mainSocketRef) (serverSettings conf)
+
   Listener.runListener appState
 
-  Admin.runAdmin appState adminSocket mainSocket (serverSettings conf)
+  -- Kick off and wait for the initial SchemaCache load before creating the
+  -- main API socket.
+  AppState.schemaCacheLoader appState
+  AppState.waitForSchemaCacheLoaded appState
+
+  mainSocket <- initServerSocket conf
+  atomicWriteIORef mainSocketRef $ Just mainSocket
 
   let app = postgrest appState (AppState.schemaCacheLoader appState)
 
-  do
-    address <- resolveSocketToAddress mainSocket
-    observer $ AppServerAddressObs address
+  address <- resolveSocketToAddress mainSocket
+  observer $ AppServerAddressObs address
 
   Warp.runSettingsSocket (serverSettings conf & setOnException onWarpException) mainSocket app
   where
@@ -258,38 +267,16 @@ addRetryHint delay response = do
 isServiceUnavailable :: Wai.Response -> Bool
 isServiceUnavailable response = Wai.responseStatus response == HTTP.status503
 
-type AppSockets = (NS.Socket, Maybe NS.Socket)
-
-initSockets :: AppConfig -> IO AppSockets
-initSockets AppConfig{..} = do
-  let
-    cfg'usp = configServerUnixSocket
-    cfg'uspm = configServerUnixSocketMode
-    cfg'host = configServerHost
-    cfg'port = configServerPort
-    cfg'adminHost = configAdminServerHost
-    cfg'adminPort = configAdminServerPort
-
-  sock <- case cfg'usp of
-    -- I'm not using `streaming-commons`' bindPath function here because it's not defined for Windows,
-    -- but we need to have runtime error if we try to use it in Windows, not compile time error
-    Just path -> createAndBindDomainSocket path cfg'uspm
-    Nothing -> do
-      (_, sock) <-
-        if cfg'port /= 0
-          then do
-            sock <- bindPortTCP cfg'port (fromString $ T.unpack cfg'host)
-            pure (cfg'port, sock)
-          else do
-            -- explicitly bind to a random port, returning bound port number
-            (num, sock) <- bindRandomPortTCP (fromString $ T.unpack cfg'host)
-            pure (num, sock)
-      pure sock
-
-  adminSock <- case cfg'adminPort of
-    Just adminPort -> do
-      adminSock <- bindPortTCP adminPort (fromString $ T.unpack cfg'adminHost)
-      pure $ Just adminSock
-    Nothing -> pure Nothing
-
-  pure (sock, adminSock)
+initServerSocket :: AppConfig -> IO NS.Socket
+initServerSocket AppConfig{..} = case configServerUnixSocket of
+  -- I'm not using `streaming-commons`' bindPath function here because it's not defined for Windows,
+  -- but we need to have runtime error if we try to use it in Windows, not compile time error
+  Just path -> createAndBindDomainSocket path configServerUnixSocketMode
+  Nothing -> bindPortTCP configServerPort (fromString $ T.unpack configServerHost)
+
+initAdminServerSocket :: AppConfig -> IO (Maybe NS.Socket)
+initAdminServerSocket AppConfig{..} =
+  traverse (`bindPortTCP` adminHost) configAdminServerPort
+  where
+    adminHost = fromString $ T.unpack configAdminServerHost
+

src/PostgREST/AppState.hs

@@ -27,6 +27,7 @@ module PostgREST.AppState
   , getObserver
   , isLoaded
   , isPending
+  , waitForSchemaCacheLoaded
   ) where
 
 import qualified Data.ByteString.Char8      as BS
@@ -394,6 +395,9 @@ markSchemaCacheLoaded = void . (`tryPutMVar` ()) . getSCStatusMVar . stateSCache
 isSchemaCacheLoaded :: AppState -> IO Bool
 isSchemaCacheLoaded = fmap not . isEmptyMVar . getSCStatusMVar . stateSCacheStatus
 
+waitForSchemaCacheLoaded :: AppState -> IO ()
+waitForSchemaCacheLoaded = atomically . void . readTMVar . stateSchemaCache
+
 -- | Reads the in-db config and reads the config file again
 -- | We don't retry reading the in-db config after it fails immediately, because it could have user errors. We just report the error and continue.
 readInDbConfig :: Bool -> AppState -> IO ()

test/io/test_io.py

@@ -6,6 +6,7 @@
 import subprocess
 import time
 import pytest
+import requests
 
 from config import CONFIGSDIR, FIXTURES, SECRET
 from util import Thread, jwtauthheader, parse_server_timings_header, relativeSeconds
@@ -1085,33 +1086,6 @@ def test_invalid_rpc_method_log_contains_role(defaultenv):
     )
 
 
-def test_empty_schema_cache_log_contains_jwt_role(defaultenv):
-    "Requests are logged with the role when the schema cache is empty on startup"
-
-    env = {
-        **defaultenv,
-        "PGRST_INTERNAL_SCHEMA_CACHE_QUERY_SLEEP": "1000",
-        "PGRST_JWT_SECRET": SECRET,
-    }
-    headers = jwtauthheader({"role": "postgrest_test_author"}, SECRET)
-
-    with run(env=env, wait_for_readiness=False) as postgrest:
-        postgrest.wait_until_scache_starts_loading()
-
-        response = postgrest.session.get("/authors_only", headers=headers)
-        assert response.status_code == 503
-
-        output = drain_stdout(postgrest)
-
-    assert any(
-        re.match(
-            r'- - postgrest_test_author \[.+\] "GET /authors_only HTTP/1.1" 503 \d+ "" "python-requests/.+"',
-            line,
-        )
-        for line in output
-    )
-
-
 def test_no_pool_connection_required_on_bad_http_logic(defaultenv):
     "no pool connection should be consumed for failing on invalid http logic"
 
@@ -1518,9 +1492,9 @@ def test_log_postgrest_host_and_port(host, defaultenv):
         if is_unix:
             re.match(r'API server listening on "/tmp/.*\.sock"', output[2])
         elif is_ipv6(host):
-            assert f"API server listening on [{host}]:{port}" in output[2]
+            assert f"API server listening on [{host}]:{port}" in output[9]
         else:  # IPv4
-            assert f"API server listening on {host}:{port}" in output[2]
+            assert f"API server listening on {host}:{port}" in output[9]
 
 
 def test_succeed_w_role_having_superuser_settings(defaultenv):
@@ -1868,29 +1842,8 @@ def test_pgrst_log_503_client_error_to_stderr(defaultenv):
         assert any(log_message in line for line in output)
 
 
-def test_log_error_when_empty_schema_cache_on_startup_to_stderr(defaultenv):
-    "Should log the 503 error message when there is an empty schema cache on startup"
-
-    env = {
-        **defaultenv,
-        "PGRST_INTERNAL_SCHEMA_CACHE_QUERY_SLEEP": "300",
-    }
-
-    with run(env=env, wait_for_readiness=False) as postgrest:
-        postgrest.wait_until_scache_starts_loading()
-
-        response = postgrest.session.get("/projects")
-        assert response.status_code == 503
-
-        output_start = postgrest.read_stdout(nlines=10)
-
-        log_err_message = '{"code":"PGRST002","details":null,"hint":null,"message":"Could not query the database for the schema cache. Retrying."}'
-
-        assert any(log_err_message in line for line in output_start)
-
-
 def test_no_double_schema_cache_reload_on_empty_schema(defaultenv):
-    "Should only load the schema cache once on a 503 error when there's an empty schema cache on startup"
+    "Should only load the schema cache once when there's an empty schema cache on startup"
 
     env = {
         **defaultenv,
@@ -1900,12 +1853,15 @@ def test_no_double_schema_cache_reload_on_empty_schema(defaultenv):
     with run(env=env, port=freeport(), wait_for_readiness=False) as postgrest:
         postgrest.wait_until_scache_starts_loading()
 
-        response = postgrest.session.get("/projects")
-        assert response.status_code == 503
+        with pytest.raises(requests.ConnectionError):
+            postgrest.session.get("/projects")
 
         # Should wait enough time to load the schema cache twice to guarantee that the test is valid
         time.sleep(1)
 
+        response = postgrest.session.get("/projects")
+        assert response.status_code == 200
+
         response = postgrest.admin.get("/metrics")
         assert response.status_code == 200
         assert 'pgrst_schema_cache_loads_total{status="SUCCESS"} 1.0' in response.text
@@ -1987,7 +1943,7 @@ def test_schema_cache_error_observation(defaultenv):
         output = postgrest.read_stdout(nlines=9)
         assert (
             "Failed to load the schema cache using db-schemas=public and db-extra-search-path=x"
-            in output[7]
+            in output[6]
         )