add: set keepalives options in listener connection string

This change adds keepalive options to listener connections string. The value of keepalives_idle is adjusted using a simple algorithm based on idle time tracking.

Author: mkleczek

src/PostgREST/App.hs

@@ -81,7 +81,7 @@ run appState = do
 
   Unix.installSignalHandlers (AppState.getMainThreadId appState) (AppState.schemaCacheLoader appState) (AppState.readInDbConfig False appState)
 
-  Listener.runListener appState
+  void $ Listener.runListener appState
 
   Admin.runAdmin appState adminSocket mainSocket (serverSettings conf)
 

src/PostgREST/Config.hs

@@ -25,6 +25,7 @@ module PostgREST.Config
   , readPGRSTEnvironment
   , toURI
   , parseSecret
+  , addConnStringOption
   , addFallbackAppName
   , addTargetSessionAttrs
   , exampleConfigFile
@@ -619,7 +620,7 @@ pgConnString conn | uriDesignator `T.isPrefixOf` conn || shortUriDesignator `T.i
 -- addFallbackAppName ver "postgresql:///postgres?host=/run/user/1000/postgrest/postgrest-with-postgresql-16-BuR/socket&user=some_protected_user&password=invalid_pass"
 -- "postgresql:///postgres?host=/run/user/1000/postgrest/postgrest-with-postgresql-16-BuR/socket&user=some_protected_user&password=invalid_pass&fallback_application_name=PostgREST%2011.1.0%20%285a04ec7%29"
 addFallbackAppName :: ByteString -> Text -> Text
-addFallbackAppName version dbUri = addConnStringOption dbUri "fallback_application_name" pgrstVer
+addFallbackAppName version = addConnStringOption "fallback_application_name" pgrstVer
   where
     pgrstVer = "PostgREST " <> T.decodeUtf8 version
 
@@ -641,10 +642,10 @@ addFallbackAppName version dbUri = addConnStringOption dbUri "fallback_applicati
 -- >>> addTargetSessionAttrs "host=localhost port=5432 dbname=postgres"
 -- "host=localhost port=5432 dbname=postgres target_session_attrs='read-write'"
 addTargetSessionAttrs :: Text -> Text
-addTargetSessionAttrs dbUri = addConnStringOption dbUri "target_session_attrs" "read-write"
+addTargetSessionAttrs = addConnStringOption "target_session_attrs" "read-write"
 
 addConnStringOption :: Text -> Text -> Text -> Text
-addConnStringOption dbUri key val = dbUri <>
+addConnStringOption key val dbUri= dbUri <>
   case pgConnString dbUri of
     Nothing  -> mempty
     Just PGKeyVal -> " " <> keyValFmt

src/PostgREST/Listener.hs

@@ -1,8 +1,11 @@
-{-# LANGUAGE LambdaCase      #-}
-{-# LANGUAGE MultiWayIf      #-}
-{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE DeriveAnyClass      #-}
+{-# LANGUAGE LambdaCase          #-}
+{-# LANGUAGE MultiWayIf          #-}
+{-# LANGUAGE RecordWildCards     #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications    #-}
 
-module PostgREST.Listener (runListener) where
+module PostgREST.Listener (runListener, runListener') where
 
 import qualified Data.ByteString.Char8 as BS
 
@@ -19,28 +22,54 @@ import qualified PostgREST.Config   as Config
 import           Control.Arrow              ((&&&))
 import           Data.Bitraversable         (bisequence)
 import           Data.Either.Combinators    (whenRight)
+import           Data.IORef                 (IORef, newIORef,
+                                             readIORef, writeIORef)
 import qualified Data.Text                  as T
+import           Data.Time                  (UTCTime, diffUTCTime,
+                                             nominalDiffTimeToSeconds)
 import qualified Database.PostgreSQL.LibPQ  as LibPQ
 import qualified Hasql.Session              as SQL
 import           PostgREST.Config.Database  (queryPgVersion)
 import           PostgREST.Config.PgVersion (pgvFullName)
 import           Protolude
+import           System.IO.Error            (isResourceVanishedError)
 
 -- | Starts the Listener in a thread
-runListener :: AppState -> IO ()
-runListener appState = do
+-- | Returns IO action to stop the listener thread.
+runListener :: AppState -> IO (IO ())
+runListener appState = runListener' appState (15 * minute) (30 * minute)
+  where
+    minute = 60
+
+data ListenerStopped = ListenerStopped deriving (Show, Exception)
+
+runListener' :: AppState -> Int -> Int -> IO (IO ())
+runListener' appState initialTcpKeepAlivesIdleSec maxTcpKeepAlivesIdleSec = do
   AppConfig{..} <- getConfig appState
-  when configDbChannelEnabled $
-    void . forkIO . void $ retryingListen appState
+  if configDbChannelEnabled then do
+    started <- newIORef Nothing
+    listenerThreadId <- forkIO . void $ retryingListen started initialTcpKeepAlivesIdleSec maxTcpKeepAlivesIdleSec False appState
+    pure $ throwTo listenerThreadId ListenerStopped
+  else
+    mempty
 
 -- | Starts a LISTEN connection and handles notifications. It recovers with exponential backoff with a cap of 32 seconds, if the LISTEN connection is lost.
 -- | This function never returns (but can throw) and return type enforces that.
-retryingListen :: AppState -> IO Void
-retryingListen appState = do
+retryingListen :: IORef (Maybe UTCTime) -> Int -> Int -> Bool -> AppState -> IO ()
+retryingListen lastActivity currentKeepalivesIdle maxKeepalivesIdle retryingOnIdleTimeout appState = do
   AppConfig{..} <- AppState.getConfig appState
   let
     dbChannel = toS configDbChannel
     onError err = do
+      -- ResourceVanished should be reported when reading from socket fails
+      -- as long as hasql-notifications does not wrap IOException in something else...
+      let resourceVanished = maybe False isResourceVanishedError (fromException @IOException err)
+      (newTcpIdle, newMaxKeepalivesIdle) <-
+        if resourceVanished then do
+          readIORef lastActivity >>= maybe (pure (currentKeepalivesIdle, maxKeepalivesIdle)) adjustTcpIdle
+        else
+          pure (currentKeepalivesIdle, maxKeepalivesIdle)
+      writeIORef lastActivity Nothing
       AppState.putIsListenerOn appState False
       observer $ DBListenFail dbChannel (Right err)
       when (isDbListenerBug err) $
@@ -55,14 +84,14 @@ retryingListen appState = do
       unless (delay == maxDelay) $
         AppState.putNextListenerDelay appState (delay * 2)
       -- loop running the listener
-      retryingListen appState
+      retryingListen lastActivity newTcpIdle newMaxKeepalivesIdle resourceVanished appState
 
   -- Execute the listener with with error handling
-  handle onError $ do
+  handle onError $ handle (\ListenerStopped -> pure ()) $ do
     -- Make sure we don't leak connections on errors
     bracket
       -- acquire connection
-      (SQL.acquire $ toUtf8 (Config.addTargetSessionAttrs $ Config.addFallbackAppName prettyVersion configDbUri))
+      (SQL.acquire $ toUtf8 (addKeepalivesOptions $ Config.addTargetSessionAttrs $ Config.addFallbackAppName prettyVersion configDbUri))
       -- release connection
       (`whenRight` releaseConnection) $
       -- use connection
@@ -82,6 +111,7 @@ retryingListen appState = do
             AppState.putNextListenerDelay appState 1
 
           observer $ DBListenStart pqHost pqPort pgFullName dbChannel
+          AppState.getTime appState >>= writeIORef lastActivity . pure
 
           -- wait for notifications
           -- this will never return, in case of an error it will throw and be caught by onError
@@ -96,15 +126,46 @@ retryingListen appState = do
     oneSecondInMicro = 1000000
     maxDelay = 32
 
-    handleNotification channel msg =
+    handleNotification channel msg = do
       if | BS.null msg            -> observer (DBListenerGotSCacheMsg channel) >> cacheReloader
          | msg == "reload schema" -> observer (DBListenerGotSCacheMsg channel) >> cacheReloader
          | msg == "reload config" -> observer (DBListenerGotConfigMsg channel) >> AppState.readInDbConfig False appState
          | otherwise              -> pure () -- Do nothing if anything else than an empty message is sent
+      AppState.getTime appState >>= writeIORef lastActivity . Just
+
 
     cacheReloader =
       AppState.schemaCacheLoader appState
 
     releaseConnection = void . forkIO . handle (observer . DBListenerConnectionCleanupFail) . SQL.release
 
     isDbListenerBug e = "could not access status of transaction" `T.isInfixOf` show e
+
+    adjustTcpIdle lastActiveTime = do
+        currentIdleSeconds <- AppState.getTime appState <&> round . nominalDiffTimeToSeconds . (`diffUTCTime` lastActiveTime)
+        let currentIdleTimeout = currentKeepalivesIdle + keepalivesInterval * keepalivesCount
+        -- if our idle time == current idle timeout setting it means
+        -- we have to make it shorter
+        if currentIdleSeconds `div` currentIdleTimeout <= 1 then
+          -- only adjust if this is the second idle timeout failure
+          -- this is to eliminate spurious adjustments (TODO rethink if it is really needed)
+          if retryingOnIdleTimeout then
+            -- try with 1/2 of current keepalive idle
+            -- remember that it is the new maximum we can try later
+            pure (max 1 $ currentKeepalivesIdle `div` 2, currentKeepalivesIdle)
+          else
+            pure (currentKeepalivesIdle, maxKeepalivesIdle)
+        else
+          pure (currentKeepalivesIdle + (maxKeepalivesIdle - currentKeepalivesIdle) `div` 2, maxKeepalivesIdle)
+
+    keepalivesInterval = max 1 $ currentKeepalivesIdle `div` (5 * keepalivesCount)
+    keepalivesCount = 5
+
+    -- (Config.addConnStringOption opt val) is an endomorphism
+    -- so it is a Monoid under function composition
+    -- Haskell is awesome
+    addKeepalivesOptions = appEndo $ foldMap (Endo . uncurry Config.addConnStringOption . fmap show) [
+        ("keepalives_count", keepalivesCount)
+      , ("keepalives_interval", keepalivesInterval)
+      , ("keepalives_idle", currentKeepalivesIdle)
+      ]

test/spec/Feature/ToxiSpec.hs

@@ -2,6 +2,7 @@
 {-# LANGUAGE FlexibleContexts    #-}
 {-# LANGUAGE MonadComprehensions #-}
 {-# LANGUAGE NamedFieldPuns      #-}
+{-# LANGUAGE NumericUnderscores #-}
 module Feature.ToxiSpec where
 
 import           Control.Monad.Trans.Control (liftBaseDiscard)
@@ -11,7 +12,10 @@ import           Protolude                   hiding (get)
 import           SpecHelper
 import           Test.Hspec                  (SpecWith, describe, it)
 import           Test.Hspec.Wai
-import           Toxiproxy                   (withDisabled)
+import           Toxiproxy                   (withDisabled, Toxic (..), Stream (..), withToxic, ToxicType (..))
+import PostgREST.Observation (Observation(..))
+import PostgREST.Listener ( runListener, runListener' )
+import qualified Data.Map as M
 
 spec :: SpecWith (SpecState, Application)
 spec = describe "Tests using Toxiproxy" $ do
@@ -32,3 +36,73 @@ spec = describe "Tests using Toxiproxy" $ do
     liftBaseDiscard (withDisabled specToxiProxy) $ do
         void $ get "/items?id=eq.5"
             `shouldRespondWith` 503
+
+  describe "Toxiproxy tests of notification listener" $ do
+    it "should start listener" $ do
+      SpecState {specAppState, specObsChan} <- getState
+      let waitFor = waitForObs specObsChan
+
+      liftIO $ withListener specAppState $
+          waitFor (1*sec) "DBListenStart" $ \x -> [ o | o@DBListenStart{} <- pure x]
+
+    it "should retry listener" $ do
+      SpecState {specAppState, specObsChan, specToxiProxy} <- getState
+      let waitFor = waitForObs specObsChan
+
+      liftIO $ bracket (
+        withDisabled specToxiProxy $ do
+          stopListener <- runListener specAppState
+          (do
+            waitFor (1*sec) "DBListenFail" $ \x -> [ o | o@DBListenFail{} <- pure x ]
+            waitFor (1*sec) "DBListenRetry" $ \x -> [ o | o@DBListenRetry{} <- pure x ])
+            `onException` stopListener
+          pure stopListener)
+        identity
+        (const $ waitFor (2*sec) "DBListenStart" $ \x -> [ o | o@DBListenStart{} <- pure x])
+
+    it "should retry listener with exponential backoff when connection broken" $ do
+      SpecState {specAppState, specObsChan, specToxiProxy} <- getState
+      let waitFor = waitForObs specObsChan
+
+      liftIO $ withListener specAppState $ do
+        waitFor (1*sec) "DBListenStart" $ \x -> [ o | o@DBListenStart{} <- pure x]
+        withDisabled specToxiProxy $ do
+          waitFor (1*sec) "DBListenFail" $ \x -> [ o | o@DBListenFail{} <- pure x ]
+          waitFor (1*sec) "DBListenRetry 1" $ \x -> [ o | o@(DBListenRetry 1) <- pure x ]
+          waitFor (2*sec) "DBListenRetry 2" $ \x -> [ o | o@(DBListenRetry 2) <- pure x ]
+          waitFor (3*sec) "DBListenRetry 4" $ \x -> [ o | o@(DBListenRetry 4) <- pure x ]
+        waitFor (5*sec) "DBListenStart after retries" $ \x -> [ o | o@DBListenStart{} <- pure x]
+        waitFor (1*sec) "SchemaCacheLoadedObs" $ \x -> [ o | o@SchemaCacheLoadedObs{} <- pure x ]
+
+    -- this scenario cannot be tested with Toxiproxy
+    -- because keepalives are handled by the kernel TCP/IP stack
+    -- left here as an example and template for a test
+    -- using some future more advanced tool
+    it "should detect broken connection using keepalives" $ do
+      pendingWith "Cannot be tested with Toxiproxy"
+      SpecState {specAppState, specObsChan, specToxiProxy} <- getState
+      let waitFor = waitForObs specObsChan
+
+      liftIO $ withKeepAliveListener specAppState $ do
+        waitFor (1*sec) "DBListenStart" $ \x -> [ o | o@DBListenStart{} <- pure x]
+        withTimedOutConnection specToxiProxy $ do
+          waitFor (10*sec) "DBListenFail" $ \x -> [ o | o@DBListenFail{} <- pure x ]
+          waitFor (1*sec) "DBListenRetry 1" $ \x -> [ o | o@(DBListenRetry 1) <- pure x ]
+        waitFor (2*sec) "DBListenStart after timeout toxic" $ \x -> [ o | o@DBListenStart{} <- pure x]
+
+  where
+    withListener appState = bracket (runListener appState) identity . const
+    withKeepAliveListener appState = bracket (runListener' appState 1 1) identity . const
+    withTimedOutConnection proxy =
+      withToxic proxy (timeoutToxic Upstream) .
+      withToxic proxy (timeoutToxic Downstream)
+    timeoutToxic stream = Toxic
+      { toxicName = case stream of
+          Upstream   -> "listener-timeout-upstream"
+          Downstream -> "listener-timeout-downstream"
+      , toxicType = Timeout
+      , toxicStream = stream
+      , toxicToxicity = 1
+      , toxicAttributes = M.fromList [("timeout", 0)]
+      }
+    sec = 1_000_000

test/spec/SpecHelper.hs

@@ -149,7 +149,7 @@ baseCfg = let secret = encodeUtf8 "reallyreallyreallyreallyverysafe" in
   , configClientErrorVerbosity      = Verbose
   , configDbAggregates              = False
   , configDbAnonRole                = Just "postgrest_test_anonymous"
-  , configDbChannel                 = mempty
+  , configDbChannel                 = "pgrst"
   , configDbChannelEnabled          = True
   , configDbExtraSearchPath         = []
   , configDbHoistedTxSettings       = ["default_transaction_isolation","plan_filter.statement_cost_limit","statement_timeout"]