ihp-pglistener (empty) → 1.0.0
raw patch · 5 files changed
+614/−0 lines, 5 filesdep +aesondep +asyncdep +base
Dependencies added: aeson, async, base, bytestring, containers, hashable, hasql, hasql-notifications, hspec, ihp-log, ihp-pglistener, safe-exceptions, string-conversions, text, unagi-chan, unordered-containers, uuid
Files
- IHP/PGListener.hs +356/−0
- LICENSE +21/−0
- Test/Main.hs +8/−0
- Test/PGListenerSpec.hs +169/−0
- ihp-pglistener.cabal +60/−0
+ IHP/PGListener.hs view
@@ -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
+ LICENSE view
@@ -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.
+ Test/Main.hs view
@@ -0,0 +1,8 @@+module Main where++import Test.Hspec+import qualified Test.PGListenerSpec++main :: IO ()+main = hspec do+ Test.PGListenerSpec.tests
+ Test/PGListenerSpec.hs view
@@ -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
+ ihp-pglistener.cabal view
@@ -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