{-# LANGUAGE OverloadedStrings #-}

-- | Handlers for the worker pool's NOTIFY channels. Each ignores a message not for this worker.
module Arbiter.Worker.ChannelHandlers
  ( handlePauseNotif
  , handleCancelNotif
  , handleCronRunNotif
  ) where

import Arbiter.Core.Listen (Notification, notificationData)
import Control.Monad (unless, when)
import Data.Aeson qualified as Aeson
import Data.Int (Int64)
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text.Encoding (decodeUtf8Lenient)
import Data.UUID (UUID)
import UnliftIO (MonadIO, atomically)
import UnliftIO.STM (TVar)
import UnliftIO.STM qualified as STM

import Arbiter.Worker.Config (WorkerConfig (..), workerStateVar, writePause)
import Arbiter.Worker.Heartbeat (PoolGuard, recheckJob)
import Arbiter.Worker.WorkerState (isShuttingDown)

-- | Decode the pause payload and, if it addresses this worker, write 'Arbiter.Worker.Config.pauseVar'.
handlePauseNotif
  :: (MonadIO m)
  => WorkerConfig n payload
  -> Notification
  -> m ()
handlePauseNotif :: forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> Notification -> m ()
handlePauseNotif WorkerConfig n payload
config Notification
notif =
  case ByteString -> Maybe PausePayload
forall a. FromJSON a => ByteString -> Maybe a
Aeson.decodeStrict (Notification -> ByteString
notificationData Notification
notif) :: Maybe PausePayload of
    Just (PausePayload UUID
wid Bool
paused) | UUID
wid UUID -> UUID -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerConfig n payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig n payload
config -> STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      stopping <- TVar WorkerState -> STM Bool
isShuttingDown (WorkerConfig n payload -> TVar WorkerState
forall (n :: * -> *) payload.
WorkerConfig n payload -> TVar WorkerState
workerStateVar WorkerConfig n payload
config)
      unless stopping $ writePause config paused
    Maybe PausePayload
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | If the cancel payload targets this worker, extend the job's batch now. The
-- extend finds the cancel and the guard stops the handler.
handleCancelNotif
  :: (MonadIO m)
  => WorkerConfig n payload
  -> PoolGuard payload
  -> Notification
  -> m ()
handleCancelNotif :: forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> PoolGuard payload -> Notification -> m ()
handleCancelNotif WorkerConfig n payload
config PoolGuard payload
guard Notification
notif =
  case ByteString -> Maybe CancelPayload
forall a. FromJSON a => ByteString -> Maybe a
Aeson.decodeStrict (Notification -> ByteString
notificationData Notification
notif) :: Maybe CancelPayload of
    Just (CancelPayload UUID
wid Int64
jid) | UUID
wid UUID -> UUID -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerConfig n payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig n payload
config -> PoolGuard payload -> Int64 -> m ()
forall (m :: * -> *) payload.
MonadIO m =>
PoolGuard payload -> Int64 -> m ()
recheckJob PoolGuard payload
guard Int64
jid
    Maybe CancelPayload
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Signal the scheduler when a run-now NOTIFY names a schedule this pool owns.
handleCronRunNotif
  :: (MonadIO m)
  => Set Text
  -- ^ This pool's own cron schedule names
  -> TVar Bool
  -> Notification
  -> m ()
handleCronRunNotif :: forall (m :: * -> *).
MonadIO m =>
Set Text -> TVar Bool -> Notification -> m ()
handleCronRunNotif Set Text
ownNames TVar Bool
runNowVar Notification
notif =
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Text -> Set Text -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (ByteString -> Text
decodeUtf8Lenient (Notification -> ByteString
notificationData Notification
notif)) Set Text
ownNames)
    (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically
    (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
STM.writeTVar TVar Bool
runNowVar Bool
True

data PausePayload = PausePayload UUID Bool

instance Aeson.FromJSON PausePayload where
  parseJSON :: Value -> Parser PausePayload
parseJSON = String
-> (Object -> Parser PausePayload) -> Value -> Parser PausePayload
forall a. String -> (Object -> Parser a) -> Value -> Parser a
Aeson.withObject String
"PausePayload" ((Object -> Parser PausePayload) -> Value -> Parser PausePayload)
-> (Object -> Parser PausePayload) -> Value -> Parser PausePayload
forall a b. (a -> b) -> a -> b
$ \Object
obj ->
    UUID -> Bool -> PausePayload
PausePayload (UUID -> Bool -> PausePayload)
-> Parser UUID -> Parser (Bool -> PausePayload)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
obj Object -> Key -> Parser UUID
forall a. FromJSON a => Object -> Key -> Parser a
Aeson..: Key
"worker_id" Parser (Bool -> PausePayload) -> Parser Bool -> Parser PausePayload
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
obj Object -> Key -> Parser Bool
forall a. FromJSON a => Object -> Key -> Parser a
Aeson..: Key
"paused"

data CancelPayload = CancelPayload UUID Int64

instance Aeson.FromJSON CancelPayload where
  parseJSON :: Value -> Parser CancelPayload
parseJSON = String
-> (Object -> Parser CancelPayload)
-> Value
-> Parser CancelPayload
forall a. String -> (Object -> Parser a) -> Value -> Parser a
Aeson.withObject String
"CancelPayload" ((Object -> Parser CancelPayload) -> Value -> Parser CancelPayload)
-> (Object -> Parser CancelPayload)
-> Value
-> Parser CancelPayload
forall a b. (a -> b) -> a -> b
$ \Object
obj ->
    UUID -> Int64 -> CancelPayload
CancelPayload (UUID -> Int64 -> CancelPayload)
-> Parser UUID -> Parser (Int64 -> CancelPayload)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
obj Object -> Key -> Parser UUID
forall a. FromJSON a => Object -> Key -> Parser a
Aeson..: Key
"worker_id" Parser (Int64 -> CancelPayload)
-> Parser Int64 -> Parser CancelPayload
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
obj Object -> Key -> Parser Int64
forall a. FromJSON a => Object -> Key -> Parser a
Aeson..: Key
"job_id"