diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,15 @@ +# consumers-2.1.0.0 (2017-09-18) + +* Added a `MonadTime` constraint to `runConsumer`. The `currentTime` + method from `MonadTime` is now used instead of `now()` in SQL. This + is mostly useful for testing, to control the flow of time in test + cases. + +* Added a new `runConsumer` variant: `runConsumerWithIdleSignal`. + When a consumer is idle, this variant signals a TMVar. This is + mostly useful for testing. + +* Added a test suite. + +# consumers-2.0.0.1 (2017-07-21) +* Now depends on 'log-base' instead of 'log'. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,6 @@ +# consumers [![Hackage version](https://img.shields.io/hackage/v/consumers.svg?label=Hackage)](https://hackage.haskell.org/package/consumers) [![Build Status](https://secure.travis-ci.org/scrive/consumers.svg?branch=master)](http://travis-ci.org/scrive/consumers) + +Library for setting up concurrent consumers of data stored inside +PostgreSQL database in a simple, declarative manner. + +See the `examples/` directory for a usage example. diff --git a/consumers.cabal b/consumers.cabal --- a/consumers.cabal +++ b/consumers.cabal @@ -1,5 +1,5 @@ name: consumers -version: 2.0.0.1 +version: 2.1.0.0 synopsis: Concurrent PostgreSQL data consumers description: Library for setting up concurrent consumers of data @@ -9,6 +9,7 @@ homepage: https://github.com/scrive/consumers license: BSD3 license-file: LICENSE +extra-source-files: CHANGELOG.md, README.md author: Scrive AB maintainer: Andrzej Rybczak , Jonathan Jouty @@ -16,7 +17,7 @@ category: Concurrency, Database build-type: Simple cabal-version: >=1.10 -tested-with: GHC == 7.8.4, GHC == 7.10.3, GHC == 8.0.1 +tested-with: GHC == 7.10.3, GHC == 8.0.2, GHC == 8.2.1 Source-repository head Type: git @@ -32,11 +33,13 @@ build-depends: base <5, containers, exceptions, + extra, hpqtypes >=1.5, lifted-base, lifted-threads, log-base >= 0.7, monad-control, + monad-time, mtl, stm, time, @@ -58,3 +61,44 @@ , TupleSections , 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 + build-depends: base, + consumers, + hpqtypes, + hpqtypes-extras, + log-base, + text, + text-show, + transformers + +test-suite consumers-test + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Test.hs + default-language: Haskell2010 + ghc-options: -Wall + build-depends: base, + consumers, + exceptions, + HUnit, + hpqtypes, + hpqtypes-extras, + log-base, + monad-control, + monad-loops, + monad-time, + mtl, + stm, + text, + text-show, + time, + transformers, + transformers-base diff --git a/example/Example.hs b/example/Example.hs new file mode 100644 --- /dev/null +++ b/example/Example.hs @@ -0,0 +1,165 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# 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 Log +import Log.Backend.StandardOutput +import System.Environment +import TextShow + +import qualified Data.Text as T + +type AppM a = DBT (LogT IO) a + +main :: IO () +main = do + args <- getArgs + let connString = case args of + [] -> defaultConnString + (cs:_) -> cs + + let connSettings = def { csConnInfo = T.pack connString } + ConnectionSource connSource = simpleSource connSettings + + withSimpleStdOutLogger $ \logger -> + runLogT "consumers-example" logger $ + runDBT connSource {- transactionSettings -} def $ do + createTables + finalize (localDomain "process" $ + runConsumer consumerConfig connSource) $ do + forM_ [(0::Int)..10] $ \_ -> do + putJob + liftIO $ threadDelay (1 * 1000000) -- 1 sec + dropTables + + where + + defaultConnString = + "postgresql://postgres@localhost/travis_ci_test" + + tables = [consumersTable, jobsTable] + -- NB: order of migrations is important. + migrations = [ createTableMigration consumersTable + , createTableMigration jobsTable ] + + createTables :: AppM () + createTables = do + migrateDatabase {- options -} [] {- extensions -} [] {- domains -} [] + tables migrations + checkDatabase {- domains -} [] tables + + dropTables :: AppM () + dropTables = do + migrateDatabase {- options -} [] {- extensions -} [] {- domains -} [] + {- tables -} [] + [ dropTableMigration jobsTable + , dropTableMigration consumersTable ] + + 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 + } + + 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')" + commit + + processJob :: (Int64, T.Text) -> AppM Result + processJob (_idx, msg) = do + logInfo_ msg + return (Ok Remove) + + handleException :: SomeException -> (Int64, T.Text) -> AppM Action + handleException exc (idx, _msg) = do + logAttention_ $ + "Job #" <> showt idx <> " failed with: " <> showt exc + return . RerunAfter $ imicroseconds 500000 + + +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 + } + ] + } + +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" + } + + +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 + } 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,6 @@ module Database.PostgreSQL.Consumers ( runConsumer + , runConsumerWithIdleSignal , module Database.PostgreSQL.Consumers.Config , module Database.PostgreSQL.Consumers.Utils ) where 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,5 +1,6 @@ module Database.PostgreSQL.Consumers.Components ( runConsumer + , runConsumerWithIdleSignal , spawnListener , spawnMonitor , spawnDispatcher @@ -12,10 +13,12 @@ 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.Function import Data.Int +import Data.Maybe import Data.Monoid import Data.Monoid.Utils import Database.PostgreSQL.PQTypes @@ -35,11 +38,33 @@ -- This function is best used in conjunction with 'finalize' to -- seamlessly handle the finalization. runConsumer - :: (MonadBaseControl IO m, MonadLog m, MonadMask m, Eq idx, Show idx, ToSQL idx) + :: ( MonadBaseControl IO m, MonadLog m, MonadMask m, Eq idx, Show idx + , ToSQL idx ) => ConsumerConfig m idx job -> ConnectionSourceM m -> m (m ()) -runConsumer cc cs = do +runConsumer cc cs = runConsumerWithMaybeIdleSignal cc cs Nothing + +runConsumerWithIdleSignal + :: ( MonadBaseControl IO m, MonadLog m, MonadMask m, Eq idx, Show idx + , ToSQL idx ) + => ConsumerConfig m idx job + -> ConnectionSourceM m + -> TMVar Bool + -> m (m ()) +runConsumerWithIdleSignal cc cs idleSignal = runConsumerWithMaybeIdleSignal cc cs (Just idleSignal) + +-- | Run the consumer and also signal whenever the consumer is waiting for +-- getNotification or threadDelay. + +runConsumerWithMaybeIdleSignal + :: ( MonadBaseControl IO m, MonadLog m, MonadMask m, Eq idx, Show idx + , ToSQL idx ) + => ConsumerConfig m idx job + -> ConnectionSourceM m + -> Maybe (TMVar Bool) + -> m (m ()) +runConsumerWithMaybeIdleSignal cc cs mIdleSignal = do semaphore <- newMVar () runningJobsInfo <- liftBase $ newTVarIO M.empty runningJobs <- liftBase $ newTVarIO 0 @@ -52,7 +77,8 @@ localData ["consumer_id" .= show cid] $ do listener <- spawnListener cc cs semaphore monitor <- localDomain "monitor" $ spawnMonitor cc cs cid - dispatcher <- localDomain "dispatcher" $ spawnDispatcher cc useSkipLocked cs cid semaphore runningJobsInfo runningJobs + dispatcher <- localDomain "dispatcher" $ spawnDispatcher cc useSkipLocked + cs cid semaphore runningJobsInfo runningJobs mIdleSignal return . localDomain "finalizer" $ do stopExecution listener stopExecution dispatcher @@ -91,17 +117,21 @@ -> ConnectionSourceM m -> MVar () -> 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 +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 where signalDispatcher = do liftBase $ tryPutMVar semaphore () @@ -120,10 +150,11 @@ -> m ThreadId 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()" + , "SET last_activity = " now , "WHERE id =" cid , " AND name =" unRawSQL ccJobsTable ] @@ -137,10 +168,11 @@ -- it was already marked as reserved by other consumer, so let's -- run it in serializable transaction. (inactiveConsumers, freedJobs) <- runDBT cs tsSerializable $ do + now <- currentTime -- Delete all inactive (assumed dead) consumers and get their ids. runSQL_ $ smconcat [ "DELETE FROM" <+> raw ccConsumersTable - , " WHERE last_activity +" iminutes 1 <+> "<= now()" + , " WHERE last_activity +" iminutes 1 <+> "<= " now , " AND name =" unRawSQL ccJobsTable , " RETURNING id::bigint" ] @@ -171,7 +203,8 @@ -- | Spawn a thread that reserves and processes jobs. spawnDispatcher - :: forall m idx job. (MonadBaseControl IO m, MonadLog m, MonadMask 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 -> Bool -> ConnectionSourceM m @@ -179,13 +212,25 @@ -> MVar () -> TVar (M.Map ThreadId idx) -> TVar Int + -> Maybe (TMVar Bool) -> m ThreadId -spawnDispatcher ConsumerConfig{..} useSkipLocked cs cid semaphore runningJobsInfo runningJobs = +spawnDispatcher ConsumerConfig{..} useSkipLocked cs cid semaphore + runningJobsInfo runningJobs mIdleSignal = forkP "dispatcher" . forever $ do void $ takeMVar semaphore - loop 1 + someJobWasProcessed <- loop 1 + if someJobWasProcessed + then setIdle False + else setIdle True where - loop :: Int -> m () + setIdle :: forall m' . (MonadBaseControl IO m') => Bool -> m' () + setIdle isIdle = case mIdleSignal of + Nothing -> return () + Just idleSignal -> atomically $ do + _ <- tryTakeTMVar idleSignal + putTMVar idleSignal isIdle + + loop :: Int -> m Bool loop limit = do (batch, batchSize) <- reserveJobs limit when (batchSize > 0) $ do @@ -201,7 +246,8 @@ atomically $ modifyTVar' runningJobs (+batchSize) let subtractJobs = atomically $ do modifyTVar' runningJobs (subtract batchSize) - void . forkP "batch processor" . (`finally` subtractJobs) . restore $ do + void . forkP "batch processor" + . (`finally` subtractJobs) . restore $ do mapM startJob batch >>= mapM joinJob >>= updateJobs when (batchSize == limit) $ do @@ -209,10 +255,13 @@ jobs <- readTVar runningJobs when (jobs >= ccMaxRunningJobs) retry return $ ccMaxRunningJobs - jobs - loop $ min maxBatchSize (2*limit) + void $ loop $ min maxBatchSize (2*limit) + return (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 @@ -220,30 +269,34 @@ , " WHEN finished_at IS NULL THEN attempts + 1" , " ELSE 1" , " END" - , "WHERE id IN (" <> reservedJobs <> ")" + , "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 where - reservedJobs :: SQL - reservedJobs = smconcat [ + reservedJobs :: UTCTime -> SQL + reservedJobs now = smconcat [ "SELECT id FROM" <+> raw ccJobsTable - -- Converting id to text and hashing it may seem silly, especially - -- when we're dealing with integers in the first place, but even in - -- such case the overhead is small enough (converting 100k integers - -- to text and hashing them takes around 15 ms on i7) to be worth - -- the generality. Note that even if IDs of two pending jobs produce - -- the same hash, it just means that in the worst case they will be - -- processed by the same consumer. + -- Converting id to text and hashing it may seem silly, + -- especially when we're dealing with integers in the + -- first place, but even in such case the overhead is + -- small enough (converting 100k integers to text and + -- hashing them takes around 15 ms on i7) to be worth the + -- generality. Note that even if IDs of two pending jobs + -- produce the same hash, it just means that in the worst + -- case they will be processed by the same consumer. , if useSkipLocked then "WHERE TRUE" - else "WHERE pg_try_advisory_xact_lock(" unRawSQL ccJobsTable <> "::regclass::integer, hashtext(id::text))" + else "WHERE pg_try_advisory_xact_lock(" + unRawSQL ccJobsTable + <> "::regclass::integer, hashtext(id::text))" , " AND reserved_by IS NULL" , " AND run_at IS NOT NULL" - , " AND run_at <= now()" + , " AND run_at <= " now , "LIMIT" limit - -- Use SKIP LOCKED if available. Otherwise utilize advisory locks. + -- Use SKIP LOCKED if available. Otherwise utilise + -- advisory locks. , if useSkipLocked then "FOR UPDATE SKIP LOCKED" else "FOR UPDATE" @@ -269,7 +322,8 @@ Right result -> return (ccJobIndex job, result) Left ex -> do action <- ccOnException ex job - logAttention "Unexpected exception caught while processing job" $ object [ + logAttention "Unexpected exception caught while processing job" $ + object [ "job_id" .= show (ccJobIndex job) , "exception" .= show ex , "action" .= show action @@ -279,6 +333,7 @@ -- | Update status of the jobs. updateJobs :: [(idx, Result)] -> m () updateJobs results = runDBT cs ts $ do + now <- currentTime runSQL_ $ smconcat [ "WITH removed AS (" , " DELETE FROM" <+> raw ccJobsTable @@ -288,19 +343,19 @@ , " reserved_by = NULL" , ", run_at = CASE" , " WHEN FALSE THEN run_at" - , smconcat $ M.foldrWithKey retryToSQL [] retries + , smconcat $ M.foldrWithKey (retryToSQL now) [] retries , " ELSE NULL" -- processed , " END" , ", finished_at = CASE" - , " WHEN id = ANY(" Array1 successes <+> ") THEN now()" + , " 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 = + retryToSQL now (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 f updates @@ -313,7 +368,8 @@ 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" + Remove -> error + "updateJobs: Remove should've been filtered out" successes = foldr step [] updates where 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 @@ -7,9 +7,14 @@ import Control.Exception (SomeException) import Data.Time -import Database.PostgreSQL.PQTypes import Prelude +import Database.PostgreSQL.PQTypes.FromRow +import Database.PostgreSQL.PQTypes.Interval +import Database.PostgreSQL.PQTypes.Notification +import Database.PostgreSQL.PQTypes.SQL +import Database.PostgreSQL.PQTypes.SQL.Raw + -- | Action to take after a job was processed. data Action = MarkProcessed @@ -24,19 +29,20 @@ -- | Config of a consumer. data ConsumerConfig m idx job = forall row. FromRow row => ConsumerConfig { --- | Name of a database table where jobs are stored. The table needs to have +-- | 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 the type +-- * 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 the type comparable with now() (TIMESTAMPTZ is recommended). +-- * 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. -- -- * finished_at - represents the time at which job processing was finished. --- Needs to be nullable, of the type you can assign now() to (TIMESTAMPTZ is +-- 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. -- @@ -70,7 +76,8 @@ -- 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. +-- | 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) @@ -94,7 +101,7 @@ , ccMaxRunningJobs :: !Int -- | Function that processes a job. , ccProcessJob :: !(job -> m Result) --- | Action taken if job processing function throws an exception. +-- | 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) 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 @@ -7,6 +7,7 @@ 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 @@ -27,18 +28,19 @@ 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. registerConsumer - :: (MonadBase IO m, MonadMask m) + :: (MonadBase IO m, MonadMask m, MonadTime m) => ConsumerConfig n idx job -> ConnectionSourceM m -> m ConsumerID registerConsumer ConsumerConfig{..} cs = runDBT cs ts $ do + now <- currentTime runSQL_ $ smconcat [ "INSERT INTO" <+> raw ccConsumersTable - , "(name, last_activity) VALUES (" unRawSQL ccJobsTable <> ", now())" + , "(name, last_activity) VALUES (" unRawSQL ccJobsTable <> ", " now <> ")" , "RETURNING id" ] fetchOne runIdentity 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 @@ -16,7 +16,8 @@ import qualified Control.Concurrent.Thread.Lifted as T import qualified Control.Exception.Lifted as E --- | Runs an action that returns a finalizer and performs it at the end. +-- | 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 diff --git a/test/Test.hs b/test/Test.hs new file mode 100644 --- /dev/null +++ b/test/Test.hs @@ -0,0 +1,227 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleContexts #-} +{-# 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.Time +import Log +import Log.Backend.StandardOutput +import System.Environment +import TextShow + +import qualified Data.Text as T +import qualified Test.HUnit as T + +data TestEnvSt = TestEnvSt { + teCurrentTime :: UTCTime + } + +type InnerTestEnv = StateT TestEnvSt (DBT (LogT IO)) + +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 #-} + +instance MonadTime TestEnv where + currentTime = gets teCurrentTime + +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) + . (runDBT connSource {- transactionSettings -} def) + . (\m' -> fst <$> (runStateT m' $ TestEnvSt $ UTCTime (ModifiedJulianDay 0) 0)) + . unTestEnv + $ m + +main :: IO () +main = void $ T.runTestTT $ T.TestCase test + +test :: IO () +test = do + args <- getArgs + let connString = case args of + [] -> defaultConnString + (cs:_) -> cs + + let connSettings = def { csConnInfo = T.pack connString } + ConnectionSource connSource = simpleSource connSettings + + withSimpleStdOutLogger $ \logger -> + runTestEnv connSource logger $ do + createTables + idleSignal <- liftIO $ atomically $ newEmptyTMVar + putJob 10 >> commit + + 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 + 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" + rowcount0 :: Int64 <- fetchOne runIdentity + -- Move time 2 hours forward + 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" + 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 () + defaultConnString = + "postgresql://postgres@localhost/travis_ci_test" + + tables = [consumersTable, jobsTable] + -- NB: order of migrations is important. + migrations = [ createTableMigration consumersTable + , createTableMigration jobsTable ] + + createTables :: TestEnv () + createTables = do + migrateDatabase {- options -} [] {- extensions -} [] {- domains -} [] + tables migrations + checkDatabase {- domains -} [] tables + + dropTables :: TestEnv () + dropTables = do + migrateDatabase {- options -} [] {- extensions -} [] {- 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 + , ccNotificationChannel = Just "consumers_test_chan" + -- 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 " + <> "(run_at, finished_at, reserved_by, attempts, countdown) " + <> "VALUES (" now <> " + interval '1 hour', NULL, NULL, 0, " countdown <> ")" + + processJob :: (Int64, Int32) -> TestEnv Result + processJob (_idx, countdown) = do + when (countdown > 0) $ do + putJob (countdown - 1) + putJob (countdown - 1) + commit + return (Ok Remove) + + handleException :: SomeException -> (Int64, Int32) -> TestEnv Action + handleException exc (idx, _countdown) = do + logAttention_ $ + "Job #" <> showt idx <> " failed with: " <> showt exc + return . 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 + } + ] + } + +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" + } + +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 + }