diff --git a/IHP/PGListener.hs b/IHP/PGListener.hs new file mode 100644 --- /dev/null +++ b/IHP/PGListener.hs @@ -0,0 +1,356 @@ +{-| +Module: IHP.PGListener +Description: Event listener handling pg_notify messages +Copyright: (c) digitally induced GmbH, 2021 + +This module is solving the problem, that previously IHP was using one database connection +per running @LISTEN ..;@ statement. A @PGListener@ provides one central object to listen on +postgres channels, without manually dealing with connection management. +-} +module IHP.PGListener +( Channel +, Callback +, Notification (..) +, Subscription (..) +, PGListener (..) +, init +, stop +, withPGListener +, subscribe +, subscribeJSON +, unsubscribe +, onReconnect +) where + +import Prelude hiding (init, show, error) +import qualified Prelude +import Data.ByteString (ByteString) +import Data.Text (Text) +import qualified Data.Text as Text +import Data.IORef +import Data.String.Conversions (cs) +import Data.UUID (UUID) +import qualified Data.UUID.V4 as UUID +import Control.Monad (forever, unless, forM_) +import Data.Maybe (fromMaybe) +import Control.Exception (SomeException, displayException, uninterruptibleMask_) +import Control.Concurrent.Async (Async, async, cancel, uninterruptibleCancel) +import Data.Function ((&)) + +import IHP.Log.Types (Logger) +import Data.Set (Set) +import qualified Data.Set as Set +import Control.Concurrent.MVar (MVar) +import qualified Control.Concurrent.MVar as MVar +import Data.HashMap.Strict as HashMap +import qualified Control.Concurrent.Async as Async +import qualified Data.List as List +import qualified Data.Aeson as Aeson +import qualified IHP.Log as Log +import qualified Control.Exception.Safe as Exception +import qualified Control.Concurrent.Chan.Unagi as Queue +import qualified Control.Concurrent + +import qualified Hasql.Connection as Hasql +import qualified Hasql.Connection.Settings as HasqlSettings +import qualified Hasql.Notifications as HasqlNotifications + +-- | Local helper: show as Text +tshow :: Prelude.Show a => a -> Text +tshow = Text.pack . Prelude.show + +-- | Wrapper to satisfy 'LoggingProvider' constraint for standalone logging +data LogContext = LogContext { logger :: !Logger } + +-- TODO: How to deal with timeout of the connection? + +-- | The channel is like the event name +-- +-- It's used in the postgres NOTIFY call: +-- +-- > NOTIFY channel [ , payload ] +-- +type Channel = ByteString + +-- | A notification received from postgres +data Notification = Notification + { notificationChannel :: !ByteString + , notificationData :: !ByteString + } deriving (Prelude.Show) + +-- | An event callback receives the notification and can do IO +type Callback = Notification -> IO () + +-- | Returned by a call to 'subscribe' +data Subscription = Subscription + { id :: !UUID + , reader :: !(Async ()) + , inChan :: !(Queue.InChan Notification) + , channel :: !Channel + } + +-- | The main datatype of the service. Keeps tracks of all channels we're watching on, as well as all open subscriptions +-- +-- Use 'init' to create a new object and 'stop' to deallocate it. +data PGListener = PGListener + { logger :: !Logger + , databaseUrl :: !ByteString + , listeningTo :: !(MVar (Set Channel)) + , listenTo :: !(MVar Channel) + , subscriptions :: !(IORef (HashMap Channel [Subscription])) + , notifyLoopAsync :: !(Async ()) + , reconnectCallbacks :: !(IORef [Hasql.Connection -> IO ()]) + } + +-- | Creates a new 'PGListener' object +-- +-- > pgListener <- PGListener.init databaseUrl logger +-- +-- This will start a new async listening for postgres notifications. This will open a dedicated +-- database connection and keep it blocked until 'stop' is called. +-- +init :: ByteString -> Logger -> IO PGListener +init databaseUrl logger = do + listeningTo <- MVar.newMVar Set.empty + subscriptions <- newIORef HashMap.empty + listenTo <- MVar.newEmptyMVar + reconnectCallbacks <- newIORef [] + + notifyLoopAsync <- async (notifyLoop logger databaseUrl listeningTo listenTo subscriptions reconnectCallbacks) + pure PGListener { logger, databaseUrl, listeningTo, subscriptions, listenTo, notifyLoopAsync, reconnectCallbacks } + +-- | Stops the database listener async and releases the database connection +-- +-- > PGListener.stop pgListener +-- +stop :: PGListener -> IO () +stop PGListener { notifyLoopAsync } = do + cancel notifyLoopAsync + +withPGListener :: ByteString -> Logger -> (PGListener -> IO a) -> IO a +withPGListener databaseUrl logger = + Exception.bracket (init databaseUrl logger) stop + +-- | After you subscribed to a channel, the provided callback will be called whenever there's a new +-- notification on the channel. +-- +-- > pgListener <- PGListener.init +-- > +-- > let callback notification = do +-- > let payload :: Text = cs (notification.notificationData) +-- > putStrLn ("Received notification: " <> payload) +-- > +-- > subscription <- pgListener |> PGListener.subscribe "my_channel" callback +-- +-- The @callback@ function will now be called whenever @NOTIFY "my_channel", "my payload"@ is executed on the postgres server. +-- +-- When the subscription is not used anymore, call 'unsubscribe' to stop the callback from being called anymore. +-- +subscribe :: Channel -> Callback -> PGListener -> IO Subscription +subscribe channel callback pgListener = do + id <- UUID.nextRandom + listenToChannelIfNeeded channel pgListener + + -- We use a queue here to guarantee that the messages are processed in the right order + -- while keeping high performance. + -- + -- A naive implementation might be just kicking of an async for each message. But in that case + -- the messages might be delivered to the final consumer out of order. + (inChan, outChan) <- Queue.newChan + + let + -- We need to log any exception, otherwise there might be silent errors + logException :: SomeException -> IO () + logException exception = logError pgListener ("Error in pg_notify handler: " <> cs (displayException exception)) + + reader <- async $ forever do + message <- Queue.readChan outChan + callback message `Exception.catch` logException + let subscription = Subscription { .. } + + modifyIORef' (pgListener.subscriptions) (HashMap.insertWith mappend channel [subscription] ) + + pure subscription + +-- | Like 'subscribe' but decodes the notification payload from JSON and passes the decoded data structure to the callback +-- +-- When JSON parsing fails, this will ignore the notification. +-- +-- > pgListener <- PGListener.init +-- > +-- > let callback (jsonObject :: Aeson.Value) = do +-- > putStrLn ("Received notification: " <> tshow jsonObject) +-- > +-- > subscription <- pgListener |> PGListener.subscribeJSON "my_json_channel" callback +-- +-- The @callback@ function will now be called whenever @NOTIFY "my_json_channel", "{\"hello\":\"world\"}"@ is executed on the postgres server. +subscribeJSON :: Aeson.FromJSON jsonValue => Channel -> (jsonValue -> IO ()) -> PGListener -> IO Subscription +subscribeJSON channel callback pgListener = subscribe channel callback' pgListener + where + callback' notification = do + let payload = (notification.notificationData) + case Aeson.decodeStrict' payload of + Just payload -> callback payload + Nothing -> logError pgListener ("PGListener.subscribeJSON: Failed to parse " <> tshow payload) + +-- | Stops the callback of a subscription from receiving further notifications +-- +-- > pgListener <- PGListener.init +-- > +-- > subscription <- pgListener |> PGListener.subscribe "my_channel" callback +-- > doSomethingExpensive +-- > pgListener |> PGListener.unsubscribe subscription +-- +unsubscribe :: Subscription -> PGListener -> IO () +unsubscribe subscription@(Subscription { .. }) pgListener = do + let + deleteById :: [Subscription] -> [Subscription] + deleteById = List.deleteBy (\a b -> a.id == b.id) subscription + modifyIORef' (pgListener.subscriptions) (HashMap.adjust deleteById channel) + uninterruptibleCancel reader + pure () + +-- | Register a callback to be called when the PGListener reconnects after a connection loss. +-- +-- The callback receives the live 'Hasql.Connection' so callers can run SQL directly +-- on the known-good connection (e.g. to recreate notification triggers). +-- +-- > PGListener.onReconnect (\connection -> do +-- > Hasql.Session.run (Hasql.Session.script triggerSQL) connection +-- > pure () +-- > ) pgListener +-- +onReconnect :: (Hasql.Connection -> IO ()) -> PGListener -> IO () +onReconnect callback pgListener = + modifyIORef' (pgListener.reconnectCallbacks) (callback :) + +-- | Runs a @LISTEN ..;@ statements on the postgres connection, if not already listening on that channel +listenToChannelIfNeeded :: Channel -> PGListener -> IO () +listenToChannelIfNeeded channel pgListener = do + listeningTo <- MVar.readMVar (pgListener.listeningTo) + let alreadyListening = channel `Set.member` listeningTo + + unless alreadyListening do + MVar.putMVar (pgListener.listenTo) channel + + + +-- | Acquires a dedicated hasql connection from the given database URL. +acquireConnection :: ByteString -> IO Hasql.Connection +acquireConnection databaseUrl = do + result <- Hasql.acquire (HasqlSettings.connectionString (cs databaseUrl)) + case result of + Right connection -> pure connection + Left err -> Prelude.error ("PGListener: Failed to connect to database: " <> Prelude.show err) + +-- | The main loop that is receiving events from the database and triggering callbacks +-- +-- Todo: What happens when the connection dies? +notifyLoop :: Logger -> ByteString -> MVar (Set Channel) -> MVar Channel -> IORef (HashMap Channel [Subscription]) -> IORef [Hasql.Connection -> IO ()] -> IO () +notifyLoop logger databaseUrl listeningToVar listenToVar subscriptions reconnectCallbacksRef = do + -- Wait until the first LISTEN is requested before opening a database connection + MVar.readMVar listenToVar + + let innerLoop = do + connection <- acquireConnection databaseUrl + let cleanup = Hasql.release connection + + flip Exception.finally cleanup do + -- If listeningTo already contains channels, this means that previously the database connection + -- died, so we're restarting here. Therefore we need to replay all LISTEN calls to restore the previous state + listeningTo <- MVar.readMVar listeningToVar + forM_ listeningTo (listenToChannel connection) + + -- Fire reconnection callbacks (non-empty listeningTo = reconnection, not initial connect) + unless (Set.null listeningTo) do + callbacks <- readIORef reconnectCallbacksRef + forM_ callbacks \callback -> + Exception.tryAny (callback connection) >>= \case + Left e -> let ?context = LogContext logger in Log.info ("PGListener reconnect callback failed: " <> displayException e) + Right _ -> pure () + + -- We use 'race' to alternate between waiting for notifications and + -- processing new LISTEN requests. This avoids a deadlock: both + -- 'waitForNotifications' and 'listen' acquire an exclusive lock on + -- the underlying libpq connection via 'Connection.use'. If we ran + -- them concurrently (as before), 'listen' would block forever + -- waiting for 'waitForNotifications' to release the lock. + -- + -- When 'race' cancels 'waitForNotifications', the lock is released, + -- allowing 'listenToChannel' to acquire it. Any buffered notifications + -- are preserved in the libpq connection and picked up when + -- 'waitForNotifications' is restarted. + let notifyAndListenLoop = do + result <- Async.race + (HasqlNotifications.waitForNotifications + (\channel payload -> uninterruptibleMask_ do + -- uninterruptibleMask_ ensures that once waitForNotifications + -- has dequeued a notification from libpq's buffer, the callback + -- runs to completion even if race cancels us with an async exception. + -- Without this, a notification could be lost: dequeued from libpq + -- but never delivered to the subscription's inChan. + let notification = Notification { notificationChannel = channel, notificationData = payload } + + allSubscriptions <- readIORef subscriptions + let channelSubscriptions = allSubscriptions + & HashMap.lookup channel + & fromMaybe [] + + forM_ channelSubscriptions \subscription -> + Queue.writeChan (subscription.inChan) notification + ) + connection + ) + (MVar.takeMVar listenToVar) + + case result of + Left () -> + -- waitForNotifications returned (connection lost) - exit to trigger retry + pure () + Right newChannel -> do + MVar.modifyMVar_ listeningToVar \currentListeningTo -> + if newChannel `Set.member` currentListeningTo + then pure currentListeningTo + else do + listenToChannel connection newChannel + pure (Set.insert newChannel currentListeningTo) + notifyAndListenLoop + + notifyAndListenLoop + + -- Initial delay (in microseconds) + let initialDelay = 500 * 1000 + -- Max delay (in microseconds) + let maxDelay = 60 * 1000 * 1000 + -- This outer loop restarts the listeners if the database connection dies (e.g. due to a timeout) + let retryLoop delay isFirstError = do + result <- Exception.tryAny innerLoop + case result of + Left error -> do + let ?context = LogContext logger + if isFirstError then do + Log.info ("PGListener is going to restart, loop failed with exception: " <> (displayException error) <> ". Retrying immediately.") + retryLoop delay False -- Retry with no delay interval on first error, but will increase delay interval in subsequent retries + else do + let increasedDelay = delay * 2 -- Double current delay + let nextDelay = min increasedDelay maxDelay -- Picks whichever delay is lowest of increasedDelay * 2 or maxDelay + Log.info ("PGListener is going to restart, loop failed with exception: " <> (displayException error) <> ". Retrying in " <> cs (printTimeToNextRetry delay) <> ".") + Control.Concurrent.threadDelay delay -- Sleep for the current delay + retryLoop nextDelay False -- Retry with longer interval + Right _ -> + retryLoop initialDelay True -- If all went well, re-run with no sleeping and reset current delay to the initial value + retryLoop initialDelay True + +printTimeToNextRetry :: Int -> Text +printTimeToNextRetry microseconds + | microseconds >= 1000000000 = tshow (microseconds `div` 1000000000) <> " min" + | microseconds >= 1000000 = tshow (microseconds `div` 1000000) <> " s" + | microseconds >= 1000 = tshow (microseconds `div` 1000) <> " ms" + | otherwise = tshow microseconds <> " µs" + +listenToChannel :: Hasql.Connection -> Channel -> IO () +listenToChannel connection channel = do + HasqlNotifications.listen connection (HasqlNotifications.toPgIdentifier (cs channel)) + +logError :: PGListener -> Text -> IO () +logError pgListener message = let ?context = LogContext pgListener.logger in Log.error message diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2020 digitally induced GmbH + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/Test/Main.hs b/Test/Main.hs new file mode 100644 --- /dev/null +++ b/Test/Main.hs @@ -0,0 +1,8 @@ +module Main where + +import Test.Hspec +import qualified Test.PGListenerSpec + +main :: IO () +main = hspec do + Test.PGListenerSpec.tests diff --git a/Test/PGListenerSpec.hs b/Test/PGListenerSpec.hs new file mode 100644 --- /dev/null +++ b/Test/PGListenerSpec.hs @@ -0,0 +1,169 @@ +{-| +Module: Test.PGListenerSpec +Copyright: (c) digitally induced GmbH, 2025 +-} +module Test.PGListenerSpec where + +import Prelude +import Test.Hspec +import Data.ByteString (ByteString) +import qualified Data.ByteString.Char8 as BS8 +import Data.IORef +import Data.String.Conversions (cs) +import Data.Function ((&)) +import Data.HashMap.Strict as HashMap +import Control.Concurrent (threadDelay) +import Control.Monad (forM_) +import qualified Control.Concurrent.Async as Async +import qualified Control.Exception.Safe as Exception +import System.Environment (lookupEnv) + +import IHP.Log.Types (Logger(..), LogLevel(..)) +import qualified IHP.PGListener as PGListener + +import qualified Hasql.Connection as Hasql +import qualified Hasql.Connection.Settings as HasqlSettings +import qualified Hasql.Session as Session + +logger :: Logger +logger = Logger + { write = \_ -> pure () + , level = Debug + , formatter = \_ _ msg -> msg + , timeCache = pure "" + , cleanup = pure () + } + +getDatabaseUrl :: IO ByteString +getDatabaseUrl = do + envUrl <- lookupEnv "DATABASE_URL" + pure (maybe "postgresql:///postgres" cs envUrl) + +acquireConnection :: ByteString -> IO Hasql.Connection +acquireConnection databaseUrl = do + result <- Hasql.acquire (HasqlSettings.connectionString (cs databaseUrl)) + case result of + Right connection -> pure connection + Left err -> error ("Test: Failed to connect to database: " <> show err) + +canConnectToPostgres :: IO Bool +canConnectToPostgres = do + connStr <- getDatabaseUrl + result <- Exception.try (Exception.bracket (acquireConnection connStr) Hasql.release (\_ -> pure ())) + case result of + Left (_ :: Exception.SomeException) -> pure False + Right _ -> pure True + +-- | Run a test that requires a real PostgreSQL connection, skipping if unavailable +withDB :: (ByteString -> IO ()) -> IO () +withDB action = do + available <- canConnectToPostgres + if available + then getDatabaseUrl >>= action + else pendingWith "PostgreSQL not available (set DATABASE_URL or start a local Postgres)" + +-- | Execute a raw SQL statement via a temporary hasql connection +execSQL :: ByteString -> ByteString -> IO () +execSQL connStr sql = Exception.bracket (acquireConnection connStr) Hasql.release \conn -> do + result <- Hasql.use conn (Session.script (cs sql)) + case result of + Right () -> pure () + Left err -> error ("SQL exec failed: " <> show err) + +tests :: Spec +tests = do + describe "IHP.PGListener" do + describe "subscribe" do + it "should add a subscriber" do + PGListener.withPGListener "" logger \pgListener -> do + subscriptionsCount <- length . concat . HashMap.elems <$> readIORef pgListener.subscriptions + subscriptionsCount `shouldBe` 0 + + let didInsertRecordCallback notification = pure () + + pgListener & PGListener.subscribe "did_insert_record" didInsertRecordCallback + + subscriptionsCount <- length . concat . HashMap.elems <$> readIORef pgListener.subscriptions + subscriptionsCount `shouldBe` 1 + + describe "unsubscribe" do + it "remove the subscription" do + PGListener.withPGListener "" logger \pgListener -> do + subscription <- pgListener & PGListener.subscribe "did_insert_record" (const (pure ())) + pgListener & PGListener.unsubscribe subscription + + subscriptionsCount <- length . concat . HashMap.elems <$> readIORef pgListener.subscriptions + + subscriptionsCount `shouldBe` 0 + + describe "multi-channel notifications" do + it "should receive notifications on multiple channels" do + withDB \connStr -> do + PGListener.withPGListener connStr logger \pgListener -> do + received1 <- newIORef ([] :: [ByteString]) + received2 <- newIORef ([] :: [ByteString]) + + _sub1 <- pgListener & PGListener.subscribe "test_channel_1" \n -> + modifyIORef' received1 (n.notificationData :) + _sub2 <- pgListener & PGListener.subscribe "test_channel_2" \n -> + modifyIORef' received2 (n.notificationData :) + + -- Allow time for LISTEN commands to be processed + threadDelay 200_000 + + -- Send notifications via a separate hasql connection + execSQL connStr "NOTIFY test_channel_1, 'hello1'" + execSQL connStr "NOTIFY test_channel_2, 'hello2'" + + -- Allow time for notifications to be delivered + threadDelay 200_000 + + r1 <- readIORef received1 + r2 <- readIORef received2 + r1 `shouldBe` ["hello1"] + r2 `shouldBe` ["hello2"] + + it "should not drop notifications when a new channel is subscribed" do + withDB \connStr -> do + PGListener.withPGListener connStr logger \pgListener -> do + received <- newIORef ([] :: [ByteString]) + + -- Subscribe to channel_1 and wait for the LISTEN to be processed + _sub1 <- pgListener & PGListener.subscribe "nodrop_1" \n -> + modifyIORef' received (n.notificationData :) + threadDelay 200_000 + + -- One thread continuously sends notifications on channel_1. + -- Concurrently, the main thread subscribes to new channels, + -- which triggers the internal race cancellation of + -- waitForNotifications. If the cancellation drops a notification + -- that was mid-delivery, the final count will be wrong. + let totalNotifications = 100 :: Int + Exception.bracket (acquireConnection connStr) Hasql.release \notifyConn -> do + Async.concurrently_ + -- Sender thread: fire notifications with small delays + -- so they arrive while waitForNotifications is active + (forM_ [1..totalNotifications] \i -> do + let payload = BS8.pack (show i) + Hasql.use notifyConn (Session.script (cs ("NOTIFY nodrop_1, '" <> payload <> "'"))) + >>= \case + Right () -> pure () + Left err -> error ("NOTIFY failed: " <> show err) + threadDelay 1_000 + ) + -- Main thread: subscribe to new channels while + -- notifications are being delivered, triggering + -- the race cancellation multiple times + (do + threadDelay 10_000 + forM_ [(2::Int)..10] \i -> do + let ch = BS8.pack ("nodrop_" <> show i) + _ <- pgListener & PGListener.subscribe ch (const (pure ())) + threadDelay 10_000 + ) + + -- Wait for all notifications to be delivered + threadDelay 500_000 + + r <- readIORef received + length r `shouldBe` totalNotifications diff --git a/ihp-pglistener.cabal b/ihp-pglistener.cabal new file mode 100644 --- /dev/null +++ b/ihp-pglistener.cabal @@ -0,0 +1,60 @@ +cabal-version: 2.2 +name: ihp-pglistener +version: 1.0.0 +synopsis: PostgreSQL LISTEN/NOTIFY channel manager for IHP +description: Manages PostgreSQL LISTEN/NOTIFY subscriptions over a single shared connection + with automatic reconnection and exponential backoff. +license: MIT +license-file: LICENSE +author: digitally induced GmbH +maintainer: hello@digitallyinduced.com +homepage: https://ihp.digitallyinduced.com/ +bug-reports: https://github.com/digitallyinduced/ihp/issues +copyright: (c) digitally induced GmbH +category: Web, Database, IHP + +common shared-properties + default-language: GHC2021 + build-depends: + base >= 4.17.0 && < 4.22 + , bytestring + , text + , string-conversions + , aeson + , uuid + , async + , hasql + , hasql-notifications + , unagi-chan + , safe-exceptions + , containers + , unordered-containers + , hashable + , ihp-log + default-extensions: + OverloadedStrings + , ImplicitParams + , DataKinds + , RecordWildCards + , OverloadedRecordDot + , DuplicateRecordFields + , BlockArguments + , LambdaCase + ghc-options: -Werror=incomplete-patterns -Werror=unused-imports -Werror=missing-fields + +library + import: shared-properties + hs-source-dirs: . + exposed-modules: + IHP.PGListener + +test-suite tests + import: shared-properties + type: exitcode-stdio-1.0 + main-is: Test/Main.hs + hs-source-dirs: . + build-depends: + ihp-pglistener + , hspec + other-modules: + Test.PGListenerSpec