fix: db-channel-enabled not reloadable on config reload

Fixes #4894.

Signed-off-by: Taimoor Zaeem <taimoorzaeem@gmail.com>

src/library/PostgREST/AppState.hs

@@ -75,6 +75,7 @@ initWithPool pool confRef loggerState metricsState observer appKiller = mdo
     <*> newIORef Nothing
     <*> newSchemaCacheStatus
     <*> newIORef False
+    <*> newIORef Nothing
     <*> makeDebouncer (retryingSchemaCacheLoad appState *> threadDelay 100000)  -- 100ms cooldown
     <*> pure confRef
     <*> mkAutoUpdate defaultUpdateSettings { updateAction = getCurrentTime }

src/library/PostgREST/AppState/Reload.hs

@@ -184,42 +184,55 @@ readInDbConfig startingUp appState@AppState{stateObserver=observer} = do
       -- entries, because they were cached using the old secret
       update (getJwtCacheState appState) newConf
 
+      -- If db-channel-enabled is changed, then reload listener
+      when (((/=) `on` configDbChannelEnabled) newConf oldConf) $ do
+        -- 1. Kill the listener thread
+        getListenerThreadId appState >>= mapM_ (`throwTo` ListenerRestart)
+        putIsListenerOn appState False
+        -- 2. Restart listener
+        runListener appState
+
       if startingUp then
         pass
       else
         observer ConfigSucceededObs
 
+
 -- | Starts the Listener in a thread
 runListener :: AppState -> IO ()
 runListener appState = do
   AppConfig{..} <- getConfig appState
   when configDbChannelEnabled $ do
     nextDelay <- newIORef 1
-    void . forkIO . void $ retryingListen appState nextDelay False
+    listenerThreadId <- forkIO . void $ retryingListen appState nextDelay False
+    putListenerThreadId appState (Just listenerThreadId)
 
 -- | 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 -> IORef Int -> Bool -> IO Void
+retryingListen :: AppState -> IORef Int -> Bool -> IO ()
 retryingListen appState nextDelay hasDbListenerBug = do
   cfg@AppConfig{..} <- getConfig appState
   let
     dbChannel = toS configDbChannel
-    onError err = do
-      putIsListenerOn appState False
-      observer $ DBListenFail dbChannel (Right err)
-      when (isDbListenerBug err) $
-        observer DBListenBugCallQueryFix
-      unless configDbPoolAutomaticRecovery $
-        killApp appState
 
-      -- retry the listener
-      delay <- readIORef nextDelay
-      observer $ DBListenRetry delay
-      threadDelay (delay * oneSecondInMicro)
-      unless (delay == maxDelay) $
-        writeIORef nextDelay (delay * 2)
-      -- loop running the listener
-      retryingListen appState nextDelay (isDbListenerBug err)
+    onError err = case fromException err of
+      Just ListenerRestart -> traverse_ killThread =<< getListenerThreadId appState
+      Nothing -> do -- for any other exception
+        putIsListenerOn appState False
+        observer $ DBListenFail dbChannel (Right err)
+        when (isDbListenerBug err) $
+          observer DBListenBugCallQueryFix
+        unless configDbPoolAutomaticRecovery $
+          killApp appState
+
+        -- retry the listener
+        delay <- readIORef nextDelay
+        observer $ DBListenRetry delay
+        threadDelay (delay * oneSecondInMicro)
+        unless (delay == maxDelay) $
+          writeIORef nextDelay (delay * 2)
+        -- loop running the listener
+        retryingListen appState nextDelay (isDbListenerBug err)
 
   -- Execute the listener with error handling
   handle onError $ do

src/library/PostgREST/AppState/Types.hs

@@ -2,6 +2,7 @@
 Module      : PostgREST.AppState.Types
 Description : AppState data type and stateful functions
 -}
+{-# LANGUAGE DeriveAnyClass #-}
 module PostgREST.AppState.Types where
 
 import qualified Hasql.Pool              as SQL
@@ -32,6 +33,8 @@ data AppState = AppState
   , stateSCacheStatus     :: SchemaCacheStatus
   -- | State of the LISTEN channel
   , stateIsListenerOn     :: IORef Bool
+  -- | Listener Thread ID
+  , stateListenerThreadId :: IORef (Maybe ThreadId)
   -- | starts the connection worker with a debounce
   , debouncedSCacheLoader :: IO ()
   -- | Config that can change at runtime
@@ -58,6 +61,12 @@ newtype SchemaCacheStatus = SchemaCacheStatus
   { getSCStatusTMVar :: TMVar Bool
   }
 
+-- |
+-- We define a custom exception and throw this on listener reload. The
+-- KillThread exception can occur in an unexpected scenario, so we should
+-- avoid using that.
+data ListenerException = ListenerRestart deriving (Show, Exception)
+
 getPgVersion :: AppState -> IO PgVersion
 getPgVersion = readIORef . statePgVersion
 
@@ -94,5 +103,11 @@ killApp = stateKillApp
 putIsListenerOn :: AppState -> Bool -> IO ()
 putIsListenerOn = atomicWriteIORef . stateIsListenerOn
 
+getListenerThreadId :: AppState -> IO (Maybe ThreadId)
+getListenerThreadId = readIORef . stateListenerThreadId
+
+putListenerThreadId :: AppState -> Maybe ThreadId -> IO ()
+putListenerThreadId = atomicWriteIORef . stateListenerThreadId
+
 getObserver :: AppState -> ObservationHandler
 getObserver = stateObserver