diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,8 @@ +# consumers-2.3.2.0 (2024-08-30) +* Use prepared queries to improve parsing and planning time. +* Prevent the consumer from crashlooping if `ccOnException` throws. +* Drop support for GHC 8.8. + # consumers-2.3.1.0 (2023-11-27) * Drop support for advisory locks. * Add support for GHC 9.6 and 9.8. diff --git a/Setup.hs b/Setup.hs deleted file mode 100644 --- a/Setup.hs +++ /dev/null @@ -1,2 +0,0 @@ -import Distribution.Simple -main = defaultMain diff --git a/consumers.cabal b/consumers.cabal --- a/consumers.cabal +++ b/consumers.cabal @@ -1,5 +1,6 @@ +cabal-version: 3.0 name: consumers -version: 2.3.1.0 +version: 2.3.2.0 synopsis: Concurrent PostgreSQL data consumers description: Library for setting up concurrent consumers of data @@ -7,32 +8,52 @@ declarative manner. homepage: https://github.com/scrive/consumers -license: BSD3 +license: BSD-3-Clause license-file: LICENSE extra-source-files: CHANGELOG.md, README.md author: Scrive AB maintainer: Andrzej Rybczak , - Jonathan Jouty , - Mikhail Glushenkov + Jonathan Jouty copyright: Scrive AB category: Concurrency, Database build-type: Simple -cabal-version: >=1.10 -tested-with: GHC ==8.8.4 || ==8.10.7 || ==9.0.2 || ==9.2.8 || ==9.4.7 || ==9.6.3 - || ==9.8.1 +tested-with: GHC == { 8.10.7, 9.0.2, 9.2.8, 9.4.8, 9.6.6, 9.8.2, 9.10.1 } -Source-repository head - Type: git - Location: https://github.com/scrive/consumers.git +bug-reports: https://github.com/scrive/consumers/issues +source-repository head + type: git + location: https://github.com/scrive/consumers.git +common language + ghc-options: -Wall -Wcompat -Werror=prepositive-qualified-module + + default-language: Haskell2010 + default-extensions: DeriveDataTypeable + , ExistentialQuantification + , FlexibleContexts + , GeneralizedNewtypeDeriving + , ImportQualifiedPost + , LambdaCase + , MultiParamTypeClasses + , OverloadedStrings + , RankNTypes + , RecordWildCards + , ScopedTypeVariables + , TupleSections + , TypeApplications + , TypeFamilies + , UndecidableInstances + library + import: language + exposed-modules: Database.PostgreSQL.Consumers, Database.PostgreSQL.Consumers.Config, Database.PostgreSQL.Consumers.Consumer, Database.PostgreSQL.Consumers.Components, Database.PostgreSQL.Consumers.Utils - build-depends: base >= 4.13 && < 5 + build-depends: base >= 4.14 && < 5 , containers >= 0.5 , exceptions >= 0.10 , hpqtypes >= 1.11 @@ -42,37 +63,17 @@ , monad-control >= 1.0 , monad-time >= 0.4 , mtl >= 2.2 + , safe-exceptions >= 0.1.7 , stm >= 2.4 + , text >= 1.2 , time >= 1.6 , transformers-base >= 0.4 hs-source-dirs: src - ghc-options: -Wall - - default-language: Haskell2010 - default-extensions: DeriveDataTypeable - , FlexibleContexts - , GeneralizedNewtypeDeriving - , LambdaCase - , NoImplicitPrelude - , OverloadedStrings - , RankNTypes - , RecordWildCards - , ScopedTypeVariables - , TupleSections - , TypeApplications - , TypeFamilies - , UndecidableInstances - test-suite consumers-example - -- Not quite a test suite, just a lazy way to disable this component - -- by default, but have Travis build it. - type: exitcode-stdio-1.0 - hs-source-dirs: example - main-is: Example.hs - default-language: Haskell2010 - ghc-options: -Wall + ghc-options: -threaded -Wall -Wcompat + build-depends: base, consumers, hpqtypes, @@ -82,12 +83,18 @@ text-show, transformers -test-suite consumers-test - type: exitcode-stdio-1.0 - hs-source-dirs: test - main-is: Test.hs + hs-source-dirs: example + default-language: Haskell2010 - ghc-options: -Wall + + type: exitcode-stdio-1.0 + main-is: Example.hs + +test-suite consumers-test + import: language + + ghc-options: -threaded + build-depends: base, consumers, exceptions, @@ -96,7 +103,6 @@ hpqtypes-extras, log-base, monad-control, - monad-loops, monad-time, mtl, stm, @@ -105,3 +111,8 @@ time, transformers, transformers-base + + hs-source-dirs: test + + type: exitcode-stdio-1.0 + main-is: Test.hs diff --git a/example/Example.hs b/example/Example.hs --- a/example/Example.hs +++ b/example/Example.hs @@ -1,196 +1,242 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ImportQualifiedPost #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} module Main where -import Database.PostgreSQL.Consumers -import Database.PostgreSQL.PQTypes -import Database.PostgreSQL.PQTypes.Checks -import Database.PostgreSQL.PQTypes.Model - import Control.Concurrent import Control.Exception import Control.Monad import Control.Monad.IO.Class import Data.Int -import Data.Monoid +import Data.Text qualified as T +import Database.PostgreSQL.Consumers +import Database.PostgreSQL.PQTypes +import Database.PostgreSQL.PQTypes.Checks +import Database.PostgreSQL.PQTypes.Model import Log import Log.Backend.StandardOutput -import Prelude import System.Environment import System.Exit import TextShow -import qualified Data.Text as T - -- | Main application monad. See the 'log-base' and the 'hpqtypes' -- packages for documentation on 'DBT' and 'LogT'. type AppM a = DBT (LogT IO) a main :: IO () main = do - connString <- getArgs >>= \case - connString : _args -> return $ T.pack connString - [] -> lookupEnv "GITHUB_ACTIONS" >>= \case - Just "true" -> return "host=postgres user=postgres password=postgres" - _ -> printUsage >> exitFailure + connString <- + getArgs >>= \case + connString : _args -> pure $ T.pack connString + [] -> + lookupEnv "GITHUB_ACTIONS" >>= \case + Just "true" -> pure "host=postgres user=postgres password=postgres" + _ -> printUsage >> exitFailure - let connSettings = defaultConnectionSettings - { csConnInfo = connString } + let connSettings = + defaultConnectionSettings + { csConnInfo = connString + } ConnectionSource connSource = simpleSource connSettings -- Monad stack initialisation. withStdOutLogger $ \logger -> - runLogT "consumers-example" logger defaultLogLevel $ - runDBT connSource defaultTransactionSettings $ do - + runLogT "consumers-example" logger defaultLogLevel + . runDBT connSource defaultTransactionSettings + $ do -- Initialise. createTables -- Create a consumer, put ten jobs into its queue, and wait -- for it to finish. 'runConsumer' returns a finaliser that is -- invoked by 'finalize' after the 'putJob' loop. - finalize (localDomain "process" $ - runConsumer consumerConfig connSource) $ - forM_ [(0::Int)..10] $ \_ -> do + finalize + ( localDomain "process" $ + runConsumer consumerConfig connSource + ) + . forM_ [(0 :: Int) .. 10] + $ \_ -> do putJob liftIO $ threadDelay (1 * 1000000) -- 1 sec -- Clean up. dropTables - - where - printUsage = do - prog <- getProgName - putStrLn $ "Usage: " <> prog <> " " + where + printUsage = do + prog <- getProgName + putStrLn $ "Usage: " <> prog <> " " - tables = [consumersTable, jobsTable] - -- NB: order of migrations is important. - migrations = [ createTableMigration consumersTable - , createTableMigration jobsTable ] + tables = [consumersTable, jobsTable] + -- NB: order of migrations is important. + migrations = + [ createTableMigration consumersTable + , createTableMigration jobsTable + ] - createTables :: AppM () - createTables = do - migrateDatabase defaultExtrasOptions - {- extensions -} [] {- composites -} [] {- domains -} [] - tables migrations - checkDatabase defaultExtrasOptions - {- composites -} [] {- domains -} [] - tables + createTables :: AppM () + createTables = do + migrateDatabase + defaultExtrasOptions + [] -- extensions + [] -- composites + [] -- domains + tables + migrations + checkDatabase + defaultExtrasOptions + [] -- composites + [] -- domains + tables - dropTables :: AppM () - dropTables = do - migrateDatabase defaultExtrasOptions - {- extensions -} [] {- composites -} [] {- domains -} [] {- tables -} [] - [ dropTableMigration jobsTable - , dropTableMigration consumersTable ] + dropTables :: AppM () + dropTables = do + migrateDatabase + defaultExtrasOptions + [] -- extensions + [] -- composites + [] -- domains + [] -- tables + [ dropTableMigration jobsTable + , dropTableMigration consumersTable + ] - -- Configuration of a consumer. See - -- 'Database.PostgreSQL.Consumers.Config.ConsumerConfig'. - consumerConfig = ConsumerConfig - { ccJobsTable = "consumers_example_jobs" - , ccConsumersTable = "consumers_example_consumers" - , ccJobSelectors = ["id", "message"] - , ccJobFetcher = id - , ccJobIndex = \(i::Int64, _msg::T.Text) -> i + -- Configuration of a consumer. See + -- 'Database.PostgreSQL.Consumers.Config.ConsumerConfig'. + consumerConfig = + ConsumerConfig + { ccJobsTable = "consumers_example_jobs" + , ccConsumersTable = "consumers_example_consumers" + , ccJobSelectors = ["id", "message"] + , ccJobFetcher = id + , ccJobIndex = \(i :: Int64, _msg :: T.Text) -> i , ccNotificationChannel = Just "consumers_example_chan" , ccNotificationTimeout = 10 * 1000000 -- 10 sec - , ccMaxRunningJobs = 1 - , ccProcessJob = processJob - , ccOnException = handleException + , ccMaxRunningJobs = 1 + , ccProcessJob = processJob + , ccOnException = handleException } - -- Add a job to the consumer's queue. - putJob :: AppM () - putJob = localDomain "put" $ do - logInfo_ "putJob" - runSQL_ $ "INSERT INTO consumers_example_jobs " + -- Add a job to the consumer's queue. + putJob :: AppM () + putJob = localDomain "put" $ do + logInfo_ "putJob" + runSQL_ $ + "INSERT INTO consumers_example_jobs " <> "(run_at, finished_at, reserved_by, attempts, message) " <> "VALUES (NOW(), NULL, NULL, 0, 'hello')" - notify "consumers_example_chan" "" - commit - - -- Invoked when a job is ready to be processed. - processJob :: (Int64, T.Text) -> AppM Result - processJob (_idx, msg) = do - logInfo_ msg - return (Ok Remove) + notify "consumers_example_chan" "" + commit - -- Invoked when 'processJob' throws an exception. Can handle - -- failure in different ways, such as: remove the job from the - -- queue, mark it as processed, or schedule it for rerun. - handleException :: SomeException -> (Int64, T.Text) -> AppM Action - handleException exc (idx, _msg) = do - logAttention_ $ - "Job #" <> showt idx <> " failed with: " <> showt exc - return . RerunAfter $ imicroseconds 500000 + -- Invoked when a job is ready to be processed. + processJob :: (Int64, T.Text) -> AppM Result + processJob (_idx, msg) = do + logInfo_ msg + pure (Ok Remove) + -- Invoked when 'processJob' throws an exception. Can handle + -- failure in different ways, such as: remove the job from the + -- queue, mark it as processed, or schedule it for rerun. + handleException :: SomeException -> (Int64, T.Text) -> AppM Action + handleException exc (idx, _msg) = do + logAttention "Job failed" $ + object + [ "job_id" .= showt idx + , "exception" .= showt exc + ] + pure . RerunAfter $ imicroseconds 500000 -- | Table where jobs are stored. See -- 'Database.PostgreSQL.Consumers.Config.ConsumerConfig'. jobsTable :: Table jobsTable = tblTable - { tblName = "consumers_example_jobs" - , tblVersion = 1 - , tblColumns = - [ tblColumn { colName = "id", colType = BigSerialT - , colNullable = False } - , tblColumn { colName = "run_at", colType = TimestampWithZoneT - , colNullable = True } - , tblColumn { colName = "finished_at", colType = TimestampWithZoneT - , colNullable = True } - , tblColumn { colName = "reserved_by", colType = BigIntT - , colNullable = True } - , tblColumn { colName = "attempts", colType = IntegerT - , colNullable = False } - - -- The only non-obligatory field: - , tblColumn { colName = "message", colType = TextT - , colNullable = False } - ] - , tblPrimaryKey = pkOnColumn "id" - , tblForeignKeys = [ - (fkOnColumn "reserved_by" "consumers_example_consumers" "id") { - fkOnDelete = ForeignKeySetNull - } - ] - } + { tblName = "consumers_example_jobs" + , tblVersion = 1 + , tblColumns = + [ tblColumn + { colName = "id" + , colType = BigSerialT + , colNullable = False + } + , tblColumn + { colName = "run_at" + , colType = TimestampWithZoneT + , colNullable = True + } + , tblColumn + { colName = "finished_at" + , colType = TimestampWithZoneT + , colNullable = True + } + , tblColumn + { colName = "reserved_by" + , colType = BigIntT + , colNullable = True + } + , tblColumn + { colName = "attempts" + , colType = IntegerT + , colNullable = False + } + , -- The only non-obligatory field: + tblColumn + { colName = "message" + , colType = TextT + , colNullable = False + } + ] + , tblPrimaryKey = pkOnColumn "id" + , tblForeignKeys = + [ (fkOnColumn "reserved_by" "consumers_example_consumers" "id") + { fkOnDelete = ForeignKeySetNull + } + ] + } -- | Table where registered consumers are stored. See -- 'Database.PostgreSQL.Consumers.Config.ConsumerConfig'. consumersTable :: Table consumersTable = tblTable - { tblName = "consumers_example_consumers" - , tblVersion = 1 - , tblColumns = - [ tblColumn { colName = "id", colType = BigSerialT - , colNullable = False } - , tblColumn { colName = "name", colType = TextT - , colNullable = False } - , tblColumn { colName = "last_activity", colType = TimestampWithZoneT - , colNullable = False } - ] - , tblPrimaryKey = pkOnColumn "id" - } - + { tblName = "consumers_example_consumers" + , tblVersion = 1 + , tblColumns = + [ tblColumn + { colName = "id" + , colType = BigSerialT + , colNullable = False + } + , tblColumn + { colName = "name" + , colType = TextT + , colNullable = False + } + , tblColumn + { colName = "last_activity" + , colType = TimestampWithZoneT + , colNullable = False + } + ] + , tblPrimaryKey = pkOnColumn "id" + } -createTableMigration :: (MonadDB m) => Table -> Migration m -createTableMigration tbl = Migration - { mgrTableName = tblName tbl - , mgrFrom = 0 - , mgrAction = StandardMigration $ do - createTable True tbl - } +createTableMigration :: MonadDB m => Table -> Migration m +createTableMigration tbl = + Migration + { mgrTableName = tblName tbl + , mgrFrom = 0 + , mgrAction = StandardMigration $ do + createTable True tbl + } dropTableMigration :: Table -> Migration m -dropTableMigration tbl = Migration - { mgrTableName = tblName tbl - , mgrFrom = 1 - , mgrAction = DropTableMigration DropTableRestrict - } +dropTableMigration tbl = + Migration + { mgrTableName = tblName tbl + , mgrFrom = 1 + , mgrAction = DropTableMigration DropTableRestrict + } diff --git a/src/Database/PostgreSQL/Consumers.hs b/src/Database/PostgreSQL/Consumers.hs --- a/src/Database/PostgreSQL/Consumers.hs +++ b/src/Database/PostgreSQL/Consumers.hs @@ -1,5 +1,5 @@ -module Database.PostgreSQL.Consumers ( - runConsumer +module Database.PostgreSQL.Consumers + ( runConsumer , runConsumerWithIdleSignal , module Database.PostgreSQL.Consumers.Config , module Database.PostgreSQL.Consumers.Utils diff --git a/src/Database/PostgreSQL/Consumers/Components.hs b/src/Database/PostgreSQL/Consumers/Components.hs --- a/src/Database/PostgreSQL/Consumers/Components.hs +++ b/src/Database/PostgreSQL/Consumers/Components.hs @@ -1,54 +1,65 @@ -module Database.PostgreSQL.Consumers.Components ( - runConsumer +module Database.PostgreSQL.Consumers.Components + ( runConsumer , runConsumerWithIdleSignal , spawnListener , spawnMonitor , spawnDispatcher ) where -import Control.Applicative import Control.Concurrent.Lifted import Control.Concurrent.STM hiding (atomically) -import Control.Exception (AsyncException(ThreadKilled)) +import Control.Concurrent.STM qualified as STM +import Control.Concurrent.Thread.Lifted qualified as T +import Control.Exception (AsyncException (ThreadKilled)) +import Control.Exception.Safe qualified as ES import Control.Monad import Control.Monad.Base import Control.Monad.Catch import Control.Monad.Time import Control.Monad.Trans import Control.Monad.Trans.Control +import Data.Foldable qualified as F import Data.Function import Data.Int -import Data.Maybe -import Data.Monoid +import Data.Map.Strict qualified as M import Data.Monoid.Utils -import Database.PostgreSQL.PQTypes -import Log -import Prelude -import qualified Control.Concurrent.STM as STM -import qualified Control.Concurrent.Thread.Lifted as T -import qualified Data.Foldable as F -import qualified Data.Map.Strict as M - import Database.PostgreSQL.Consumers.Config import Database.PostgreSQL.Consumers.Consumer import Database.PostgreSQL.Consumers.Utils +import Database.PostgreSQL.PQTypes +import Log --- | Run the consumer. The purpose of the returned monadic --- action is to wait for currently processed jobs and clean up. --- This function is best used in conjunction with 'finalize' to --- seamlessly handle the finalization. +-- | Run the consumer. The purpose of the returned monadic action is to wait for +-- currently processed jobs and clean up. This function is best used in +-- conjunction with 'finalize' to seamlessly handle the finalization. runConsumer - :: ( MonadBaseControl IO m, MonadLog m, MonadMask m, MonadTime m, Eq idx, Show idx - , FromSQL idx, ToSQL idx ) + :: ( MonadBaseControl IO m + , MonadLog m + , MonadMask m + , MonadTime m + , Eq idx + , Show idx + , FromSQL idx + , ToSQL idx + ) => ConsumerConfig m idx job + -- ^ The consumer. -> ConnectionSourceM m -> m (m ()) runConsumer cc cs = runConsumerWithMaybeIdleSignal cc cs Nothing runConsumerWithIdleSignal - :: ( MonadBaseControl IO m, MonadLog m, MonadMask m, MonadTime m, Eq idx, Show idx - , FromSQL idx, ToSQL idx ) + :: ( MonadBaseControl IO m + , MonadLog m + , MonadMask m + , MonadTime m + , Eq idx + , Show idx + , FromSQL idx + , ToSQL idx + ) => ConsumerConfig m idx job + -- ^ The consumer. -> ConnectionSourceM m -> TMVar Bool -> m (m ()) @@ -57,16 +68,23 @@ -- | Run the consumer and also signal whenever the consumer is waiting for -- getNotification or threadDelay. runConsumerWithMaybeIdleSignal - :: ( MonadBaseControl IO m, MonadLog m, MonadMask m, MonadTime m, Eq idx, Show idx - , FromSQL idx, ToSQL idx ) + :: ( MonadBaseControl IO m + , MonadLog m + , MonadMask m + , MonadTime m + , Eq idx + , Show idx + , FromSQL idx + , ToSQL idx + ) => ConsumerConfig m idx job -> ConnectionSourceM m -> Maybe (TMVar Bool) -> m (m ()) -runConsumerWithMaybeIdleSignal cc cs mIdleSignal +runConsumerWithMaybeIdleSignal cc0 cs mIdleSignal | ccMaxRunningJobs cc < 1 = do logInfo_ "ccMaxRunningJobs < 1, not starting the consumer" - return $ return () + pure $ pure () | otherwise = do semaphore <- newMVar () runningJobsInfo <- liftBase $ newTVarIO M.empty @@ -74,7 +92,7 @@ skipLockedTest :: Either DBException () <- try . runDBT cs defaultTransactionSettings $ - runSQL_ "SELECT TRUE FOR UPDATE SKIP LOCKED" + runSQL_ "SELECT TRUE FOR UPDATE SKIP LOCKED" -- If we can't lock rows using 'skip locked' throw an exception either (const $ error "PostgreSQL version with support for SKIP LOCKED is required") pure skipLockedTest @@ -82,40 +100,75 @@ localData ["consumer_id" .= show cid] $ do listener <- spawnListener cc cs semaphore monitor <- localDomain "monitor" $ spawnMonitor cc cs cid - dispatcher <- localDomain "dispatcher" $ spawnDispatcher cc - cs cid semaphore runningJobsInfo runningJobs mIdleSignal - return . localDomain "finalizer" $ do + dispatcher <- + localDomain "dispatcher" $ + spawnDispatcher + cc + cs + cid + semaphore + runningJobsInfo + runningJobs + mIdleSignal + pure . localDomain "finalizer" $ do stopExecution listener stopExecution dispatcher waitForRunningJobs runningJobsInfo runningJobs stopExecution monitor unregisterConsumer cc cs cid where + cc = + cc0 + { ccOnException = \ex job -> do + -- Let asynchronous exceptions through (StopExecution in particular). + ccOnException cc0 ex job `ES.catchAny` \handlerEx -> do + -- Arbitrary delay, but better than letting exceptions from the + -- handler through and potentially crashlooping the consumer: + -- + -- 1. A job J fails with an exception, ccOnException is called and + -- it throws an exception. + -- + -- 2. The consumer goes down, J is now stuck. + -- + -- 3. The consumer is restarted, it tries to clean up stuck jobs + -- (which include J), the cleanup code calls ccOnException on J + -- and if it throws again, we're back to (2). + let action = RerunAfter $ idays 1 + logAttention "ccOnException threw an exception" $ + object + [ "job_id" .= show (ccJobIndex cc0 job) + , "exception" .= show handlerEx + , "action" .= show action + ] + pure action + } + waitForRunningJobs runningJobsInfo runningJobs = do initialJobs <- liftBase $ readTVarIO runningJobsInfo (`fix` initialJobs) $ \loop jobsInfo -> do -- If jobs are still running, display info about them. - when (not $ M.null jobsInfo) $ do - logInfo "Waiting for running jobs" $ object [ - "job_id" .= showJobsInfo jobsInfo - ] + unless (M.null jobsInfo) $ do + logInfo "Waiting for running jobs" $ + object + [ "job_ids" .= showJobsInfo jobsInfo + ] join . atomically $ do jobs <- readTVar runningJobs if jobs == 0 - then return $ return () + then pure $ pure () else do newJobsInfo <- readTVar runningJobsInfo - -- If jobs info didn't change, wait for it to change. - -- Otherwise loop so it either displays the new info - -- or exits if there are no jobs running anymore. - if (newJobsInfo == jobsInfo) + -- If jobs info didn't change, wait for it to change. Otherwise + -- loop so it either displays the new info or exits if there are + -- no jobs running anymore. + if newJobsInfo == jobsInfo then retry - else return $ loop newJobsInfo + else pure $ loop newJobsInfo where showJobsInfo = M.foldr (\idx acc -> show idx : acc) [] --- | Spawn a thread that generates signals for the --- dispatcher to probe the database for incoming jobs. +-- | Spawn a thread that generates signals for the dispatcher to probe the +-- database for incoming jobs. spawnListener :: (MonadBaseControl IO m, MonadMask m) => ConsumerConfig m idx job @@ -124,50 +177,62 @@ -> m ThreadId spawnListener cc cs semaphore = forkP "listener" $ - case ccNotificationChannel cc of - Just chan -> - runDBT cs noTs . bracket_ (listen chan) (unlisten chan) - . forever $ do - -- If there are many notifications, we need to collect them - -- as soon as possible, because they are stored in memory by - -- libpq. They are also not squashed, so we perform the - -- squashing ourselves with the help of MVar (). - void . getNotification $ ccNotificationTimeout cc - lift signalDispatcher - Nothing -> forever $ do - liftBase . threadDelay $ ccNotificationTimeout cc - signalDispatcher + case ccNotificationChannel cc of + Just chan -> + runDBT cs noTs + . bracket_ (listen chan) (unlisten chan) + . forever + $ do + -- If there are many notifications, we need to collect them as soon + -- as possible, because they are stored in memory by libpq. They are + -- also not squashed, so we perform the squashing ourselves with the + -- help of MVar (). + void . getNotification $ ccNotificationTimeout cc + lift signalDispatcher + Nothing -> forever $ do + liftBase . threadDelay $ ccNotificationTimeout cc + signalDispatcher where signalDispatcher = do liftBase $ tryPutMVar semaphore () - noTs = defaultTransactionSettings { - tsAutoTransaction = False - } + noTs = + defaultTransactionSettings + { tsAutoTransaction = False + } --- | Spawn a thread that monitors working consumers --- for activity and periodically updates its own. +-- | Spawn a thread that monitors working consumers for activity and +-- periodically updates its own. spawnMonitor - :: forall m idx job. (MonadBaseControl IO m, MonadLog m, MonadMask m, MonadTime m, - Show idx, FromSQL idx, ToSQL idx) + :: forall m idx job + . ( MonadBaseControl IO m + , MonadLog m + , MonadMask m + , MonadTime m + , Show idx + , FromSQL idx + , ToSQL idx + ) => ConsumerConfig m idx job -> ConnectionSourceM m -> ConsumerID -> m ThreadId -spawnMonitor ConsumerConfig{..} cs cid = forkP "monitor" . forever $ do +spawnMonitor ConsumerConfig {..} cs cid = forkP "monitor" . forever $ do runDBT cs ts $ do now <- currentTime -- Update last_activity of the consumer. - ok <- runSQL01 $ smconcat [ - "UPDATE" <+> raw ccConsumersTable - , "SET last_activity = " now - , "WHERE id =" cid - , " AND name =" unRawSQL ccJobsTable - ] + ok <- + runPreparedSQL01 (preparedSqlName "setActivity" ccConsumersTable) $ + smconcat + [ "UPDATE" <+> raw ccConsumersTable + , "SET last_activity = " now + , "WHERE id =" cid + , " AND name =" unRawSQL ccJobsTable + ] if ok then logInfo_ "Activity of the consumer updated" else do - logInfo_ $ "Consumer is not registered" + logInfo_ "Consumer is not registered" throwM ThreadKilled (inactiveConsumers, freedJobs) <- runDBT cs ts $ do now <- currentTime @@ -175,50 +240,61 @@ -- delete them here, because if the coresponding reserved_by column in the -- jobs table has an IMMEDIATE foreign key with the ON DELETE SET NULL -- property, we will not be able to determine stuck jobs in the next step. - runSQL_ $ smconcat - [ "SELECT id::bigint" - , "FROM" <+> raw ccConsumersTable - , "WHERE last_activity +" iminutes 1 <+> "<= " now - , " AND name =" unRawSQL ccJobsTable - , "FOR UPDATE SKIP LOCKED" - ] + runPreparedSQL_ (preparedSqlName "reserveConsumers" ccConsumersTable) $ + smconcat + [ "SELECT id::bigint" + , "FROM" <+> raw ccConsumersTable + , "WHERE last_activity +" iminutes 1 <+> "<= " now + , " AND name =" unRawSQL ccJobsTable + , "FOR UPDATE SKIP LOCKED" + ] fetchMany (runIdentity @Int64) >>= \case [] -> pure (0, []) inactive -> do -- Fetch all stuck jobs and run ccOnException on them to determine -- actions. This is necessary e.g. to be able to apply exponential -- backoff to them correctly. - runSQL_ $ smconcat - [ "SELECT" <+> mintercalate ", " ccJobSelectors - , "FROM" <+> raw ccJobsTable - , "WHERE reserved_by = ANY(" Array1 inactive <+> ")" - , "FOR UPDATE SKIP LOCKED" - ] + runPreparedSQL_ (preparedSqlName "findStuck" ccJobsTable) $ + smconcat + [ "SELECT" <+> mintercalate ", " ccJobSelectors + , "FROM" <+> raw ccJobsTable + , "WHERE reserved_by = ANY(" Array1 inactive <+> ")" + , "FOR UPDATE SKIP LOCKED" + ] stuckJobs <- fetchMany ccJobFetcher unless (null stuckJobs) $ do results <- forM stuckJobs $ \job -> do action <- lift $ ccOnException (toException ThreadKilled) job pure (ccJobIndex job, Failed action) - runSQL_ $ updateJobsQuery ccJobsTable results now - runSQL_ $ smconcat - [ "DELETE FROM" <+> raw ccConsumersTable - , "WHERE id = ANY(" Array1 inactive <+> ")" - ] + runPreparedSQL_ (preparedSqlName "updateJobs" ccJobsTable) $ updateJobsQuery ccJobsTable results now + runPreparedSQL_ (preparedSqlName "removeInactive" ccConsumersTable) $ + smconcat + [ "DELETE FROM" <+> raw ccConsumersTable + , "WHERE id = ANY(" Array1 inactive <+> ")" + ] pure (length inactive, map ccJobIndex stuckJobs) when (inactiveConsumers > 0) $ do - logInfo "Unregistered inactive consumers" $ object [ - "inactive_consumers" .= inactiveConsumers - ] + logInfo "Unregistered inactive consumers" $ + object + [ "inactive_consumers" .= inactiveConsumers + ] unless (null freedJobs) $ do - logInfo "Freed locked jobs" $ object [ - "freed_jobs" .= map show freedJobs - ] + logInfo "Freed locked jobs" $ + object + [ "freed_jobs" .= map show freedJobs + ] liftBase . threadDelay $ 30 * 1000000 -- wait 30 seconds -- | Spawn a thread that reserves and processes jobs. spawnDispatcher - :: forall m idx job. ( MonadBaseControl IO m, MonadLog m, MonadMask m, MonadTime m - , Show idx, ToSQL idx ) + :: forall m idx job + . ( MonadBaseControl IO m + , MonadLog m + , MonadMask m + , MonadTime m + , Show idx + , ToSQL idx + ) => ConsumerConfig m idx job -> ConnectionSourceM m -> ConsumerID @@ -227,8 +303,7 @@ -> TVar Int -> Maybe (TMVar Bool) -> m ThreadId -spawnDispatcher ConsumerConfig{..} cs cid semaphore - runningJobsInfo runningJobs mIdleSignal = +spawnDispatcher ConsumerConfig {..} cs cid semaphore runningJobsInfo runningJobs mIdleSignal = forkP "dispatcher" . forever $ do void $ takeMVar semaphore someJobWasProcessed <- loop 1 @@ -236,9 +311,9 @@ then setIdle False else setIdle True where - setIdle :: forall m' . (MonadBaseControl IO m') => Bool -> m' () + setIdle :: forall m'. MonadBaseControl IO m' => Bool -> m' () setIdle isIdle = case mIdleSignal of - Nothing -> return () + Nothing -> pure () Just idleSignal -> atomically $ do _ <- tryTakeTMVar idleSignal putTMVar idleSignal isIdle @@ -247,88 +322,95 @@ loop limit = do (batch, batchSize) <- reserveJobs limit when (batchSize > 0) $ do - logInfo "Processing batch" $ object [ - "batch_size" .= batchSize - ] - -- Update runningJobs before forking so that we can - -- adjust maxBatchSize appropriately later. We also - -- need to mask asynchronous exceptions here as we - -- rely on correct value of runningJobs to perform - -- graceful termination. + logInfo "Processing batch" $ + object + [ "batch_size" .= batchSize + ] + -- Update runningJobs before forking so that we can adjust + -- maxBatchSize appropriately later. We also need to mask asynchronous + -- exceptions here as we rely on correct value of runningJobs to + -- perform graceful termination. mask $ \restore -> do - atomically $ modifyTVar' runningJobs (+batchSize) + atomically $ modifyTVar' runningJobs (+ batchSize) let subtractJobs = atomically $ do modifyTVar' runningJobs (subtract batchSize) - void . forkP "batch processor" - . (`finally` subtractJobs) . restore $ do - mapM startJob batch >>= mapM joinJob >>= updateJobs + void + . forkP "batch processor" + . (`finally` subtractJobs) + . restore + $ do + mapM startJob batch >>= mapM joinJob >>= updateJobs when (batchSize == limit) $ do maxBatchSize <- atomically $ do jobs <- readTVar runningJobs when (jobs >= ccMaxRunningJobs) retry - return $ ccMaxRunningJobs - jobs - void $ loop $ min maxBatchSize (2*limit) + pure $ ccMaxRunningJobs - jobs + void . loop $ min maxBatchSize (2 * limit) - return (batchSize > 0) + pure (batchSize > 0) reserveJobs :: Int -> m ([job], Int) reserveJobs limit = runDBT cs ts $ do now <- currentTime - n <- runSQL $ smconcat [ - "UPDATE" <+> raw ccJobsTable <+> "SET" - , " reserved_by =" cid - , ", attempts = CASE" - , " WHEN finished_at IS NULL THEN attempts + 1" - , " ELSE 1" - , " END" - , "WHERE id IN (" <> reservedJobs now <> ")" - , "RETURNING" <+> mintercalate ", " ccJobSelectors - ] + n <- + runPreparedSQL (preparedSqlName "setReservation" ccJobsTable) $ + smconcat + [ "UPDATE" <+> raw ccJobsTable <+> "SET" + , " reserved_by =" cid + , ", attempts = CASE" + , " WHEN finished_at IS NULL THEN attempts + 1" + , " ELSE 1" + , " END" + , "WHERE id IN (" <> reservedJobs now <> ")" + , "RETURNING" <+> mintercalate ", " ccJobSelectors + ] -- Decode lazily as we want the transaction to be as short as possible. - (, n) . F.toList . fmap ccJobFetcher <$> queryResult + (,n) . F.toList . fmap ccJobFetcher <$> queryResult where reservedJobs :: UTCTime -> SQL - reservedJobs now = smconcat [ - "SELECT id FROM" <+> raw ccJobsTable - , "WHERE" - , " reserved_by IS NULL" - , " AND run_at IS NOT NULL" - , " AND run_at <= " now - , " ORDER BY run_at" - , "LIMIT" limit - , "FOR UPDATE SKIP LOCKED" - ] + reservedJobs now = + smconcat + [ "SELECT id FROM" <+> raw ccJobsTable + , "WHERE" + , " reserved_by IS NULL" + , " AND run_at IS NOT NULL" + , " AND run_at <= " now + , " ORDER BY run_at" + , "LIMIT" limit + , "FOR UPDATE SKIP LOCKED" + ] - -- | Spawn each job in a separate thread. + -- Spawn each job in a separate thread. startJob :: job -> m (job, m (T.Result Result)) startJob job = do (_, joinFork) <- mask $ \restore -> T.fork $ do tid <- myThreadId bracket_ (registerJob tid) (unregisterJob tid) . restore $ do ccProcessJob job - return (job, joinFork) + pure (job, joinFork) where registerJob tid = atomically $ do modifyTVar' runningJobsInfo . M.insert tid $ ccJobIndex job unregisterJob tid = atomically $ do - modifyTVar' runningJobsInfo $ M.delete tid + modifyTVar' runningJobsInfo $ M.delete tid - -- | Wait for all the jobs and collect their results. + -- Wait for all the jobs and collect their results. joinJob :: (job, m (T.Result Result)) -> m (idx, Result) - joinJob (job, joinFork) = joinFork >>= \eres -> case eres of - Right result -> return (ccJobIndex job, result) - Left ex -> do - action <- ccOnException ex job - logAttention "Unexpected exception caught while processing job" $ - object [ - "job_id" .= show (ccJobIndex job) - , "exception" .= show ex - , "action" .= show action - ] - return (ccJobIndex job, Failed action) + joinJob (job, joinFork) = + joinFork >>= \case + Right result -> pure (ccJobIndex job, result) + Left ex -> do + action <- ccOnException ex job + logAttention "Unexpected exception caught while processing job" $ + object + [ "job_id" .= show (ccJobIndex job) + , "exception" .= show ex + , "action" .= show action + ] + pure (ccJobIndex job, Failed action) - -- | Update status of the jobs. + -- Update status of the jobs. updateJobs :: [(idx, Result)] -> m () updateJobs results = runDBT cs ts $ do now <- currentTime @@ -343,66 +425,68 @@ -> [(idx, Result)] -> UTCTime -> SQL -updateJobsQuery jobsTable results now = smconcat - [ "WITH removed AS (" - , " DELETE FROM" <+> raw jobsTable - , " WHERE id = ANY(" Array1 deletes <+> ")" - , ")" - , "UPDATE" <+> raw jobsTable <+> "SET" - , " reserved_by = NULL" - , ", run_at = CASE" - , " WHEN FALSE THEN run_at" - , smconcat $ M.foldrWithKey retryToSQL [] retries - , " ELSE NULL" -- processed - , " END" - , ", finished_at = CASE" - , " WHEN id = ANY(" Array1 successes <+> ") THEN " now - , " ELSE NULL" - , " END" - , "WHERE id = ANY(" Array1 (map fst updates) <+> ")" - ] +updateJobsQuery jobsTable results now = + smconcat + [ "WITH removed AS (" + , " DELETE FROM" <+> raw jobsTable + , " WHERE id = ANY(" Array1 deletes <+> ")" + , ")" + , "UPDATE" <+> raw jobsTable <+> "SET" + , " reserved_by = NULL" + , ", run_at = CASE" + , " WHEN FALSE THEN run_at" + , smconcat $ M.foldrWithKey retryToSQL [] retries + , " ELSE NULL" -- processed + , " END" + , ", finished_at = CASE" + , " WHEN id = ANY(" Array1 successes <+> ") THEN " now + , " ELSE NULL" + , " END" + , "WHERE id = ANY(" Array1 (map fst updates) <+> ")" + ] where retryToSQL (Left int) ids = ("WHEN id = ANY(" Array1 ids <+> ") THEN " now <> " +" int :) retryToSQL (Right time) ids = ("WHEN id = ANY(" Array1 ids <+> ") THEN" time :) - retries = foldr step M.empty $ map getAction updates + retries = foldr (step . getAction) M.empty updates where getAction (idx, result) = case result of - Ok action -> (idx, action) + Ok action -> (idx, action) Failed action -> (idx, action) step (idx, action) iretries = case action of - MarkProcessed -> iretries + MarkProcessed -> iretries RerunAfter int -> M.insertWith (++) (Left int) [idx] iretries - RerunAt time -> M.insertWith (++) (Right time) [idx] iretries - Remove -> error "updateJobs: Remove should've been filtered out" + RerunAt time -> M.insertWith (++) (Right time) [idx] iretries + Remove -> error "updateJobs: Remove should've been filtered out" successes = foldr step [] updates where - step (idx, Ok _) acc = idx : acc - step (_, Failed _) acc = acc + step (idx, Ok _) acc = idx : acc + step (_, Failed _) acc = acc (deletes, updates) = foldr step ([], []) results where step job@(idx, result) (ideletes, iupdates) = case result of - Ok Remove -> (idx : ideletes, iupdates) + Ok Remove -> (idx : ideletes, iupdates) Failed Remove -> (idx : ideletes, iupdates) - _ -> (ideletes, job : iupdates) - + _ -> (ideletes, job : iupdates) ts :: TransactionSettings -ts = defaultTransactionSettings { - -- PostgreSQL doesn't seem to handle very high amount of - -- concurrent transactions that modify multiple rows in - -- the same table well (see updateJobs) and sometimes (very - -- rarely though) ends up in a deadlock. It doesn't matter - -- much though, we just restart the transaction in such case. - tsRestartPredicate = Just . RestartPredicate - $ \e _ -> qeErrorCode e == DeadlockDetected - || qeErrorCode e == SerializationFailure -} +ts = + defaultTransactionSettings + { -- PostgreSQL doesn't seem to handle very high amount of concurrent + -- transactions that modify multiple rows in the same table well (see + -- updateJobs) and sometimes (very rarely though) ends up in a + -- deadlock. It doesn't matter much though, we just restart the + -- transaction in such case. + tsRestartPredicate = Just . RestartPredicate $ + \e _ -> + qeErrorCode e == DeadlockDetected + || qeErrorCode e == SerializationFailure + } atomically :: MonadBase IO m => STM a -> m a atomically = liftBase . STM.atomically diff --git a/src/Database/PostgreSQL/Consumers/Config.hs b/src/Database/PostgreSQL/Consumers/Config.hs --- a/src/Database/PostgreSQL/Consumers/Config.hs +++ b/src/Database/PostgreSQL/Consumers/Config.hs @@ -1,14 +1,11 @@ -{-# LANGUAGE ExistentialQuantification #-} -module Database.PostgreSQL.Consumers.Config ( - Action(..) - , Result(..) - , ConsumerConfig(..) +module Database.PostgreSQL.Consumers.Config + ( Action (..) + , Result (..) + , ConsumerConfig (..) ) where import Control.Exception (SomeException) import Data.Time -import Prelude - import Database.PostgreSQL.PQTypes.FromRow import Database.PostgreSQL.PQTypes.Interval import Database.PostgreSQL.PQTypes.Notification @@ -21,101 +18,99 @@ | RerunAfter Interval | RerunAt UTCTime | Remove - deriving (Eq, Ord, Show) + deriving (Eq, Ord, Show) -- | Result of processing a job. data Result = Ok Action | Failed Action deriving (Eq, Ord, Show) -- | Config of a consumer. -data ConsumerConfig m idx job = forall row. FromRow row => ConsumerConfig { --- | Name of the database table where jobs are stored. The table needs to have --- the following columns in order to be suitable for acting as a job queue: --- --- * __id__ - represents ID of the job. Needs to be a primary key of a type --- convertible to text, not nullable. --- --- * __run_at__ - represents the time at which the job will be --- processed. Needs to be nullable, of a type comparable with @now()@ --- (TIMESTAMPTZ is recommended). --- --- Note: a job with run_at set to NULL is never picked for processing. Useful --- for storing already processed/expired jobs for debugging purposes. --- --- It's highly recommended to have an index on this column. --- --- * __finished_at__ - represents the time at which job processing was finished. --- Needs to be nullable, of a type you can assign @now()@ to (TIMESTAMPTZ is --- recommended). NULL means that the job was either never processed or that --- it was started and failed at least once. --- --- * __reserved_by__ - represents ID of the consumer that currently processes the --- job. Needs to be nullable, of the type corresponding to id in the table --- 'ccConsumersTable'. It's recommended (though not neccessary) to make it --- a foreign key referencing id in 'ccConsumersTable' with ON DELETE SET NULL. --- --- * __attempts__ - represents number of job processing attempts made so far. Needs --- to be not nullable, of type INTEGER. Initial value of a fresh job should be --- 0, therefore it makes sense to make the column default to 0. --- - ccJobsTable :: !(RawSQL ()) --- | Name of a database table where registered consumers are stored. The table --- itself needs to have the following columns: --- --- * __id__ - represents ID of a consumer. Needs to be a primary key of the type --- SERIAL or BIGSERIAL (recommended). --- --- * __name__ - represents jobs table of the consumer. Needs to be not nullable, --- of type TEXT. Allows for tracking consumers of multiple queues with one --- table. Set to 'ccJobsTable'. --- --- * __last_activity__ - represents the last registered activity of the consumer. --- It's updated periodically by all currently running consumers every 30 --- seconds to prove that they are indeed running. They also check for the --- registered consumers that didn't update their status for a minute. If any --- such consumers are found, they are presumed to be not working and all the --- jobs reserved by them are released. This prevents the situation where --- a consumer with reserved jobs silently fails (e.g. because of a hard crash) --- and these jobs stay locked forever, yet are never processed. --- -, ccConsumersTable :: !(RawSQL ()) --- | Fields needed to be selected from the jobs table in order to --- assemble a job. -, ccJobSelectors :: ![SQL] --- | Function that transforms the list of fields into a job. -, ccJobFetcher :: !(row -> job) --- | Selector for taking out job ID from the job object. -, ccJobIndex :: !(job -> idx) --- | Notification channel used for listening for incoming jobs. --- Whenever the consumer receives a notification, it checks the database --- for any pending jobs (@'run_at <= NOW()'@) and runs them all. --- If set to 'Nothing', no listening is performed and 'ccNotificationTimeout' --- should be set to a positive number, otherwise no jobs would be ever run. --- 'ccNotificationChannel' and 'ccNotificationTimeout' can be combined. --- The consumer will check for pending jobs either when notification is received or --- no notification is received for 'ccNotificationTimeout' microseconds since --- the last check. -, ccNotificationChannel :: !(Maybe Channel) --- | Timeout of checking for any pending jobs (@'run_at <= NOW()'@), in --- microseconds. The consumer checks the database for any pending jobs after --- 'ccNotificationTimeout' microseconds since the last check was performed, --- runs them until all pending jobs are processed and after that, the cycle --- repeats. Note that even if 'ccNotificationChannel' is 'Just', you need to set --- 'ccNotificationTimeout' to a reasonable number if the jobs you process may fail --- and are retried later, as there is no way to signal with a notification that --- a job will need to be performed e.g. in 5 minutes. However, if --- 'ccNotificationChannel' is 'Just' and jobs are never retried, you can set it --- to -1, then listening will never timeout. Otherwise it needs to be a positive --- number. -, ccNotificationTimeout :: !Int --- | Maximum amount of jobs that can be processed in parallel. -, ccMaxRunningJobs :: !Int --- | Function that processes a job. It's recommended to process each --- job in a separate DB transaction, otherwise you'll have to remember --- to commit your changes to the database manually. -, ccProcessJob :: !(job -> m Result) --- | Action taken if a job processing function throws an exception. --- Note that if this action throws an exception, the consumer goes --- down, so it's best to ensure that it doesn't throw. -, ccOnException :: !(SomeException -> job -> m Action) -} +data ConsumerConfig m idx job = forall row. FromRow row => ConsumerConfig + { ccJobsTable :: !(RawSQL ()) + -- ^ Name of the database table where jobs are stored. The table needs to have + -- the following columns in order to be suitable for acting as a job queue: + -- + -- * __id__ - represents ID of the job. Needs to be a primary key of a type + -- convertible to text, not nullable. + -- + -- * __run_at__ - represents the time at which the job will be + -- processed. Needs to be nullable, of a type comparable with @now()@ + -- (TIMESTAMPTZ is recommended). + -- + -- Note: a job with run_at set to NULL is never picked for processing. Useful + -- for storing already processed/expired jobs for debugging purposes. + -- + -- It's highly recommended to have an index on this column. + -- + -- * __finished_at__ - represents the time at which job processing was + -- finished. Needs to be nullable, of a type you can assign @now()@ to + -- (TIMESTAMPTZ is recommended). NULL means that the job was either never + -- processed or that it was started and failed at least once. + -- + -- * __reserved_by__ - represents ID of the consumer that currently processes + -- the job. Needs to be nullable, of the type corresponding to id in the table + -- 'ccConsumersTable'. It's recommended (though not neccessary) to make it a + -- foreign key referencing id in 'ccConsumersTable' with ON DELETE SET NULL. + -- + -- * __attempts__ - represents number of job processing attempts made so + -- far. Needs to be not nullable, of type INTEGER. Initial value of a fresh + -- job should be 0, therefore it makes sense to make the column default to 0. + , ccConsumersTable :: !(RawSQL ()) + -- ^ Name of a database table where registered consumers are stored. The table + -- itself needs to have the following columns: + -- + -- * __id__ - represents ID of a consumer. Needs to be a primary key of the + -- type SERIAL or BIGSERIAL (recommended). + -- + -- * __name__ - represents jobs table of the consumer. Needs to be not + -- nullable, of type TEXT. Allows for tracking consumers of multiple queues + -- with one table. Set to 'ccJobsTable'. + -- + -- * __last_activity__ - represents the last registered activity of the + -- consumer. It's updated periodically by all currently running consumers + -- every 30 seconds to prove that they are indeed running. They also check for + -- the registered consumers that didn't update their status for a minute. If + -- any such consumers are found, they are presumed to be not working and all + -- the jobs reserved by them are released. This prevents the situation where a + -- consumer with reserved jobs silently fails (e.g. because of a hard crash) + -- and these jobs stay locked forever, yet are never processed. + , ccJobSelectors :: ![SQL] + -- ^ Fields needed to be selected from the jobs table in order to assemble a + -- job. + , ccJobFetcher :: !(row -> job) + -- ^ Function that transforms the list of fields into a job. + , ccJobIndex :: !(job -> idx) + -- ^ Selector for taking out job ID from the job object. + , ccNotificationChannel :: !(Maybe Channel) + -- ^ Notification channel used for listening for incoming jobs. Whenever the + -- consumer receives a notification, it checks the database for any pending + -- jobs (@'run_at <= NOW()'@) and runs them all. If set to 'Nothing', no + -- listening is performed and 'ccNotificationTimeout' should be set to a + -- positive number, otherwise no jobs would be ever run. + -- 'ccNotificationChannel' and 'ccNotificationTimeout' can be combined. The + -- consumer will check for pending jobs either when notification is received + -- or no notification is received for 'ccNotificationTimeout' microseconds + -- since the last check. + , ccNotificationTimeout :: !Int + -- ^ Timeout of checking for any pending jobs (@'run_at <= NOW()'@), in + -- microseconds. The consumer checks the database for any pending jobs after + -- 'ccNotificationTimeout' microseconds since the last check was performed, + -- runs them until all pending jobs are processed and after that, the cycle + -- repeats. Note that even if 'ccNotificationChannel' is 'Just', you need to + -- set 'ccNotificationTimeout' to a reasonable number if the jobs you process + -- may fail and are retried later, as there is no way to signal with a + -- notification that a job will need to be performed e.g. in 5 + -- minutes. However, if 'ccNotificationChannel' is 'Just' and jobs are never + -- retried, you can set it to -1, then listening will never timeout. Otherwise + -- it needs to be a positive number. + , ccMaxRunningJobs :: !Int + -- ^ Maximum amount of jobs that can be processed in parallel. + , ccProcessJob :: !(job -> m Result) + -- ^ Function that processes a job. It's recommended to process each job in a + -- separate DB transaction, otherwise you'll have to remember to commit your + -- changes to the database manually. + , ccOnException :: !(SomeException -> job -> m Action) + -- ^ Action taken if a job processing function throws an exception. For + -- robustness it's best to ensure that it doesn't throw. If it does, the + -- exception will be logged and the job in question postponed by a day. + } diff --git a/src/Database/PostgreSQL/Consumers/Consumer.hs b/src/Database/PostgreSQL/Consumers/Consumer.hs --- a/src/Database/PostgreSQL/Consumers/Consumer.hs +++ b/src/Database/PostgreSQL/Consumers/Consumer.hs @@ -1,21 +1,17 @@ -{-# LANGUAGE TypeApplications #-} -module Database.PostgreSQL.Consumers.Consumer ( - ConsumerID +module Database.PostgreSQL.Consumers.Consumer + ( ConsumerID , registerConsumer , unregisterConsumer ) where -import Control.Applicative import Control.Monad.Base import Control.Monad.Catch import Control.Monad.Time import Data.Int -import Data.Monoid import Data.Monoid.Utils -import Database.PostgreSQL.PQTypes -import Prelude - import Database.PostgreSQL.Consumers.Config +import Database.PostgreSQL.Consumers.Utils +import Database.PostgreSQL.PQTypes -- | ID of a consumer. newtype ConsumerID = ConsumerID Int64 @@ -33,25 +29,27 @@ instance Show ConsumerID where showsPrec p (ConsumerID n) = showsPrec p n --- | Register consumer in the consumers table, --- so that it can reserve jobs using acquired ID. +-- | Register consumer in the consumers table, so that it can reserve jobs using +-- acquired ID. registerConsumer :: (MonadBase IO m, MonadMask m, MonadTime m) => ConsumerConfig n idx job -> ConnectionSourceM m -> m ConsumerID -registerConsumer ConsumerConfig{..} cs = runDBT cs ts $ do +registerConsumer ConsumerConfig {..} cs = runDBT cs ts $ do now <- currentTime - runSQL_ $ smconcat [ - "INSERT INTO" <+> raw ccConsumersTable - , "(name, last_activity) VALUES (" unRawSQL ccJobsTable <> ", " now <> ")" - , "RETURNING id" - ] + runPreparedSQL_ (preparedSqlName "registerConsumer" ccConsumersTable) $ + smconcat + [ "INSERT INTO" <+> raw ccConsumersTable + , "(name, last_activity) VALUES (" unRawSQL ccJobsTable <> ", " now <> ")" + , "RETURNING id" + ] fetchOne runIdentity where - ts = defaultTransactionSettings { - tsAutoTransaction = False - } + ts = + defaultTransactionSettings + { tsAutoTransaction = False + } -- | Unregister consumer with a given ID. unregisterConsumer @@ -60,21 +58,24 @@ -> ConnectionSourceM m -> ConsumerID -> m () -unregisterConsumer ConsumerConfig{..} cs wid = runDBT cs ts $ do - -- Free tasks manually in case there is no - -- foreign key constraint on reserved_by, - runSQL_ $ smconcat [ - "UPDATE" <+> raw ccJobsTable - , " SET reserved_by = NULL" - , " WHERE reserved_by =" wid - ] - runSQL_ $ smconcat [ - "DELETE FROM " <+> raw ccConsumersTable - , "WHERE id =" wid - , " AND name =" unRawSQL ccJobsTable - ] +unregisterConsumer ConsumerConfig {..} cs wid = runDBT cs ts $ do + -- Free tasks manually in case there is no foreign key constraint on + -- reserved_by, + runPreparedSQL_ (preparedSqlName "deregisterJobs" ccJobsTable) $ + smconcat + [ "UPDATE" <+> raw ccJobsTable + , " SET reserved_by = NULL" + , " WHERE reserved_by =" wid + ] + runPreparedSQL_ (preparedSqlName "removeConsumers" ccConsumersTable) $ + smconcat + [ "DELETE FROM " <+> raw ccConsumersTable + , "WHERE id =" wid + , " AND name =" unRawSQL ccJobsTable + ] where - ts = defaultTransactionSettings { - tsRestartPredicate = Just . RestartPredicate - $ \e _ -> qeErrorCode e == DeadlockDetected - } + ts = + defaultTransactionSettings + { tsRestartPredicate = Just . RestartPredicate $ + \e _ -> qeErrorCode e == DeadlockDetected + } diff --git a/src/Database/PostgreSQL/Consumers/Utils.hs b/src/Database/PostgreSQL/Consumers/Utils.hs --- a/src/Database/PostgreSQL/Consumers/Utils.hs +++ b/src/Database/PostgreSQL/Consumers/Utils.hs @@ -1,43 +1,51 @@ -module Database.PostgreSQL.Consumers.Utils ( - finalize - , ThrownFrom(..) +module Database.PostgreSQL.Consumers.Utils + ( finalize + , ThrownFrom (..) , stopExecution , forkP , gforkP + , preparedSqlName ) where import Control.Concurrent.Lifted +import Control.Concurrent.Thread.Group.Lifted qualified as TG +import Control.Concurrent.Thread.Lifted qualified as T +import Control.Exception.Lifted qualified as E import Control.Monad.Base import Control.Monad.Catch import Control.Monad.Trans.Control +import Data.Maybe +import Data.Text qualified as T import Data.Typeable -import Prelude -import qualified Control.Concurrent.Thread.Group.Lifted as TG -import qualified Control.Concurrent.Thread.Lifted as T -import qualified Control.Exception.Lifted as E +import Database.PostgreSQL.PQTypes.Class +import Database.PostgreSQL.PQTypes.SQL.Raw --- | Run an action 'm' that returns a finalizer and perform the --- returned finalizer after the action 'action' completes. +-- | Run an action 'm' that returns a finalizer and perform the returned +-- finalizer after the action 'action' completes. finalize :: (MonadMask m, MonadBase IO m) => m (m ()) -> m a -> m a finalize m action = do finalizer <- newEmptyMVar - flip finally (tryTakeMVar finalizer >>= maybe (return ()) id) $ do + flip finally (tryTakeMVar finalizer >>= fromMaybe (pure ())) $ do putMVar finalizer =<< m action ---------------------------------------- -- | Exception thrown to a thread to stop its execution. --- All exceptions other than 'StopExecution' thrown to --- threads spawned by 'forkP' and 'gforkP' are propagated --- back to the parent thread. +-- +-- All exceptions other than 'StopExecution' thrown to threads spawned by +-- 'forkP' and 'gforkP' are propagated back to the parent thread. data StopExecution = StopExecution deriving (Show, Typeable) -instance Exception StopExecution +instance Exception StopExecution where + toException = E.asyncExceptionToException + fromException = E.asyncExceptionFromException + -- | Exception thrown from a child thread. data ThrownFrom = ThrownFrom String SomeException deriving (Show, Typeable) + instance Exception ThrownFrom -- | Stop execution of a thread. @@ -46,30 +54,36 @@ ---------------------------------------- --- | Modified version of 'fork' that propagates --- thrown exceptions to the parent thread. +-- | Modified version of 'fork' that propagates thrown exceptions to the parent +-- thread. forkP :: MonadBaseControl IO m => String -> m () -> m ThreadId forkP = forkImpl fork --- | Modified version of 'TG.fork' that propagates --- thrown exceptions to the parent thread. -gforkP :: MonadBaseControl IO m - => TG.ThreadGroup - -> String - -> m () - -> m (ThreadId, m (T.Result ())) +-- | Modified version of 'TG.fork' that propagates thrown exceptions to the +-- parent thread. +gforkP + :: MonadBaseControl IO m + => TG.ThreadGroup + -> String + -> m () + -> m (ThreadId, m (T.Result ())) gforkP = forkImpl . TG.fork ---------------------------------------- -forkImpl :: MonadBaseControl IO m - => (m () -> m a) - -> String - -> m () - -> m a +forkImpl + :: MonadBaseControl IO m + => (m () -> m a) + -> String + -> m () + -> m a forkImpl ffork tname m = E.mask $ \release -> do parent <- myThreadId - ffork $ release m `E.catches` [ - E.Handler $ \StopExecution -> return () - , E.Handler $ (throwTo parent . ThrownFrom tname) - ] + ffork $ + release m + `E.catches` [ E.Handler $ \StopExecution -> pure () + , E.Handler $ throwTo parent . ThrownFrom tname + ] + +preparedSqlName :: T.Text -> RawSQL () -> QueryName +preparedSqlName baseName tableName = QueryName . T.take 63 $ baseName <> "$" <> unRawSQL tableName diff --git a/test/Test.hs b/test/Test.hs --- a/test/Test.hs +++ b/test/Test.hs @@ -1,240 +1,281 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} - module Main where -import Database.PostgreSQL.Consumers -import Database.PostgreSQL.PQTypes -import Database.PostgreSQL.PQTypes.Checks -import Database.PostgreSQL.PQTypes.Model - import Control.Concurrent.STM import Control.Exception import Control.Monad import Control.Monad.Base import Control.Monad.Catch import Control.Monad.IO.Class -import Control.Monad.Loops import Control.Monad.State.Strict import Control.Monad.Time import Control.Monad.Trans.Control import Data.Int -import Data.Monoid +import Data.Text qualified as T import Data.Time +import Database.PostgreSQL.Consumers +import Database.PostgreSQL.PQTypes +import Database.PostgreSQL.PQTypes.Checks +import Database.PostgreSQL.PQTypes.Model import Log import Log.Backend.StandardOutput -import Prelude import System.Environment import System.Exit +import Test.HUnit qualified as T import TextShow -import qualified Data.Text as T -import qualified Test.HUnit as T - -data TestEnvSt = TestEnvSt { - teCurrentTime :: UTCTime +data TestEnvSt = TestEnvSt + { teCurrentTime :: UTCTime , teMonotonicTime :: Double } type InnerTestEnv = StateT TestEnvSt (DBT (LogT IO)) -newtype TestEnv a = TestEnv { unTestEnv :: InnerTestEnv a } +newtype TestEnv a = TestEnv {unTestEnv :: InnerTestEnv a} deriving (Applicative, Functor, Monad, MonadLog, MonadDB, MonadThrow, MonadCatch, MonadMask, MonadIO, MonadBase IO, MonadState TestEnvSt) instance MonadBaseControl IO TestEnv where type StM TestEnv a = StM InnerTestEnv a - liftBaseWith f = TestEnv $ liftBaseWith $ \run -> f $ run . unTestEnv - restoreM = TestEnv . restoreM - {-# INLINE liftBaseWith #-} - {-# INLINE restoreM #-} + liftBaseWith f = TestEnv $ liftBaseWith (\run -> f $ run . unTestEnv) + restoreM = TestEnv . restoreM instance MonadTime TestEnv where currentTime = gets teCurrentTime monotonicTime = gets teMonotonicTime -modifyTestTime :: (MonadState TestEnvSt m) => (UTCTime -> UTCTime) -> m () -modifyTestTime modtime = modify (\te -> te { teCurrentTime = modtime . teCurrentTime $ te }) +modifyTestTime :: MonadState TestEnvSt m => (UTCTime -> UTCTime) -> m () +modifyTestTime modtime = modify (\te -> te {teCurrentTime = modtime . teCurrentTime $ te}) runTestEnv :: ConnectionSourceM (LogT IO) -> Logger -> TestEnv a -> IO a -runTestEnv connSource logger m = - (runLogT "consumers-test" logger defaultLogLevel) - . (runDBT connSource defaultTransactionSettings) - . (\m' -> fst <$> (runStateT m' $ TestEnvSt (UTCTime (ModifiedJulianDay 0) 0) 0)) - . unTestEnv - $ m +runTestEnv connSource logger = + runLogT "consumers-test" logger defaultLogLevel + . runDBT connSource defaultTransactionSettings + . (\m' -> fst <$> runStateT m' (TestEnvSt (UTCTime (ModifiedJulianDay 0) 0) 0)) + . unTestEnv main :: IO () -main = void $ T.runTestTT $ T.TestCase test +main = void . T.runTestTT $ T.TestCase test test :: IO () test = do - connString <- getArgs >>= \case - connString : _args -> return $ T.pack connString - [] -> lookupEnv "GITHUB_ACTIONS" >>= \case - Just "true" -> return "host=postgres user=postgres password=postgres" - _ -> printUsage >> exitFailure + connString <- + getArgs >>= \case + connString : _args -> pure $ T.pack connString + [] -> + lookupEnv "GITHUB_ACTIONS" >>= \case + Just "true" -> pure "host=postgres user=postgres password=postgres" + _ -> printUsage >> exitFailure - let connSettings = defaultConnectionSettings - { csConnInfo = connString } + let connSettings = + defaultConnectionSettings + { csConnInfo = connString + } ConnectionSource connSource = simpleSource connSettings withStdOutLogger $ \logger -> runTestEnv connSource logger $ do createTables - idleSignal <- liftIO $ atomically $ newEmptyTMVar + idleSignal <- liftIO newEmptyTMVarIO putJob 10 >> commit - forM_ [1..10::Int] $ \_ -> do + forM_ [1 .. 10 :: Int] $ \_ -> do -- Move time forward 2hours, because jobs are scheduled 1 hour into future - modifyTestTime $ addUTCTime (2*60*60) - finalize (localDomain "process" $ - runConsumerWithIdleSignal consumerConfig connSource idleSignal) $ do - waitUntilTrue idleSignal + modifyTestTime $ addUTCTime (2 * 60 * 60) + finalize + ( localDomain "process" $ + runConsumerWithIdleSignal consumerConfig connSource idleSignal + ) + $ do + waitUntilTrue idleSignal currentTime >>= (logInfo_ . T.pack . ("current time: " ++) . show) -- Each job creates 2 new jobs, so there should be 1024 jobs in table. - runSQL_ $ "SELECT COUNT(*) from consumers_test_jobs" + runSQL_ "SELECT COUNT(*) from consumers_test_jobs" rowcount0 :: Int64 <- fetchOne runIdentity -- Move time 2 hours forward - modifyTestTime $ addUTCTime (2*60*60) - finalize (localDomain "process" $ - runConsumerWithIdleSignal consumerConfig connSource idleSignal) $ do - waitUntilTrue idleSignal + modifyTestTime $ addUTCTime (2 * 60 * 60) + finalize + ( localDomain "process" $ + runConsumerWithIdleSignal consumerConfig connSource idleSignal + ) + $ do + waitUntilTrue idleSignal -- Jobs are designed to double only 10 times, so there should be no jobs left now. - runSQL_ $ "SELECT COUNT(*) from consumers_test_jobs" + runSQL_ "SELECT COUNT(*) from consumers_test_jobs" rowcount1 :: Int64 <- fetchOne runIdentity liftIO $ T.assertEqual "Number of jobs in table after 10 steps is 1024" 1024 rowcount0 liftIO $ T.assertEqual "Number of jobs in table after 11 steps is 0" 0 rowcount1 dropTables - where - waitUntilTrue tmvar = whileM_ (not <$> (liftIO $ atomically $ takeTMVar tmvar)) $ return () + where + waitUntilTrue tmvar = liftIO . atomically $ do + takeTMVar tmvar >>= \case + True -> pure () + False -> retry - printUsage = do - prog <- getProgName - putStrLn $ "Usage: " <> prog <> " " + printUsage = do + prog <- getProgName + putStrLn $ "Usage: " <> prog <> " " - tables = [consumersTable, jobsTable] - -- NB: order of migrations is important. - migrations = [ createTableMigration consumersTable - , createTableMigration jobsTable ] + tables = [consumersTable, jobsTable] + -- NB: order of migrations is important. + migrations = + [ createTableMigration consumersTable + , createTableMigration jobsTable + ] - createTables :: TestEnv () - createTables = do - migrateDatabase defaultExtrasOptions - {- extensions -} [] {- composites -} [] {- domains -} [] - tables migrations - checkDatabase defaultExtrasOptions - {- composites -} [] {- domains -} [] - tables + createTables :: TestEnv () + createTables = do + migrateDatabase + defaultExtrasOptions + [] -- extensions + [] -- composites + [] -- domains + tables + migrations + checkDatabase + defaultExtrasOptions + [] -- composites + [] -- domains + tables - dropTables :: TestEnv () - dropTables = do - migrateDatabase defaultExtrasOptions - {- extensions -} [] {- composites -} [] {- domains -} [] {- tables -} [] - [ dropTableMigration jobsTable - , dropTableMigration consumersTable ] + dropTables :: TestEnv () + dropTables = do + migrateDatabase + defaultExtrasOptions + [] -- extensions + [] -- composites + [] -- domains + [] -- tables + [ dropTableMigration jobsTable + , dropTableMigration consumersTable + ] - consumerConfig = ConsumerConfig - { ccJobsTable = "consumers_test_jobs" - , ccConsumersTable = "consumers_test_consumers" - , ccJobSelectors = ["id", "countdown"] - , ccJobFetcher = id - , ccJobIndex = \(i::Int64, _::Int32) -> i + consumerConfig = + ConsumerConfig + { ccJobsTable = "consumers_test_jobs" + , ccConsumersTable = "consumers_test_consumers" + , ccJobSelectors = ["id", "countdown"] + , ccJobFetcher = id + , ccJobIndex = \(i :: Int64, _ :: Int32) -> i , ccNotificationChannel = Just "consumers_test_chan" - -- select some small timeout - , ccNotificationTimeout = 100 * 1000 -- 100 msec - , ccMaxRunningJobs = 20 - , ccProcessJob = processJob - , ccOnException = handleException + , -- select some small timeout + ccNotificationTimeout = 100 * 1000 -- 100 msec + , ccMaxRunningJobs = 20 + , ccProcessJob = processJob + , ccOnException = handleException } - putJob :: Int32 -> TestEnv () - putJob countdown = localDomain "put" $ do - now <- currentTime - runSQL_ $ "INSERT INTO consumers_test_jobs " + putJob :: Int32 -> TestEnv () + putJob countdown = localDomain "put" $ do + now <- currentTime + runSQL_ $ + "INSERT INTO consumers_test_jobs " <> "(run_at, finished_at, reserved_by, attempts, countdown) " - <> "VALUES (" now <> " + interval '1 hour', NULL, NULL, 0, " countdown <> ")" - notify "consumers_test_chan" "" - - processJob :: (Int64, Int32) -> TestEnv Result - processJob (_idx, countdown) = do - when (countdown > 0) $ do - putJob (countdown - 1) - putJob (countdown - 1) - commit - return (Ok Remove) + <> "VALUES (" now + <> " + interval '1 hour', NULL, NULL, 0, " countdown + <> ")" + notify "consumers_test_chan" "" - handleException :: SomeException -> (Int64, Int32) -> TestEnv Action - handleException exc (idx, _countdown) = do - logAttention_ $ - "Job #" <> showt idx <> " failed with: " <> showt exc - return . RerunAfter $ imicroseconds 500000 + processJob :: (Int64, Int32) -> TestEnv Result + processJob (_idx, countdown) = do + when (countdown > 0) $ do + putJob (countdown - 1) + putJob (countdown - 1) + commit + pure (Ok Remove) + handleException :: SomeException -> (Int64, Int32) -> TestEnv Action + handleException exc (idx, _countdown) = do + logAttention "Job failed" $ + object + [ "job_id" .= showt idx + , "exception" .= showt exc + ] + pure . RerunAfter $ imicroseconds 500000 jobsTable :: Table jobsTable = tblTable - { tblName = "consumers_test_jobs" - , tblVersion = 1 - , tblColumns = - [ tblColumn { colName = "id", colType = BigSerialT - , colNullable = False } - , tblColumn { colName = "run_at", colType = TimestampWithZoneT - , colNullable = True } - , tblColumn { colName = "finished_at", colType = TimestampWithZoneT - , colNullable = True } - , tblColumn { colName = "reserved_by", colType = BigIntT - , colNullable = True } - , tblColumn { colName = "attempts", colType = IntegerT - , colNullable = False } - - -- The only non-obligatory field: - , tblColumn { colName = "countdown", colType = IntegerT - , colNullable = False } - ] - , tblPrimaryKey = pkOnColumn "id" - , tblForeignKeys = [ - (fkOnColumn "reserved_by" "consumers_test_consumers" "id") { - fkOnDelete = ForeignKeySetNull - } - ] - } + { tblName = "consumers_test_jobs" + , tblVersion = 1 + , tblColumns = + [ tblColumn + { colName = "id" + , colType = BigSerialT + , colNullable = False + } + , tblColumn + { colName = "run_at" + , colType = TimestampWithZoneT + , colNullable = True + } + , tblColumn + { colName = "finished_at" + , colType = TimestampWithZoneT + , colNullable = True + } + , tblColumn + { colName = "reserved_by" + , colType = BigIntT + , colNullable = True + } + , tblColumn + { colName = "attempts" + , colType = IntegerT + , colNullable = False + } + , -- The only non-obligatory field: + tblColumn + { colName = "countdown" + , colType = IntegerT + , colNullable = False + } + ] + , tblPrimaryKey = pkOnColumn "id" + , tblForeignKeys = + [ (fkOnColumn "reserved_by" "consumers_test_consumers" "id") + { fkOnDelete = ForeignKeySetNull + } + ] + } consumersTable :: Table consumersTable = tblTable - { tblName = "consumers_test_consumers" - , tblVersion = 1 - , tblColumns = - [ tblColumn { colName = "id", colType = BigSerialT - , colNullable = False } - , tblColumn { colName = "name", colType = TextT - , colNullable = False } - , tblColumn { colName = "last_activity", colType = TimestampWithZoneT - , colNullable = False } - ] - , tblPrimaryKey = pkOnColumn "id" - } + { tblName = "consumers_test_consumers" + , tblVersion = 1 + , tblColumns = + [ tblColumn + { colName = "id" + , colType = BigSerialT + , colNullable = False + } + , tblColumn + { colName = "name" + , colType = TextT + , colNullable = False + } + , tblColumn + { colName = "last_activity" + , colType = TimestampWithZoneT + , colNullable = False + } + ] + , tblPrimaryKey = pkOnColumn "id" + } -createTableMigration :: (MonadDB m) => Table -> Migration m -createTableMigration tbl = Migration - { mgrTableName = tblName tbl - , mgrFrom = 0 - , mgrAction = StandardMigration $ do - createTable True tbl - } +createTableMigration :: MonadDB m => Table -> Migration m +createTableMigration tbl = + Migration + { mgrTableName = tblName tbl + , mgrFrom = 0 + , mgrAction = StandardMigration $ do + createTable True tbl + } dropTableMigration :: Table -> Migration m -dropTableMigration tbl = Migration - { mgrTableName = tblName tbl - , mgrFrom = 1 - , mgrAction = DropTableMigration DropTableRestrict - } +dropTableMigration tbl = + Migration + { mgrTableName = tblName tbl + , mgrFrom = 1 + , mgrAction = DropTableMigration DropTableRestrict + }