{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-x-partial #-}

-- | Worker lifecycle tests, instantiated for each backend.
module Arbiter.Test.Worker.Lifecycle
  ( lifecycleSpec
  ) where

import Arbiter.Core.Codec (Col (..), pval)
import Arbiter.Core.CronSchedule qualified as CS
import Arbiter.Core.Exceptions (throwBranchCancel, throwRetryable, throwTreeCancel)
import Arbiter.Core.HighLevel (QueueOperation, RegistryAdmissionPolicies)
import Arbiter.Core.HighLevel qualified as HL
import Arbiter.Core.Job.DLQ qualified as DLQ
import Arbiter.Core.Job.Schema qualified as Schema
import Arbiter.Core.Job.Types
  ( JobRead
  , ObservabilityHooks (..)
  , Stored
  , attempts
  , claimSeq
  , claimedBy
  , decodeStored
  , defaultJob
  , defaultObservabilityHooks
  , payload
  , primaryKey
  , setGroupKey
  , setMaxAttempts
  )
import Arbiter.Core.JobTree ((<~~))
import Arbiter.Core.JobTree qualified as JT
import Arbiter.Core.MonadArbiter (MonadArbiter, RegistryOf, ResultOf, withDbTransaction)
import Arbiter.Core.Operations qualified as Ops
import Arbiter.Core.QueueRegistry (RegistryTables)
import Arbiter.Core.Queues qualified as Q
import Arbiter.Core.Worker qualified as WR
import Arbiter.Worker
  ( MaintenancePace (..)
  , WorkerState (..)
  , mergedChildResults
  , runMaintenancePass
  , runWorkerPool
  )
import Arbiter.Worker.Config
  ( MaintenanceOp (..)
  , WorkerConfig (..)
  , ackAll
  , ackAllWith
  , ackWith
  , batchedWorkerConfig
  , getListenerReady
  , getWorkerState
  , nack
  , shutdownWorkerPool
  , transactionalWorkerConfig
  )
import Arbiter.Worker.Cron (OverlapPolicy (AllowOverlap), cronJob)
import Arbiter.Worker.Logger (silentLogConfig)
import Arbiter.Worker.Reaper (runReaperOp)
import Control.Concurrent (threadDelay)
import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import Control.Exception (SomeException, finally, throwIO, try)
import Control.Monad (replicateM, void, when)
import Control.Monad.IO.Class (liftIO)
import Data.Aeson qualified as Aeson
import Data.ByteString (ByteString)
import Data.ByteString.Char8 qualified as BSC
import Data.ByteString.Lazy qualified as BL
import Data.Either (isRight)
import Data.Foldable (for_, toList, traverse_)
import Data.IORef (atomicModifyIORef', newIORef, readIORef, writeIORef)
import Data.Int (Int32, Int64)
import Data.List (find)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map.Strict qualified as Map
import Data.Maybe (catMaybes, fromMaybe, isJust, isNothing)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Time (NominalDiffTime, diffUTCTime, getCurrentTime)
import Data.UUID (UUID)
import Data.UUID.V4 qualified as UUID
import Database.PostgreSQL.Simple (Only (..), query)
import Database.PostgreSQL.Simple qualified as PG
import Database.PostgreSQL.Simple.Notification (Notification (..), getNotification)
import System.Directory qualified as Dir
import System.Timeout (timeout)
import Test.Hspec
  ( Expectation
  , Spec
  , before
  , describe
  , expectationFailure
  , it
  , shouldBe
  , shouldContain
  , shouldMatchList
  , shouldNotContain
  , shouldReturn
  , shouldSatisfy
  )
import UnliftIO.Async (withAsync)
import UnliftIO.Async qualified as Async

import Arbiter.Test.Poll (bumpRef, waitUntil, withLinkedAsync)
import Arbiter.Test.Setup (addQueueTable, execStatement, execute_, toQuery, withConn)
import Arbiter.Test.Worker.Backend (TestBackend (..))
import Arbiter.Test.Worker.Rows (reclaimJob, updateJob)

-- | Worker lifecycle suite. The queue under test declares @Maybe [Text]@ as its result type.
lifecycleSpec
  :: forall payload m env
   . ( Eq payload
     , QueueOperation m payload
     , RegistryAdmissionPolicies (RegistryOf m)
     , RegistryTables (RegistryOf m)
     , ResultOf m payload ~ Maybe [Text]
     , Show payload
     )
  => TestBackend payload m env
  -> Spec
lifecycleSpec :: forall payload (m :: * -> *) env.
(Eq payload, QueueOperation m payload,
 RegistryAdmissionPolicies (RegistryOf m),
 RegistryTables (RegistryOf m), ResultOf m payload ~ Maybe [Text],
 Show payload) =>
TestBackend payload m env -> Spec
lifecycleSpec TestBackend {Text
schema :: Text
schema :: forall payload (m :: * -> *) env. TestBackend payload m env -> Text
schema, Text
table :: Text
table :: forall payload (m :: * -> *) env. TestBackend payload m env -> Text
table, ByteString
connStr :: ByteString
connStr :: forall payload (m :: * -> *) env.
TestBackend payload m env -> ByteString
connStr, Text -> payload
mkSimple :: Text -> payload
mkSimple :: forall payload (m :: * -> *) env.
TestBackend payload m env -> Text -> payload
mkSimple, IO env
mkEnv :: IO env
mkEnv :: forall payload (m :: * -> *) env.
TestBackend payload m env -> IO env
mkEnv, env -> env
pollOnly :: env -> env
pollOnly :: forall payload (m :: * -> *) env.
TestBackend payload m env -> env -> env
pollOnly, (JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler :: (JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler :: forall payload (m :: * -> *) env.
TestBackend payload m env
-> (JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler, Text -> m ()
runCommand :: Text -> m ()
runCommand :: forall payload (m :: * -> *) env.
TestBackend payload m env -> Text -> m ()
runCommand, forall a. env -> m a -> IO a
runM :: forall a. env -> m a -> IO a
runM :: forall payload (m :: * -> *) env.
TestBackend payload m env -> forall a. env -> m a -> IO a
runM} =
  IO env -> SpecWith env -> Spec
forall a. IO a -> SpecWith a -> Spec
before IO env
mkEnv (SpecWith env -> Spec) -> SpecWith env -> Spec
forall a b. (a -> b) -> a -> b
$ do
    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Maintenance pass" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      let everyPass :: MaintenancePace
everyPass = MaintenancePace {paceWindow :: NominalDiffTime
paceWindow = NominalDiffTime
0, paceSparseWindow :: NominalDiffTime
paceSparseWindow = NominalDiffTime
0, paceBucketIdle :: NominalDiffTime
paceBucketIdle = NominalDiffTime
0}
          dlqPayloads :: env -> IO [payload]
dlqPayloads env
env = (DLQJob payload -> payload) -> [DLQJob payload] -> [payload]
forall a b. (a -> b) -> [a] -> [b]
map (JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload (JobRead payload -> payload)
-> (DLQJob payload -> JobRead payload) -> DLQJob payload -> payload
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DLQJob payload -> JobRead payload
forall payload. DLQJob payload -> JobSnapshot payload
DLQ.jobSnapshot) ([DLQJob payload] -> [payload])
-> IO [DLQJob payload] -> IO [payload]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> env -> m [DLQJob payload] -> IO [DLQJob payload]
forall a. env -> m a -> IO a
runM env
env (forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> m [DLQJob payload]
HL.listDLQJobs @payload Int
10 Int
0)
          deadLetterAged :: env -> Text -> IO ()
deadLetterAged env
env Text
name = do
            Just inserted <- env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
name)))
            [claimed] <- runM env (HL.claimNextVisibleJobs 1 60) :: IO [JobRead payload]
            void $ runM env (HL.moveToDLQ "boom" claimed)
            runM env $
              runCommand
                ( "UPDATE "
                    <> Schema.jobQueueDLQTable schema table
                    <> " SET failed_at = NOW() - interval '2 hours' WHERE job_id = "
                    <> T.pack (show (primaryKey inserted))
                )

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"deletes a flagged job whose lease lapsed" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        Just inserted <- env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"orphan-flagged")))
        [_] <- runM env (HL.claimNextVisibleJobs 1 60) :: IO [JobRead payload]
        runM env (HL.forceCancelJob @payload (primaryKey inserted)) `shouldReturn` 1
        Just flaggedRow <- runM env (HL.getJobById @payload (primaryKey inserted))
        void $ runM env (HL.setVisibilityTimeout 0 flaggedRow)

        failed <- runM env (runMaintenancePass silentLogConfig (\MaintenanceOp
_ Int64
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) everyPass 5)
        failed `shouldBe` []
        runM env (HL.getJobById @payload (primaryKey inserted)) `shouldReturn` Nothing

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"purges DLQ jobs past the queue's retention" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        env -> Text -> IO ()
deadLetterAged env
env Text
"dlq-aged"
        env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (Text -> Maybe NominalDiffTime -> m ()
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe NominalDiffTime -> m ()
HL.setQueueDLQRetention Text
table (NominalDiffTime -> Maybe NominalDiffTime
forall a. a -> Maybe a
Just NominalDiffTime
3600))
        purgedRef <- Int64 -> IO (IORef Int64)
forall a. a -> IO (IORef a)
newIORef (Int64
0 :: Int64)
        let report MaintenanceOp
operation Int64
count = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (MaintenanceOp
operation MaintenanceOp -> MaintenanceOp -> Bool
forall a. Eq a => a -> a -> Bool
== MaintenanceOp
PurgeDLQ) (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef Int64 -> Int64 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int64
purgedRef Int64
count))

        failed <- runM env (runMaintenancePass silentLogConfig report everyPass 5)
        failed `shouldBe` []
        readIORef purgedRef `shouldReturn` 1
        dlqPayloads env `shouldReturn` []

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"keeps DLQ jobs without a retention" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        env -> Text -> IO ()
deadLetterAged env
env Text
"dlq-kept"
        failed <- env -> m [MaintenanceOp] -> IO [MaintenanceOp]
forall a. env -> m a -> IO a
runM env
env (LogConfig
-> (MaintenanceOp -> Int64 -> m ())
-> MaintenancePace
-> NominalDiffTime
-> m [MaintenanceOp]
forall (m :: * -> *).
(RegistryAdmissionPolicies (RegistryOf m), MonadArbiter m,
 RegistryTables (RegistryOf m)) =>
LogConfig
-> (MaintenanceOp -> Int64 -> m ())
-> MaintenancePace
-> NominalDiffTime
-> m [MaintenanceOp]
runMaintenancePass LogConfig
silentLogConfig (\MaintenanceOp
_ Int64
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) MaintenancePace
everyPass NominalDiffTime
5)
        failed `shouldBe` []
        dlqPayloads env `shouldReturn` [mkSimple "dlq-kept"]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"covers a queue outside the caller's registry" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let other :: Text
other = Text
table Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_ext"
            ownDLQ :: Text
ownDLQ = Text -> Text -> Text
Schema.jobQueueDLQTable Text
schema Text
table
            otherDLQ :: Text
otherDLQ = Text -> Text -> Text
Schema.jobQueueDLQTable Text
schema Text
other
        ByteString -> Text -> Text -> Bool -> IO ()
addQueueTable ByteString
connStr Text
schema Text
other Bool
False
        env -> Text -> IO ()
deadLetterAged env
env Text
"dlq-elsewhere"
        env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
          Text -> m ()
runCommand (Text
"INSERT INTO " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
otherDLQ Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" SELECT * FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
ownDLQ)
          Text -> m ()
runCommand (Text
"DELETE FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
ownDLQ)
          Text -> Maybe NominalDiffTime -> m ()
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe NominalDiffTime -> m ()
HL.setQueueDLQRetention Text
other (NominalDiffTime -> Maybe NominalDiffTime
forall a. a -> Maybe a
Just NominalDiffTime
3600)

        failed <- env -> m [MaintenanceOp] -> IO [MaintenanceOp]
forall a. env -> m a -> IO a
runM env
env (LogConfig
-> (MaintenanceOp -> Int64 -> m ())
-> MaintenancePace
-> NominalDiffTime
-> m [MaintenanceOp]
forall (m :: * -> *).
(RegistryAdmissionPolicies (RegistryOf m), MonadArbiter m,
 RegistryTables (RegistryOf m)) =>
LogConfig
-> (MaintenanceOp -> Int64 -> m ())
-> MaintenancePace
-> NominalDiffTime
-> m [MaintenanceOp]
runMaintenancePass LogConfig
silentLogConfig (\MaintenanceOp
_ Int64
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) MaintenancePace
everyPass NominalDiffTime
5)
        failed `shouldBe` []
        remaining <- withConn connStr $ \Connection
conn -> Connection -> Query -> () -> IO [Only Int64]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
conn (Text -> Query
toQuery (Text
"SELECT count(*) FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
otherDLQ)) ()
        remaining `shouldBe` [Only (0 :: Int64)]

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Reaper op bounding" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"completes an op longer than the timeout when each statement is within it" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let sleep :: m ()
sleep = Text -> m ()
runCommand Text
"DO $$ BEGIN PERFORM pg_sleep(0.4); END $$"
        result <-
          env -> m (Maybe Int) -> IO (Maybe Int)
forall a. env -> m a -> IO a
runM env
env (m (Maybe Int) -> IO (Maybe Int))
-> m (Maybe Int) -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$
            LogConfig
-> NominalDiffTime
-> Text
-> NominalDiffTime
-> m Int
-> m (Maybe Int)
forall (m :: * -> *) a.
MonadArbiter m =>
LogConfig
-> NominalDiffTime -> Text -> NominalDiffTime -> m a -> m (Maybe a)
runReaperOp LogConfig
silentLogConfig NominalDiffTime
1 Text
"test-reaper-slow-op" NominalDiffTime
0 (m Int -> m (Maybe Int)) -> m Int -> m (Maybe Int)
forall a b. (a -> b) -> a -> b
$ do
              m ()
sleep
              m ()
sleep
              m ()
sleep
              Int -> m Int
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
42 :: Int)
        result `shouldBe` Just 42
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"aborts a stuck statement at the timeout without killing the caller" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        result <-
          env -> m (Maybe ()) -> IO (Maybe ())
forall a. env -> m a -> IO a
runM env
env
            (m (Maybe ()) -> IO (Maybe ())) -> m (Maybe ()) -> IO (Maybe ())
forall a b. (a -> b) -> a -> b
$ LogConfig
-> NominalDiffTime
-> Text
-> NominalDiffTime
-> m ()
-> m (Maybe ())
forall (m :: * -> *) a.
MonadArbiter m =>
LogConfig
-> NominalDiffTime -> Text -> NominalDiffTime -> m a -> m (Maybe a)
runReaperOp LogConfig
silentLogConfig NominalDiffTime
0.5 Text
"test-reaper-stuck-op" NominalDiffTime
0
            (m () -> m (Maybe ())) -> m () -> m (Maybe ())
forall a b. (a -> b) -> a -> b
$ Text -> m ()
runCommand Text
"DO $$ BEGIN PERFORM pg_sleep(5); END $$"
        result `shouldBe` Nothing

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Transactional Atomicity" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"rolls back user operations when handler fails" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> IO () -> IO ()
forall a. IO a -> IO a
withOpsTable (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
job = do
              Text -> Int64 -> m ()
forall (n :: * -> *). MonadArbiter n => Text -> Int64 -> n ()
recordOp Text
schema (JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job)
              Text -> m ()
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwRetryable Text
"Simulated failure"

        IO (Maybe (JobRead payload)) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void
          (IO (Maybe (JobRead payload)) -> IO ())
-> IO (Maybe (JobRead payload)) -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env
          (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob
          (JobWrite payload -> m (Maybe (JobRead payload)))
-> JobWrite payload -> m (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ Maybe Int32 -> JobWrite payload -> JobWrite payload
forall payload. Maybe Int32 -> JobWrite payload -> JobWrite payload
setMaxAttempts (Int32 -> Maybe Int32
forall a. a -> Maybe a
Just Int32
1)
          (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"g1")
          (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"WillFail")

        config :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
10 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))

        withLinkedAsync
          (runM env $ runWorkerPool config {workerCount = 1, pollInterval = 0.1})
          $ \Async ()
_ -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
              dlqJobs <- env -> m [DLQJob payload] -> IO [DLQJob payload]
forall a. env -> m a -> IO a
runM env
env (m [DLQJob payload] -> IO [DLQJob payload])
-> m [DLQJob payload] -> IO [DLQJob payload]
forall a b. (a -> b) -> a -> b
$ Int -> Int -> m [DLQJob payload]
forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> m [DLQJob payload]
HL.listDLQJobs Int
10 Int
0 :: IO [DLQ.DLQJob payload]
              pure (length dlqJobs == 1)

            dlqJobs <- env -> m [DLQJob payload] -> IO [DLQJob payload]
forall a. env -> m a -> IO a
runM env
env (m [DLQJob payload] -> IO [DLQJob payload])
-> m [DLQJob payload] -> IO [DLQJob payload]
forall a b. (a -> b) -> a -> b
$ Int -> Int -> m [DLQJob payload]
forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> m [DLQJob payload]
HL.listDLQJobs Int
10 Int
0 :: IO [DLQ.DLQJob payload]
            length dlqJobs `shouldBe` 1

            count <- opsCount
            count `shouldBe` 0

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"commits user operations when handler succeeds" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> IO () -> IO ()
forall a. IO a -> IO a
withOpsTable (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
job = Text -> Int64 -> m ()
forall (n :: * -> *). MonadArbiter n => Text -> Int64 -> n ()
recordOp Text
schema (JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job)

        IO (Maybe (JobRead payload)) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe (JobRead payload)) -> IO ())
-> IO (Maybe (JobRead payload)) -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob (JobWrite payload -> m (Maybe (JobRead payload)))
-> JobWrite payload -> m (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"g1") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"WillSucceed")

        config :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
