More tests, refactorings

bin/simple-worker/Main.hs

@@ -14,6 +14,7 @@ import Async.Worker.Types (State(..), PerformAction, getJob, formatStr, TimeoutS
 import Control.Applicative ((<|>))
 import Control.Concurrent (forkIO, threadDelay)
 import Control.Exception (Exception, throwIO)
+import Control.Monad (void)
 import Data.Aeson (FromJSON(..), ToJSON(..), object, (.=), (.:), withObject, withText)
 import Data.Text qualified as T
 import Database.PostgreSQL.Simple qualified as PSQL
@@ -75,7 +76,7 @@ main :: IO ()
 main = do
   let connInfo = PSQL.defaultConnectInfo { PSQL.connectUser = "postgres"
                                          , PSQL.connectDatabase = "postgres" }
-  let brokerInitParams = PGMQBrokerInitParams connInfo :: BrokerInitParams PGMQBroker (Job Message)
+  let brokerInitParams = PGMQBrokerInitParams connInfo 10 :: BrokerInitParams PGMQBroker (Job Message)
 
   let queue = "simple_worker"
   
@@ -108,17 +109,17 @@ main = do
   let mkJob msg = mkDefaultSendJob' broker queue msg
     
   mapM_ (\idx -> do
-    sendJob' $ mkJob $ Ping
-    sendJob' $ mkJob $ Wait 1
-    sendJob' $ mkJob $ Echo $ "hello " <> show idx
-    sendJob' $ mkJob $ Error $ "error " <> show idx
+    void $ sendJob' $ mkJob $ Ping
+    void $ sendJob' $ mkJob $ Wait 1
+    void $ sendJob' $ mkJob $ Echo $ "hello " <> show idx
+    void $ sendJob' $ mkJob $ Error $ "error " <> show idx
     ) tasksLst
 
   -- a job that will timeout
   let timedOut =
         (mkDefaultSendJob broker queue (Wait 5) 1)
           { toStrat = TSRepeatNElseArchive 3 }
-  sendJob' timedOut
+  void $ sendJob' timedOut
     
   threadDelay (10*second)
 

cabal.project

@@ -8,6 +8,6 @@ packages:
 source-repository-package
     type: git
     location: https://gitlab.iscpif.fr/gargantext/haskell-pgmq
-    tag: 268398735f61008af099918a24b3fb57f9533ba3
+    tag: fcb7d4fb811e5b7239078b48268c469c8d28fdf9
     
 tests: true

haskell-bee.cabal

@@ -83,9 +83,11 @@ library
                     , haskell-pgmq >= 0.1.0.0 && < 0.2
                     , hedis >= 0.15.2 && < 0.16
                     , mtl >= 2.2 && < 2.4
+                    , postgresql-libpq >= 0.10 && < 0.11
                     , postgresql-simple >= 0.6 && < 0.8
                     , safe >= 0.3 && < 0.4
                     , safe-exceptions >= 0.1.7 && < 0.2
+                    , stm >= 2.5.3 && < 3
                     , text >= 1.2 && < 2.2
                     , time >= 1.10 && < 1.15
                     , units >= 2.4 && < 2.5
@@ -204,4 +206,4 @@ test-suite test-integration
       OverloadedStrings
       RecordWildCards
 
-    ghc-options: -threaded
+    ghc-options: -threaded -fprof-auto

src/Async/Worker/Broker/PGMQ.hs

@@ -19,7 +19,11 @@ module Async.Worker.Broker.PGMQ
 where
 
 import Async.Worker.Broker.Types (HasBroker(..), SerializableMessage)
+import Control.Concurrent (threadDelay)
+import Control.Concurrent.MVar (withMVar)
+import Database.PostgreSQL.LibPQ qualified as LibPQ
 import Database.PostgreSQL.Simple qualified as PSQL
+import Database.PostgreSQL.Simple.Internal qualified as PSQLInternal
 import Database.PGMQ.Simple qualified as PGMQ
 import Database.PGMQ.Types qualified as PGMQ
 
@@ -30,22 +34,30 @@ instance (SerializableMessage a, Show a) => HasBroker PGMQBroker a where
   data Broker PGMQBroker a =
     PGMQBroker' {
         conn :: PSQL.Connection
+      , defaultVt :: PGMQ.VisibilityTimeout
       }
   data BrokerMessage PGMQBroker a = PGMQBM (PGMQ.Message a)
     deriving (Show)
   data Message PGMQBroker a = PGMQM a
   data MessageId PGMQBroker = PGMQMid Int
     deriving (Eq, Show)
