{-# OPTIONS_HADDOCK not-home #-}

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- Batch queue with in-flight counts and a slot-availability signal.
module Arbiter.Worker.WorkQueue
  ( WorkQueue
  , newWorkQueue
  , pushWork
  , popWork
  , finishWork
  , awaitFinished
  , queuedCount
  , busyCount
  , inFlight
  ) where

import Control.Monad.IO.Class (MonadIO)
import Data.Foldable (traverse_)
import UnliftIO (MonadUnliftIO, mask_)
import UnliftIO.Chan (Chan, newChan, readChan, writeChan)
import UnliftIO.STM (STM, TVar, atomically, checkSTM, modifyTVar', newTVarIO, readTVar, writeTVar)

-- | The channel, the batches queued on it, the batches a worker holds, and
-- whether a slot freed since the dispatcher last looked.
data WorkQueue a = WorkQueue
  { forall a. WorkQueue a -> Chan a
wqChan :: Chan a
  , forall a. WorkQueue a -> TVar Int
wqQueued :: TVar Int
  , forall a. WorkQueue a -> TVar Int
wqBusy :: TVar Int
  , forall a. WorkQueue a -> TVar Bool
wqFinished :: TVar Bool
  }

-- | An empty queue with no busy slots.
newWorkQueue :: (MonadIO m) => m (WorkQueue a)
newWorkQueue :: forall (m :: * -> *) a. MonadIO m => m (WorkQueue a)
newWorkQueue = Chan a -> TVar Int -> TVar Int -> TVar Bool -> WorkQueue a
forall a.
Chan a -> TVar Int -> TVar Int -> TVar Bool -> WorkQueue a
WorkQueue (Chan a -> TVar Int -> TVar Int -> TVar Bool -> WorkQueue a)
-> m (Chan a)
-> m (TVar Int -> TVar Int -> TVar Bool -> WorkQueue a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Chan a)
forall (m :: * -> *) a. MonadIO m => m (Chan a)
newChan m (TVar Int -> TVar Int -> TVar Bool -> WorkQueue a)
-> m (TVar Int) -> m (TVar Int -> TVar Bool -> WorkQueue a)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> m (TVar Int)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Int
0 m (TVar Int -> TVar Bool -> WorkQueue a)
-> m (TVar Int) -> m (TVar Bool -> WorkQueue a)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> m (TVar Int)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Int
0 m (TVar Bool -> WorkQueue a) -> m (TVar Bool) -> m (WorkQueue a)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> m (TVar Bool)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Bool
False

-- | Enqueue in order.
pushWork :: (MonadUnliftIO m) => WorkQueue a -> [a] -> m ()
pushWork :: forall (m :: * -> *) a.
MonadUnliftIO m =>
WorkQueue a -> [a] -> m ()
pushWork WorkQueue a
_ [] = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
pushWork WorkQueue a
queue [a]
items = m () -> m ()
forall (m :: * -> *) a. MonadUnliftIO m => m a -> m a
mask_ (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (WorkQueue a -> TVar Int
forall a. WorkQueue a -> TVar Int
wqQueued WorkQueue a
queue) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
items))
  (a -> m ()) -> [a] -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Chan a -> a -> m ()
forall (m :: * -> *) a. MonadIO m => Chan a -> a -> m ()
writeChan (WorkQueue a -> Chan a
forall a. WorkQueue a -> Chan a
wqChan WorkQueue a
queue)) [a]
items

-- | Take the next item and transfer its count from queued to busy without interruption.
popWork :: (MonadUnliftIO m) => WorkQueue a -> m a
popWork :: forall (m :: * -> *) a. MonadUnliftIO m => WorkQueue a -> m a
popWork WorkQueue a
queue = m a -> m a
forall (m :: * -> *) a. MonadUnliftIO m => m a -> m a
mask_ (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$ do
  item <- Chan a -> m a
forall (m :: * -> *) a. MonadIO m => Chan a -> m a
readChan (WorkQueue a -> Chan a
forall a. WorkQueue a -> Chan a
wqChan WorkQueue a
queue)
  item <$ atomically (modifyTVar' (wqQueued queue) (subtract 1) *> modifyTVar' (wqBusy queue) (+ 1))

-- | Free the busy slot a 'popWork' took and signal the dispatcher.
finishWork :: WorkQueue a -> STM ()
finishWork :: forall a. WorkQueue a -> STM ()
finishWork WorkQueue a
queue = TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (WorkQueue a -> TVar Int
forall a. WorkQueue a -> TVar Int
wqBusy WorkQueue a
queue) (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) STM () -> STM () -> STM ()
forall a b. STM a -> STM b -> STM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (WorkQueue a -> TVar Bool
forall a. WorkQueue a -> TVar Bool
wqFinished WorkQueue a
queue) Bool
True

-- | Block until a slot freed, consuming the signal.
awaitFinished :: WorkQueue a -> STM ()
awaitFinished :: forall a. WorkQueue a -> STM ()
awaitFinished WorkQueue a
queue = TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar (WorkQueue a -> TVar Bool
forall a. WorkQueue a -> TVar Bool
wqFinished WorkQueue a
queue) STM Bool -> (Bool -> STM ()) -> STM ()
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
checkSTM STM () -> STM () -> STM ()
forall a b. STM a -> STM b -> STM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (WorkQueue a -> TVar Bool
forall a. WorkQueue a -> TVar Bool
wqFinished WorkQueue a
queue) Bool
False

-- | Batches queued and not yet taken by a worker.
queuedCount :: WorkQueue a -> STM Int
queuedCount :: forall a. WorkQueue a -> STM Int
queuedCount = TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar (TVar Int -> STM Int)
-> (WorkQueue a -> TVar Int) -> WorkQueue a -> STM Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WorkQueue a -> TVar Int
forall a. WorkQueue a -> TVar Int
wqQueued

-- | Batches a worker holds.
busyCount :: WorkQueue a -> STM Int
busyCount :: forall a. WorkQueue a -> STM Int
busyCount = TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar (TVar Int -> STM Int)
-> (WorkQueue a -> TVar Int) -> WorkQueue a -> STM Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WorkQueue a -> TVar Int
forall a. WorkQueue a -> TVar Int
wqBusy

-- | Batches queued or held by a worker.
inFlight :: WorkQueue a -> STM Int
inFlight :: forall a. WorkQueue a -> STM Int
inFlight WorkQueue a
queue = Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) (Int -> Int -> Int) -> STM Int -> STM (Int -> Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkQueue a -> STM Int
forall a. WorkQueue a -> STM Int
queuedCount WorkQueue a
queue STM (Int -> Int) -> STM Int -> STM Int
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> WorkQueue a -> STM Int
forall a. WorkQueue a -> STM Int
busyCount WorkQueue a
queue