10 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))

        withLinkedAsync
          (runM env $ runWorkerPool config {workerCount = 1, pollInterval = 0.1})
          $ \Async ()
_ -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1) (Int -> Bool) -> IO Int -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Int
opsCount

            count <- IO Int
opsCount
            count `shouldBe` 1

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"manual commit inside handler persists despite subsequent failure" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> IO () -> IO ()
forall a. IO a -> IO a
withOpsTable (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
job = do
              Text -> Int64 -> m ()
forall (n :: * -> *). MonadArbiter n => Text -> Int64 -> n ()
recordOp Text
schema (JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job)
              -- Manual commit violates the worker transaction boundary.
              Text -> m ()
runCommand Text
"COMMIT"
              Text -> m ()
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwRetryable Text
"Simulated failure after commit"

        IO (Maybe (JobRead payload)) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void
          (IO (Maybe (JobRead payload)) -> IO ())
-> IO (Maybe (JobRead payload)) -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env
          (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob
          (JobWrite payload -> m (Maybe (JobRead payload)))
-> JobWrite payload -> m (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ Maybe Int32 -> JobWrite payload -> JobWrite payload
forall payload. Maybe Int32 -> JobWrite payload -> JobWrite payload
setMaxAttempts (Int32 -> Maybe Int32
forall a. a -> Maybe a
Just Int32
1)
          (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"g1")
          (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"ManualCommit")

        config :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
10 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))

        withLinkedAsync
          (runM env $ runWorkerPool config {workerCount = 1, pollInterval = 0.1})
          $ \Async ()
_ -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
              dlqJobs <- env -> m [DLQJob payload] -> IO [DLQJob payload]
forall a. env -> m a -> IO a
runM env
env (m [DLQJob payload] -> IO [DLQJob payload])
-> m [DLQJob payload] -> IO [DLQJob payload]
forall a b. (a -> b) -> a -> b
$ Int -> Int -> m [DLQJob payload]
forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> m [DLQJob payload]
HL.listDLQJobs Int
10 Int
0 :: IO [DLQ.DLQJob payload]
              pure (length dlqJobs == 1)

            dlqJobs <- env -> m [DLQJob payload] -> IO [DLQJob payload]
forall a. env -> m a -> IO a
runM env
env (m [DLQJob payload] -> IO [DLQJob payload])
-> m [DLQJob payload] -> IO [DLQJob payload]
forall a b. (a -> b) -> a -> b
$ Int -> Int -> m [DLQJob payload]
forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> m [DLQJob payload]
HL.listDLQJobs Int
10 Int
0 :: IO [DLQ.DLQJob payload]
            length dlqJobs `shouldBe` 1

            -- User's manual commit survives despite handler failure
            count <- opsCount
            count `shouldBe` 1

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Graceful Shutdown" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"graceful shutdown waits for in-flight jobs to complete" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        -- Track job completion
        completedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        startedRef <- newIORef False

        let handler :: JobRead payload -> m ()
            handler JobRead payload
_job = do
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> (Bool -> (Bool, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Bool
startedRef ((Bool -> (Bool, ())) -> IO ()) -> (Bool -> (Bool, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Bool
_ -> (Bool
True, ())
              -- Simulate long-running job
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay Int
inFlightJobMicros
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> (Bool -> (Bool, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Bool
completedRef ((Bool -> (Bool, ())) -> IO ()) -> (Bool -> (Bool, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Bool
_ -> (Bool
True, ())

        -- Insert a job
        let job = Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"g1") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"LongJob")
        void $ runM env $ HL.insertJob job

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler (noResult handler))

        let configWithTimeout =
              WorkerConfig m payload
config
                { workerCount = 1
                , pollInterval = 0.1
                , gracefulShutdownTimeout = Just 10 -- 10 second timeout (plenty of time)
                }

        -- Start worker and wait for job to start processing
        withLinkedAsync (runM env $ runWorkerPool configWithTimeout) $ \Async ()
worker -> do
          -- Wait for job to start
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
startedRef

          -- Trigger shutdown while job is running
          WorkerConfig m payload -> IO ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m ()
shutdownWorkerPool WorkerConfig m payload
configWithTimeout

          -- Wait for worker to exit (should complete job first)
          Async () -> IO ()
forall (m :: * -> *) a. MonadIO m => Async a -> m a
Async.wait Async ()
worker

          -- Verify job completed despite shutdown
          completed <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
completedRef
          completed `shouldBe` True

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"graceful shutdown times out if jobs take too long" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        -- Track if job started and completed
        startedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        completedRef <- newIORef False

        let handler :: JobRead payload -> m ()
            handler JobRead payload
_job = do
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> (Bool -> (Bool, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Bool
startedRef ((Bool -> (Bool, ())) -> IO ()) -> (Bool -> (Bool, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Bool
_ -> (Bool
True, ())
              -- Very long running job that exceeds timeout
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay Int
overrunJobMicros
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> (Bool -> (Bool, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Bool
completedRef ((Bool -> (Bool, ())) -> IO ()) -> (Bool -> (Bool, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Bool
_ -> (Bool
True, ())

        -- Insert a job
        let job = Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"g1") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"VeryLongJob")
        void $ runM env $ HL.insertJob job

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler (noResult handler))

        let configWithShortTimeout =
              WorkerConfig m payload
config
                { workerCount = 1
                , pollInterval = 0.05 -- Faster polling for test
                , gracefulShutdownTimeout = Just 1 -- Only 1 second timeout
                }

        withLinkedAsync (runM env $ runWorkerPool configWithShortTimeout) $ \Async ()
worker -> do
          -- Wait for the job to start processing
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
startedRef

          -- Measure the shutdown duration.
          startTime <- IO UTCTime -> IO UTCTime
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UTCTime
getCurrentTime
          shutdownWorkerPool configWithShortTimeout
          Async.wait worker
          endTime <- liftIO getCurrentTime

          let elapsed = UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
endTime UTCTime
startTime
          -- Shutdown takes about 1s, the graceful timeout.
          elapsed `shouldSatisfy` (< 5)

        -- The job did not complete. The timeout cancelled it.
        completed <- readIORef completedRef
        completed `shouldBe` False

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"graceful shutdown hands a claim in flight to a worker" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        ranRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        let handler :: JobRead payload -> m ()
            handler JobRead payload
_job = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
ranRef Bool
True
            idleCron = (FilePath -> CronJob payload)
-> (CronJob payload -> CronJob payload)
-> Either FilePath (CronJob payload)
-> CronJob payload
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either FilePath -> CronJob payload
forall a. HasCallStack => FilePath -> a
error CronJob payload -> CronJob payload
forall a. a -> a
id (Text
-> Text
-> OverlapPolicy
-> (TickKind -> UTCTime -> JobWrite payload)
-> Either FilePath (CronJob payload)
forall payload.
Text
-> Text
-> OverlapPolicy
-> (TickKind -> UTCTime -> JobWrite payload)
-> Either FilePath (CronJob payload)
cronJob Text
"shutdown-idle" Text
"0 3 1 1 *" OverlapPolicy
AllowOverlap (\TickKind
_ UTCTime
_ -> payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"cron")))
            queueTable = Text -> Text -> Text
Schema.jobQueueTable Text
schema Text
table
        void $ runM env $ HL.insertJob (defaultJob (mkSimple "claimed-at-shutdown"))

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler (noResult handler))
        let cfg = WorkerConfig m payload
config {workerCount = 1, gracefulShutdownTimeout = Just 10, cronJobs = [idleCron]}
            claimWaiting =
              ByteString -> (Connection -> IO Bool) -> IO Bool
forall a. ByteString -> (Connection -> IO a) -> IO a
withConn ByteString
connStr ((Connection -> IO Bool) -> IO Bool)
-> (Connection -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> do
                [Only waiting] <-
                  Connection -> Query -> () -> IO [Only Int64]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query
                    Connection
conn
                    Query
"SELECT count(*) FROM pg_stat_activity WHERE wait_event_type = 'Lock' AND query LIKE '%group_candidates%'"
                    ()
                pure (waiting > (0 :: Int64))

        withConn connStr $ \Connection
lockConn -> do
          Connection -> IO ()
PG.begin Connection
lockConn
          Connection -> Text -> IO ()
execute_ Connection
lockConn (Text
"LOCK TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
queueTable Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" IN ACCESS EXCLUSIVE MODE")
          IO () -> (Async () -> IO ()) -> IO ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withLinkedAsync (env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> m ()
forall payload (m :: * -> *).
(RegistryAdmissionPolicies (RegistryOf m),
 EncodeJobResult (ResultOf m payload), QueueOperation m payload,
 RegistryTables (RegistryOf m)) =>
WorkerConfig m payload -> m ()
runWorkerPool WorkerConfig m payload
cfg) ((Async () -> IO ()) -> IO ()) -> (Async () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Async ()
worker -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 IO Bool
claimWaiting
            WorkerConfig m payload -> IO ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m ()
shutdownWorkerPool WorkerConfig m payload
cfg
            Int -> IO ()
threadDelay Int
shutdownSignalMicros
            Connection -> IO ()
PG.commit Connection
lockConn
            Async () -> IO ()
forall (m :: * -> *) a. MonadIO m => Async a -> m a
Async.wait Async ()
worker

        readIORef ranRef `shouldReturn` True
        remaining <- withConn connStr $ \Connection
conn ->
          Connection -> Query -> () -> IO [Only Int64]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
conn (Text -> Query
toQuery (Text
"SELECT count(*) FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
queueTable)) ()
        remaining `shouldBe` [Only (0 :: Int64)]

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Liveness Probe" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"creates a health check file when liveness is enabled" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay Int
shortHandlerMicros

        -- Get system temp directory and create liveness file path
        tmpDir <- IO FilePath
Dir.getTemporaryDirectory
        let livenessPath = FilePath
tmpDir FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
"/arbiter-test-liveness"

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler (noResult handler))
        let configWithLiveness =
              WorkerConfig m payload
config
                { livenessFile = Just livenessPath
                , workerCount = 1
                , pollInterval = 0.1
                }

        withLinkedAsync (runM env $ runWorkerPool configWithLiveness) $ \Async ()
worker -> do
          -- Wait for liveness probe to create the file
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> IO Bool
Dir.doesFileExist FilePath
livenessPath

          -- Check that the specific liveness file was created
          exists <- FilePath -> IO Bool
Dir.doesFileExist FilePath
livenessPath
          exists `shouldBe` True

          -- Shutdown and verify cleanup
          shutdownWorkerPool configWithLiveness
          _ <- Async.waitCatch worker
          waitUntil 5_000 $ not <$> Dir.doesFileExist livenessPath
          cleaned <- Dir.doesFileExist livenessPath
          cleaned `shouldBe` False

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Batched Job Mode" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"ackAll skips a job reclaimed mid-batch" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        successRef <- [payload] -> IO (IORef [payload])
forall a. a -> IO (IORef a)
newIORef ([] :: [payload])
        let hooks =
              ObservabilityHooks m payload
forall (m :: * -> *) payload.
Applicative m =>
ObservabilityHooks m payload
defaultObservabilityHooks
                { onJobSuccess = \JobRead payload
job UTCTime
_ UTCTime
_ -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef [payload] -> ([payload] -> ([payload], ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef [payload]
successRef (([payload] -> ([payload], ())) -> IO ())
-> ([payload] -> ([payload], ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \[payload]
seen -> (JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> [payload] -> [payload]
forall a. a -> [a] -> [a]
: [payload]
seen, ())
                }
        let batchHandler NonEmpty (JobRead payload)
jobs BatchCallbacks m payload (Maybe [Text])
cbs = do
              let batchJobs :: [JobRead payload]
batchJobs = NonEmpty (JobRead payload) -> [JobRead payload]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty (JobRead payload)
jobs
              -- Simulate a reclaim of "ca-stolen". A claim bumps both counters. The bulk ack skips it.
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$
                (JobRead payload -> IO ()) -> [JobRead payload] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_
                  (\JobRead payload
job -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"ca-stolen") (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Text -> Text -> Int64 -> IO ()
reclaimJob ByteString
connStr Text
schema Text
table (JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job))
                  [JobRead payload]
batchJobs
              BatchCallbacks m payload (Maybe [Text])
-> [JobRead payload] -> m ()
forall (m :: * -> *) payload result.
BatchCallbacks m payload result -> [JobRead payload] -> m ()
ackAll BatchCallbacks m payload (Maybe [Text])
cbs [JobRead payload]
batchJobs
        let jobs =
              [ Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"ca") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"ca-keep1")
              , Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"ca") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"ca-stolen")
              , Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"ca") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"ca-keep2")
              ]
        void $ runM env $ HL.insertJobsBatch jobs
        config :: WorkerConfig m payload <- batchedWorkerConfig 1 10 batchHandler
        withLinkedAsync (runM env $ runWorkerPool (config {pollInterval = 0.05, observabilityHooks = hooks})) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2) (Int -> Bool) -> ([payload] -> Int) -> [payload] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [payload] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([payload] -> Bool) -> IO [payload] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [payload] -> IO [payload]
forall a. IORef a -> IO a
readIORef IORef [payload]
successRef
          successes <- IORef [payload] -> IO [payload]
forall a. IORef a -> IO a
readIORef IORef [payload]
successRef
          -- onJobSuccess fired only for the survivors. The reclaimed job was skipped.
          successes `shouldMatchList` [mkSimple "ca-keep1", mkSimple "ca-keep2"]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"a tree cancel after a bulk ack also cancels the job the ack found held elsewhere" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let batchHandler :: NonEmpty (JobRead payload)
-> BatchCallbacks m payload (Maybe [Text]) -> m ()
batchHandler NonEmpty (JobRead payload)
jobs BatchCallbacks m payload (Maybe [Text])
cbs = do
              let stolen :: [JobRead payload]
stolen = (JobRead payload -> Bool) -> [JobRead payload] -> [JobRead payload]
forall a. (a -> Bool) -> [a] -> [a]
filter ((payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"tc-stolen") (payload -> Bool)
-> (JobRead payload -> payload) -> JobRead payload -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload) (NonEmpty (JobRead payload) -> [JobRead payload]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty (JobRead payload)
jobs)
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ (JobRead payload -> IO ()) -> [JobRead payload] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (ByteString -> Text -> Text -> Int64 -> IO ()
reclaimJob ByteString
connStr Text
schema Text
table (Int64 -> IO ())
-> (JobRead payload -> Int64) -> JobRead payload -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey) [JobRead payload]
stolen
              BatchCallbacks m payload (Maybe [Text])
-> [JobRead payload] -> m ()
forall (m :: * -> *) payload result.
BatchCallbacks m payload result -> [JobRead payload] -> m ()
ackAll BatchCallbacks m payload (Maybe [Text])
cbs [JobRead payload]
stolen
              Text -> m ()
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwTreeCancel Text
"cancel the batch's trees"
        let jobs :: [JobWrite payload]
jobs =
              [ Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"tc") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"tc-other")
              , Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"tc") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"tc-stolen")
              ]
        IO [JobRead payload] -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO [JobRead payload] -> IO ()) -> IO [JobRead payload] -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m [JobRead payload] -> IO [JobRead payload]
forall a. env -> m a -> IO a
runM env
env (m [JobRead payload] -> IO [JobRead payload])
-> m [JobRead payload] -> IO [JobRead payload]
forall a b. (a -> b) -> a -> b
$ [JobWrite payload] -> m [JobRead payload]
forall payload (m :: * -> *).
QueueOperation m payload =>
[JobWrite payload] -> m [JobRead payload]
HL.insertJobsBatch [JobWrite payload]
jobs
        config :: WorkerConfig m payload <- Int
-> Int
-> (NonEmpty (JobRead payload)
    -> BatchCallbacks m payload (ResultOf m payload) -> m ())
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> Int
-> (NonEmpty (JobRead payload)
    -> BatchCallbacks n payload (ResultOf n payload) -> n ())
-> m (WorkerConfig n payload)
batchedWorkerConfig Int
1 Int
10 NonEmpty (JobRead payload)
-> BatchCallbacks m payload (Maybe [Text]) -> m ()
NonEmpty (JobRead payload)
-> BatchCallbacks m payload (ResultOf m payload) -> m ()
batchHandler
        withLinkedAsync (runM env $ runWorkerPool (config {pollInterval = 0.05})) $ \Async ()
_ ->
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64
0) (Int64 -> Bool) -> IO Int64 -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (forall payload (m :: * -> *). QueueOperation m payload => m Int64
HL.countJobs @payload)

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"rollup ackAllWith stores each job's result for the parent" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        finalRef <- [Text] -> IO (IORef [Text])
forall a. a -> IO (IORef a)
newIORef ([] :: [Text])
        let resultFor :: payload -> Maybe [Text]
            resultFor payload
task = payload -> [(payload, [Text])] -> Maybe [Text]
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup payload
task [(Text -> payload
mkSimple Text
"rb-ca", [Text
"alpha"]), (Text -> payload
mkSimple Text
"rb-cb", [Text
"beta"])]
            isReducer payload
task = payload
task payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"rb-reducer"
            handler NonEmpty (JobRead payload)
jobs BatchCallbacks m payload (Maybe [Text])
cbs =
              if (JobRead payload -> Bool) -> NonEmpty (JobRead payload) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (payload -> Bool
isReducer (payload -> Bool)
-> (JobRead payload -> payload) -> JobRead payload -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload) NonEmpty (JobRead payload)
jobs
                then
                  (JobRead payload -> m ()) -> [JobRead payload] -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_
                    ( \JobRead payload
parent -> do
                        (merged, _dlq) <- JobRead payload -> m (ResultOf m payload, Map Int64 Text)
forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m,
 Monoid (ResultOf m payload)) =>
JobRead payload -> m (ResultOf m payload, Map Int64 Text)
mergedChildResults JobRead payload
parent
                        liftIO $ atomicModifyIORef' finalRef $ \[Text]
_ -> ([Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Text]
merged, ())
                        ackWith cbs parent merged
                    )
                    (NonEmpty (JobRead payload) -> [JobRead payload]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty (JobRead payload)
jobs)
                else BatchCallbacks m payload (Maybe [Text])
