{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}

-- | Reclaimed-job, stolen-job, and handler-exception tests, instantiated for each backend.
module Arbiter.Test.Worker.Reclaim (reclaimSpec) where

import Arbiter.Core.HighLevel (QueueOperation, RegistryAdmissionPolicies)
import Arbiter.Core.HighLevel qualified as HL
import Arbiter.Core.Job.Types
  ( ObservabilityHooks (..)
  , defaultJob
  , defaultObservabilityHooks
  , payload
  , primaryKey
  , setMaxAttempts
  )
import Arbiter.Core.MonadArbiter (RegistryOf, ResultOf)
import Arbiter.Core.QueueRegistry (RegistryTables)
import Arbiter.Worker (runWorkerPool)
import Arbiter.Worker.Config (WorkerConfig (..), transactionalWorkerConfig)
import Control.Concurrent (threadDelay)
import Control.Monad (void, when)
import Control.Monad.IO.Class (liftIO)
import Data.Foldable (traverse_)
import Data.IORef (atomicModifyIORef', newIORef, readIORef)
import Test.Hspec

import Arbiter.Test.Poll (bumpRef, waitUntil, withLinkedAsync)
import Arbiter.Test.Worker.Backend (TestBackend (..))
import Arbiter.Test.Worker.Rows (reclaimJob)

-- | A pause for the worker to settle the job and run its hooks.
settleMicros :: Int
settleMicros :: Int
settleMicros = Int
500_000

-- | How long the handler runs after its job is reclaimed.
afterReclaimMicros :: Int
afterReclaimMicros :: Int
afterReclaimMicros = Int
5_000_000

-- | The tag on the failing payloads. The handler fails each payload with this tag.
alwaysFailing :: Int
alwaysFailing :: Int
alwaysFailing = Int
999

-- | Reclaim, heartbeat theft, and worker-loop exception safety suite. The queue
-- under test declares @()@ as its result type.
reclaimSpec
  :: forall payload m env
   . ( Eq payload
     , QueueOperation m payload
     , RegistryAdmissionPolicies (RegistryOf m)
     , RegistryTables (RegistryOf m)
     , ResultOf m payload ~ ()
     )
  => TestBackend payload m env
  -> Spec
reclaimSpec :: forall payload (m :: * -> *) env.
(Eq payload, QueueOperation m payload,
 RegistryAdmissionPolicies (RegistryOf m),
 RegistryTables (RegistryOf m), ResultOf m payload ~ ()) =>
TestBackend payload m env -> Spec
reclaimSpec 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, Int -> payload
mkFailing :: Int -> payload
mkFailing :: forall payload (m :: * -> *) env.
TestBackend payload m env -> Int -> payload
mkFailing, IO env
mkEnv :: IO env
mkEnv :: forall payload (m :: * -> *) env.
TestBackend payload m env -> IO env
mkEnv, (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, 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
    String -> SpecWith env -> SpecWith env
forall a. HasCallStack => String -> SpecWith a -> SpecWith a
describe String
"Job Reclaim During Processing" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"gracefully skips retry when job is reclaimed by another worker" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        failureCalls <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)
        successCalls <- newIORef (0 :: Int)
        unavailableCalls <- newIORef (0 :: Int)
        handlerCompleted <- newIORef False

        let hooks =
              ObservabilityHooks m payload
forall (m :: * -> *) payload.
Applicative m =>
ObservabilityHooks m payload
defaultObservabilityHooks
                { onJobSuccess = \JobRead payload
_ StartTime
_ StartTime
_ -> 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
successCalls
                , onJobFailure = \JobRead payload
_ Text
_ StartTime
_ StartTime
_ -> 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
failureCalls
                , onJobUnavailable = \JobRead payload
_ Text
_ -> 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
unavailableCalls
                }

        Just inserted <- runM env $ HL.insertJob (defaultJob (mkSimple "slow"))
        let jobId = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
inserted

        let jobHandler 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
$ do
              ByteString -> Text -> Text -> Int64 -> IO ()
reclaimJob ByteString
connStr Text
schema Text
table Int64
jobId
              IORef Bool -> (Bool -> (Bool, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Bool
handlerCompleted (\Bool
_ -> (Bool
True, ()))

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler jobHandler)
        let configWithHooks =
              WorkerConfig m payload
config
                { observabilityHooks = hooks
                , pollInterval = 0.1
                }

        withLinkedAsync
          (runM env $ runWorkerPool configWithHooks)
          $ \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
$ IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
handlerCompleted
            Int -> IO ()
threadDelay Int
settleMicros

        failureCount <- readIORef failureCalls
        successCount <- readIORef successCalls
        failureCount `shouldBe` 0
        successCount `shouldBe` 0

        unavailableCount <- readIORef unavailableCalls
        unavailableCount `shouldBe` 1

        allJobs <- runM env $ HL.listJobs @payload 10 0
        map primaryKey allJobs `shouldContain` [jobId]

      String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"onJobFailure fires when handler throws a retryable exception" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        failureCalls <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)
        successCalls <- newIORef (0 :: Int)

        let hooks =
              ObservabilityHooks m payload
forall (m :: * -> *) payload.
Applicative m =>
ObservabilityHooks m payload
defaultObservabilityHooks
                { onJobSuccess = \JobRead payload
_ StartTime
_ StartTime
_ -> 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
successCalls
                , onJobFailure = \JobRead payload
_ Text
_ StartTime
_ StartTime
_ -> 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
failureCalls
                }

        void
          $ runM env
          $ HL.insertJob (setMaxAttempts (Just 1) $ defaultJob (mkSimple "will-fail"))

        let jobHandler p
_job = String -> a
forall a. HasCallStack => String -> a
error String
"intentional failure"

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler jobHandler)
        let configWithHooks =
              WorkerConfig m payload
