Queue.hs

{-# OPTIONS_HADDOCK hide, not-home #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
-- |
-- Module      : Control.Scheduler.Queue
-- Copyright   : (c) Alexey Kuleshevich 2018-2019
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Control.Scheduler.Queue
  (  -- * Job queue
    Job(Job_)
  , mkJob
  , Queue(..)
  , JQueue(..)
  , WorkerId(..)
  , newJQueue
  , pushJQueue
  , popJQueue
  , clearPendingJQueue
  , readResults
  , blockPopJQueue
  , unblockPopJQueue
  ) where

import Control.Concurrent.MVar
import Control.Monad (join)
import Control.Monad.IO.Unlift
import Data.Atomics (atomicModifyIORefCAS, atomicModifyIORefCAS_)
import Data.Maybe
import Data.IORef

-- | A blocking unbounded queue that keeps the jobs in FIFO order and the results IORefs
-- in reversed
data Queue m a = Queue
  { qQueue   :: ![Job m a]
  , qStack   :: ![Job m a]
  , qResults :: ![IORef (Maybe a)]
  , qBaton   :: {-# UNPACK #-}!(MVar ())
  }


-- | A unique id for the worker in the `Control.Scheduler.Scheduler` context. It will
-- always be a number from @0@ up to, but not including, the number of workers a scheduler
-- has, which in turn can always be determined with `Control.Scheduler.numWorkers` function.
--
-- @since 1.4.0
newtype WorkerId = WorkerId
  { getWorkerId :: Int
  } deriving (Show, Read, Eq, Ord, Enum, Bounded, Num, Real, Integral)


popQueue :: Queue m a -> Maybe (Job m a, Queue m a)
popQueue queue =
  case qQueue queue of
    x:xs -> Just (x, queue {qQueue = xs})
    [] ->
      case reverse (qStack queue) of
        []   -> Nothing
        y:ys -> Just (y, queue {qQueue = ys, qStack = []})
{-# INLINEABLE popQueue #-}

data Job m a
  = Job {-# UNPACK #-} !(IORef (Maybe a)) (WorkerId -> m ())
  | Job_ (WorkerId -> m ())


mkJob :: MonadIO m => ((a -> m ()) -> WorkerId -> m ()) -> m (Job m a)
mkJob action = do
  resRef <- liftIO $ newIORef Nothing
  return $ Job resRef (action (liftIO . writeIORef resRef . Just))
{-# INLINEABLE mkJob #-}

data JQueue m a =
  JQueue
    { jqQueueRef :: {-# UNPACK #-}!(IORef (Queue m a))
    , jqLock     :: {-# UNPACK #-}!(MVar ())
    }

newJQueue :: MonadIO m => m (JQueue m a)
newJQueue =
  liftIO $ do
    newLock <- newEmptyMVar
    newBaton <- newEmptyMVar
    queueRef <- newIORef (Queue [] [] [] newBaton)
    return $ JQueue queueRef newLock

-- | Pushes an item onto a queue and returns the previous count.
pushJQueue :: MonadIO m => JQueue m a -> Job m a -> m ()
pushJQueue (JQueue jQueueRef _) job =
  liftIO $ do
    newBaton <- newEmptyMVar
    join $
      atomicModifyIORefCAS jQueueRef $ \queue@Queue {qStack, qResults, qBaton} ->
        ( queue
            { qResults =
                case job of
                  Job resRef _ -> resRef : qResults
                  _ -> qResults
            , qStack = job : qStack
            , qBaton = newBaton
            }
        , putMVar qBaton ())
{-# INLINEABLE pushJQueue #-}

-- | Pops an item from the queue. The job returns the total job counts that is still left
-- in the queue
popJQueue :: MonadUnliftIO m => JQueue m a -> m (WorkerId -> m ())
popJQueue (JQueue jQueueRef lock) = liftIO inner
  where
    inner = do
      readMVar lock
      join $
        atomicModifyIORefCAS jQueueRef $ \queue@Queue {qBaton} ->
          case popQueue queue of
            Nothing -> (queue, readMVar qBaton >> inner)
            Just (job, newQueue) ->
              ( newQueue
              , case job of
                  Job _ action -> return action
                  Job_ action_ -> return action_)
{-# INLINEABLE popJQueue #-}

unblockPopJQueue :: MonadIO m => JQueue m a -> m ()
unblockPopJQueue (JQueue _ lock) = liftIO $ putMVar lock ()
{-# INLINEABLE unblockPopJQueue #-}

blockPopJQueue :: MonadIO m => JQueue m a -> m ()
blockPopJQueue (JQueue _ lock) = liftIO $ takeMVar lock
{-# INLINEABLE blockPopJQueue #-}

-- | Clears any jobs that haven't been started yet. Returns the number of jobs that are
-- still in progress and have not been yet been completed.
clearPendingJQueue :: MonadIO m => JQueue m a -> m ()
clearPendingJQueue (JQueue queueRef _) =
  liftIO $ atomicModifyIORefCAS_ queueRef $ \queue -> (queue {qQueue = [], qStack = []})
{-# INLINEABLE clearPendingJQueue #-}


-- | Extracts all results available up to now, the uncomputed ones are discarded. This
-- also has an affect of resetting the total job count to zero.
readResults :: MonadIO m => JQueue m a -> m [a]
readResults (JQueue jQueueRef _) =
  liftIO $ do
    results <-
      atomicModifyIORefCAS jQueueRef $ \queue ->
        (queue {qQueue = [], qStack = [], qResults = []}, qResults queue)
    catMaybes <$> mapM readIORef results
{-# INLINEABLE readResults #-}