-> [(JobRead payload, Maybe [Text])] -> m ()
forall (m :: * -> *) payload result.
BatchCallbacks m payload result
-> [(JobRead payload, result)] -> m ()
ackAllWith BatchCallbacks m payload (Maybe [Text])
cbs ((JobRead payload -> (JobRead payload, Maybe [Text]))
-> [JobRead payload] -> [(JobRead payload, Maybe [Text])]
forall a b. (a -> b) -> [a] -> [b]
map (\JobRead payload
job -> (JobRead payload
job, payload -> Maybe [Text]
resultFor (JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job))) (NonEmpty (JobRead payload) -> [JobRead payload]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty (JobRead payload)
jobs))
        runM env
          $ void
          $ HL.insertJobTree
          $ defaultJob (mkSimple "rb-reducer")
            <~~ (defaultJob (mkSimple "rb-ca") :| [defaultJob (mkSimple "rb-cb")])
        config :: WorkerConfig m payload <- batchedWorkerConfig 1 10 handler
        withLinkedAsync (runM env $ runWorkerPool (config {pollInterval = 0.1})) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2) (Int -> Bool) -> ([Text] -> Int) -> [Text] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Text] -> Bool) -> IO [Text] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalRef
          final <- IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalRef
          final `shouldMatchList` ["alpha", "beta"]

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Fan-out/fan-in with rollup" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"worker auto-appends handler results; finalizer reads merged state" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        finalResultRef <- [Text] -> IO (IORef [Text])
forall a. a -> IO (IORef a)
newIORef ([] :: [Text])

        let handler JobRead payload
job
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"mapper-a" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"sales", Text
"growth"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"mapper-b" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"revenue"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"mapper-c" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"forecast", Text
"trend"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"reducer" = do
                  (merged, _dlq) <- JobRead payload -> m (ResultOf m payload, Map Int64 Text)
forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m,
 Monoid (ResultOf m payload)) =>
JobRead payload -> m (ResultOf m payload, Map Int64 Text)
mergedChildResults JobRead payload
job
                  liftIO $ atomicModifyIORef' finalResultRef $ \[Text]
_ -> ([Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Text]
merged, ())
                  pure merged
              | Bool
otherwise = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [Text]
forall a. Maybe a
Nothing

        -- Insert the rollup tree
        runM env
          $ void
          $ HL.insertJobTree
          $ defaultJob (mkSimple "reducer")
            <~~ ( defaultJob (mkSimple "mapper-a")
                    :| [ defaultJob (mkSimple "mapper-b")
                       , defaultJob (mkSimple "mapper-c")
                       ]
                )

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler handler)

        withLinkedAsync
          ( runM env $
              runWorkerPool
                ( config
                    { workerCount = 3
                    , pollInterval = 0.1
                    }
                )
          )
          $ \Async ()
_ -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
5) (Int -> Bool) -> ([Text] -> Int) -> [Text] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Text] -> Bool) -> IO [Text] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalResultRef

            -- The reducer handler should have received all 5 words
            finalResult <- IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalResultRef
            length finalResult `shouldBe` 5
            finalResult `shouldMatchList` ["sales", "growth", "revenue", "forecast", "trend"]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"nested rollup: section finalizers explicitly merge and propagate" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        --   root (rollup) ← merges section results
        --   ├── section-1 (rollup) ← merges mapper results, returns merged
        --   │   ├── mapper-1a  → ["sales", "growth"]
        --   │   └── mapper-1b  → ["revenue"]
        --   └── section-2 (rollup) ← merges mapper results, returns merged
        --       ├── mapper-2a  → ["forecast"]
        --       └── mapper-2b  → ["trend"]
        finalResultRef <- [Text] -> IO (IORef [Text])
forall a. a -> IO (IORef a)
newIORef ([] :: [Text])

        let handler JobRead payload
job
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"mapper-1a" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"sales", Text
"growth"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"mapper-1b" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"revenue"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"mapper-2a" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"forecast"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"mapper-2b" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"trend"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"section-1" = (Maybe [Text], Map Int64 Text) -> Maybe [Text]
forall a b. (a, b) -> a
fst ((Maybe [Text], Map Int64 Text) -> Maybe [Text])
-> m (Maybe [Text], Map Int64 Text) -> m (Maybe [Text])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> JobRead payload -> m (ResultOf m payload, Map Int64 Text)
forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m,
 Monoid (ResultOf m payload)) =>
JobRead payload -> m (ResultOf m payload, Map Int64 Text)
mergedChildResults JobRead payload
job
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"section-2" = (Maybe [Text], Map Int64 Text) -> Maybe [Text]
forall a b. (a, b) -> a
fst ((Maybe [Text], Map Int64 Text) -> Maybe [Text])
-> m (Maybe [Text], Map Int64 Text) -> m (Maybe [Text])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> JobRead payload -> m (ResultOf m payload, Map Int64 Text)
forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m,
 Monoid (ResultOf m payload)) =>
JobRead payload -> m (ResultOf m payload, Map Int64 Text)
mergedChildResults JobRead payload
job
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"root" = do
                  (merged, _dlq) <- JobRead payload -> m (ResultOf m payload, Map Int64 Text)
forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m,
 Monoid (ResultOf m payload)) =>
JobRead payload -> m (ResultOf m payload, Map Int64 Text)
mergedChildResults JobRead payload
job
                  liftIO $ atomicModifyIORef' finalResultRef $ \[Text]
_ -> ([Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Text]
merged, ())
                  pure merged
              | Bool
otherwise = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [Text]
forall a. Maybe a
Nothing

        runM env
          $ void
          $ HL.insertJobTree
          $ JT.rollup (defaultJob (mkSimple "root"))
          $ ( defaultJob (mkSimple "section-1")
                <~~ (defaultJob (mkSimple "mapper-1a") :| [defaultJob (mkSimple "mapper-1b")])
            )
            :| [ defaultJob (mkSimple "section-2")
                   <~~ (defaultJob (mkSimple "mapper-2a") :| [defaultJob (mkSimple "mapper-2b")])
               ]

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler handler)

        withLinkedAsync
          ( runM env $
              runWorkerPool
                ( config
                    { workerCount = 3
                    , pollInterval = 0.1
                    }
                )
          )
          $ \Async ()
_ -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
15_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
5) (Int -> Bool) -> ([Text] -> Int) -> [Text] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Text] -> Bool) -> IO [Text] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalResultRef

            -- Root handler should receive all 5 words merged from both sections
            finalResult <- IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalResultRef
            length finalResult `shouldBe` 5
            finalResult
              `shouldMatchList` ["sales", "growth", "revenue", "forecast", "trend"]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"completing a rollup parent stores its result and acks" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        finalResultRef <- [Text] -> IO (IORef [Text])
forall a. a -> IO (IORef a)
newIORef ([] :: [Text])

        let handler (JobRead payload
job :| [JobRead payload]
_) BatchCallbacks m payload (Maybe [Text])
cbs
              -- ackWith stores the child result for the parent and acks it
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"child-a" = BatchCallbacks m payload (Maybe [Text])
-> JobRead payload -> Maybe [Text] -> m ()
forall (m :: * -> *) payload result.
BatchCallbacks m payload result
-> JobRead payload -> result -> m ()
ackWith BatchCallbacks m payload (Maybe [Text])
cbs JobRead payload
job ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"alpha"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"child-b" = BatchCallbacks m payload (Maybe [Text])
-> JobRead payload -> Maybe [Text] -> m ()
forall (m :: * -> *) payload result.
BatchCallbacks m payload result
-> JobRead payload -> result -> m ()
ackWith BatchCallbacks m payload (Maybe [Text])
cbs JobRead payload
job ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"beta", Text
"gamma"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"manual-reducer" = do
                  (merged, _dlq) <- JobRead payload -> m (ResultOf m payload, Map Int64 Text)
forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m,
 Monoid (ResultOf m payload)) =>
JobRead payload -> m (ResultOf m payload, Map Int64 Text)
mergedChildResults JobRead payload
job
                  liftIO $ atomicModifyIORef' finalResultRef $ \[Text]
_ -> ([Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Text]
merged, ())
                  ackWith cbs job merged
              | Bool
otherwise = BatchCallbacks m payload (Maybe [Text])
-> JobRead payload -> Maybe [Text] -> m ()
forall (m :: * -> *) payload result.
BatchCallbacks m payload result
-> JobRead payload -> result -> m ()
ackWith BatchCallbacks m payload (Maybe [Text])
cbs JobRead payload
job Maybe [Text]
forall a. Maybe a
Nothing

        -- Insert the rollup tree
        runM env
          $ void
          $ HL.insertJobTree
          $ defaultJob (mkSimple "manual-reducer")
            <~~ ( defaultJob (mkSimple "child-a")
                    :| [defaultJob (mkSimple "child-b")]
                )

        config :: WorkerConfig m payload <- batchedWorkerConfig 3 1 handler

        withLinkedAsync
          ( runM env $
              runWorkerPool (config {pollInterval = 0.1})
          )
          $ \Async ()
_ -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3) (Int -> Bool) -> ([Text] -> Int) -> [Text] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Text] -> Bool) -> IO [Text] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalResultRef

            -- Finalizer should have received child results
            finalResult <- IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalResultRef
            length finalResult `shouldBe` 3
            finalResult `shouldMatchList` ["alpha", "beta", "gamma"]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"batched mode: two rollup parents claimed in one batch each get their own child results" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        -- Two ungrouped rollup parents are batched together once their children
        -- complete. Each parent receives its own merged children, keyed by primary
        -- key. No group keys are needed.
        receivedRef <- Map Text [Text] -> IO (IORef (Map Text [Text]))
forall a. a -> IO (IORef a)
newIORef (Map Text [Text]
forall k a. Map k a
Map.empty :: Map.Map Text [Text])
        batchSizeRef <- newIORef (0 :: Int)

        let reducers = [Text
"reducer-1", Text
"reducer-2"] :: [Text]
            reducerName payload
task = (Text -> Bool) -> [Text] -> Maybe Text
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== payload
task) (payload -> Bool) -> (Text -> payload) -> Text -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> payload
mkSimple) [Text]
reducers
            childResult payload
task =
              payload -> [(payload, [Text])] -> Maybe [Text]
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup
                payload
task
                [ (Text -> payload
mkSimple Text
"child-1a", [Text
"a1"])
                , (Text -> payload
mkSimple Text
"child-1b", [Text
"b1"])
                , (Text -> payload
mkSimple Text
"child-2a", [Text
"a2"])
                , (Text -> payload
mkSimple Text
"child-2b", [Text
"b2"])
                ]
            handler NonEmpty (JobRead payload)
jobs BatchCallbacks m payload (Maybe [Text])
cbs = do
              let reducerCount :: Int
reducerCount = [JobRead payload] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((JobRead payload -> Bool) -> [JobRead payload] -> [JobRead payload]
forall a. (a -> Bool) -> [a] -> [a]
filter (Maybe Text -> Bool
forall a. Maybe a -> Bool
isJust (Maybe Text -> Bool)
-> (JobRead payload -> Maybe Text) -> JobRead payload -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. payload -> Maybe Text
reducerName (payload -> Maybe Text)
-> (JobRead payload -> payload) -> JobRead payload -> Maybe Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload) (NonEmpty (JobRead payload) -> [JobRead payload]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty (JobRead payload)
jobs))
              Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
reducerCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
                IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$
                  IORef Int -> (Int -> (Int, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
batchSizeRef ((Int -> (Int, ())) -> IO ()) -> (Int -> (Int, ())) -> IO ()
forall a b. (a -> b) -> a -> b
$
                    \Int
largest -> (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
largest Int
reducerCount, ())
              NonEmpty (JobRead payload) -> (JobRead payload -> m ()) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ NonEmpty (JobRead payload)
jobs ((JobRead payload -> m ()) -> m ())
-> (JobRead payload -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \JobRead payload
job -> case payload -> Maybe Text
reducerName (JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job) of
                Just Text
name -> do
                  (merged, _dlq) <- JobRead payload -> m (ResultOf m payload, Map Int64 Text)
forall (m :: * -> *) payload.
(FromJSON (ResultOf m payload), MonadArbiter m,
 Monoid (ResultOf m payload)) =>
JobRead payload -> m (ResultOf m payload, Map Int64 Text)
mergedChildResults JobRead payload
job
                  liftIO $ atomicModifyIORef' receivedRef $ \Map Text [Text]
collected -> (Text -> [Text] -> Map Text [Text] -> Map Text [Text]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
name ([Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Text]
merged) Map Text [Text]
collected, ())
                  ackWith cbs job merged
                Maybe Text
Nothing -> BatchCallbacks m payload (Maybe [Text])
-> JobRead payload -> Maybe [Text] -> m ()
forall (m :: * -> *) payload result.
BatchCallbacks m payload result
-> JobRead payload -> result -> m ()
ackWith BatchCallbacks m payload (Maybe [Text])
cbs JobRead payload
job (payload -> Maybe [Text]
childResult (JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job))

        -- Two independent rollup trees, all ungrouped. The four children drain in
        -- one ungrouped batch, then both parents unblock and batch together.
        runM env
          $ void
          $ HL.insertJobTree
          $ defaultJob (mkSimple "reducer-1")
            <~~ (defaultJob (mkSimple "child-1a") :| [defaultJob (mkSimple "child-1b")])
        runM env
          $ void
          $ HL.insertJobTree
          $ defaultJob (mkSimple "reducer-2")
            <~~ (defaultJob (mkSimple "child-2a") :| [defaultJob (mkSimple "child-2b")])

        config :: WorkerConfig m payload <- batchedWorkerConfig 1 10 handler

        withLinkedAsync
          ( runM env $
              runWorkerPool (config {pollInterval = 0.1})
          )
          $ \Async ()
_ -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2) (Int -> Bool)
-> (Map Text [Text] -> Int) -> Map Text [Text] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Text [Text] -> Int
forall k a. Map k a -> Int
Map.size (Map Text [Text] -> Bool) -> IO (Map Text [Text]) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef (Map Text [Text]) -> IO (Map Text [Text])
forall a. IORef a -> IO a
readIORef IORef (Map Text [Text])
receivedRef

            received <- IORef (Map Text [Text]) -> IO (Map Text [Text])
forall a. IORef a -> IO a
readIORef IORef (Map Text [Text])
receivedRef
            Map.findWithDefault [] "reducer-1" received `shouldMatchList` ["a1", "b1"]
            Map.findWithDefault [] "reducer-2" received `shouldMatchList` ["a2", "b2"]

            -- Both parents were handled together in a single batch.
            batchSize <- readIORef batchSizeRef
            batchSize `shouldBe` 2

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"DLQ snapshot round-trip preserves child results" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        attemptRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)
        finalResultRef <- newIORef ([] :: [Text])

        let handler JobRead payload
job
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"dlq-child-a" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"x"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"dlq-child-b" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"y", Text
"z"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"dlq-reducer" = do
                  attempt <- IO Int -> m Int
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> m Int) -> IO Int -> m Int
forall a b. (a -> b) -> a -> b
$ IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
attemptRef ((Int -> (Int, Int)) -> IO Int) -> (Int -> (Int, Int)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Int
count -> (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                  if attempt == 1
                    then throwRetryable "Intentional failure on first attempt"
                    else do
                      (merged, _dlq) <- mergedChildResults job
                      liftIO $ atomicModifyIORef' finalResultRef $ \[Text]
_ -> ([Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Text]
merged, ())
                      pure merged
              | Bool
otherwise = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [Text]
forall a. Maybe a
Nothing

        -- Insert the rollup tree
        Right (_parent :| _children) <-
          runM env
            $ HL.insertJobTree
            $ (setMaxAttempts (Just 1) $ defaultJob (mkSimple "dlq-reducer"))
              <~~ ( defaultJob (mkSimple "dlq-child-a")
                      :| [defaultJob (mkSimple "dlq-child-b")]
                  )

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler handler)

        let cfg =
              WorkerConfig m payload
config
                { workerCount = 3
                , pollInterval = 0.1
                }

        -- Phase 1: Run workers - children succeed, reducer fails → DLQ
        withLinkedAsync (runM env $ runWorkerPool cfg) $ \Async ()
_ ->
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            dlqJobs <- env -> m [DLQJob payload] -> IO [DLQJob payload]
forall a. env -> m a -> IO a
runM env
env (m [DLQJob payload] -> IO [DLQJob payload])
-> m [DLQJob payload] -> IO [DLQJob payload]
forall a b. (a -> b) -> a -> b
$ forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> m [DLQJob payload]
HL.listDLQJobs @payload Int
10 Int
0
            pure $ any (\DLQJob payload
dlqJob -> JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload (DLQJob payload -> JobRead payload
forall payload. DLQJob payload -> JobSnapshot payload
DLQ.jobSnapshot DLQJob payload
dlqJob) payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"dlq-reducer") dlqJobs

        -- Verify reducer is in DLQ with snapshot
        dlqJobs <- runM env $ HL.listDLQJobs @payload 10 0
        let reducerDlq = (DLQJob payload -> Bool)
-> [DLQJob payload] -> Maybe (DLQJob payload)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\DLQJob payload
dlqJob -> JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload (DLQJob payload -> JobRead payload
forall payload. DLQJob payload -> JobSnapshot payload
DLQ.jobSnapshot DLQJob payload
dlqJob) payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"dlq-reducer") [DLQJob payload]
dlqJobs

        -- Phase 2: Retry from DLQ - reducer should see preserved results from snapshot
        dlqId <- maybe (fail "dlq-reducer is not in the DLQ") (pure . DLQ.dlqPrimaryKey) reducerDlq
        mRetried <- runM env $ HL.retryFromDLQ @payload dlqId
        case mRetried of
          Maybe (JobRead payload)
