consumers 2.0.0.1 → 2.1.0.0

raw patch · 10 files changed

+580/−56 lines, 10 filesdep +HUnitdep +consumersdep +extradep ~basedep ~hpqtypesdep ~log-basePVP ok

version bump matches the API change (PVP)

Dependencies added: HUnit, consumers, extra, hpqtypes-extras, monad-loops, monad-time, text, text-show, transformers

Dependency ranges changed: base, hpqtypes, log-base

API changes (from Hackage documentation)

Files

• CHANGELOG.md +15/−0
• README.md +6/−0
• consumers.cabal +46/−2
• example/Example.hs +165/−0
• src/Database/PostgreSQL/Consumers.hs +1/−0
• src/Database/PostgreSQL/Consumers/Components.hs +98/−42
• src/Database/PostgreSQL/Consumers/Config.hs +15/−8
• src/Database/PostgreSQL/Consumers/Consumer.hs +5/−3
• src/Database/PostgreSQL/Consumers/Utils.hs +2/−1
• test/Test.hs +227/−0

@@ -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'.

@@ -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.

@@ -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 <andrzej@rybczak.net>,                     Jonathan Jouty <jonathan@scrive.com>@@ -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

@@ -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+  }

@@ -1,5 +1,6 @@ module Database.PostgreSQL.Consumers (     runConsumer+  , runConsumerWithIdleSignal   , module Database.PostgreSQL.Consumers.Config   , module Database.PostgreSQL.Consumers.Utils   ) where

@@ -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

@@ -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)

@@ -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

@@ -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

@@ -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+  }