config
                { observabilityHooks = hooks
                , pollInterval = 0.1
                }

        withLinkedAsync
          (runM env $ runWorkerPool configWithHooks)
          $ \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
$ (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
<$> IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
failureCalls

        failureCount <- readIORef failureCalls
        successCount <- readIORef successCalls
        failureCount `shouldBe` 1
        successCount `shouldBe` 0

    String -> SpecWith env -> SpecWith env
forall a. HasCallStack => String -> SpecWith a -> SpecWith a
describe String
"Heartbeat Stolen Job Detection" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"heartbeat cancels handler via race when job is stolen mid-processing" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        handlerCompleted <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
        unavailableCalls <- newIORef (0 :: Int)

        let hooks =
              ObservabilityHooks m payload
forall (m :: * -> *) payload.
Applicative m =>
ObservabilityHooks m payload
defaultObservabilityHooks
                { onJobUnavailable = \JobRead payload
_ Text
_ -> 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
unavailableCalls
                }

        Just inserted <- runM env $ HL.insertJob (defaultJob (mkSimple "slow"))
        let jobId = JobRead payload -> Int64
forall payload key q insertedAt adm.
Job payload key q insertedAt adm -> key
primaryKey JobRead payload
inserted

        let jobHandler 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
$ do
              ByteString -> Text -> Text -> Int64 -> IO ()
reclaimJob ByteString
connStr Text
schema Text
table Int64
jobId
              Int -> IO ()
threadDelay Int
afterReclaimMicros
              IORef Bool -> (Bool -> (Bool, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Bool
handlerCompleted (\Bool
_ -> (Bool
True, ()))

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler jobHandler)
        let configWithHooks =
              WorkerConfig m payload
config
                { observabilityHooks = hooks
                , pollInterval = 0.1
                , visibilityTimeout = 2
                , jobHeartbeatInterval = 1
                }

        withLinkedAsync
          (runM env $ runWorkerPool configWithHooks)
          $ \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
$ (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
<$> IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
unavailableCalls

        completed <- readIORef handlerCompleted
        completed `shouldBe` False

        allJobs <- runM env $ HL.listJobs @payload 10 0
        map primaryKey allJobs `shouldContain` [jobId]

    String -> SpecWith env -> SpecWith env
forall a. HasCallStack => String -> SpecWith a -> SpecWith a
describe String
"Worker Loop Exception Safety" (SpecWith env -> SpecWith env) -> SpecWith env -> SpecWith env
forall a b. (a -> b) -> a -> b
$ do
      String -> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"continues processing after handler exceptions" ((env -> IO ()) -> SpecWith (Arg (env -> IO ())))
-> (env -> IO ()) -> SpecWith (Arg (env -> IO ()))
forall a b. (a -> b) -> a -> b
$ \env
env -> do
        processedCount <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)

        let hooks =
              ObservabilityHooks m payload
forall (m :: * -> *) payload.
Applicative m =>
ObservabilityHooks m payload
defaultObservabilityHooks
                { onJobSuccess = \JobRead payload
_ StartTime
_ StartTime
_ -> 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
processedCount
                , onJobFailure = \JobRead payload
_ Text
_ StartTime
_ StartTime
_ -> 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
processedCount
                }

        runM env $
          traverse_
            (void . HL.insertJob . setMaxAttempts (Just 1) . defaultJob . mkFailing)
            [alwaysFailing, alwaysFailing, alwaysFailing]

        let jobHandler JobRead payload
job = Bool -> m () -> m ()
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
== Int -> payload
mkFailing Int
alwaysFailing) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> m ()
forall a. HasCallStack => String -> a
error String
"Failing task"

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler jobHandler)
        let configWithHooks =
              WorkerConfig m payload
config
                { observabilityHooks = hooks
                , workerCount = 1
                , pollInterval = 0.1
                }

        withLinkedAsync
          (runM env $ runWorkerPool configWithHooks)
          $ \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
$ (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
3) (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
processedCount

        processed <- readIORef processedCount
        processed `shouldBe` 3