Nothing -> HasCallStack => FilePath -> IO ()
FilePath -> IO ()
expectationFailure FilePath
"retryFromDLQ returned Nothing"
          Just JobRead payload
retried -> JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
retried payload -> payload -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Text -> payload
mkSimple Text
"dlq-reducer"

        withLinkedAsync (runM env $ runWorkerPool cfg) $ \Async ()
_ ->
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> ([Text] -> Bool) -> [Text] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Text] -> Bool) -> IO [Text] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalResultRef

        -- The retried reducer should have received the preserved child results
        finalResult <- readIORef finalResultRef
        length finalResult `shouldBe` 3
        finalResult `shouldMatchList` ["x", "y", "z"]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"DLQ tree recovery: both parent and child in DLQ" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        attemptRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)
        finalResultRef <- newIORef ([] :: [Text])

        let handler JobRead payload
job
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"recover-child-ok" = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
"alpha"])
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"recover-child-fail" = Text -> m (Maybe [Text])
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwRetryable Text
"Permanent child failure"
              | JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"recover-reducer" = do
                  attempt <- IO Int -> m Int
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> m Int) -> IO Int -> m Int
forall a b. (a -> b) -> a -> b
$ IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
attemptRef ((Int -> (Int, Int)) -> IO Int) -> (Int -> (Int, Int)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Int
count -> (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                  if attempt == 1
                    then throwRetryable "Reducer fails first time"
                    else do
                      (merged, _dlq) <- mergedChildResults job
                      liftIO $ atomicModifyIORef' finalResultRef $ \[Text]
_ -> ([Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Text]
merged, ())
                      pure merged
              | Bool
otherwise = Maybe [Text] -> m (Maybe [Text])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [Text]
forall a. Maybe a
Nothing

        -- Insert rollup tree: reducer + 2 children
        Right (_parent :| _children) <-
          runM env
            $ HL.insertJobTree
            $ (setMaxAttempts (Just 1) $ defaultJob (mkSimple "recover-reducer"))
              <~~ ( defaultJob (mkSimple "recover-child-ok")
                      :| [setMaxAttempts (Just 1) $ defaultJob (mkSimple "recover-child-fail")]
                  )

        -- Phase 1: Worker runs - child-ok succeeds, child-fail DLQs, reducer wakes, reducer DLQs
        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler handler)

        let cfg =
              WorkerConfig m payload
config
                { workerCount = 3
                , pollInterval = 0.1
                }

        withLinkedAsync (runM env $ runWorkerPool cfg) $ \Async ()
_ ->
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
15_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            dlqJobs <- env -> m [DLQJob payload] -> IO [DLQJob payload]
forall a. env -> m a -> IO a
runM env
env (m [DLQJob payload] -> IO [DLQJob payload])
-> m [DLQJob payload] -> IO [DLQJob payload]
forall a b. (a -> b) -> a -> b
$ forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> m [DLQJob payload]
HL.listDLQJobs @payload Int
10 Int
0
            pure (length dlqJobs == 2)

        -- Both child-fail and reducer should be in DLQ
        dlqJobs <- runM env $ HL.listDLQJobs @payload 10 0
        let dlqPayloads = (DLQJob payload -> payload) -> [DLQJob payload] -> [payload]
forall a b. (a -> b) -> [a] -> [b]
map (JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload (JobRead payload -> payload)
-> (DLQJob payload -> JobRead payload) -> DLQJob payload -> payload
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DLQJob payload -> JobRead payload
forall payload. DLQJob payload -> JobSnapshot payload
DLQ.jobSnapshot) [DLQJob payload]
dlqJobs
        dlqPayloads `shouldContain` [mkSimple "recover-child-fail"]
        dlqPayloads `shouldContain` [mkSimple "recover-reducer"]

        -- Phase 2: Retry child-fail from DLQ → auto-retries reducer (suspended)
        childDlq <-
          maybe (fail "recover-child-fail is not in the DLQ") pure $
            find (\DLQJob payload
dlqJob -> JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload (DLQJob payload -> JobRead payload
forall payload. DLQJob payload -> JobSnapshot payload
DLQ.jobSnapshot DLQJob payload
dlqJob) payload -> payload -> Bool
forall a. Eq a => a -> a -> Bool
== Text -> payload
mkSimple Text
"recover-child-fail") dlqJobs
        mRetried <- runM env $ HL.retryFromDLQ @payload (DLQ.dlqPrimaryKey childDlq)
        case mRetried of
          Maybe (JobRead payload)
Nothing -> HasCallStack => FilePath -> IO ()
FilePath -> IO ()
expectationFailure FilePath
"retryFromDLQ returned Nothing"
          Just JobRead payload
retried -> JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
retried payload -> payload -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Text -> payload
mkSimple Text
"recover-child-fail"

        -- Phase 3: Run workers again - child-fail still fails, goes back to DLQ,
        -- but reducer wakes with partial results from snapshot
        withLinkedAsync (runM env $ runWorkerPool cfg) $ \Async ()
_ ->
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
15_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> ([Text] -> Bool) -> [Text] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Text] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Text] -> Bool) -> IO [Text] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [Text] -> IO [Text]
forall a. IORef a -> IO a
readIORef IORef [Text]
finalResultRef

        -- The retried reducer (second attempt) should have received at least child-ok's result
        finalResult <- readIORef finalResultRef
        finalResult `shouldBe` ["alpha"]

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Worker Registry" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"registers, stamps claimed_by, and reconciles pause from the registry" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        processedRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)
        let handler :: JobRead payload -> m ()
            handler JobRead payload
_job =
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Int -> IO ()
forall a. Num a => IORef a -> IO ()
bumpRef IORef Int
processedRef

        baseConfig :: WorkerConfig m payload <- transactionalWorkerConfig 2 (mkHandler (noResult handler))
        let config = WorkerConfig m payload
baseConfig {workerCount = 2, pollInterval = 0.1}
            wid = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
config

        void $ runM env $ HL.insertJob (defaultJob (mkSimple "first"))

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1) (Int -> Bool) -> IO Int -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
processedRef

          rows <- env -> m [WorkerRow] -> IO [WorkerRow]
forall a. env -> m a -> IO a
runM env
env (m [WorkerRow] -> IO [WorkerRow])
-> m [WorkerRow] -> IO [WorkerRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
Ops.listWorkers Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table) Maybe NominalDiffTime
forall a. Maybe a
Nothing
          map WR.workerId rows `shouldContain` [wid]
          map WR.queueName rows `shouldContain` [table]

          -- Pause the worker before the next claim.
          void $ runM env $ Ops.setWorkerPaused schema wid True
          waitUntil 5_000 $ (== Paused) <$> getWorkerState config

          -- Insert and manually claim to read back claimed_by.
          void $ runM env $ HL.insertJob (defaultJob (mkSimple "attribution"))
          claimed <-
            runM env $
              Ops.claimNextVisibleJobsAs @_ @payload schema table 1 60 wid
          case claimed of
            (JobRead payload
claimedJob : [JobRead payload]
_) -> JobRead payload -> Maybe UUID
forall payload q insertedAt adm.
Job payload Int64 q insertedAt adm -> Maybe UUID
claimedBy JobRead payload
claimedJob Maybe UUID -> Maybe UUID -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` UUID -> Maybe UUID
forall a. a -> Maybe a
Just UUID
wid
            [] -> HasCallStack => FilePath -> IO ()
FilePath -> IO ()
expectationFailure FilePath
"expected a job to be claimable for the claimed_by assertion"

          void $ runM env $ Ops.setWorkerPaused schema wid False
          waitUntil 5_000 $ (== Running) <$> getWorkerState config

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"re-registers if the registry row is swept out from under it" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

        baseConfig :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
1 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))
        let config =
              WorkerConfig m payload
baseConfig
                { workerCount = 1
                , pollInterval = 0.1
                , workerHeartbeatInterval = 0.2
                }
            wid = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
config

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            rows <- env -> m [WorkerRow] -> IO [WorkerRow]
forall a. env -> m a -> IO a
runM env
env (m [WorkerRow] -> IO [WorkerRow])
-> m [WorkerRow] -> IO [WorkerRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
Ops.listWorkers Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table) Maybe NominalDiffTime
forall a. Maybe a
Nothing
            pure $ wid `elem` map WR.workerId rows

          _ <- env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> UUID -> m Int64
forall (m :: * -> *). MonadArbiter m => Text -> UUID -> m Int64
Ops.deregisterWorker Text
schema UUID
wid
          rowsAfterDelete <- runM env $ Ops.listWorkers schema (Just table) Nothing
          map WR.workerId rowsAfterDelete `shouldNotContain` [wid]

          -- Insert a job. The dispatcher then signals the heartbeat.
          void $ runM env $ HL.insertJob (defaultJob (mkSimple "wake"))

          waitUntil 5_000 $ do
            rows <- runM env $ Ops.listWorkers schema (Just table) Nothing
            pure $ wid `elem` map WR.workerId rows

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"does not re-register a swept row while it drains" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        startedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        let handler :: JobRead payload -> m ()
            handler JobRead payload
_job = do
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
startedRef Bool
True
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay Int
drainingHandlerMicros

        baseConfig :: WorkerConfig m payload <- transactionalWorkerConfig 1 (mkHandler (noResult handler))
        let config =
              WorkerConfig m payload
baseConfig
                { workerCount = 1
                , pollInterval = 0.1
                , jobHeartbeatInterval = 0.2
                , workerHeartbeatInterval = 0.2
                , gracefulShutdownTimeout = Just 10
                }
            wid = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
config
            registered = do
              rows <- env -> m [WorkerRow] -> IO [WorkerRow]
forall a. env -> m a -> IO a
runM env
env (m [WorkerRow] -> IO [WorkerRow])
-> m [WorkerRow] -> IO [WorkerRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
Ops.listWorkers Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table) Maybe NominalDiffTime
forall a. Maybe a
Nothing
              pure $ wid `elem` map WR.workerId rows

        void $ runM env $ HL.insertJob (defaultJob (mkSimple "draining"))
        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
worker -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
startedRef
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 IO Bool
registered
          WorkerConfig m payload -> IO ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m ()
shutdownWorkerPool WorkerConfig m payload
config
          _ <- env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> UUID -> m Int64
forall (m :: * -> *). MonadArbiter m => Text -> UUID -> m Int64
Ops.deregisterWorker Text
schema UUID
wid
          -- The job's lease extensions keep the registry heartbeat ticking.
          threadDelay drainTicksMicros
          registered `shouldReturn` False
          Async.wait worker

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"paused worker keeps heartbeating and survives the sweeper" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

        baseConfig :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
1 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))
        let config =
              WorkerConfig m payload
baseConfig
                { workerCount = 1
                , pollInterval = 0.1
                , workerHeartbeatInterval = 0.2
                , workerStaleThreshold = 1
                }
            wid = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
config

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            rows <- env -> m [WorkerRow] -> IO [WorkerRow]
forall a. env -> m a -> IO a
runM env
env (m [WorkerRow] -> IO [WorkerRow])
-> m [WorkerRow] -> IO [WorkerRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
Ops.listWorkers Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table) Maybe NominalDiffTime
forall a. Maybe a
Nothing
            pure $ wid `elem` map WR.workerId rows

          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> UUID -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> UUID -> Bool -> m Int64
Ops.setWorkerPaused Text
schema UUID
wid Bool
True
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Paused) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config

          -- Sit past stale_threshold_secs while paused.
          Int -> IO ()
threadDelay Int
pastStaleWhilePausedMicros
          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> m Int64
forall (m :: * -> *). MonadArbiter m => Text -> m Int64
Ops.sweepStaleWorkers Text
schema

          rows <- env -> m [WorkerRow] -> IO [WorkerRow]
forall a. env -> m a -> IO a
runM env
env (m [WorkerRow] -> IO [WorkerRow])
-> m [WorkerRow] -> IO [WorkerRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
Ops.listWorkers Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table) Maybe NominalDiffTime
forall a. Maybe a
Nothing
          map WR.workerId rows `shouldContain` [wid]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"re-registers with the queue's pause state" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        baseConfig :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
1 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))
        let config = WorkerConfig m payload
baseConfig {workerCount = 1, pollInterval = 5.0, workerHeartbeatInterval = 2.0}
            wid = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
config
            registered = do
              rows <- env -> m [WorkerRow] -> IO [WorkerRow]
forall a. env -> m a -> IO a
runM env
env (m [WorkerRow] -> IO [WorkerRow])
-> m [WorkerRow] -> IO [WorkerRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
Ops.listWorkers Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table) Maybe NominalDiffTime
forall a. Maybe a
Nothing
              pure $ wid `elem` map WR.workerId rows

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Running) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> IO Bool
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m Bool
getListenerReady WorkerConfig m payload
config
          -- Let the tick the startup claim signalled pass.
          Int -> IO ()
threadDelay Int
startupTickMicros

          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> UUID -> m Int64
forall (m :: * -> *). MonadArbiter m => Text -> UUID -> m Int64
Ops.deregisterWorker Text
schema UUID
wid
          -- The pause fans out per registry row. A worker without one hears nothing.
          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Bool -> m Int64
Ops.setQueuePaused Text
schema Text
table Bool
True
          Int -> IO ()
threadDelay Int
pauseLandMicros
          WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config IO WorkerState -> WorkerState -> IO ()
forall a. (HasCallStack, Show a, Eq a) => IO a -> a -> IO ()
`shouldReturn` WorkerState
Running

          -- A claim signals the heartbeat, whose tick finds the row gone and re-registers.
          IO (Maybe (JobRead payload)) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe (JobRead payload)) -> IO ())
-> IO (Maybe (JobRead payload)) -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"wake"))
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 IO Bool
registered
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
500 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Paused) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config
          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Bool -> m Int64
Ops.setQueuePaused Text
schema Text
table Bool
False

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"starts paused when the registry insert fails, then registers on a later heartbeat" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        baseConfig :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
1 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))
        let config = WorkerConfig m payload
baseConfig {workerCount = 1, pollInterval = 0.1, workerHeartbeatInterval = 0.2}
            wid = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
config
            registry = Text
schema Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
".arbiter_workers"
            hidden = Text
schema Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
".arbiter_workers_hidden"

        withConn connStr $ \Connection
conn -> do
          Connection -> Text -> IO ()
execute_ Connection
conn (Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
registry Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" RENAME TO arbiter_workers_hidden")
          let restore :: IO ()
restore = Connection -> Text -> IO ()
execute_ Connection
conn (Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
hidden Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" RENAME TO arbiter_workers")
          IO () -> (Async () -> IO ()) -> IO ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withLinkedAsync (env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> m ()
forall payload (m :: * -> *).
(RegistryAdmissionPolicies (RegistryOf m),
 EncodeJobResult (ResultOf m payload), QueueOperation m payload,
 RegistryTables (RegistryOf m)) =>
WorkerConfig m payload -> m ()
runWorkerPool WorkerConfig m payload
config) ((Async () -> IO ()) -> IO ()) -> (Async () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Async ()
_ -> do
            (Int -> IO ()
threadDelay Int
registrationAttemptMicros IO () -> IO WorkerState -> IO WorkerState
forall a b. IO a -> IO b -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config) IO WorkerState -> IO () -> IO WorkerState
forall a b. IO a -> IO b -> IO a
`finally` IO ()
restore IO WorkerState -> (WorkerState -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (WorkerState -> WorkerState -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` WorkerState
Paused)
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Running) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config
            rows <- env -> m [WorkerRow] -> IO [WorkerRow]
forall a. env -> m a -> IO a
runM env
env (m [WorkerRow] -> IO [WorkerRow])
-> m [WorkerRow] -> IO [WorkerRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
Ops.listWorkers Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table) Maybe NominalDiffTime
forall a. Maybe a
Nothing
            map WR.workerId rows `shouldContain` [wid]

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Queue pause" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"stamps paused_at on first pause and clears it on resume" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> m Int64
forall (m :: * -> *). MonadArbiter m => Text -> Text -> m Int64
Ops.ensureQueue Text
schema Text
table
        IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Bool -> m Int64
Ops.setQueuePaused Text
schema Text
table Bool
True
        Just row1 <- env -> m (Maybe QueueRow) -> IO (Maybe QueueRow)
forall a. env -> m a -> IO a
runM env
env (m (Maybe QueueRow) -> IO (Maybe QueueRow))
-> m (Maybe QueueRow) -> IO (Maybe QueueRow)
forall a b. (a -> b) -> a -> b
$ Text -> Text -> m (Maybe QueueRow)
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> m (Maybe QueueRow)
Ops.getQueue Text
schema Text
table
        Q.paused row1 `shouldBe` True
        Q.pausedAt row1 `shouldSatisfy` isJust

        void $ runM env $ Ops.setQueuePaused schema table False
        Just row2 <- runM env $ Ops.getQueue schema table
        Q.paused row2 `shouldBe` False
        Q.pausedAt row2 `shouldBe` Nothing

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"preserves paused_at on idempotent re-pause" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> m Int64
forall (m :: * -> *). MonadArbiter m => Text -> Text -> m Int64
Ops.ensureQueue Text
schema Text
table
        IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Bool -> m Int64
Ops.setQueuePaused Text
schema Text
table Bool
True
        Just first <- env -> m (Maybe QueueRow) -> IO (Maybe QueueRow)
forall a. env -> m a -> IO a
runM env
env (m (Maybe QueueRow) -> IO (Maybe QueueRow))
-> m (Maybe QueueRow) -> IO (Maybe QueueRow)
forall a b. (a -> b) -> a -> b
$ Text -> Text -> m (Maybe QueueRow)
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> m (Maybe QueueRow)
Ops.getQueue Text
schema Text
table
        let original = QueueRow -> Maybe UTCTime
