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

-- | Connection failure recovery tests, instantiated for each backend.
module Arbiter.Test.Worker.ConnectionRecovery (connectionRecoverySpec) where

import Arbiter.Core.Codec (Col (CInt4), col)
import Arbiter.Core.HighLevel (QueueOperation, RegistryAdmissionPolicies)
import Arbiter.Core.HighLevel qualified as HL
import Arbiter.Core.Job.Types
  ( defaultJob
  , payload
  , setGroupKey
  , setMaxAttempts
  )
import Arbiter.Core.MonadArbiter (RegistryOf, ResultOf, withDbTransaction)
import Arbiter.Core.QueueRegistry (RegistryTables)
import Arbiter.Worker (runWorkerPool)
import Arbiter.Worker.BackoffStrategy (Jitter (NoJitter))
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 Data.Text qualified as T
import Test.Hspec
import UnliftIO (bracket)
import UnliftIO.Async (withAsync)

import Arbiter.Test.Poll (bumpRef, waitUntil)
import Arbiter.Test.Setup (execQuery, terminateBackendsMatching, terminatePid)
import Arbiter.Test.Worker.Backend (TestBackend (..))

-- | A pause for the workers to reconnect after the kill.
reconnectMicros :: Int
reconnectMicros :: Int
reconnectMicros = Int
7_000_000

-- | Connection recovery suite. The queue under test declares @()@ as its result type.
-- It terminates other connections on the database whose query names the schema.
connectionRecoverySpec
  :: forall payload m env
   . ( Eq payload
     , QueueOperation m payload
     , RegistryAdmissionPolicies (RegistryOf m)
     , RegistryTables (RegistryOf m)
     , ResultOf m payload ~ ()
     )
  => TestBackend payload m env
  -> Spec
connectionRecoverySpec :: forall payload (m :: * -> *) env.
(Eq payload, QueueOperation m payload,
 RegistryAdmissionPolicies (RegistryOf m),
 RegistryTables (RegistryOf m), ResultOf m payload ~ ()) =>
TestBackend payload m env -> Spec
connectionRecoverySpec TestBackend {Text
schema :: Text
schema :: forall payload (m :: * -> *) env. TestBackend payload m env -> Text
schema, 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
mkFreshEnv :: IO env
mkFreshEnv :: forall payload (m :: * -> *) env.
TestBackend payload m env -> IO env
mkFreshEnv, env -> IO ()
destroyEnv :: env -> IO ()
destroyEnv :: forall payload (m :: * -> *) env.
TestBackend payload m env -> env -> IO ()
destroyEnv, (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} =
  (((env, env) -> IO ()) -> IO ()) -> SpecWith (env, env) -> Spec
forall a. (ActionWith a -> IO ()) -> SpecWith a -> Spec
around (IO (env, env)
-> ((env, env) -> IO ()) -> ((env, env) -> IO ()) -> IO ()
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket ((,) (env -> env -> (env, env)) -> IO env -> IO (env -> (env, env))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO env
mkFreshEnv IO (env -> (env, env)) -> IO env -> IO (env, env)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO env
mkFreshEnv) (\(env
env, env
spare) -> env -> IO ()
destroyEnv env
env IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> env -> IO ()
destroyEnv env
spare)) (SpecWith (env, env) -> Spec) -> SpecWith (env, env) -> Spec
forall a b. (a -> b) -> a -> b
$
    String -> SpecWith (env, env) -> SpecWith (env, env)
forall a. HasCallStack => String -> SpecWith a -> SpecWith a
describe String
"Connection Recovery" (SpecWith (env, env) -> SpecWith (env, env))
-> SpecWith (env, env) -> SpecWith (env, env)
forall a b. (a -> b) -> a -> b
$ do
      String
-> ((env, env) -> IO ()) -> SpecWith (Arg ((env, env) -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"processes jobs inserted before and after a connection kill" (((env, env) -> IO ()) -> SpecWith (Arg ((env, env) -> IO ())))
-> ((env, env) -> IO ()) -> SpecWith (Arg ((env, env) -> IO ()))
forall a b. (a -> b) -> a -> b
$ \(env
env, env
spare) -> do
        completedRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)

        let 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
completedRef

        traverse_
          ( \Int
jobIndex ->
              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 (Maybe (JobRead payload)) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void
                (m (Maybe (JobRead payload)) -> m ())
-> m (Maybe (JobRead payload)) -> m ()
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 (String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ String
"Pre-kill " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
jobIndex))
          )
          [1 :: Int .. 3]

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler handler)
        let workerConfig =
              WorkerConfig m payload
