refactor: Provide API to stop listener programatically

This change introduces possibility to stop listener thread programatically. It is a prerequisite for HSpec tests starting the listener as they have to be able to stop it when cleaning up.

Author: mkleczek

postgrest.cabal

@@ -310,6 +310,7 @@ test-suite observability
   main-is:            Main.hs
   other-modules:      ObsHelper
                       Observation.JwtCache
+                      Observation.ListenerSpec
                       Observation.MetricsSpec
                       Observation.SchemaCacheSpec
   build-depends:      base              >= 4.9 && < 4.22

src/PostgREST/App.hs

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

src/PostgREST/Listener.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE LambdaCase      #-}
+{-# LANGUAGE DeriveAnyClass  #-}
 {-# LANGUAGE MultiWayIf      #-}
 {-# LANGUAGE RecordWildCards #-}
 
@@ -24,24 +24,30 @@ import qualified Hasql.Session              as SQL
 import           PostgREST.Config.Database  (queryPgVersion)
 import           PostgREST.Config.PgVersion (pgvFullName)
 import           Protolude
+import System.Mem.Weak (deRefWeak)
+import Data.Monoid.Extra (mwhen)
+
+data ListenerStopped = ListenerStopped deriving (Show, Exception)
+newtype ListenerConnectionError = ListenerConnectionError SQL.ConnectionError deriving (Show, Exception)
 
 -- | Starts the Listener in a thread
-runListener :: AppState -> IO ()
+runListener :: AppState -> IO (IO ())
 runListener appState = do
   AppConfig{..} <- getConfig appState
-  when configDbChannelEnabled $
-    void . forkIO . void $ retryingListen appState False
+  mwhen configDbChannelEnabled $ do
+    listenerThreadId <- mkWeakThreadId =<< forkIO (retryingListen appState False)
+    pure $ deRefWeak listenerThreadId >>= foldMap (`throwTo` ListenerStopped)
 
 -- | 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 -> Bool -> IO Void
+-- | This function returns upon receiving of ListenerStopped async exception.
+retryingListen :: AppState -> Bool -> IO ()
 retryingListen appState hasDbListenerBug = do
   cfg@AppConfig{..} <- AppState.getConfig appState
   let
     dbChannel = toS configDbChannel
     onError err = do
       AppState.putIsListenerOn appState False
-      observer $ DBListenFail dbChannel (Right err)
+      observer $ DBListenFail dbChannel err
       when (isDbListenerBug err) $
         observer DBListenBugCallQueryFix
       unless configDbPoolAutomaticRecovery $
@@ -54,24 +60,24 @@ retryingListen appState hasDbListenerBug = do
       unless (delay == maxDelay) $
         AppState.putNextListenerDelay appState (delay * 2)
       -- loop running the listener
-      retryingListen appState (isDbListenerBug err)
+      pure $ retryingListen appState (isDbListenerBug err)
 
   -- Execute the listener with with error handling
-  handle onError $ do
-    -- Make sure we don't leak connections on errors
+  join $ handle onError $ handle (\ListenerStopped -> mempty) $
     bracket
       -- acquire connection
       (SQL.acquire $
         Config.toConnectionSettings Config.addTargetSessionAttrs cfg)
       -- release connection
       (`whenRight` releaseConnection) $
       -- use connection
-      \case
-        Right db -> do
+      either (throwIO . ListenerConnectionError) $ \db -> do
           (pqHost, pqPort) <- SQL.withLibPQConnection db $ bisequence . (LibPQ.host &&& LibPQ.port)
           pgFullName <- SQL.run queryPgVersion db >>= either throwIO (pure . pgvFullName)
           when hasDbListenerBug $ SQL.run callNotifQueryUsage db >>= either throwIO pure
           SQL.listen db $ SQL.toPgIdentifier dbChannel
+          when hasDbListenerBug $ SQL.run callNotifQueryUsage db >>= either throwIO pure
+          SQL.listen db $ SQL.toPgIdentifier dbChannel
 
           AppState.putIsListenerOn appState True
 
@@ -88,9 +94,6 @@ retryingListen appState hasDbListenerBug = do
           -- this will never return, in case of an error it will throw and be caught by onError
           forever $ SQL.waitForNotifications handleNotification db
 
-        Left err -> do
-          observer $ DBListenFail dbChannel (Left err)
-          exitFailure
   where
     observer = AppState.getObserver appState
     mainThreadId = AppState.getMainThreadId appState

src/PostgREST/Logger.hs

@@ -40,7 +40,6 @@ import PostgREST.SchemaCache (queryTimingsWLabels)
 
 import qualified Data.ByteString.Lazy       as LBS
 import qualified Data.Text                  as T
-import qualified Hasql.Connection           as SQL
 import qualified Hasql.Pool                 as SQL
 import qualified Hasql.Pool.Observation     as SQL
 import           Numeric                    (showFFloat)
@@ -185,7 +184,7 @@ observationMessages = \case
     pure $ "Listener connected to " <> fullName <> " on " <> show (fold $ host <> fmap (":" <>) port) <> " and listening for database notifications on the " <> show channel <> " channel"
   DBListenFail channel listenErr ->
     pure $ "Failed listening for database notifications on the " <> show channel <> " channel. " <>
-      either showListenerConnError showListenerException listenErr
+      showListenerException listenErr
   DBListenRetry delay ->
     pure $ "Retrying listening for database notifications in " <> (show delay::Text) <> " seconds..."
   DBListenBugCallQueryFix ->
@@ -249,10 +248,6 @@ observationMessages = \case
 
     jsonMessage err = T.decodeUtf8 . LBS.toStrict . Error.errorPayload Verbose $ Error.PgError False err
 
-
-    showListenerConnError :: SQL.ConnectionError -> Text
-    showListenerConnError = maybe "Connection error" (showOnSingleLine '\t' . T.decodeUtf8)
-
     showListenerException :: SomeException -> Text
     showListenerException = showOnSingleLine '\t' . show
 

src/PostgREST/Observation.hs

@@ -12,7 +12,6 @@ module PostgREST.Observation
   , ObservationHandler
   ) where
 
-import qualified Hasql.Connection           as SQL
 import qualified Hasql.Pool                 as SQL
 import qualified Hasql.Pool.Observation     as SQL
 import           Network.HTTP.Types.Status  (Status)
@@ -36,7 +35,7 @@ data Observation
   | SchemaCacheLoadedObs Double Text
   | ConnectionRetryObs Int
   | DBListenStart (Maybe ByteString) (Maybe ByteString) Text Text -- host, port, version string, channel
-  | DBListenFail Text (Either SQL.ConnectionError SomeException)
+  | DBListenFail Text SomeException
   | DBListenRetry Int
   | DBListenBugCallQueryFix
   | DBListenerGotSCacheMsg ByteString

test/observability/Main.hs

@@ -18,6 +18,7 @@ import           PostgREST.SchemaCache     (querySchemaCache)
 import qualified Observation.JwtCache
 import qualified Observation.MetricsSpec
 
+import qualified Observation.ListenerSpec
 import qualified Observation.SchemaCacheSpec
 import           ObsHelper
 import           PostgREST.Observation       (Observation (HasqlPoolObs))
@@ -68,6 +69,8 @@ main = do
       describe "Feature.MetricsSpec" Observation.MetricsSpec.spec
     before (initApp baseSchemaCache testCfg) $
       describe "Feature.SchemaCacheSpec" Observation.SchemaCacheSpec.spec
+    before (initApp baseSchemaCache testCfg) $
+      describe "Observation.ListenerSpec" Observation.ListenerSpec.spec
 
   where
     loadSCache pool conf =

test/observability/ObsHelper.hs

@@ -74,7 +74,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"]

test/observability/Observation/ListenerSpec.hs

@@ -0,0 +1,31 @@
+{-# LANGUAGE DataKinds           #-}
+{-# LANGUAGE MonadComprehensions #-}
+{-# LANGUAGE NamedFieldPuns      #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+module Observation.ListenerSpec where
+
+import Network.Wai (Application)
+
+import ObsHelper
+
+import qualified PostgREST.Listener    as Listener
+import           PostgREST.Observation
+
+import Test.Hspec     (SpecWith, describe, it)
+import Test.Hspec.Wai (getState)
+
+import Protolude
+
+spec :: SpecWith (SpecState, Application)
+spec = describe "Listener tests" $ do
+
+  it "Should start the listener" $ do
+    SpecState{specAppState = appState, specObsChan} <- getState
+    let waitFor = waitForObs specObsChan
+
+    liftIO $ withListener appState $ do
+      waitFor (1 * sec) "DBListenStart" $ \x -> [ o | o@(DBListenStart {}) <- pure x ]
+
+  where
+    sec = 1000000
+    withListener appState = bracket (Listener.runListener appState) identity . const