Q.pausedAt QueueRow
first
        original `shouldSatisfy` isJust

        threadDelay clockAdvanceMicros
        void $ runM env $ Ops.setQueuePaused schema table True
        Just second <- runM env $ Ops.getQueue schema table
        Q.pausedAt second `shouldBe` original

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"lists workers filtered by liveness across all queues" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        liveWid <- IO UUID -> IO UUID
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
UUID.nextRandom
        staleWid <- liftIO UUID.nextRandom
        void
          $ runM env
          $ Ops.registerWorker schema liveWid table Nothing (Just testWorkerThreads) farStaleAfter Nothing
        void
          $ runM env
          $ Ops.registerWorker schema staleWid table Nothing (Just testWorkerThreads) farStaleAfter Nothing

        -- Age both rows past the query threshold, then bump only the live row.
        threadDelay pastListThresholdMicros
        void $ runM env $ Ops.heartbeatWorker schema liveWid

        -- No liveness filter returns both rows.
        allRows <- runM env $ Ops.listWorkers schema Nothing Nothing
        let allIds = (WorkerRow -> UUID) -> [WorkerRow] -> [UUID]
forall a b. (a -> b) -> [a] -> [b]
map WorkerRow -> UUID
WR.workerId [WorkerRow]
allRows
        allIds `shouldSatisfy` (liveWid `elem`)
        allIds `shouldSatisfy` (staleWid `elem`)

        -- Queueless live filter at 1s threshold keeps the freshly-heartbeated row only.
        liveOnly <- runM env $ Ops.listWorkers schema Nothing (Just 1)
        let liveIds = (WorkerRow -> UUID) -> [WorkerRow] -> [UUID]
forall a b. (a -> b) -> [a] -> [b]
map WorkerRow -> UUID
WR.workerId [WorkerRow]
liveOnly
        liveIds `shouldSatisfy` (liveWid `elem`)
        liveIds `shouldSatisfy` (staleWid `notElem`)

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"propagates queue pause to local pauseVar via heartbeat reconcile" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        baseConfig :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
1 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))
        let config = WorkerConfig m payload
baseConfig {workerCount = 1, pollInterval = 0.1}

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Running) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config

          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Bool -> m Int64
Ops.setQueuePaused Text
schema Text
table Bool
True
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Paused) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config

          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Bool -> m Int64
Ops.setQueuePaused Text
schema Text
table Bool
False
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Running) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"propagates queue pause via NOTIFY at steady state under one pollInterval" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        baseConfig :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
1 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))
        let config = WorkerConfig m payload
baseConfig {workerCount = 1, pollInterval = 5.0}

            -- Steady-state toggles complete via NOTIFY. The next heartbeat tick
            -- is one pollInterval away.
            timed Bool
paused = do
              let expected :: WorkerState
expected = if Bool
paused then WorkerState
Paused else WorkerState
Running
              start <- IO UTCTime
getCurrentTime
              void $ runM env $ Ops.setQueuePaused schema table paused
              waitUntil 5_000 $ (== expected) <$> getWorkerState config
              elapsed <- (`diffUTCTime` start) <$> getCurrentTime
              elapsed `shouldSatisfy` (< 1.0)

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Running) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> IO Bool
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m Bool
getListenerReady WorkerConfig m payload
config
          Bool -> IO ()
timed Bool
True
          Bool -> IO ()
timed Bool
False
          Bool -> IO ()
timed Bool
True
          Bool -> IO ()
timed Bool
False

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"keeps a pause NOTIFY that lands while a heartbeat reading is in flight" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        baseConfig :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
1 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))
        let config = WorkerConfig m payload
baseConfig {workerCount = 1, pollInterval = 5.0, workerHeartbeatInterval = 1.0}

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Running) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> IO Bool
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m Bool
getListenerReady WorkerConfig m payload
config

          -- Holding the worker's registry row blocks the pool's next heartbeat.
          -- Its reading of the queue's pause state predates the pause below.
          released <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
          let holdRow = env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ m () -> m ()
forall a. m a -> m a
forall (m :: * -> *) a. MonadArbiter m => m a -> m a
withDbTransaction (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
                m (Maybe Bool) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (Maybe Bool) -> m ()) -> m (Maybe Bool) -> m ()
forall a b. (a -> b) -> a -> b
$ Text -> UUID -> m (Maybe Bool)
forall (m :: * -> *).
MonadArbiter m =>
Text -> UUID -> m (Maybe Bool)
Ops.heartbeatWorker Text
schema (WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
config)
                IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
released
          withAsync holdRow $ \Async ()
_ -> do
            Int -> IO ()
threadDelay Int
heartbeatBlockMicros
            IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Bool -> m Int64
Ops.setQueuePaused Text
schema Text
table Bool
True
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Paused) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config
            MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
released ()
            -- The blocked heartbeat completes here, one interval before the next.
            Int -> IO ()
threadDelay Int
heartbeatLandMicros
            WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config IO WorkerState -> WorkerState -> IO ()
forall a. (HasCallStack, Show a, Eq a) => IO a -> a -> IO ()
`shouldReturn` WorkerState
Paused

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"claims immediately on unpause without waiting another poll cycle" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        processedRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)
        let handler :: JobRead payload -> m ()
            handler JobRead payload
_job =
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Int -> IO ()
forall a. Num a => IORef a -> IO ()
bumpRef IORef Int
processedRef
        baseConfig :: WorkerConfig m payload <- transactionalWorkerConfig 1 (mkHandler (noResult handler))
        let config = WorkerConfig m payload
baseConfig {workerCount = 1, pollInterval = 2.0}

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Running) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> IO Bool
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m Bool
getListenerReady WorkerConfig m payload
config

          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Bool -> m Int64
Ops.setQueuePaused Text
schema Text
table Bool
True
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
Paused) (WorkerState -> Bool) -> IO WorkerState -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkerConfig m payload -> IO WorkerState
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m WorkerState
getWorkerState WorkerConfig m payload
config

          IO (Maybe (JobRead payload)) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe (JobRead payload)) -> IO ())
-> IO (Maybe (JobRead payload)) -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"post-resume"))

          start <- IO UTCTime
getCurrentTime
          void $ runM env $ Ops.setQueuePaused schema table False
          waitUntil 10_000 $ (>= 1) <$> readIORef processedRef
          elapsed <- (`diffUTCTime` start) <$> getCurrentTime
          elapsed `shouldSatisfy` (< 3.0)

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"setWorkerPaused only targets the addressed worker" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        baseA :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
1 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))
        baseB :: WorkerConfig m payload <- transactionalWorkerConfig 1 (mkHandler (noResult handler))
        let cfgA = WorkerConfig m payload
baseA {workerCount = 1, pollInterval = 5.0}
            cfgB = WorkerConfig m payload
baseB {workerCount = 1, pollInterval = 5.0}
            widA = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
cfgA
            widB = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
cfgB

        withLinkedAsync (runM env $ runWorkerPool cfgA) $ \Async ()
_ ->
          IO () -> (Async () -> IO ()) -> IO ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withLinkedAsync (env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> m ()
forall payload (m :: * -> *).
(RegistryAdmissionPolicies (RegistryOf m),
 EncodeJobResult (ResultOf m payload), QueueOperation m payload,
 RegistryTables (RegistryOf m)) =>
WorkerConfig m payload -> m ()
runWorkerPool WorkerConfig m payload
cfgB) ((Async () -> IO ()) -> IO ()) -> (Async () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Async ()
_ -> do
            -- Wait on the registry rows. getWorkerState reads only TVars.
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
              rows <- env -> m [WorkerRow] -> IO [WorkerRow]
forall a. env -> m a -> IO a
runM env
env (m [WorkerRow] -> IO [WorkerRow])
-> m [WorkerRow] -> IO [WorkerRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
Ops.listWorkers Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table) Maybe NominalDiffTime
forall a. Maybe a
Nothing
              let ids = (WorkerRow -> UUID) -> [WorkerRow] -> [UUID]
forall a b. (a -> b) -> [a] -> [b]
map WorkerRow -> UUID
WR.workerId [WorkerRow]
rows
              pure (widA `elem` ids && widB `elem` ids)
            -- And on subscription, before the pause NOTIFY is sent.
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> IO Bool
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m Bool
getListenerReady WorkerConfig m payload
cfgA
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> IO Bool
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m Bool
getListenerReady WorkerConfig m payload
cfgB

            start <- IO UTCTime
getCurrentTime
            void $ runM env $ Ops.setWorkerPaused schema widA True
            waitUntil 5_000 $ (== Paused) <$> getWorkerState cfgA
            elapsed <- (`diffUTCTime` start) <$> getCurrentTime
            elapsed `shouldSatisfy` (< 1.0)

            stateB <- getWorkerState cfgB
            stateB `shouldBe` Running

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"setQueuePaused fans out to every worker in the queue" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let handler :: JobRead payload -> m ()
            handler :: JobRead payload -> m ()
handler JobRead payload
_job = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        baseA :: WorkerConfig m payload <- Int
-> JobHandler m payload (ResultOf m payload)
-> IO (WorkerConfig m payload)
forall (n :: * -> *) (m :: * -> *) payload.
(MonadArbiter n, MonadIO m) =>
Int
-> JobHandler n payload (ResultOf n payload)
-> m (WorkerConfig n payload)
transactionalWorkerConfig Int
1 ((JobRead payload -> m (ResultOf m payload))
-> JobHandler m payload (ResultOf m payload)
mkHandler ((JobRead payload -> m ()) -> JobRead payload -> m (Maybe [Text])
forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult JobRead payload -> m ()
handler))
        baseB :: WorkerConfig m payload <- transactionalWorkerConfig 1 (mkHandler (noResult handler))
        let cfgA = WorkerConfig m payload
baseA {workerCount = 1, pollInterval = 5.0}
            cfgB = WorkerConfig m payload
baseB {workerCount = 1, pollInterval = 5.0}
            widA = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
cfgA
            widB = WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
cfgB

        withLinkedAsync (runM env $ runWorkerPool cfgA) $ \Async ()
_ ->
          IO () -> (Async () -> IO ()) -> IO ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withLinkedAsync (env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> m ()
forall payload (m :: * -> *).
(RegistryAdmissionPolicies (RegistryOf m),
 EncodeJobResult (ResultOf m payload), QueueOperation m payload,
 RegistryTables (RegistryOf m)) =>
WorkerConfig m payload -> m ()
runWorkerPool WorkerConfig m payload
cfgB) ((Async () -> IO ()) -> IO ()) -> (Async () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Async ()
_ -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
              rows <- env -> m [WorkerRow] -> IO [WorkerRow]
forall a. env -> m a -> IO a
runM env
env (m [WorkerRow] -> IO [WorkerRow])
-> m [WorkerRow] -> IO [WorkerRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> Maybe NominalDiffTime -> m [WorkerRow]
Ops.listWorkers Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table) Maybe NominalDiffTime
forall a. Maybe a
Nothing
              let ids = (WorkerRow -> UUID) -> [WorkerRow] -> [UUID]
forall a b. (a -> b) -> [a] -> [b]
map WorkerRow -> UUID
WR.workerId [WorkerRow]
rows
              pure (widA `elem` ids && widB `elem` ids)
            -- And on subscription, before the pause NOTIFY is sent.
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> IO Bool
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m Bool
getListenerReady WorkerConfig m payload
cfgA
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> IO Bool
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m Bool
getListenerReady WorkerConfig m payload
cfgB

            start <- IO UTCTime
getCurrentTime
            void $ runM env $ Ops.setQueuePaused schema table True
            waitUntil 5_000 $ (== Paused) <$> getWorkerState cfgA
            waitUntil 5_000 $ (== Paused) <$> getWorkerState cfgB
            elapsed <- (`diffUTCTime` start) <$> getCurrentTime
            elapsed `shouldSatisfy` (< 1.0)

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Force cancel" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"interrupts a long-running handler and removes the job" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        startedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        completedRef <- newIORef False
        let handler :: JobRead payload -> m ()
            handler JobRead payload
_job = do
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
startedRef Bool
True
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay Int
stuckHandlerMicros
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
completedRef Bool
True

        baseConfig :: WorkerConfig m payload <- transactionalWorkerConfig 1 (mkHandler (noResult handler))
        let config = WorkerConfig m payload
baseConfig {workerCount = 1, pollInterval = 0.2}

        Just job <- runM env $ HL.insertJob (defaultJob (mkSimple "long"))

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
startedRef

          start <- IO UTCTime
getCurrentTime
          cancelled <- runM env $ Ops.forceCancelJob schema table (primaryKey job)
          cancelled `shouldBe` 1

          -- Handler should be interrupted well before its 30s sleep finishes.
          waitUntil 5_000 $ do
            mJob <- runM env $ HL.getJobById @payload (primaryKey job)
            pure (isNothing mJob)
          elapsed <- (`diffUTCTime` start) <$> getCurrentTime
          elapsed `shouldSatisfy` (< 3.0)

          -- The handler did not run to completion.
          completed <- readIORef completedRef
          completed `shouldBe` False

          -- The cancel produced no DLQ entry.
          dlqJobs <- runM env $ HL.listDLQJobs 10 0 :: IO [DLQ.DLQJob payload]
          dlqJobs `shouldBe` []

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"interrupts a CPU-bound handler (no DB I/O)" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        -- The handler runs a tight IORef-bumping loop with no blocking I/O. The
        -- test samples the counter twice after the force-cancel. The loop has no
        -- interruptible point. A masked child keeps incrementing.
        startedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        counterRef <- newIORef (0 :: Int)
        let handler :: JobRead payload -> m ()
            handler JobRead payload
_job = do
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
startedRef Bool
True
              let go :: IO ()
go = do
                    IORef Int -> IO ()
forall a. Num a => IORef a -> IO ()
bumpRef IORef Int
counterRef
                    IO ()
go
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
go

        baseConfig :: WorkerConfig m payload <- transactionalWorkerConfig 1 (mkHandler (noResult handler))
        let config = WorkerConfig m payload
baseConfig {workerCount = 1, pollInterval = 0.2}

        Just job <- runM env $ HL.insertJob (defaultJob (mkSimple "cpu"))

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
startedRef
          cancelled <- env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Int64 -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Int64 -> m Int64
Ops.forceCancelJob Text
schema Text
table (JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job)
          cancelled `shouldBe` 1
          -- Let cancellation propagate.
          threadDelay cancelPropagateMicros
          countBefore <- readIORef counterRef
          threadDelay counterWindowMicros
          countAfter <- readIORef counterRef
          -- The counter freezes. A live handler bumps it millions of times in 500ms.
          countAfter `shouldBe` countBefore

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"cancelling one job of a batch interrupts the whole batch handler" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        -- A batch runs in a single handler thread. All its job ids point at the
        -- same async. Targeting one job throws into that thread and tears down
        -- the in-flight batch.
        startedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        completedRef <- newIORef False
        let batchHandler NonEmpty (JobRead payload)
_jobs BatchCallbacks m payload (Maybe [Text])
_cbs = do
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
startedRef Bool
True
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay Int
stuckHandlerMicros
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
completedRef Bool
True

        let jobs =
              [ Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"bc") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"bc-1")
              , Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"bc") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"bc-2")
              ]
        inserted <- runM env $ HL.insertJobsBatch jobs
        let firstId = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey ([JobRead payload] -> JobRead payload
forall a. HasCallStack => [a] -> a
head [JobRead payload]
inserted)

        config :: WorkerConfig m payload <- batchedWorkerConfig 1 10 batchHandler
        threadDelay beforeStartMicros

        withLinkedAsync (runM env $ runWorkerPool config {pollInterval = 0.1}) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
startedRef

          start <- IO UTCTime
getCurrentTime
          -- Cancel only the first job. The whole batch thread unwinds.
          cancelled <- runM env $ Ops.forceCancelJob schema table firstId
          cancelled `shouldBe` 1
          waitUntil 5_000 $ do
            mJob <- runM env $ HL.getJobById @payload firstId
            pure (isNothing mJob)
          elapsed <- (`diffUTCTime` start) <$> getCurrentTime
          elapsed `shouldSatisfy` (< 3.0)

          -- The handler was interrupted.
          readIORef completedRef `shouldReturn` False
          -- No DLQ entries from the cancel.
          dlqJobs <- runM env $ HL.listDLQJobs 10 0 :: IO [DLQ.DLQJob payload]
          dlqJobs `shouldBe` []

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"interrupts a running handler in poll-only mode via the flag" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let pollEnv :: env
pollEnv = env -> env
pollOnly env
env
        -- With no listener, the heartbeat polls cancel_requested_at and throws
        -- into the handler.
        startedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        completedRef <- newIORef False
        let handler :: JobRead payload -> m ()
            handler JobRead payload
_job = do
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
startedRef Bool
True
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay Int
stuckHandlerMicros
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
completedRef Bool
True

        baseConfig :: WorkerConfig m payload <- transactionalWorkerConfig 1 (mkHandler (noResult handler))
        let config =
              WorkerConfig m payload
baseConfig
                { workerCount = 1
                , pollInterval = 0.2
                , jobHeartbeatInterval = 0.3
                , visibilityTimeout = 3
                }

        Just job <- runM env $ HL.insertJob (defaultJob (mkSimple "poll-cancel"))

        withLinkedAsync (runM pollEnv $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
startedRef

          start <- IO UTCTime
getCurrentTime
          cancelled <- runM env $ Ops.forceCancelJob schema table (primaryKey job)
          cancelled `shouldBe` 1

          waitUntil 5_000 $ do
            mJob <- runM env $ HL.getJobById @payload (primaryKey job)
            pure (isNothing mJob)
          elapsed <- (`diffUTCTime` start) <$> getCurrentTime
          elapsed `shouldSatisfy` (< 3.0)

          readIORef completedRef `shouldReturn` False
          dlqJobs <- runM env $ HL.listDLQJobs 10 0 :: IO [DLQ.DLQJob payload]
          dlqJobs `shouldBe` []

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"flags a job that is claimed concurrently with the force-cancel" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        Just job <- env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"concurrent-claim"))
        let jid = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job
            claimSql =
              Text -> Query