-  data BrokerInitParams PGMQBroker a = PGMQBrokerInitParams PSQL.ConnectInfo
+  data BrokerInitParams PGMQBroker a = PGMQBrokerInitParams PSQL.ConnectInfo PGMQ.VisibilityTimeout
 
   messageId (PGMQBM (PGMQ.Message { msgId })) = PGMQMid msgId
   getMessage (PGMQBM (PGMQ.Message { message })) = PGMQM message
   toMessage message = PGMQM message
   toA (PGMQM message) = message
-  initBroker (PGMQBrokerInitParams connInfo) = do
+  initBroker (PGMQBrokerInitParams connInfo defaultVt) = do
     conn <- PSQL.connect connInfo
+    -- PGMQ is quite verbose because of initialization. We can disable
+    -- notices
+    -- https://hackage.haskell.org/package/postgresql-simple-0.7.0.0/docs/src/Database.PostgreSQL.Simple.Internal.html#Connection
+    -- https://hackage.haskell.org/package/postgresql-libpq-0.10.1.0/docs/Database-PostgreSQL-LibPQ.html#g:13
+    -- https://www.postgresql.org/docs/current/libpq-notice-processing.html
+    withMVar (PSQLInternal.connectionHandle conn) $ \c -> do
+      LibPQ.disableNoticeReporting c
     PGMQ.initialize conn
