Fix dropping schema cache reload notifications

* test: bad schema reload * refactor: DRY using the "extra" lib * refactor: move worker funtions inside AppState * Also rename Workers.hs to Admin.hs
2023-06-02 12:22:17 -05:00
parent 14be3fb671
commit a852b766eb
15 changed files with 460 additions and 409 deletions
@@ -1,9 +1,10 @@
+{-# LANGUAGE LambdaCase      #-}
+{-# LANGUAGE NamedFieldPuns  #-}
 {-# LANGUAGE RecordWildCards #-}

 module PostgREST.AppState
  ( AppState
  , destroy
-  , flushPool
  , getConfig
  , getSchemaCache
  , getIsListenerOn
@@ -11,40 +12,50 @@ module PostgREST.AppState
  , getPgVersion
  , getRetryNextIn
  , getTime
-  , getWorkerSem
  , init
  , initWithPool
  , logWithZTime
-  , logPgrstError
-  , putConfig
  , putSchemaCache
-  , putIsListenerOn
  , putPgVersion
-  , putRetryNextIn
-  , signalListener
  , usePool
-  , waitListener
  , debounceLogAcquisitionTimeout
+  , loadSchemaCache
+  , reReadConfig
+  , connectionWorker
+  , runListener
  ) where

-import qualified Data.ByteString.Lazy as LBS
-import qualified Data.Text.Encoding   as T
-import qualified Hasql.Pool           as SQL
-import qualified Hasql.Session        as SQL
-import qualified PostgREST.Error      as Error
+import qualified Data.ByteString            as BS
+import qualified Data.ByteString.Lazy       as LBS
+import qualified Data.Text.Encoding         as T
+import           Hasql.Connection           (acquire)
+import qualified Hasql.Notifications        as SQL
+import qualified Hasql.Pool                 as SQL
+import qualified Hasql.Session              as SQL
+import qualified Hasql.Transaction.Sessions as SQL
+import qualified PostgREST.Error            as Error

 import Control.AutoUpdate (defaultUpdateSettings, mkAutoUpdate,
                           updateAction)
 import Control.Debounce
+import Control.Retry      (RetryStatus, capDelay, exponentialBackoff,
+                           retrying, rsPreviousDelay)
 import Data.IORef         (IORef, atomicWriteIORef, newIORef,
                           readIORef)
 import Data.Time          (ZonedTime, defaultTimeLocale, formatTime,
                           getZonedTime)
 import Data.Time.Clock    (UTCTime, getCurrentTime)

-import PostgREST.Config           (AppConfig (..))
-import PostgREST.Config.PgVersion (PgVersion (..), minimumPgVersion)
-import PostgREST.SchemaCache      (SchemaCache)
+import PostgREST.Config                  (AppConfig (..),
+                                          readAppConfig)
+import PostgREST.Config.Database         (queryDbSettings,
+                                          queryPgVersion,
+                                          queryRoleSettings)
+import PostgREST.Config.PgVersion        (PgVersion (..),
+                                          minimumPgVersion)
+import PostgREST.SchemaCache             (SchemaCache,
+                                          querySchemaCache)
+import PostgREST.SchemaCache.Identifiers (dumpQi)

 import Protolude

@@ -56,8 +67,8 @@ data AppState = AppState
  , statePgVersion                :: IORef PgVersion
  -- | No schema cache at the start. Will be filled in by the connectionWorker
  , stateSchemaCache              :: IORef (Maybe SchemaCache)
-  -- | Binary semaphore to make sure just one connectionWorker can run at a time
-  , stateWorkerSem                :: MVar ()
+  -- | starts the connection worker with a debounce
+  , debouncedConnectionWorker     :: IO ()
  -- | Binary semaphore used to sync the listener(NOTIFY reload) with the connectionWorker.
  , stateListener                 :: MVar ()
  -- | State of the LISTEN channel, used for the admin server checks
@@ -86,7 +97,7 @@ initWithPool pool conf = do
  appState <- AppState pool
    <$> newIORef minimumPgVersion -- assume we're in a supported version when starting, this will be corrected on a later step
    <*> newIORef Nothing
-    <*> newEmptyMVar
+    <*> pure (pure ())
    <*> newEmptyMVar
    <*> newIORef False
    <*> newIORef conf
@@ -96,7 +107,8 @@ initWithPool pool conf = do
    <*> newIORef 0
    <*> pure (pure ())

-  deb <-
+
+  debLogTimeout <-
    let oneSecond = 1000000 in
    mkDebounce defaultDebounceSettings
       { debounceAction = logPgrstError appState SQL.AcquisitionTimeoutUsageError
@@ -104,7 +116,15 @@ initWithPool pool conf = do
       , debounceEdge = leadingEdge -- logs at the start and the end
       }

-  return appState { debounceLogAcquisitionTimeout = deb }
+  debWorker <-
+    let decisecond = 100000 in
+    mkDebounce defaultDebounceSettings
+       { debounceAction = internalConnectionWorker appState
+       , debounceFreq = decisecond
+       , debounceEdge = leadingEdge -- runs the worker at the start and the end
+       }
+
+  return appState { debounceLogAcquisitionTimeout = debLogTimeout, debouncedConnectionWorker = debWorker }

 destroy :: AppState -> IO ()
 destroy = destroyPool
@@ -143,8 +163,8 @@ getSchemaCache = readIORef . stateSchemaCache
 putSchemaCache :: AppState -> Maybe SchemaCache -> IO ()
 putSchemaCache appState = atomicWriteIORef (stateSchemaCache appState)

-getWorkerSem :: AppState -> MVar ()
-getWorkerSem = stateWorkerSem
+connectionWorker :: AppState -> IO ()
+connectionWorker = debouncedConnectionWorker

 getRetryNextIn :: AppState -> IO Int
 getRetryNextIn = readIORef . stateRetryNextIn
@@ -189,3 +209,236 @@ getIsListenerOn = readIORef . stateIsListenerOn

 putIsListenerOn :: AppState -> Bool -> IO ()
 putIsListenerOn = atomicWriteIORef . stateIsListenerOn
+
+-- | Schema cache status
+data SCacheStatus
+  = SCLoaded
+  | SCOnRetry
+  | SCFatalFail
+
+-- | Load the SchemaCache by using a connection from the pool.
+loadSchemaCache :: AppState -> IO SCacheStatus
+loadSchemaCache appState = do
+  conf@AppConfig{..} <- getConfig appState
+  result <-
+    let transaction = if configDbPreparedStatements then SQL.transaction else SQL.unpreparedTransaction in
+    usePool appState . transaction SQL.ReadCommitted SQL.Read $
+      querySchemaCache conf
+  case result of
+    Left e -> do
+      case Error.checkIsFatal e of
+        Just hint -> do
+          logWithZTime appState "A fatal error ocurred when loading the schema cache"
+          logPgrstError appState e
+          logWithZTime appState hint
+          return SCFatalFail
+        Nothing -> do
+          putSchemaCache appState Nothing
+          logWithZTime appState "An error ocurred when loading the schema cache"
+          logPgrstError appState e
+          return SCOnRetry
+
+    Right sCache -> do
+      putSchemaCache appState (Just sCache)
+      logWithZTime appState "Schema cache loaded"
+      return SCLoaded
+
+-- | Current database connection status data ConnectionStatus
+data ConnectionStatus
+  = NotConnected
+  | Connected PgVersion
+  | FatalConnectionError Text
+  deriving (Eq)
+
+-- | The purpose of this worker is to obtain a healthy connection to pg and an
+-- up-to-date schema cache(SchemaCache).  This method is meant to be called
+-- multiple times by the same thread, but does nothing if the previous
+-- invocation has not terminated. In all cases this method does not halt the
+-- calling thread, the work is performed in a separate thread.
+--
+-- Background thread that does the following :
+--  1. Tries to connect to pg server and will keep trying until success.
+--  2. Checks if the pg version is supported and if it's not it kills the main
+--     program.
+--  3. Obtains the sCache. If this fails, it goes back to 1.
+internalConnectionWorker :: AppState -> IO ()
+internalConnectionWorker appState = work
+  where
+    work = do
+      AppConfig{..} <- getConfig appState
+      logWithZTime appState "Attempting to connect to the database..."
+      connected <- establishConnection appState
+      case connected of
+        FatalConnectionError reason ->
+          -- Fatal error when connecting
+          logWithZTime appState reason >> killThread (getMainThreadId appState)
+        NotConnected ->
+          -- Unreachable because establishConnection will keep trying to connect
+          return ()
+        Connected actualPgVersion -> do
+          -- Procede with initialization
+          putPgVersion appState actualPgVersion
+          when configDbChannelEnabled $
+            signalListener appState
+          logWithZTime appState "Connection successful"
+          -- this could be fail because the connection drops, but the loadSchemaCache will pick the error and retry again
+          -- We cannot retry after it fails immediately, because db-pre-config could have user errors. We just log the error and continue.
+          when configDbConfig $ reReadConfig False appState
+          scStatus <- loadSchemaCache appState
+          case scStatus of
+            SCLoaded ->
+              -- do nothing and proceed if the load was successful
+              return ()
+            SCOnRetry ->
+              -- retry reloading the schema cache
+              work
+            SCFatalFail ->
+              -- die if our schema cache query has an error
+              killThread $ getMainThreadId appState
+
+-- | Repeatedly flush the pool, and check if a connection from the
+-- pool allows access to the PostgreSQL database.
+--
+-- Releasing the pool is key for rapid recovery. Otherwise, the pool
+-- timeout would have to be reached for new healthy connections to be acquired.
+-- Which might not happen if the server is busy with requests. No idle
+-- connection, no pool timeout.
+--
+-- The connection tries are capped, but if the connection times out no error is
+-- thrown, just 'False' is returned.
+establishConnection :: AppState -> IO ConnectionStatus
+establishConnection appState =
+  retrying retrySettings shouldRetry $
+    const $ flushPool appState >> getConnectionStatus
+  where
+    retrySettings = capDelay delayMicroseconds $ exponentialBackoff backoffMicroseconds
+    delayMicroseconds = 32000000 -- 32 seconds
+    backoffMicroseconds = 1000000 -- 1 second
+
+    getConnectionStatus :: IO ConnectionStatus
+    getConnectionStatus = do
+      pgVersion <- usePool appState $ queryPgVersion False -- No need to prepare the query here, as the connection might not be established
+      case pgVersion of
+        Left e -> do
+          logPgrstError appState e
+          case Error.checkIsFatal e of
+            Just reason ->
+              return $ FatalConnectionError reason
+            Nothing ->
+              return NotConnected
+        Right version ->
+          if version < minimumPgVersion then
+            return . FatalConnectionError $
+              "Cannot run in this PostgreSQL version, PostgREST needs at least "
+              <> pgvName minimumPgVersion
+          else
+            return . Connected  $ version
+
+    shouldRetry :: RetryStatus -> ConnectionStatus -> IO Bool
+    shouldRetry rs isConnSucc = do
+      let
+        delay = fromMaybe 0 (rsPreviousDelay rs) `div` backoffMicroseconds
+        itShould = NotConnected == isConnSucc
+      when itShould . logWithZTime appState $
+        "Attempting to reconnect to the database in "
+        <> (show delay::Text)
+        <> " seconds..."
+      when itShould $ putRetryNextIn appState delay
+      return itShould
+
+-- | Re-reads the config plus config options from the db
+reReadConfig :: Bool -> AppState -> IO ()
+reReadConfig startingUp appState = do
+  AppConfig{..} <- getConfig appState
+  dbSettings <-
+    if configDbConfig then do
+      qDbSettings <- usePool appState $ queryDbSettings (dumpQi <$> configDbPreConfig) configDbPreparedStatements
+      case qDbSettings of
+        Left e -> do
+          logWithZTime appState
+            "An error ocurred when trying to query database settings for the config parameters"
+          case Error.checkIsFatal e of
+            Just hint -> do
+              logPgrstError appState e
+              logWithZTime appState hint
+              killThread (getMainThreadId appState)
+            Nothing -> do
+              logPgrstError appState e
+          pure mempty
+        Right x -> pure x
+    else
+      pure mempty
+  roleSettings <-
+    if configDbConfig then do
+      rSettings <- usePool appState $ queryRoleSettings configDbPreparedStatements
+      case rSettings of
+        Left e -> do
+          logWithZTime appState "An error ocurred when trying to query the role settings"
+          logPgrstError appState e
+          pure mempty
+        Right x -> pure x
+    else
+      pure mempty
+  readAppConfig dbSettings configFilePath (Just configDbUri) roleSettings >>= \case
+    Left err   ->
+      if startingUp then
+        panic err -- die on invalid config if the program is starting up
+      else
+        logWithZTime appState $ "Failed reloading config: " <> err
+    Right newConf -> do
+      putConfig appState newConf
+      if startingUp then
+        pass
+      else
+        logWithZTime appState "Config reloaded"
+
+
+runListener :: AppConfig -> AppState -> IO ()
+runListener AppConfig{configDbChannelEnabled} appState =
+  when configDbChannelEnabled $ listener appState
+
+-- | Starts a dedicated pg connection to LISTEN for notifications.  When a
+-- NOTIFY <db-channel> - with an empty payload - is done, it refills the schema
+-- cache.  It uses the connectionWorker in case the LISTEN connection dies.
+listener :: AppState -> IO ()
+listener appState = do
+  AppConfig{..} <- getConfig appState
+  let dbChannel = toS configDbChannel
+
+  -- The listener has to wait for a signal from the connectionWorker.
+  -- This is because when the connection to the db is lost, the listener also
+  -- tries to recover the connection, but not with the same pace as the connectionWorker.
+  -- Not waiting makes stderr quickly fill with connection retries messages from the listener.
+  waitListener appState
+
+  -- forkFinally allows to detect if the thread dies
+  void . flip forkFinally (handleFinally dbChannel) $ do
+    dbOrError <- acquire $ toUtf8 configDbUri
+    case dbOrError of
+      Right db -> do
+        logWithZTime appState $ "Listening for notifications on the " <> dbChannel <> " channel"
+        putIsListenerOn appState True
+        SQL.listen db $ SQL.toPgIdentifier dbChannel
+        SQL.waitForNotifications handleNotification db
+      _ ->
+        die $ "Could not listen for notifications on the " <> dbChannel <> " channel"
+  where
+    handleFinally dbChannel _ = do
+      -- if the thread dies, we try to recover
+      logWithZTime appState $ "Retrying listening for notifications on the " <> dbChannel <> " channel.."
+      putIsListenerOn appState False
+      -- assume the pool connection was also lost, call the connection worker
+      connectionWorker appState
+      -- retry the listener
+      listener appState
+
+    handleNotification _ msg
+      | BS.null msg            = cacheReloader
+      | msg == "reload schema" = cacheReloader
+      | msg == "reload config" = reReadConfig False appState
+      | otherwise              = pure () -- Do nothing if anything else than an empty message is sent
+
+    cacheReloader =
+      -- reloads the schema cache + restarts pool connections
+      -- it's necessary to restart the pg connections because they cache the pg catalog(see #2620)
+      connectionWorker appState