toQuery (Text -> Query) -> Text -> Query
forall a b. (a -> b) -> a -> b
$
                Text
"UPDATE "
                  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
Schema.jobQueueTable Text
schema Text
table
                  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" SET claimed_by = '00000000-0000-0000-0000-000000000abc'::uuid"
                  Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
", not_visible_until = NOW() + interval '60 second', attempts = attempts + 1, claim_seq = claim_seq + 1 WHERE id = ?"

        cancelledCount <- withConn connStr $ \Connection
connB -> do
          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
PG.execute_ Connection
connB Query
"BEGIN"
          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> Only Int64 -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
PG.execute Connection
connB Query
claimSql (Int64 -> Only Int64
forall a. a -> Only a
Only Int64
jid)
          IO Int64 -> (Async Int64 -> IO Int64) -> IO Int64
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withAsync (env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Int64 -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Int64 -> m Int64
Ops.forceCancelJob Text
schema Text
table Int64
jid) ((Async Int64 -> IO Int64) -> IO Int64)
-> (Async Int64 -> IO Int64) -> IO Int64
forall a b. (a -> b) -> a -> b
$ \Async Int64
cancelAsync -> do
            Int -> IO ()
threadDelay Int
lockBlockMicros
            IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
PG.execute_ Connection
connB Query
"COMMIT"
            Async Int64 -> IO Int64
forall (m :: * -> *) a. MonadIO m => Async a -> m a
Async.wait Async Int64
cancelAsync

        cancelledCount `shouldBe` 1
        [Only flagged] <-
          withConn connStr $ \Connection
conn ->
            Connection -> Query -> Only Int64 -> IO [Only Bool]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
PG.query
              Connection
conn
              ( Text -> Query
toQuery (Text -> Query) -> Text -> Query
forall a b. (a -> b) -> a -> b
$
                  Text
"SELECT cancel_requested_at IS NOT NULL FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
Schema.jobQueueTable Text
schema Text
table Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" WHERE id = ?"
              )
              (Int64 -> Only Int64
forall a. a -> Only a
Only Int64
jid)
              :: IO [Only Bool]
        flagged `shouldBe` True

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"notifies the worker when force-cancel deletes its lease-lapsed claimed job" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        wid <- IO UUID
UUID.nextRandom
        Just job <- runM env $ HL.insertJob (defaultJob (mkSimple "lapsed-cancel"))
        let jid = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job
        claimed <- runM env (HL.claimNextVisibleJobsAs 1 60 wid) :: IO [JobRead payload]
        length claimed `shouldBe` 1

        updateJob connStr schema table "not_visible_until = NOW() - interval '1 second'" jid

        mNotif <- withConn connStr $ \Connection
lconn -> do
          let chan :: Text
chan = Text -> Text -> Text
Schema.cancelNotifyChannel Text
schema Text
table
          Connection -> Text -> IO ()
execute_ Connection
lconn (Text
"LISTEN \"" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
chan Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\"")

          cancelled <- env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Int64 -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Int64 -> m Int64
Ops.forceCancelJob Text
schema Text
table Int64
jid
          cancelled `shouldBe` 1

          runM env (HL.getJobById @payload jid)
            >>= (`shouldSatisfy` isNothing)

          timeout notifyTimeoutMicros (getNotification lconn)
        case mNotif of
          Maybe Notification
Nothing -> HasCallStack => FilePath -> IO ()
FilePath -> IO ()
expectationFailure FilePath
"expected a cancel NOTIFY for the deleted lease-lapsed job"
          Just Notification
notif -> Notification -> ByteString
notificationData Notification
notif ByteString -> (ByteString -> Bool) -> IO ()
forall a. (HasCallStack, Show a) => a -> (a -> Bool) -> IO ()
`shouldSatisfy` ByteString -> ByteString -> Bool
BSC.isInfixOf (FilePath -> ByteString
BSC.pack (Int64 -> FilePath
forall a. Show a => a -> FilePath
show Int64
jid))

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"does not deadlock against a concurrent ack of the last child" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        Right (parent :| [child]) <-
          env
-> m (Either Text (NonEmpty (JobRead payload)))
-> IO (Either Text (NonEmpty (JobRead payload)))
forall a. env -> m a -> IO a
runM env
env
            (m (Either Text (NonEmpty (JobRead payload)))
 -> IO (Either Text (NonEmpty (JobRead payload))))
-> m (Either Text (NonEmpty (JobRead payload)))
-> IO (Either Text (NonEmpty (JobRead payload)))
forall a b. (a -> b) -> a -> b
$ JobTree payload -> m (Either Text (NonEmpty (JobRead payload)))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobTree payload -> m (Either Text (NonEmpty (JobRead payload)))
HL.insertJobTree
            (JobTree payload -> m (Either Text (NonEmpty (JobRead payload))))
-> JobTree payload -> m (Either Text (NonEmpty (JobRead payload)))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> NonEmpty (JobTree payload) -> JobTree payload
forall payload.
JobWrite payload -> NonEmpty (JobTree payload) -> JobTree payload
JT.rollup
              (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"dl-parent"))
              (JobWrite payload -> JobTree payload
forall payload. JobWrite payload -> JobTree payload
JT.leaf (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"dl-child")) JobTree payload -> [JobTree payload] -> NonEmpty (JobTree payload)
forall a. a -> [a] -> NonEmpty a
:| [])
        let pid = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
parent
            cid = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
child

        (efc, epA) <- withConn connStr $ \Connection
connA -> do
          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
PG.execute_ Connection
connA Query
"BEGIN"
          IO (Either SomeException [Only Int64]) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either SomeException [Only Int64]) -> IO ())
-> IO (Either SomeException [Only Int64]) -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Int64 -> IO (Either SomeException [Only Int64])
lockRow Connection
connA Int64
cid
          IO (Either SomeException Int64)
-> (Async (Either SomeException Int64)
    -> IO
         (Either SomeException Int64, Either SomeException [Only Int64]))
-> IO
     (Either SomeException Int64, Either SomeException [Only Int64])
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withAsync (IO Int64 -> IO (Either SomeException Int64)
forall e a. Exception e => IO a -> IO (Either e a)
try (env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Int64 -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Int64 -> m Int64
Ops.forceCancelJob Text
schema Text
table Int64
pid) :: IO (Either SomeException Int64)) ((Async (Either SomeException Int64)
  -> IO
       (Either SomeException Int64, Either SomeException [Only Int64]))
 -> IO
      (Either SomeException Int64, Either SomeException [Only Int64]))
-> (Async (Either SomeException Int64)
    -> IO
         (Either SomeException Int64, Either SomeException [Only Int64]))
-> IO
     (Either SomeException Int64, Either SomeException [Only Int64])
forall a b. (a -> b) -> a -> b
$ \Async (Either SomeException Int64)
cancelAsync -> do
            Int -> IO ()
threadDelay Int
lockBlockMicros
            epA <- Connection -> Int64 -> IO (Either SomeException [Only Int64])
lockRow Connection
connA Int64
pid
            void (try (PG.execute_ connA "COMMIT") :: IO (Either SomeException Int64))
            efc <- Async.wait cancelAsync
            pure (efc, epA)

        efc `shouldSatisfy` isRight
        epA `shouldSatisfy` isRight

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"settles a failed batch children-first, so a tree lock cannot deadlock it" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        startedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        goVar <- newEmptyMVar
        let batchHandler NonEmpty (JobRead payload)
_jobs BatchCallbacks m payload (Maybe [Text])
_cbs = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
              IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
startedRef Bool
True
              MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
goVar
              IOError -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (FilePath -> IOError
userError FilePath
"dlb-boom")
        let jobs =
              [ Maybe Int32 -> JobWrite payload -> JobWrite payload
forall payload. Maybe Int32 -> JobWrite payload -> JobWrite payload
setMaxAttempts (Int32 -> Maybe Int32
forall a. a -> Maybe a
Just Int32
1) (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"dlb") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"dlb-1")
              , Maybe Int32 -> JobWrite payload -> JobWrite payload
forall payload. Maybe Int32 -> JobWrite payload -> JobWrite payload
setMaxAttempts (Int32 -> Maybe Int32
forall a. a -> Maybe a
Just Int32
1) (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"dlb") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"dlb-2")
              ]
        inserted <- runM env $ HL.insertJobsBatch jobs
        let ids = (JobRead payload -> Int64) -> [JobRead payload] -> [Int64]
forall a b. (a -> b) -> [a] -> [b]
map JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey [JobRead payload]
inserted
        config :: WorkerConfig m payload <- batchedWorkerConfig 1 10 batchHandler

        withLinkedAsync
          (runM env $ runWorkerPool config {pollInterval = 0.05, jobHeartbeatInterval = 30, visibilityTimeout = 60})
          $ \Async ()
_ -> do
            HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
startedRef

            eLo <- ByteString
-> (Connection -> IO (Either SomeException [Only Int64]))
-> IO (Either SomeException [Only Int64])
forall a. ByteString -> (Connection -> IO a) -> IO a
withConn ByteString
connStr ((Connection -> IO (Either SomeException [Only Int64]))
 -> IO (Either SomeException [Only Int64]))
-> (Connection -> IO (Either SomeException [Only Int64]))
-> IO (Either SomeException [Only Int64])
forall a b. (a -> b) -> a -> b
$ \Connection
connA -> do
              IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
PG.execute_ Connection
connA Query
"BEGIN"
              -- Hold the batch's higher id, the row a force-cancel over the tree takes first.
              IO (Either SomeException [Only Int64]) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either SomeException [Only Int64]) -> IO ())
-> IO (Either SomeException [Only Int64]) -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Int64 -> IO (Either SomeException [Only Int64])
lockRow Connection
connA ([Int64] -> Int64
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Int64]
ids)
              MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
goVar ()
              Int -> IO ()
threadDelay Int
lockBlockMicros
              -- The failure transaction does not yet hold the lower id.
              eLo <- Connection -> Int64 -> IO (Either SomeException [Only Int64])
lockRow Connection
connA ([Int64] -> Int64
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Int64]
ids)
              void (try (PG.execute_ connA "COMMIT") :: IO (Either SomeException Int64))
              pure eLo

            eLo `shouldSatisfy` isRight
            waitUntil 10_000 $ do
              dlqJobs <- runM env $ HL.listDLQJobs 10 0 :: IO [DLQ.DLQJob payload]
              pure (length dlqJobs == 2)

      [Bool] -> (Bool -> SpecWith env) -> SpecWith env
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ [Bool
False, Bool
True] ((Bool -> SpecWith env) -> SpecWith env)
-> (Bool -> SpecWith env) -> SpecWith env
forall a b. (a -> b) -> a -> b
$ \Bool
nested ->
        FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it
          (FilePath
"takes branch-cancel advisory locks before settlement row locks (" FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> (if Bool
nested then FilePath
"nested" else FilePath
"standalone") FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> FilePath
")")
          ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
            (jid, lockId) <-
              if Bool
nested
                then do
                  Right (root :| [_, child]) <-
                    env
-> m (Either Text (NonEmpty (JobRead payload)))
-> IO (Either Text (NonEmpty (JobRead payload)))
forall a. env -> m a -> IO a
runM env
env
                      (m (Either Text (NonEmpty (JobRead payload)))
 -> IO (Either Text (NonEmpty (JobRead payload))))
-> m (Either Text (NonEmpty (JobRead payload)))
-> IO (Either Text (NonEmpty (JobRead payload)))
forall a b. (a -> b) -> a -> b
$ JobTree payload -> m (Either Text (NonEmpty (JobRead payload)))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobTree payload -> m (Either Text (NonEmpty (JobRead payload)))
HL.insertJobTree
                      (JobTree payload -> m (Either Text (NonEmpty (JobRead payload))))
-> JobTree payload -> m (Either Text (NonEmpty (JobRead payload)))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> NonEmpty (JobTree payload) -> JobTree payload
forall payload.
JobWrite payload -> NonEmpty (JobTree payload) -> JobTree payload
JT.rollup
                        (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"branch-grandparent"))
                        ( JobWrite payload -> NonEmpty (JobTree payload) -> JobTree payload
forall payload.
JobWrite payload -> NonEmpty (JobTree payload) -> JobTree payload
JT.rollup
                            (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"branch-parent"))
                            (JobWrite payload -> JobTree payload
forall payload. JobWrite payload -> JobTree payload
JT.leaf (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"branch-child")) JobTree payload -> [JobTree payload] -> NonEmpty (JobTree payload)
forall a. a -> [a] -> NonEmpty a
:| [])
                            JobTree payload -> [JobTree payload] -> NonEmpty (JobTree payload)
forall a. a -> [a] -> NonEmpty a
:| []
                        )
                  pure (primaryKey child, primaryKey root)
                else do
                  Just job <- env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"branch-lock-order"))
                  pure (primaryKey job, primaryKey job)
            let handler p
_jobs p
_callbacks = Text -> m a
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwBranchCancel Text
"cancel branch"
            base :: WorkerConfig m payload <- batchedWorkerConfig 1 1 handler
            let config = WorkerConfig m payload
base {pollInterval = 0.05, logConfig = silentLogConfig}
            withConn connStr $ \Connection
conn -> do
              IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
PG.execute_ Connection
conn Query
"BEGIN"
              (IO () -> IO () -> IO ()) -> IO () -> IO () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
finally (IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
PG.execute_ Connection
conn Query
"ROLLBACK") (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                -- An external cancel takes this advisory lock before touching the row.
                IO [Only Text] -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void
                  ( Connection -> Query -> (Text, Int64) -> IO [Only Text]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
PG.query
                      Connection
conn
                      Query
"SELECT pg_advisory_xact_lock(hashtextextended(?, ?))::text"
                      (Text
schema Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
table, Int64
lockId)
                      :: IO [Only Text]
                  )
                IO () -> (Async () -> IO ()) -> IO ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withLinkedAsync (env -> m () -> IO ()
forall a. env -> m a -> IO a
runM (env -> env
pollOnly env
env) (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ WorkerConfig m payload -> m ()
forall payload (m :: * -> *).
(RegistryAdmissionPolicies (RegistryOf m),
 EncodeJobResult (ResultOf m payload), QueueOperation m payload,
 RegistryTables (RegistryOf m)) =>
WorkerConfig m payload -> m ()
runWorkerPool WorkerConfig m payload
config) ((Async () -> IO ()) -> IO ()) -> (Async () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Async ()
_ -> do
                  -- Wait until settlement requests the advisory lock we hold.
                  HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                    blocked <-
                      Connection -> Query -> IO [Only Bool]
forall r. FromRow r => Connection -> Query -> IO [r]
PG.query_
                        Connection
conn
                        Query
"SELECT EXISTS (SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND NOT granted AND pg_backend_pid() = ANY(pg_blocking_pids(pid)))"
                        :: IO [Only Bool]
                    pure (blocked == [Only True])
                  -- Settlement must not already hold this row. NOWAIT makes an
                  -- inverted order fail deterministically instead of relying on a victim.
                  result <-
                    IO [Only Int64] -> IO (Either SomeException [Only Int64])
forall e a. Exception e => IO a -> IO (Either e a)
try
                      ( Connection -> Query -> Only Int64 -> IO [Only Int64]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
PG.query
                          Connection
conn
                          (Text -> Query
toQuery (Text -> Query) -> Text -> Query
forall a b. (a -> b) -> a -> b
$ Text
"SELECT id FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
Schema.jobQueueTable Text
schema Text
table Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" WHERE id = ? FOR UPDATE NOWAIT")
                          (Int64 -> Only Int64
forall a. a -> Only a
Only Int64
jid)
                          :: IO [Only Int64]
                      )
                      :: IO (Either SomeException [Only Int64])
                  void $ PG.execute_ conn "ROLLBACK"
                  -- Keep the cleanup transaction valid after releasing the held locks.
                  void $ PG.execute_ conn "BEGIN"
                  result `shouldSatisfy` isRight
                  waitUntil 5_000 $ isNothing <$> runM env (HL.getJobById @payload jid)

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"deletes a nacked job outright on force-cancel" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        -- A nack releases the claim, so the force-cancel deletes the row outright.
        nackedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        let jobs =
              [ Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"fcn") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"fcn-1")
              , Maybe Text -> JobWrite payload -> JobWrite payload
forall payload. Maybe Text -> JobWrite payload -> JobWrite payload
setGroupKey (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"fcn") (JobWrite payload -> JobWrite payload)
-> JobWrite payload -> JobWrite payload
forall a b. (a -> b) -> a -> b
$ payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"fcn-2")
              ]
        inserted <- runM env $ HL.insertJobsBatch jobs
        let firstId = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey ([JobRead payload] -> JobRead payload
forall a. HasCallStack => [a] -> a
head [JobRead payload]
inserted)
            batchHandler NonEmpty (JobRead payload)
batch BatchCallbacks m payload (Maybe [Text])
cbs = do
              (JobRead payload -> m ()) -> NonEmpty (JobRead payload) -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (\JobRead payload
job -> Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64
firstId) (BatchCallbacks m payload (Maybe [Text]) -> JobRead payload -> m ()
forall (m :: * -> *) payload result.
BatchCallbacks m payload result -> JobRead payload -> m ()
nack BatchCallbacks m payload (Maybe [Text])
cbs JobRead payload
job)) NonEmpty (JobRead payload)
batch
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
nackedRef Bool
True
              IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay Int
stuckHandlerMicros

        baseConfig :: WorkerConfig m payload <- batchedWorkerConfig 1 10 batchHandler
        let config = WorkerConfig m payload
baseConfig {pollInterval = 0.1, jobHeartbeatInterval = 0.3, visibilityTimeout = 60}

        withLinkedAsync (runM env $ runWorkerPool config) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
