-- |
-- Makes sure that all dispatched works are done.
-- Notify when all dispatched works are done.
module Server.ResponseController
  ( ResponseController
  , new
  , dispatch
  , setCheckpointAndWait
  ) where

import           Control.Concurrent
import           Control.Concurrent.SizedChan
import           Control.Monad                  ( void
                                                , when
                                                )
import           Data.IORef

data ResponseController = ResponseController
  { -- | The number of work dispatched
    ResponseController -> IORef Int
dispatchedCount :: IORef Int
  ,
    -- | The number of work completed
    ResponseController -> IORef Int
completedCount  :: IORef Int
  ,
    -- | A channel of "Checkpoints" to be met
    ResponseController -> SizedChan Checkpoint
checkpointChan  :: SizedChan Checkpoint
  }

-- | An "Checkpoint" is just a number with a callback, the callback will be invoked once the number is "met"
type Checkpoint = (Int, () -> IO ())

-- | Constructs a new ResponseController
new :: IO ResponseController
new :: IO ResponseController
new = IORef Int
-> IORef Int -> SizedChan Checkpoint -> ResponseController
ResponseController forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> IO (IORef a)
newIORef Int
0 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. a -> IO (IORef a)
newIORef Int
0 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. IO (SizedChan a)
newSizedChan

-- | Returns a callback, invoked the callback to signal completion.
-- This function and the returned callback are both non-blocking.
dispatch :: ResponseController -> IO (() -> IO ())
dispatch :: ResponseController -> IO (() -> IO ())
dispatch ResponseController
controller = do
  -- bump `dispatchedCount`
  forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (ResponseController -> IORef Int
dispatchedCount ResponseController
controller) forall a. Enum a => a -> a
succ
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ \() -> do
    -- work completed, bump `completedCount`
    forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (ResponseController -> IORef Int
completedCount ResponseController
controller) forall a. Enum a => a -> a
succ

    -- see if there's any Checkpoint
    Maybe Checkpoint
result <- forall a. SizedChan a -> IO (Maybe a)
tryPeekSizedChan (ResponseController -> SizedChan Checkpoint
checkpointChan ResponseController
controller)
    case Maybe Checkpoint
result of
      -- no checkpoints, do nothing
      Maybe Checkpoint
Nothing                     -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
      -- a checkpoint is set!
      Just (Int
dispatched, () -> IO ()
callback) -> do
        Int
completed <- forall a. IORef a -> IO a
readIORef (ResponseController -> IORef Int
completedCount ResponseController
controller)
        -- see if the checkpoint is met
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
dispatched forall a. Eq a => a -> a -> Bool
== Int
completed) forall a b. (a -> b) -> a -> b
$ do
          -- invoke the callback and remove the checkpoint
          () -> IO ()
callback ()
          forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. SizedChan a -> IO a
readSizedChan (ResponseController -> SizedChan Checkpoint
checkpointChan ResponseController
controller)

-- | Expects a callback, which will be invoked once all works dispatched BEFORE have been completed
-- This function is non-blocking
setCheckpoint :: ResponseController -> (() -> IO ()) -> IO ()
setCheckpoint :: ResponseController -> (() -> IO ()) -> IO ()
setCheckpoint ResponseController
controller () -> IO ()
callback = do
  Int
dispatched <- forall a. IORef a -> IO a
readIORef (ResponseController -> IORef Int
dispatchedCount ResponseController
controller)
  Int
completed  <- forall a. IORef a -> IO a
readIORef (ResponseController -> IORef Int
completedCount ResponseController
controller)
  -- see if the previously dispatched works have been completed
  if Int
dispatched forall a. Eq a => a -> a -> Bool
== Int
completed
    then () -> IO ()
callback ()
    else do
      -- constructs a Checkpoint from `dispatchedCount`
      let checkpoint :: Checkpoint
checkpoint = (Int
dispatched, () -> IO ()
callback)
      -- write it to the channel
      forall a. SizedChan a -> a -> IO ()
writeSizedChan (ResponseController -> SizedChan Checkpoint
checkpointChan ResponseController
controller) Checkpoint
checkpoint

-- | The blocking version of `setCheckpoint`
setCheckpointAndWait :: ResponseController -> IO ()
setCheckpointAndWait :: ResponseController -> IO ()
setCheckpointAndWait ResponseController
controller = do
  MVar ()
mvar <- forall a. IO (MVar a)
newEmptyMVar
  ResponseController -> (() -> IO ()) -> IO ()
setCheckpoint ResponseController
controller (forall a. MVar a -> a -> IO ()
putMVar MVar ()
mvar)
  forall a. MVar a -> IO a
takeMVar MVar ()
mvar