-    pure $ PGMQBroker' { conn }
+    pure $ PGMQBroker' { conn, defaultVt }
   deinitBroker (PGMQBroker' { conn }) = PSQL.close conn
   
   createQueue (PGMQBroker' { conn }) queue = do
@@ -54,17 +66,45 @@ instance (SerializableMessage a, Show a) => HasBroker PGMQBroker a where
   dropQueue (PGMQBroker' { conn }) queue = do
     PGMQ.dropQueue conn queue
 
-  readMessageWaiting q@(PGMQBroker' { conn }) queue = loop
+  readMessageWaiting q@(PGMQBroker' { conn, defaultVt }) queue = loop
     where
       -- loop :: PGMQ.SerializableMessage a => IO (BrokerMessage PGMQBroker' a)
       loop = do
-        mMsg <- PGMQ.readMessageWithPoll conn queue 10 5 100
+        -- NOTE readMessageWithPoll is not thread-safe, i.e. the
+        -- blocking is outside of GHC (in PostgreSQL itself) and we
+        -- can't reliably use it in a highly concurrent situation.
+        
+        -- mMsg <- PGMQ.readMessageWithPoll conn queue 10 5 100
+        mMsg <- PGMQ.readMessage conn queue defaultVt
         case mMsg of
-          Nothing -> readMessageWaiting q queue
-          Just msg -> return $ PGMQBM msg
+          Nothing -> do
+            -- wait a bit, then retry
+            threadDelay (50 * 1000)
+            readMessageWaiting q queue
+          Just msg -> do
+            -- TODO! we want to set message visibility timeout so that other workers don't start this job
+            return $ PGMQBM msg
+
+  popMessageWaiting q@(PGMQBroker' { conn }) queue = loop
+    where
+      -- loop :: PGMQ.SerializableMessage a => IO (BrokerMessage PGMQBroker' a)
+      loop = do
+        -- mMsg <- PGMQ.readMessageWithPoll conn queue 10 5 100
+        mMsg <- PGMQ.popMessage conn queue
+        case mMsg of
+          Nothing -> do
+            -- wait a bit, then retry
+            threadDelay (50 * 1000)
+            popMessageWaiting q queue
+          Just msg -> do
+            -- TODO! we want to set message visibility timeout so that other workers don't start this job
+            return $ PGMQBM msg
+
+  setMessageTimeout (PGMQBroker' { conn }) queue (PGMQMid msgId) timeoutS =
+    PGMQ.setMessageVt conn queue msgId timeoutS
 
   sendMessage (PGMQBroker' { conn }) queue (PGMQM message) =
-    PGMQ.sendMessage conn queue message 0
+    PGMQMid <$> PGMQ.sendMessage conn queue message 0
 
   deleteMessage (PGMQBroker' { conn }) queue (PGMQMid msgId) = do
     PGMQ.deleteMessage conn queue msgId
@@ -73,10 +113,14 @@ instance (SerializableMessage a, Show a) => HasBroker PGMQBroker a where
     PGMQ.archiveMessage conn queue msgId
 
   getQueueSize (PGMQBroker' { conn }) queue = do
-    mMetrics <- PGMQ.getMetrics conn queue
-    case mMetrics of
-      Nothing -> return 0
-      Just (PGMQ.Metrics { queueLength }) -> return queueLength
+    -- NOTE: pgmq.metrics is NOT a proper way to deal with messages
+    -- that have vt in the future
+    -- (c.f. https://github.com/tembo-io/pgmq/issues/301)
+    -- mMetrics <- PGMQ.getMetrics conn queue
+    -- case mMetrics of
+    --   Nothing -> return 0
+    --   Just (PGMQ.Metrics { queueLength }) -> return queueLength
+    PGMQ.queueAvailableLength conn queue
 
   getArchivedMessage (PGMQBroker' { conn }) queue (PGMQMid msgId) = do
     mMsg <- PGMQ.readMessageFromArchive conn queue msgId

src/Async/Worker/Broker/Redis.hs

@@ -12,10 +12,13 @@ https://redis.io/glossary/redis-queue/
 
 The design is as follows:
 - for each queue we have an 'id counter'
-- each queue is represented as a list of message ids
+- each queue is represented as a set of message ids
 - each message is stored under unique key, derived from its id
 - the above allows us to have an archive with messages
 - deleting a message means removing it's unique key from Redis
+
+The queue itself is a list, the archive is a set (so that we can use
+SISMEMBER).
 -}
 
 
@@ -31,7 +34,7 @@ module Async.Worker.Broker.Redis
 where
 
 import Async.Worker.Broker.Types (HasBroker(..), Queue, SerializableMessage)
-import Control.Concurrent (threadDelay)
+-- import Control.Concurrent (threadDelay)
 import Control.Monad (void)
 import Data.Aeson qualified as Aeson
 import Data.Aeson (FromJSON(..), ToJSON(..), (.:), (.=), withObject, object)
@@ -55,8 +58,6 @@ instance (SerializableMessage a, Show a) => HasBroker RedisBroker a where
     deriving (Eq, Show)
   data BrokerInitParams RedisBroker a = RedisBrokerInitParams Redis.ConnectInfo
 
-  -- We're using simple QUEUE so we don't care about message id as we
-  -- won't be deleting/archiving the messages
   messageId (RedisBM (RedisWithMsgId { rmidId })) = RedisMid rmidId
   getMessage (RedisBM (RedisWithMsgId { rmida })) = RedisM rmida
   toMessage message = RedisM message
@@ -76,23 +77,45 @@ instance (SerializableMessage a, Show a) => HasBroker RedisBroker a where
     let queueK = queueKey queue
     void $ Redis.runRedis conn $ Redis.del [queueK]
 
-  readMessageWaiting b@(RedisBroker' { conn }) queue = loop
+  -- TODO This is simplified
+  readMessageWaiting = popMessageWaiting
+
+  popMessageWaiting b@(RedisBroker' { conn }) queue = loop
     where
       queueK = queueKey queue
       loop = do
-        eMsgId <- Redis.runRedis conn $ Redis.spop queueK
-        case eMsgId of
+        -- 0 means block indefinitely
+        -- https://redis.io/docs/latest/commands/blpop/
+        eData <- Redis.runRedis conn $ Redis.blpop [queueK] 0
+        case eData of
           Left _ -> undefined
-          Right Nothing -> do
-            threadDelay 100
-            readMessageWaiting b queue
-          Right (Just msgIdBS) -> case bsToId msgIdBS of
+          Right Nothing -> undefined
+          Right (Just (_queueK, msgIdBS)) -> case bsToId msgIdBS of
             Nothing -> undefined
             Just msgId -> do
               mMsg <- getRedisMessage b queue msgId
-              case mMsg of
-                Nothing -> undefined
-                Just msg -> return msg
+              maybe undefined return mMsg
+
+  -- popMessageWaiting b@(RedisBroker' { conn }) queue = loop
+  --   where
+  --     queueK = queueKey queue
+  --     loop = do
+  --       eMsgId <- Redis.runRedis conn $ Redis.spop queueK
+  --       case eMsgId of
+  --         Left _ -> undefined
+  --         Right Nothing -> do
+  --           threadDelay (10*1000)
+  --           popMessageWaiting b queue
+  --         Right (Just msgIdBS) -> case bsToId msgIdBS of
+  --           Nothing -> undefined
+  --           Just msgId -> do
+  --             mMsg <- getRedisMessage b queue msgId
+  --             case mMsg of
+  --               Nothing -> undefined
+  --               Just msg -> return msg
+                
+  setMessageTimeout _broker _queue _msgId _timeoutS =
+    pure ()
 
   sendMessage b@(RedisBroker' { conn }) queue (RedisM message) = do
     mId <- nextId b queue
@@ -104,34 +127,44 @@ instance (SerializableMessage a, Show a) => HasBroker RedisBroker a where
         let msgK = messageKey queue msgId
         let queueK = queueKey queue
         void $ Redis.runRedis conn $ do
+          -- write the message itself under unique key
           _ <- Redis.set msgK (BSL.toStrict $ Aeson.encode m)
-          Redis.sadd queueK [idToBS msgId]
+          -- add message id to the list
+          -- Redis.sadd queueK [idToBS msgId]
+          Redis.lpush queueK [idToBS msgId]
+        return msgId
 
   -- deleteMessage (RedisBroker' { conn }) queue (RedisMid msgId) = do
   deleteMessage (RedisBroker' { conn }) queue msgId = do
     let queueK = queueKey queue
-    void $ Redis.runRedis conn $ Redis.srem queueK [idToBS msgId]
     let messageK = messageKey queue msgId
-    void $ Redis.runRedis conn $ Redis.del [messageK]
+    -- void $ Redis.runRedis conn $ Redis.srem queueK [idToBS msgId]
+    void $ Redis.runRedis conn $ do
+      _ <- Redis.lrem queueK 1 (idToBS msgId)
+      Redis.del [messageK]
 
   -- archiveMessage (RedisBroker' { conn }) queue (RedisMid msgId) = do
   archiveMessage (RedisBroker' { conn }) queue msgId = do
     let queueK = queueKey queue
     let archiveK = archiveKey queue
-    eMove <- Redis.runRedis conn $ Redis.smove queueK archiveK (idToBS msgId)
-    case eMove of
-      Left _ -> undefined
-      Right True -> return ()
-      Right False -> do
-        -- OK so the queue might not have the id, we just add it to archive to make sure
-        void $ Redis.runRedis conn $ Redis.sadd archiveK [idToBS msgId]
+    void $ Redis.runRedis conn $ do
+      _ <- Redis.lrem queueK 1 (idToBS msgId)
+      Redis.sadd archiveK [idToBS msgId]
+    -- eMove <- Redis.runRedis conn $ Redis.smove queueK archiveK (idToBS msgId)
+    -- case eMove of
+    --   Left _ -> undefined
+    --   Right True -> return ()
+    --   Right False -> do
+    --     -- OK so the queue might not have the id, we just add it to archive to make sure
+    --     void $ Redis.runRedis conn $ Redis.sadd archiveK [idToBS msgId]
 
   getQueueSize (RedisBroker' { conn }) queue = do
     let queueK = queueKey queue
-    eLen <- Redis.runRedis conn $ Redis.scard queueK
+    -- eLen <- Redis.runRedis conn $ Redis.scard queueK
+    eLen <- Redis.runRedis conn $ Redis.llen queueK
     case eLen of
       Right len -> return $ fromIntegral len
-      Left _ -> return 0
+      Left _ -> undefined
 
   getArchivedMessage b@(RedisBroker' { conn }) queue msgId = do
     let archiveK = archiveKey queue
@@ -169,7 +202,7 @@ nextId (RedisBroker' { conn }) queue = do
 
 -- | Key under which a message is stored
 messageKey :: Queue -> MessageId RedisBroker -> BS.ByteString
-messageKey queue (RedisMid msgId) = BS.pack $ beePrefix <> "queue-" <> queue <> "-message-" <> show msgId
+messageKey queue (RedisMid msgId) = queueKey queue <> BS.pack ("-message-" <> show msgId)
 
 getRedisMessage :: FromJSON a
                => Broker RedisBroker a

src/Async/Worker/Broker/Types.hs

@@ -19,6 +19,7 @@ Broker typeclass definition.
     
 module Async.Worker.Broker.Types
   ( Queue
+  , TimeoutS
   -- * Main broker typeclass
   -- $broker
   , HasBroker(..)
@@ -32,6 +33,7 @@ import Data.Typeable (Typeable)
 
 
 type Queue = String
+type TimeoutS = Int  -- timeout for a message, in seconds
 
 
 {- $broker
@@ -117,11 +119,23 @@ class (
   {-| Drop queue -}
   dropQueue :: Broker b a -> Queue -> IO ()
 
-  {-| Read message, waiting for it if not present -}
+  {-| Read message from queue, waiting for it if not present (this leaves
+  the message in queue, you need to use 'setMessageTimeout' to prevent
+  other workers from seeing this message). -}
   readMessageWaiting :: Broker b a -> Queue -> IO (BrokerMessage b a)
 
+  {-| Pop message from queue, waiting for it if not present -}
+  popMessageWaiting :: Broker b a -> Queue -> IO (BrokerMessage b a)
+
+  {-| We sometimes need a way to tell the broker that a message shouldn't
+  be visible for given amount of time (e.g. 'visibility timeout'
+  setting in PGMQ). The broker operates only on 'a' level and isn't
+  aware of 'Job' with its 'JobMetadata'. Hence, it's the worker's
+  responsibility to properly set timeout after message is read. -}
+  setMessageTimeout :: Broker b a -> Queue -> MessageId b -> TimeoutS -> IO ()
+
   {-| Send message -}
-  sendMessage :: Broker b a -> Queue -> Message b a -> IO ()
+  sendMessage :: Broker b a -> Queue -> Message b a -> IO (MessageId b)
 
   {-| Delete message -}
   deleteMessage :: Broker b a -> Queue -> MessageId b -> IO ()

src/Async/Worker/Types.hs

@@ -41,13 +41,12 @@ module Async.Worker.Types
   -- * Other useful types and functions
   , HasWorkerBroker
   , formatStr
-  , JobTimeout(..)
-  , JobException(..) )
+  , JobTimeout(..) )
 where
 
 import Async.Worker.Broker.Types (Broker, BrokerMessage, HasBroker, Queue)
 import Control.Applicative ((<|>))
-import Control.Exception.Safe (Exception, SomeException)
+import Control.Exception.Safe (Exception)
 import Data.Aeson (FromJSON(..), ToJSON(..), object, (.=), (.:), withObject, withText)
 import Data.Text qualified as T
 import Data.Typeable (Typeable)
@@ -156,7 +155,12 @@ data JobMetadata =
               , timeout          :: Timeout
               -- | Read count so we know how many times this message
               -- was processed
-              , readCount        :: ReadCount }
+              , readCount        :: ReadCount
+              -- | A worker might have processed a task and be
+              -- killed. If 'resendWhenWorkerKilled' is 'True', this
+              -- job will be resent to broker and picked up
+              -- later. Otherwise it will be discarded.
+              , resendWhenWorkerKilled :: Bool }
   deriving (Eq, Show)
 instance ToJSON JobMetadata where
   toJSON (JobMetadata { .. }) =
@@ -166,6 +170,7 @@ instance ToJSON JobMetadata where
         , "tstrat" .= timeoutStrategy
         , "timeout" .= timeout
         , "readCount" .= readCount
+        , "resendWhenWorkerKilled" .= resendWhenWorkerKilled
         ]
 instance FromJSON JobMetadata where
   parseJSON = withObject "JobMetadata" $ \o -> do
@@ -174,6 +179,7 @@ instance FromJSON JobMetadata where
     timeoutStrategy <- o .: "tstrat"
     timeout <- o .: "timeout"
     readCount <- o .: "readCount"
+    resendWhenWorkerKilled <- o .: "resendWhenWorkerKilled"
     return $ JobMetadata { .. }
 
 -- | For a typical 'Job' it's probably sane to just archive it no
@@ -184,7 +190,8 @@ defaultMetadata =
               , errorStrategy = ESArchive
               , timeoutStrategy = TSArchive
               , timeout = 10
-              , readCount = 0 }
+              , readCount = 0
+              , resendWhenWorkerKilled = True }
     
 -- | Worker 'Job' is 'a' (defining action to call via 'performAction')
 -- together with associated 'JobMetadata'.
@@ -277,10 +284,3 @@ data JobTimeout b a =
              , jtTimeout   :: Timeout }
 deriving instance (HasWorkerBroker b a) => Show (JobTimeout b a)
 instance (HasWorkerBroker b a) => Exception (JobTimeout b a)
-
--- | An exception, thrown when job ends with error
-data JobException b a =
-  JobException { jeBMessage :: BrokerMessage b (Job a)
-               , jeException :: SomeException }
-deriving instance (HasWorkerBroker b a) => Show (JobException b a)
-instance (HasWorkerBroker b a) => Exception (JobException b a)

tests/Test/Integration/Broker.hs

@@ -20,7 +20,8 @@ import Data.Aeson (ToJSON(..), FromJSON(..), withText)
 import Data.Maybe (isJust)
 import Data.Text qualified as T
 import Test.Hspec
-import Test.Integration.Utils (getPSQLEnvConnectInfo, getRedisEnvConnectInfo, randomQueueName, waitUntil)
+import Test.Integration.Utils (defaultPGMQVt, getPSQLEnvConnectInfo, getRedisEnvConnectInfo, randomQueueName, waitUntil)
+import Test.RandomStrings (randomASCII, randomString, onlyAlphaNum)
 
 
 data TestEnv b =
@@ -72,28 +73,56 @@ brokerTests bInitParams =
   parallel $ around (withBroker bInitParams) $ describe "Broker tests" $ do
     it "can send and receive a message" $ \(TestEnv { broker, queue }) -> do
       let msg = Message { text = "test" }
-      BT.sendMessage broker queue (BT.toMessage msg)
+      msgId <- BT.sendMessage broker queue (BT.toMessage msg)
       msg2 <- BT.readMessageWaiting broker queue
       -- putStrLn $ "[messageId] " <> show (BT.messageId msg2)
       msg `shouldBe` BT.toA (BT.getMessage msg2)
+      msgId `shouldBe` BT.messageId msg2
 
     it "can send, archive and read message from archive" $ \(TestEnv { broker, queue }) -> do
       let msg = Message { text = "test" }
-      BT.sendMessage broker queue (BT.toMessage msg)
+      msgId <- BT.sendMessage broker queue (BT.toMessage msg)
       msg2 <- BT.readMessageWaiting broker queue
-      let msgId = BT.messageId msg2
+      msgId `shouldBe` BT.messageId msg2
       BT.archiveMessage broker queue msgId
-      putStrLn $ "Reading msg " <> show msgId <> " from archive queue " <> queue
       -- It might take some time to archive a message so we wait a bit
       waitUntil (isJust <$> BT.getArchivedMessage broker queue msgId) 200
       msgArchive <- BT.getArchivedMessage broker queue msgId
       let msgIdArchive = BT.messageId <$> msgArchive
       msgIdArchive `shouldBe` Just msgId
-      
+
+    it "returns correct message id when sending message to broker" $ \(TestEnv { broker, queue }) -> do
+      let iter = [1..20] :: [Int]  -- number of steps
+      mapM_ (\_i -> do
+                -- Generate random strings and make sure that the
+                -- message ids we get from sendMessage match our data
+                text <- randomString (onlyAlphaNum randomASCII) 20
+                let msg = Message { text }
+                msgId <- BT.sendMessage broker queue (BT.toMessage msg)
+                bMsg <- BT.readMessageWaiting broker queue
+                msg `shouldBe` BT.toA (BT.getMessage bMsg)
+                msgId `shouldBe` BT.messageId bMsg
+                BT.deleteMessage broker queue msgId
+                ) iter
+
+    it "preserves msgId when archiving a message" $ \(TestEnv { broker, queue }) -> do
+      let iter = [1..20] :: [Int]  -- number of steps
+      mapM_ (\_i -> do
+                -- Generate random strings and make sure that the
+                -- message ids we get from sendMessage match our data
+                text <- randomString (onlyAlphaNum randomASCII) 20
+                let msg = Message { text }
+                msgId <- BT.sendMessage broker queue (BT.toMessage msg)
+                BT.archiveMessage broker queue msgId
+                msgArchive <- BT.getArchivedMessage broker queue msgId
+                Just msg `shouldBe` (BT.toA . BT.getMessage <$> msgArchive)
+                ) iter
+
 
 pgmqBrokerInitParams :: IO (BT.BrokerInitParams PGMQ.PGMQBroker Message)
 pgmqBrokerInitParams = do
-  PGMQ.PGMQBrokerInitParams <$> getPSQLEnvConnectInfo
+  conn <- getPSQLEnvConnectInfo
+  return $ PGMQ.PGMQBrokerInitParams conn defaultPGMQVt
 
 redisBrokerInitParams :: IO (BT.BrokerInitParams Redis.RedisBroker Message)
 redisBrokerInitParams = do

tests/Test/Integration/Utils.hs

 module Test.Integration.Utils
-  ( getPSQLEnvConnectInfo
+  ( defaultPGMQVt
+  , getPSQLEnvConnectInfo
   , getRedisEnvConnectInfo
   , randomQueueName
   , waitUntil
   , waitUntilTVarEq
-  , waitUntilTVarPred )
+  , waitUntilTVarPred
+  , waitUntilQueueSizeIs
+  , waitUntilQueueEmpty )
 where
 
 import Async.Worker.Broker qualified as B
@@ -20,6 +23,26 @@ import Test.Hspec (expectationFailure, shouldBe, shouldSatisfy, Expectation, Has
 import Test.RandomStrings (randomASCII, randomString, onlyLower)
 
 
+
+-- | Timeout for 'wait' jobs, in ms.
+newtype TimeoutMs = TimeoutMs Int
+  deriving (Eq, Show, Num, Integral, Real, Enum, Ord)
+
+
+-- | Visibility timeout is a very important parameter for PGMQ. It is
+-- mainly used when reading a job: it specifies for how many seconds
+-- this job should be invisible for other workers. We need more tests
+-- and setting this correctly, preferably in accordance with
+-- 'Job.timeout'. Issue is that at the broker level we don't know
+-- anything about 'Job'...
+--
+-- The lower the value, the more probable that some other worker will
+-- pick up the same job at about the same time (before broker marks it
+-- as invisible).
+defaultPGMQVt :: Int
+defaultPGMQVt = 1
+
+
 -- | PSQL connect info that is fetched from env
 getPSQLEnvConnectInfo :: IO PSQL.ConnectInfo
 getPSQLEnvConnectInfo = do
@@ -46,10 +69,14 @@ randomQueueName prefix = do
   postfix <- randomString (onlyLower randomASCII) 10
   return $ prefix <> "_" <> postfix
 
+
+waitThreadDelay :: Int
+waitThreadDelay = 50 * 1000
+
 -- | Given a predicate IO action, test it for given number of
 -- milliseconds or fail
-waitUntil :: HasCallStack => IO Bool -> Int -> Expectation
-waitUntil pred' timeoutMs = do
+waitUntil :: HasCallStack => IO Bool -> TimeoutMs -> Expectation
+waitUntil pred' (TimeoutMs timeoutMs) = do
   _mTimeout <- Timeout.timeout (timeoutMs * 1000) performTest
   -- shortcut for testing mTimeout
   p <- pred'
@@ -61,12 +88,12 @@ waitUntil pred' timeoutMs = do
       if p
         then return ()
         else do
-          threadDelay 50
+          threadDelay waitThreadDelay
           performTest
 
 -- | Similar to 'waitUntil' but specialized to 'TVar' equality checking
-waitUntilTVarEq :: (HasCallStack, Show a, Eq a) => TVar a -> a -> Int -> Expectation
-waitUntilTVarEq tvar expected timeoutMs = do
+waitUntilTVarEq :: (HasCallStack, Show a, Eq a) => TVar a -> a -> TimeoutMs -> Expectation
+waitUntilTVarEq tvar expected (TimeoutMs timeoutMs) = do
   _mTimeout <- Timeout.timeout (timeoutMs * 1000) performTest
   -- shortcut for testing mTimeout
   val <- readTVarIO tvar
@@ -78,12 +105,12 @@ waitUntilTVarEq tvar expected timeoutMs = do
       if val == expected
         then return ()
         else do
-          threadDelay 50
+          threadDelay waitThreadDelay
           performTest
 
 -- | Similar to 'waitUntilTVarEq' but with predicate checking
-waitUntilTVarPred :: (HasCallStack, Show a, Eq a) => TVar a -> (a -> Bool) -> Int -> Expectation
-waitUntilTVarPred tvar predicate timeoutMs = do
+waitUntilTVarPred :: (HasCallStack, Show a, Eq a) => TVar a -> (a -> Bool) -> TimeoutMs -> Expectation
+waitUntilTVarPred tvar predicate (TimeoutMs timeoutMs) = do
   _mTimeout <- Timeout.timeout (timeoutMs * 1000) performTest
   -- shortcut for testing mTimeout
   val <- readTVarIO tvar
@@ -95,5 +122,24 @@ waitUntilTVarPred tvar predicate timeoutMs = do
       if predicate val
         then return ()
         else do
-          threadDelay 50
+          threadDelay waitThreadDelay
           performTest
+
+waitUntilQueueSizeIs :: (B.HasBroker b a) => B.Broker b a -> B.Queue -> Int -> TimeoutMs -> Expectation
+waitUntilQueueSizeIs b queue size (TimeoutMs timeoutMs) = do
+  _mTimeout <- Timeout.timeout (timeoutMs * 1000) performTest
+
+  qSize <- B.getQueueSize b queue
+  qSize `shouldBe` size
+
+  where
+    performTest = do
+      qSize <- B.getQueueSize b queue
+      if qSize == size
+        then return ()
+        else do
+          threadDelay waitThreadDelay
+          performTest
+
+waitUntilQueueEmpty :: (B.HasBroker b a) => B.Broker b a -> B.Queue -> TimeoutMs -> Expectation
+waitUntilQueueEmpty b queue timeoutMs = waitUntilQueueSizeIs b queue 0 timeoutMs

tests/integration-tests.hs

 module Main where
 
 import Test.Integration.Broker (brokerTests, pgmqBrokerInitParams, redisBrokerInitParams)
-import Test.Integration.Worker (workerTests, pgmqWorkerBrokerInitParams, redisWorkerBrokerInitParams)
+import Test.Integration.Worker (workerTests, multiWorkerTests, pgmqWorkerBrokerInitParams, redisWorkerBrokerInitParams)
 import Test.Tasty
 import Test.Tasty.Hspec
 
@@ -13,18 +13,22 @@ main = do
   pgmqBrokerSpec <- testSpec "brokerTests (pgmq)" (brokerTests pgmqBInitParams)
   pgmqWBInitParams <- pgmqWorkerBrokerInitParams
   pgmqWorkerSpec <- testSpec "workerTests (pgmq)" (workerTests pgmqWBInitParams)
+  pgmqMultiWorkerSpec <- testSpec "multiWorkerTests (pgmq)" (multiWorkerTests pgmqWBInitParams 5)
 
   redisBInitParams <- redisBrokerInitParams
   redisBrokerSpec <- testSpec "brokerTests (redis)" (brokerTests redisBInitParams)
   redisWBInitParams <- redisWorkerBrokerInitParams
   redisWorkerSpec <- testSpec "workerTests (redis)" (workerTests redisWBInitParams)
+  redisMultiWorkerSpec <- testSpec "multiWorkerTests (redis)" (multiWorkerTests redisWBInitParams 5)
 
   defaultMain $ testGroup "integration tests"
     [
       pgmqBrokerSpec
     , pgmqWorkerSpec
+    , pgmqMultiWorkerSpec
     
     , redisBrokerSpec
     , redisWorkerSpec
+    , redisMultiWorkerSpec
     ]
 

tests/unit-tests.hs

@@ -45,6 +45,7 @@ instance QC.Arbitrary WT.JobMetadata where
     timeoutStrategy <- arbitrary
     timeout <- arbitrary
     readCount <- arbitrary
+    resendWhenWorkerKilled <- arbitrary
     return $ WT.JobMetadata { .. }
    
 aesonPropJobMetadataTests = testGroup "Aeson WT.JobMetadata (de-)serialization tests" $