config
                { workerCount = 1
                , pollInterval = 0.2
                , jitter = NoJitter
                }

        withAsync (runM env $ runWorkerPool workerConfig) $ \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) -> 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
completedRef

          preKillCompleted <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
completedRef
          preKillCompleted `shouldBe` 3

          terminateBackendsMatching connStr ("%" <> schema <> "%")

          threadDelay reconnectMicros

          insertJobsDirect
            spare
            [ mkSimple "Post-kill 1"
            , mkSimple "Post-kill 2"
            , mkSimple "Post-kill 3"
            ]

          waitUntil 10_000 $ (== 6) <$> readIORef completedRef

          completed <- readIORef completedRef
          completed `shouldBe` 6

      String
-> ((env, env) -> IO ()) -> SpecWith (Arg ((env, env) -> IO ()))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
it String
"retries job after worker ack fails due to terminated connection" (((env, env) -> IO ()) -> SpecWith (Arg ((env, env) -> IO ())))
-> ((env, env) -> IO ()) -> SpecWith (Arg ((env, env) -> IO ()))
forall a b. (a -> b) -> a -> b
$ \(env
env, env
_) -> do
        completedRef <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)
        killOnceRef <- newIORef True

        let handler JobRead payload
job = do
              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
== Text -> payload
mkSimple Text
"slow") (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
                shouldKill <- IO Bool -> m Bool
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> m Bool) -> IO Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ IORef Bool -> (Bool -> (Bool, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Bool
killOnceRef ((Bool -> (Bool, Bool)) -> IO Bool)
-> (Bool -> (Bool, Bool)) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Bool
armed -> (Bool
False, Bool
armed)
                when shouldKill $ do
                  myPid <- execQuery "SELECT pg_backend_pid() AS pid" [] (col "pid" CInt4)
                  liftIO $ traverse_ (terminatePid connStr) myPid
              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
completedRef

        runM env
          $ void
          $ HL.insertJob
          $ setMaxAttempts (Just 3)
          $ setGroupKey (Just "g1")
          $ defaultJob (mkSimple "slow")

        config :: WorkerConfig m payload <- transactionalWorkerConfig 10 (mkHandler handler)
        let workerConfig =
              WorkerConfig m payload
config
                { workerCount = 1
                , pollInterval = 0.2
                , jitter = NoJitter
                }

        withAsync (runM env $ runWorkerPool workerConfig) $ \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
2) (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
completedRef

          completed <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
completedRef
          completed `shouldBe` 2

          let remainingCount =
                [JobRead payload] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([JobRead payload] -> Int) -> IO [JobRead payload] -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> env -> m [JobRead payload] -> IO [JobRead payload]
forall a. env -> m a -> IO a
runM env
env (forall payload (m :: * -> *).
QueueOperation m payload =>
Int -> Int -> m [JobRead payload]
HL.listJobs @payload Int
10 Int
0)
          waitUntil 10_000 $ (== 0) <$> remainingCount

          remaining <- remainingCount
          remaining `shouldBe` 0
  where
    insertJobsDirect :: env -> [payload] -> IO ()
    insertJobsDirect :: env -> [payload] -> IO ()
insertJobsDirect env
spare [payload]
payloads =
      env -> m () -> IO ()
forall a. env -> m a -> IO a
runM env
spare
        (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
$ (payload -> m ()) -> [payload] -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (m (Maybe (JobRead payload)) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (Maybe (JobRead payload)) -> m ())
-> (payload -> m (Maybe (JobRead payload))) -> payload -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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)))
-> (payload -> JobWrite payload)
-> payload
-> m (Maybe (JobRead payload))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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)
-> (payload -> JobWrite payload) -> payload -> JobWrite payload
forall b c a. (b -> c) -> (a -> b) -> a -> c
. payload -> JobWrite payload
forall payload. payload -> JobWrite payload
defaultJob) [payload]
payloads