module Arbiter.Worker.WorkerState
  ( WorkerState (..)
  , newWorkerState
  , signalShutdown
  ) where

import Control.Concurrent.STM (TVar)
import Control.Concurrent.STM qualified as STM

-- | Create a new worker state initialized to 'Running'.
newWorkerState :: IO (TVar WorkerState)
newWorkerState :: IO (TVar WorkerState)
newWorkerState = WorkerState -> IO (TVar WorkerState)
forall a. a -> IO (TVar a)
STM.newTVarIO WorkerState
Running

-- | Signal graceful shutdown on a worker state.
--
-- Workers will stop claiming new jobs, finish in-flight work, then exit.
signalShutdown :: TVar WorkerState -> IO ()
signalShutdown :: TVar WorkerState -> IO ()
signalShutdown TVar WorkerState
st = STM () -> IO ()
forall a. STM a -> IO a
STM.atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar WorkerState -> WorkerState -> STM ()
forall a. TVar a -> a -> STM ()
STM.writeTVar TVar WorkerState
st WorkerState
ShuttingDown

-- | State for worker pool coordination.
--
-- Controls whether workers claim new jobs:
--
--   * 'Running': Normal operation, claim jobs continuously
--   * 'Paused': Stop claiming new jobs, finish in-flight jobs, wait for resume
--   * 'ShuttingDown': Stop claiming new jobs, finish in-flight jobs, then exit
data WorkerState
  = -- | Normal operation
    Running
  | -- | Paused (stop claiming, finish in-flight, wait for resume)
    Paused
  | -- | Graceful shutdown in progress (drain and exit)
    ShuttingDown
  deriving stock (WorkerState -> WorkerState -> Bool
(WorkerState -> WorkerState -> Bool)
-> (WorkerState -> WorkerState -> Bool) -> Eq WorkerState
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: WorkerState -> WorkerState -> Bool
== :: WorkerState -> WorkerState -> Bool
$c/= :: WorkerState -> WorkerState -> Bool
/= :: WorkerState -> WorkerState -> Bool
Eq, Int -> WorkerState -> ShowS
[WorkerState] -> ShowS
WorkerState -> String
(Int -> WorkerState -> ShowS)
-> (WorkerState -> String)
-> ([WorkerState] -> ShowS)
-> Show WorkerState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> WorkerState -> ShowS
showsPrec :: Int -> WorkerState -> ShowS
$cshow :: WorkerState -> String
show :: WorkerState -> String
$cshowList :: [WorkerState] -> ShowS
showList :: [WorkerState] -> ShowS
Show)