nackedRef

          env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (Text -> Text -> Int64 -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Int64 -> m Int64
Ops.forceCancelJob Text
schema Text
table Int64
firstId) IO Int64 -> Int64 -> IO ()
forall a. (HasCallStack, Show a, Eq a) => IO a -> a -> IO ()
`shouldReturn` Int64
1

          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
10_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            mJob <- env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ forall payload (m :: * -> *).
QueueOperation m payload =>
Int64 -> m (Maybe (JobRead payload))
HL.getJobById @payload Int64
firstId
            pure (isNothing mJob)

          dlqJobs <- env -> m [DLQJob payload] -> IO [DLQJob payload]
forall a. env -> m a -> IO a
runM env
env (m [DLQJob payload] -> IO [DLQJob payload])
-> m [DLQJob payload] -> IO [DLQJob payload]
forall a b. (a -> b) -> a -> b
$ Int -> Int -> m [DLQJob payload]
forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> m [DLQJob payload]
HL.listDLQJobs Int
10 Int
0 :: IO [DLQ.DLQJob payload]
          dlqJobs `shouldBe` []

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"refuses a stale worker's ack after a reclaim and nack restored attempts" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        -- A nack restores the attempt it consumed. The attempts value repeats across claims.
        staleWorker <- IO UUID
UUID.nextRandom
        holdingWorker <- UUID.nextRandom
        Just job <- runM env $ HL.insertJob (defaultJob (mkSimple "aba"))
        let jid = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
job

        [stale] <- runM env (HL.claimNextVisibleJobsAs 1 60 staleWorker) :: IO [JobRead payload]
        primaryKey stale `shouldBe` jid
        updateJob connStr schema table "not_visible_until = NOW() - interval '1 second'" jid

        [held] <- runM env (HL.claimNextVisibleJobsAs 1 60 holdingWorker) :: IO [JobRead payload]
        primaryKey held `shouldBe` jid
        runM env (HL.nackJob held) `shouldReturn` 1

        -- The nack put attempts back to what the stale worker recorded. An
        -- attempts-keyed predicate would match here.
        reread <- runM env $ HL.getJobById @payload jid
        fmap attempts reread `shouldBe` Just (attempts stale)

        -- Every finalize the stale worker can still issue matches no row.
        runM env (HL.setVisibilityTimeoutBatch 60 [stale])
          `shouldReturn` [HL.JobReclaimed jid (claimSeq stale) (claimSeq held)]
        runM env (HL.nackJob stale) `shouldReturn` 0
        runM env (HL.ackJob stale) `shouldReturn` 0
        runM env (HL.getJobById @payload jid) >>= (`shouldSatisfy` isJust)

        -- The holding worker still owns it and can finish.
        runM env (HL.ackJob held) `shouldReturn` 1

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"does not deadlock a tree cancel against a concurrent lock walk" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        Right (parent :| [child]) <-
          env
-> m (Either Text (NonEmpty (JobRead payload)))
-> IO (Either Text (NonEmpty (JobRead payload)))
forall a. env -> m a -> IO a
runM env
env
            (m (Either Text (NonEmpty (JobRead payload)))
 -> IO (Either Text (NonEmpty (JobRead payload))))
-> m (Either Text (NonEmpty (JobRead payload)))
-> IO (Either Text (NonEmpty (JobRead payload)))
forall a b. (a -> b) -> a -> b
$ JobTree payload -> m (Either Text (NonEmpty (JobRead payload)))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobTree payload -> m (Either Text (NonEmpty (JobRead payload)))
HL.insertJobTree
            (JobTree payload -> m (Either Text (NonEmpty (JobRead payload))))
-> JobTree payload -> m (Either Text (NonEmpty (JobRead payload)))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> NonEmpty (JobTree payload) -> JobTree payload
forall payload.
JobWrite payload -> NonEmpty (JobTree payload) -> JobTree payload
JT.rollup
              (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"tc-parent"))
              (JobWrite payload -> JobTree payload
forall payload. JobWrite payload -> JobTree payload
JT.leaf (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"tc-child")) JobTree payload -> [JobTree payload] -> NonEmpty (JobTree payload)
forall a. a -> [a] -> NonEmpty a
:| [])
        let pid = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
parent
            cid = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
child

        (etc, epA) <- withConn connStr $ \Connection
connA -> do
          IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
PG.execute_ Connection
connA Query
"BEGIN"
          -- Hold the child, the row a tree cancel takes first.
          IO (Either SomeException [Only Int64]) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either SomeException [Only Int64]) -> IO ())
-> IO (Either SomeException [Only Int64]) -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Int64 -> IO (Either SomeException [Only Int64])
lockRow Connection
connA Int64
cid
          IO (Either SomeException [Int64])
-> (Async (Either SomeException [Int64])
    -> IO
         (Either SomeException [Int64], Either SomeException [Only Int64]))
-> IO
     (Either SomeException [Int64], Either SomeException [Only Int64])
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withAsync (IO [Int64] -> IO (Either SomeException [Int64])
forall e a. Exception e => IO a -> IO (Either e a)
try (env -> m [Int64] -> IO [Int64]
forall a. env -> m a -> IO a
runM env
env (m [Int64] -> IO [Int64]) -> m [Int64] -> IO [Int64]
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Int64 -> m [Int64]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> Int64 -> m [Int64]
Ops.cancelJobTree Text
schema Text
table Int64
cid) :: IO (Either SomeException [Int64])) ((Async (Either SomeException [Int64])
  -> IO
       (Either SomeException [Int64], Either SomeException [Only Int64]))
 -> IO
      (Either SomeException [Int64], Either SomeException [Only Int64]))
-> (Async (Either SomeException [Int64])
    -> IO
         (Either SomeException [Int64], Either SomeException [Only Int64]))
-> IO
     (Either SomeException [Int64], Either SomeException [Only Int64])
forall a b. (a -> b) -> a -> b
$ \Async (Either SomeException [Int64])
cancelAsync -> do
            Int -> IO ()
threadDelay Int
lockBlockMicros
            epA <- Connection -> Int64 -> IO (Either SomeException [Only Int64])
lockRow Connection
connA Int64
pid
            void (try (PG.execute_ connA "COMMIT") :: IO (Either SomeException Int64))
            etc <- Async.wait cancelAsync
            pure (etc, epA)

        etc `shouldSatisfy` isRight
        epA `shouldSatisfy` isRight

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Sweeper" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"deletes a stale unpaused worker row" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        wid <- IO UUID -> IO UUID
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
UUID.nextRandom
        void
          $ runM env
          $ Ops.registerWorker schema wid table Nothing (Just testWorkerThreads) briefStaleAfter Nothing
        threadDelay pastStaleMicros
        swept <- runM env $ Ops.sweepStaleWorkers schema
        swept `shouldSatisfy` (>= 1)
        rows <- runM env $ Ops.listWorkers schema (Just table) Nothing
        map WR.workerId rows `shouldNotContain` [wid]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"deletes a stale paused worker row" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        wid <- IO UUID -> IO UUID
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
UUID.nextRandom
        void
          $ runM env
          $ Ops.registerWorker schema wid table Nothing (Just testWorkerThreads) briefStaleAfter Nothing
        void $ runM env $ Ops.setWorkerPaused schema wid True
        threadDelay pastStaleMicros
        swept <- runM env $ Ops.sweepStaleWorkers schema
        swept `shouldSatisfy` (>= 1)
        rows <- runM env $ Ops.listWorkers schema (Just table) Nothing
        map WR.workerId rows `shouldNotContain` [wid]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"deletes a stale shutting-down worker row" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        wid <- IO UUID -> IO UUID
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
UUID.nextRandom
        void
          $ runM env
          $ Ops.registerWorker schema wid table Nothing (Just testWorkerThreads) briefStaleAfter Nothing
        void $ runM env $ Ops.markWorkerShuttingDown schema wid
        threadDelay pastStaleMicros
        swept <- runM env $ Ops.sweepStaleWorkers schema
        swept `shouldSatisfy` (>= 1)
        rows <- runM env $ Ops.listWorkers schema (Just table) Nothing
        map WR.workerId rows `shouldNotContain` [wid]

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"preserves paused state across re-registration" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        wid <- IO UUID -> IO UUID
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
UUID.nextRandom
        void
          $ runM env
          $ Ops.registerWorker schema wid table Nothing (Just testWorkerThreads) farStaleAfter Nothing
        void $ runM env $ Ops.setWorkerPaused schema wid True
        -- Re-register with different metadata to confirm upsert touches the row.
        void
          $ runM env
          $ Ops.registerWorker schema wid table (Just "fresh-host") (Just testWorkerThreads) farStaleAfter Nothing
        rows <- runM env $ Ops.listWorkers schema (Just table) Nothing
        soleWorkerRow wid rows $ \WorkerRow
row -> do
          WorkerRow -> Bool
WR.paused WorkerRow
row Bool -> Bool -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Bool
True
          WorkerRow -> Maybe Text
WR.hostName WorkerRow
row Maybe Text -> Maybe Text -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"fresh-host"

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Worker health" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"reports a freshly-registered worker as live" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        wid <- IO UUID -> IO UUID
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
UUID.nextRandom
        void
          $ runM env
          $ Ops.registerWorker schema wid table Nothing (Just testWorkerThreads) farStaleAfter Nothing
        rows <- runM env $ Ops.listWorkers schema (Just table) Nothing
        soleWorkerRow wid rows $ \WorkerRow
row -> WorkerRow -> WorkerHealth
WR.health WorkerRow
row WorkerHealth -> WorkerHealth -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` WorkerHealth
WR.Live

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"reports a worker past its stale threshold as stale" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        wid <- IO UUID -> IO UUID
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
UUID.nextRandom
        void
          $ runM env
          $ Ops.registerWorker schema wid table Nothing (Just testWorkerThreads) briefStaleAfter Nothing
        threadDelay pastStaleMicros
        rows <- runM env $ Ops.listWorkers schema (Just table) Nothing
        soleWorkerRow wid rows $ \WorkerRow
row -> WorkerRow -> WorkerHealth
WR.health WorkerRow
row WorkerHealth -> WorkerHealth -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` WorkerHealth
WR.Stale

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"reports a fresh shutting-down worker as draining" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        wid <- IO UUID -> IO UUID
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UUID
UUID.nextRandom
        void
          $ runM env
          $ Ops.registerWorker schema wid table Nothing (Just testWorkerThreads) farStaleAfter Nothing
        void $ runM env $ Ops.markWorkerShuttingDown schema wid
        rows <- runM env $ Ops.listWorkers schema (Just table) Nothing
        soleWorkerRow wid rows $ \WorkerRow
