-- |
-- 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
    dispatchedCount :: IORef Int
  ,
    -- | The number of work completed
    completedCount  :: IORef Int
  ,
    -- | A channel of "Checkpoints" to be met
    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 = ResponseController <$> newIORef 0 <*> newIORef 0 <*> 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 controller = do
  -- bump `dispatchedCount`
  modifyIORef' (dispatchedCount controller) succ
  return $ \() -> do
    -- work completed, bump `completedCount`
    modifyIORef' (completedCount controller) succ

    -- see if there's any Checkpoint
    result <- tryPeekSizedChan (checkpointChan controller)
    case result of
      -- no checkpoints, do nothing
      Nothing                     -> return ()
      -- a checkpoint is set!
      Just (dispatched, callback) -> do
        completed <- readIORef (completedCount controller)
        -- see if the checkpoint is met
        when (dispatched == completed) $ do
          -- invoke the callback and remove the checkpoint
          callback ()
          void $ readSizedChan (checkpointChan 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 controller callback = do
  dispatched <- readIORef (dispatchedCount controller)
  completed  <- readIORef (completedCount controller)
  -- see if the previously dispatched works have been completed
  if dispatched == completed
    then callback ()
    else do
      -- constructs a Checkpoint from `dispatchedCount`
      let checkpoint = (dispatched, callback)
      -- write it to the channel
      writeSizedChan (checkpointChan controller) checkpoint

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