row -> WorkerRow -> WorkerHealth
WR.health WorkerRow
row WorkerHealth -> WorkerHealth -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` WorkerHealth
WR.Draining

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Cron schedule defaults" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"leaves an unchanged schedule's updated_at alone" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let upsert :: Text -> m Int64
upsert Text
expr = Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
Ops.upsertCronDefault Text
schema Text
"cron-steady" Text
table Text
expr Text
"AllowOverlap" Maybe Text
forall a. Maybe a
Nothing Bool
True
            readBack :: IO (Maybe CronScheduleRow)
readBack = env -> m (Maybe CronScheduleRow) -> IO (Maybe CronScheduleRow)
forall a. env -> m a -> IO a
runM env
env (m (Maybe CronScheduleRow) -> IO (Maybe CronScheduleRow))
-> m (Maybe CronScheduleRow) -> IO (Maybe CronScheduleRow)
forall a b. (a -> b) -> a -> b
$ Text -> Text -> m (Maybe CronScheduleRow)
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> m (Maybe CronScheduleRow)
Ops.getCronScheduleByName Text
schema Text
"cron-steady"
        IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env (Text -> m Int64
upsert Text
"* * * * *")
        Just first <- IO (Maybe CronScheduleRow)
readBack
        void $ runM env (upsert "* * * * *")
        Just second <- readBack
        CS.updatedAt second `shouldBe` CS.updatedAt first
        void $ runM env (upsert "*/5 * * * *")
        Just third <- readBack
        CS.updatedAt third `shouldSatisfy` (> CS.updatedAt first)

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"registers a schedule with initiallyEnabled False as disabled" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void
          (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env
          (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
Ops.upsertCronDefault Text
schema Text
"cron-suspended" Text
table Text
"* * * * *" Text
"AllowOverlap" Maybe Text
forall a. Maybe a
Nothing Bool
False
        Just CS.CronScheduleRow {CS.enabled = isEnabled} <-
          env -> m (Maybe CronScheduleRow) -> IO (Maybe CronScheduleRow)
forall a. env -> m a -> IO a
runM env
env (m (Maybe CronScheduleRow) -> IO (Maybe CronScheduleRow))
-> m (Maybe CronScheduleRow) -> IO (Maybe CronScheduleRow)
forall a b. (a -> b) -> a -> b
$ Text -> Text -> m (Maybe CronScheduleRow)
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> m (Maybe CronScheduleRow)
Ops.getCronScheduleByName Text
schema Text
"cron-suspended"
        isEnabled `shouldBe` False

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"leaves an existing row's enabled state alone on re-upsert" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let readBack :: IO (Maybe CronScheduleRow)
readBack = env -> m (Maybe CronScheduleRow) -> IO (Maybe CronScheduleRow)
forall a. env -> m a -> IO a
runM env
env (m (Maybe CronScheduleRow) -> IO (Maybe CronScheduleRow))
-> m (Maybe CronScheduleRow) -> IO (Maybe CronScheduleRow)
forall a b. (a -> b) -> a -> b
$ Text -> Text -> m (Maybe CronScheduleRow)
forall (m :: * -> *).
MonadArbiter m =>
Text -> Text -> m (Maybe CronScheduleRow)
Ops.getCronScheduleByName Text
schema Text
"cron-resurrected"
        IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void
          (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env
          (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
Ops.upsertCronDefault Text
schema Text
"cron-resurrected" Text
table Text
"* * * * *" Text
"AllowOverlap" Maybe Text
forall a. Maybe a
Nothing Bool
True
        IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void
          (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ env -> m Int64 -> IO Int64
forall a. env -> m a -> IO a
runM env
env
          (m Int64 -> IO Int64) -> m Int64 -> IO Int64
forall a b. (a -> b) -> a -> b
$ Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
Ops.upsertCronDefault Text
schema Text
"cron-resurrected" Text
table Text
"*/5 * * * *" Text
"AllowOverlap" Maybe Text
forall a. Maybe a
Nothing Bool
False
        Just CS.CronScheduleRow {CS.enabled = isEnabled} <- IO (Maybe CronScheduleRow)
readBack
        isEnabled `shouldBe` True

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Cron queue filter" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"filters cron schedules by queue" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        let otherQueue :: Text
otherQueue = Text
"other_queue"
        env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
env (m () -> IO ()) -> m () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
          m Int64 -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Int64 -> m ()) -> m Int64 -> m ()
forall a b. (a -> b) -> a -> b
$ Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
Ops.upsertCronDefault Text
schema Text
"cron-here" Text
table Text
"* * * * *" Text
"AllowOverlap" Maybe Text
forall a. Maybe a
Nothing Bool
True
          m Int64 -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Int64 -> m ()) -> m Int64 -> m ()
forall a b. (a -> b) -> a -> b
$ Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
forall (m :: * -> *).
MonadArbiter m =>
Text
-> Text -> Text -> Text -> Text -> Maybe Text -> Bool -> m Int64
Ops.upsertCronDefault Text
schema Text
"cron-elsewhere" Text
otherQueue Text
"* * * * *" Text
"AllowOverlap" Maybe Text
forall a. Maybe a
Nothing Bool
True

        hereOnly <- env -> m [CronScheduleRow] -> IO [CronScheduleRow]
forall a. env -> m a -> IO a
runM env
env (m [CronScheduleRow] -> IO [CronScheduleRow])
-> m [CronScheduleRow] -> IO [CronScheduleRow]
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text -> m [CronScheduleRow]
forall (m :: * -> *).
MonadArbiter m =>
Text -> Maybe Text -> m [CronScheduleRow]
Ops.listCronSchedules Text
schema (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
table)
        map CS.name hereOnly `shouldContain` ["cron-here"]
        map CS.name hereOnly `shouldNotContain` ["cron-elsewhere"]

        elsewhereOnly <- runM env $ Ops.listCronSchedules schema (Just otherQueue)
        map CS.name elsewhereOnly `shouldContain` ["cron-elsewhere"]
        map CS.name elsewhereOnly `shouldNotContain` ["cron-here"]

        all_ <- runM env $ Ops.listCronSchedules schema Nothing
        map CS.name all_ `shouldSatisfy` (\[Text]
names -> Text
"cron-here" Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
names Bool -> Bool -> Bool
&& Text
"cron-elsewhere" Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
names)

    FilePath -> SpecWith env -> SpecWith env
forall a. HasCallStack => FilePath -> SpecWith a -> SpecWith a
describe FilePath
"Undecodable payload" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"dead-letters a row its payload type rejects and runs the rest of the claim" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        completedRef <- [payload] -> IO (IORef [payload])
forall a. a -> IO (IORef a)
newIORef []
        config :: WorkerConfig m payload <-
          transactionalWorkerConfig 10 $
            mkHandler (noResult (\JobRead payload
job -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef [payload] -> ([payload] -> ([payload], ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef [payload]
completedRef (\[payload]
seen -> (JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> [payload] -> [payload]
forall a. a -> [a] -> [a]
: [payload]
seen, ()))))
        Just poison <- runM env $ HL.insertJob (defaultJob (mkSimple "poison"))
        void $ runM env $ HL.insertJob (defaultJob (mkSimple "sibling"))
        updateJob connStr schema table "payload = '{\"bogus\": 1}'" (primaryKey poison)

        withAsync (runM env $ runWorkerPool config {workerCount = 2, pollInterval = 0.1, visibilityTimeout = 60}) $ \Async ()
_ -> do
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ ([payload] -> [payload] -> Bool
forall a. Eq a => a -> a -> Bool
== [Text -> payload
mkSimple Text
"sibling"]) ([payload] -> Bool) -> IO [payload] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [payload] -> IO [payload]
forall a. IORef a -> IO a
readIORef IORef [payload]
completedRef
          dlq <-
            ByteString
-> (Connection -> IO [(Int64, Text)]) -> IO [(Int64, Text)]
forall a. ByteString -> (Connection -> IO a) -> IO a
withConn ByteString
connStr ((Connection -> IO [(Int64, Text)]) -> IO [(Int64, Text)])
-> (Connection -> IO [(Int64, Text)]) -> IO [(Int64, Text)]
forall a b. (a -> b) -> a -> b
$ \Connection
conn ->
              Connection -> Query -> IO [(Int64, Text)]
forall r. FromRow r => Connection -> Query -> IO [r]
PG.query_ Connection
conn (Text -> Query
toQuery (Text -> Query) -> Text -> Query
forall a b. (a -> b) -> a -> b
$ Text
"SELECT job_id, last_error FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
Schema.jobQueueDLQTable Text
schema Text
table)
                :: IO [(Int64, Text)]
          map fst dlq `shouldBe` [primaryKey poison]
          map snd dlq `shouldSatisfy` all (T.isPrefixOf "Failed to decode job payload")

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"claims again after an all-poison claim rather than waiting for the next poll" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        completedRef <- [payload] -> IO (IORef [payload])
forall a. a -> IO (IORef a)
newIORef []
        config :: WorkerConfig m payload <-
          transactionalWorkerConfig 10 $
            mkHandler (noResult (\JobRead payload
job -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IORef [payload] -> ([payload] -> ([payload], ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef [payload]
completedRef (\[payload]
seen -> (JobRead payload -> payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead payload
job payload -> [payload] -> [payload]
forall a. a -> [a] -> [a]
: [payload]
seen, ()))))
        poison <- replicateM 5 $ runM env $ HL.insertJob (defaultJob (mkSimple "poison"))
        void $ runM env $ HL.insertJob (defaultJob (mkSimple "sibling"))
        void $
          withConn connStr $ \Connection
conn ->
            Connection -> Query -> Only (In [Int64]) -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
PG.execute
              Connection
conn
              ( Text -> Query
toQuery (Text -> Query) -> Text -> Query
forall a b. (a -> b) -> a -> b
$
                  Text
"UPDATE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
Schema.jobQueueTable Text
schema Text
table Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" SET payload = '{\"bogus\": 1}' WHERE id IN ?"
              )
              (In [Int64] -> Only (In [Int64])
forall a. a -> Only a
Only ([Int64] -> In [Int64]
forall a. a -> In a
PG.In ((JobRead payload -> Int64) -> [JobRead payload] -> [Int64]
forall a b. (a -> b) -> [a] -> [b]
map JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey ([Maybe (JobRead payload)] -> [JobRead payload]
forall a. [Maybe a] -> [a]
catMaybes [Maybe (JobRead payload)]
poison))))

        withAsync (runM env $ runWorkerPool config {workerCount = 2, pollInterval = 30, visibilityTimeout = 60}) $ \Async ()
_ ->
          HasCallStack => Int -> IO Bool -> IO ()
Int -> IO Bool -> IO ()
waitUntil Int
5_000 (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ ([payload] -> [payload] -> Bool
forall a. Eq a => a -> a -> Bool
== [Text -> payload
mkSimple Text
"sibling"]) ([payload] -> Bool) -> IO [payload] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef [payload] -> IO [payload]
forall a. IORef a -> IO a
readIORef IORef [payload]
completedRef

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"keeps the decoded jobs of a claim when a rejected row cannot be dead-lettered" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        Just poison <- env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"poison"))
        void $ runM env $ HL.insertJob (defaultJob (mkSimple "sibling"))
        let dlqTbl = Text -> Text -> Text
Schema.jobQueueDLQTable Text
schema Text
table
            poisonId = FilePath -> Text
T.pack (Int64 -> FilePath
forall a. Show a => a -> FilePath
show (JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
poison))
        updateJob connStr schema table "payload = '{\"bogus\": 1}'" (primaryKey poison)
        claimed <-
          withConn connStr $ \Connection
conn -> do
            Connection -> Text -> IO ()
execute_ Connection
conn (Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
dlqTbl Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD CONSTRAINT reject_poison CHECK (job_id <> " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
poisonId Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")")
            (env -> m [JobRead payload] -> IO [JobRead payload]
forall a. env -> m a -> IO a
runM env
env (Int -> NominalDiffTime -> m [JobRead payload]
forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> NominalDiffTime -> m [JobRead payload]
HL.claimNextVisibleJobs Int
2 NominalDiffTime
60) :: IO [JobRead payload])
              IO [JobRead payload] -> IO () -> IO [JobRead payload]
forall a b. IO a -> IO b -> IO a
`finally` Connection -> Text -> IO ()
execute_ Connection
conn (Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
dlqTbl Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" DROP CONSTRAINT reject_poison")
        map payload claimed `shouldBe` [mkSimple "sibling"]
        runM env (Ops.listDLQJobs schema table 100 0 :: m [DLQ.DLQJob (Stored payload)]) >>= (`shouldBe` [])
        Just held <- runM env (Ops.getJobById @_ @payload schema table (primaryKey poison))
        claimedBy held `shouldSatisfy` isJust

      FilePath -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
FilePath -> a -> SpecWith (Arg a)
it FilePath
"lists and encodes a row its payload type rejects, decoding only on request" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        Just poison <- env -> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a. env -> m a -> IO a
runM env
env (m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload)))
-> m (Maybe (JobRead payload)) -> IO (Maybe (JobRead payload))
forall a b. (a -> b) -> a -> b
$ JobWrite payload -> m (Maybe (JobRead payload))
forall payload (m :: * -> *).
QueueOperation m payload =>
JobWrite payload -> m (Maybe (JobRead payload))
HL.insertJob (payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob (Text -> payload
mkSimple Text
"poison"))
        updateJob connStr schema table "payload = '{\"bogus\": 1}'" (primaryKey poison)
        moved <- runM env $ Ops.moveToDLQ Ops.TakeLocks schema table "poison" poison
        moved `shouldBe` 1

        [entry] <- runM env (Ops.listDLQJobs schema table 100 0) :: IO [DLQ.DLQJob (Stored payload)]
        let row = DLQJob (Stored payload) -> JobRead (Stored payload)
forall payload. DLQJob payload -> JobSnapshot payload
DLQ.jobSnapshot DLQJob (Stored payload)
entry
        primaryKey row `shouldBe` primaryKey poison
        -- JSONB spaces its output. A payload spliced from the row keeps that spacing.
        BL.toStrict (Aeson.encode entry) `shouldSatisfy` BSC.isInfixOf "\"payload\":{\"bogus\": 1}"
        Aeson.toJSON (payload row) `shouldBe` Aeson.object ["bogus" Aeson..= (1 :: Int)]
        decodeStored (payload row) `shouldSatisfy` either (T.isPrefixOf "Failed to decode job payload") (const False)
        case Aeson.eitherDecode (Aeson.encode row) :: Either String (JobRead (Stored payload)) of
          Left FilePath
err -> HasCallStack => FilePath -> IO ()
FilePath -> IO ()
expectationFailure FilePath
err
          Right JobRead (Stored payload)
decoded -> Stored payload -> Value
forall a. ToJSON a => a -> Value
Aeson.toJSON (JobRead (Stored payload) -> Stored payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead (Stored payload)
decoded) Value -> Value -> IO ()
forall a. (HasCallStack, Show a, Eq a) => a -> a -> IO ()
`shouldBe` Stored payload -> Value
forall a. ToJSON a => a -> Value
Aeson.toJSON (JobRead (Stored payload) -> Stored payload
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> payload
payload JobRead (Stored payload)
row)
  where
    lockRow :: Connection -> Int64 -> IO (Either SomeException [Only Int64])
lockRow = Text
-> Text
-> Connection
-> Int64
-> IO (Either SomeException [Only Int64])
lockJobRow Text
schema Text
table
    withOpsTable :: IO a -> IO a
    withOpsTable :: forall a. IO a -> IO a
withOpsTable = ByteString -> Text -> IO a -> IO a
forall a. ByteString -> Text -> IO a -> IO a
withTestOpsTable ByteString
connStr Text
schema
    opsCount :: IO Int
opsCount = ByteString -> Text -> IO Int
queryOpsCount ByteString
connStr Text
schema

-- | A handler run shorter than the graceful-shutdown timeout.
inFlightJobMicros :: Int
inFlightJobMicros :: Int
inFlightJobMicros = Int
2_000_000

-- | A handler run longer than the graceful-shutdown timeout.
overrunJobMicros :: Int
overrunJobMicros :: Int
overrunJobMicros = Int
10_000_000

-- | A pause for the shutdown signal to reach the dispatcher.
shutdownSignalMicros :: Int
shutdownSignalMicros :: Int
shutdownSignalMicros = Int
500_000

-- | A short handler run.
shortHandlerMicros :: Int
shortHandlerMicros :: Int
shortHandlerMicros = Int
500_000

-- | A handler run that covers the drain checks.
drainingHandlerMicros :: Int
drainingHandlerMicros :: Int
drainingHandlerMicros = Int
2_500_000

-- | Several heartbeat ticks during the drain.
drainTicksMicros :: Int
drainTicksMicros :: Int
drainTicksMicros = Int
1_500_000

-- | Longer than the stale threshold, spent paused.
pastStaleWhilePausedMicros :: Int
pastStaleWhilePausedMicros :: Int
pastStaleWhilePausedMicros = Int
2_000_000

-- | A pause for the tick that the startup claim signalled to pass.
startupTickMicros :: Int
startupTickMicros :: Int
startupTickMicros = Int
2_500_000

-- | A pause for a pause signal to land.
pauseLandMicros :: Int
pauseLandMicros :: Int
pauseLandMicros = Int
300_000

-- | A pause for the failed registration attempt.
registrationAttemptMicros :: Int
registrationAttemptMicros :: Int
registrationAttemptMicros = Int
500_000

-- | Long enough for @NOW()@ to differ.
clockAdvanceMicros :: Int
clockAdvanceMicros :: Int
clockAdvanceMicros = Int
1_100_000

-- | Longer than the liveness threshold of the worker list query.
pastListThresholdMicros :: Int
pastListThresholdMicros :: Int
pastListThresholdMicros = Int
1_200_000

-- | Longer than one heartbeat interval, so the next heartbeat blocks.
heartbeatBlockMicros :: Int
heartbeatBlockMicros :: Int
heartbeatBlockMicros = Int
1_500_000

-- | A pause for the blocked heartbeat to complete.
heartbeatLandMicros :: Int
heartbeatLandMicros :: Int
heartbeatLandMicros = Int
500_000

-- | A handler run longer than any test.
stuckHandlerMicros :: Int
stuckHandlerMicros :: Int
stuckHandlerMicros = Int
30_000_000

-- | A pause for a cancel to reach the handler.
cancelPropagateMicros :: Int
cancelPropagateMicros :: Int
cancelPropagateMicros = Int
500_000

-- | The window in which a live handler would bump the counter.
counterWindowMicros :: Int
counterWindowMicros :: Int
counterWindowMicros = Int
500_000

-- | A pause after the inserts, before the worker starts.
beforeStartMicros :: Int
beforeStartMicros :: Int
beforeStartMicros = Int
100_000

-- | A pause for the concurrent statement to block on the held lock.
lockBlockMicros :: Int
lockBlockMicros :: Int
lockBlockMicros = Int
300_000

-- | The most time to wait for a cancel notification.
notifyTimeoutMicros :: Int
notifyTimeoutMicros :: Int
notifyTimeoutMicros = Int
2_000_000

-- | Longer than 'briefStaleAfter'.
pastStaleMicros :: Int
pastStaleMicros :: Int
pastStaleMicros = Int
1_500_000

-- | A stale threshold that a test passes by sleeping 'pastStaleMicros'.
briefStaleAfter :: NominalDiffTime
briefStaleAfter :: NominalDiffTime
briefStaleAfter = NominalDiffTime
1

-- | A stale threshold that no test reaches.
farStaleAfter :: NominalDiffTime
farStaleAfter :: NominalDiffTime
farStaleAfter = NominalDiffTime
300

-- | Thread count a test worker registers with.
testWorkerThreads :: Int32
testWorkerThreads :: Int32
testWorkerThreads = Int32
1

-- | Run the check on the one row for this worker. Fails on zero or several.
soleWorkerRow :: UUID -> [WR.WorkerRow] -> (WR.WorkerRow -> Expectation) -> Expectation
soleWorkerRow :: UUID -> [WorkerRow] -> (WorkerRow -> IO ()) -> IO ()
soleWorkerRow UUID
wid [WorkerRow]
rows WorkerRow -> IO ()
check = case (WorkerRow -> Bool) -> [WorkerRow] -> [WorkerRow]
forall a. (a -> Bool) -> [a] -> [a]
filter ((UUID -> UUID -> Bool
forall a. Eq a => a -> a -> Bool
== UUID
wid) (UUID -> Bool) -> (WorkerRow -> UUID) -> WorkerRow -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WorkerRow -> UUID
WR.workerId) [WorkerRow]
rows of
  [WorkerRow
row] -> WorkerRow -> IO ()
check WorkerRow
row
  [WorkerRow]
_ -> HasCallStack => FilePath -> IO ()
FilePath -> IO ()
expectationFailure FilePath
"expected exactly one row for the worker"

noResult :: (Functor n) => (j -> n ()) -> j -> n (Maybe [Text])
noResult :: forall (n :: * -> *) j.
Functor n =>
(j -> n ()) -> j -> n (Maybe [Text])
noResult j -> n ()
handler j
job = Maybe [Text]
forall a. Maybe a
Nothing Maybe [Text] -> n () -> n (Maybe [Text])
forall a b. a -> n b -> n a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ j -> n ()
handler j
job

opsTable :: Text -> Text
opsTable :: Text -> Text
opsTable Text
schema = Text
schema Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
".test_operations"

-- | Record an operation for the job on the handler's connection.
recordOp :: (MonadArbiter n) => Text -> Int64 -> n ()
recordOp :: forall (n :: * -> *). MonadArbiter n => Text -> Int64 -> n ()
recordOp Text
schema Int64
jobId =
  n Int64 -> n ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (n Int64 -> n ()) -> n Int64 -> n ()
forall a b. (a -> b) -> a -> b
$
    Text -> Params -> n Int64
forall (m :: * -> *). MonadArbiter m => Text -> Params -> m Int64
execStatement
      (Text
"INSERT INTO " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
opsTable Text
schema Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" (job_id, operation) VALUES (?, ?)")
      [Col Int64 -> Int64 -> SomeParam
forall a. Col a -> a -> SomeParam
pval Col Int64
CInt8 Int64
jobId, Col Text -> Text -> SomeParam
forall a. Col a -> a -> SomeParam
pval Col Text
CText Text
"processed"]

-- | Take a job's row lock on a connection of the test's own.
lockJobRow :: Text -> Text -> PG.Connection -> Int64 -> IO (Either SomeException [Only Int64])
lockJobRow :: Text
-> Text
-> Connection
-> Int64
-> IO (Either SomeException [Only Int64])
lockJobRow Text
schema Text
table Connection
conn Int64
jobId = IO [Only Int64] -> IO (Either SomeException [Only Int64])
forall e a. Exception e => IO a -> IO (Either e a)
try (Connection -> Query -> Only Int64 -> IO [Only Int64]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
PG.query Connection
conn Query
lockSql (Int64 -> Only Int64
forall a. a -> Only a
Only Int64
jobId))
  where
    lockSql :: Query
lockSql = Text -> Query
toQuery (Text -> Query) -> Text -> Query
forall a b. (a -> b) -> a -> b
$ Text
"SELECT id FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
Schema.jobQueueTable Text
schema Text
table Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" WHERE id = ? FOR UPDATE"

withTestOpsTable :: ByteString -> Text -> IO a -> IO a
withTestOpsTable :: forall a. ByteString -> Text -> IO a -> IO a
withTestOpsTable ByteString
connStr Text
schema IO a
action = do
  ByteString -> (Connection -> IO ()) -> IO ()
forall a. ByteString -> (Connection -> IO a) -> IO a
withConn ByteString
connStr ((Connection -> IO ()) -> IO ()) -> (Connection -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> do
    Connection -> Text -> IO ()
execute_ Connection
conn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"CREATE TABLE IF NOT EXISTS " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
opsTable Text
schema Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" (job_id INT, operation TEXT)"
    Connection -> Text -> IO ()
execute_ Connection
conn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"TRUNCATE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
opsTable Text
schema
  result <- IO a
action
  withConn connStr $ \Connection
conn ->
    Connection -> Text -> IO ()
execute_ Connection
conn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"DROP TABLE IF EXISTS " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
opsTable Text
schema
  pure result

queryOpsCount :: ByteString -> Text -> IO Int
queryOpsCount :: ByteString -> Text -> IO Int
queryOpsCount ByteString
connStr Text
schema =
  ByteString -> (Connection -> IO Int) -> IO Int
forall a. ByteString -> (Connection -> IO a) -> IO a
withConn ByteString
connStr ((Connection -> IO Int) -> IO Int)
-> (Connection -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Connection
conn -> do
    [Only count] <-
      Connection -> Query -> Only Text -> IO [Only Int]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query
        Connection
conn
        (Text -> Query
toQuery (Text -> Query) -> Text -> Query
forall a b. (a -> b) -> a -> b
$ Text
"SELECT COUNT(*) FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
opsTable Text
schema Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" WHERE operation = ?")
        (Text -> Only Text
forall a. a -> Only a
Only (Text
"processed" :: Text))
        :: IO [Only Int]
    pure count