JobControl.hs

-----------------------------------------------------------------------------
-- |
-- Module      :  Distribution.Client.JobControl
-- Copyright   :  (c) Duncan Coutts 2012
-- License     :  BSD-like
--
-- Maintainer  :  cabal-devel@haskell.org
-- Stability   :  provisional
-- Portability :  portable
--
-- A job control concurrency abstraction
-----------------------------------------------------------------------------
module Distribution.Client.JobControl (
    JobControl,
    newSerialJobControl,
    newParallelJobControl,
    spawnJob,
    collectJob,
    remainingJobs,
    cancelJobs,

    JobLimit,
    newJobLimit,
    withJobLimit,

    Lock,
    newLock,
    criticalSection
  ) where

import Control.Monad
import Control.Concurrent (forkIO)
import Control.Concurrent.MVar
import Control.Concurrent.STM (STM, atomically)
import Control.Concurrent.STM.TVar
import Control.Concurrent.STM.TChan
import Control.Exception (SomeException, bracket_, throwIO, try)
import Distribution.Client.Compat.Semaphore


-- | A simple concurrency abstraction. Jobs can be spawned and can complete
-- in any order. This allows both serial and parallel implementations.
--
data JobControl m a = JobControl {
       -- | Add a new job to the pool of jobs
       spawnJob    :: m a -> m (),

       -- | Wait until one job is complete
       collectJob  :: m a,

       -- | Returns True if there are any outstanding jobs
       -- (ie spawned but yet to be collected)
       remainingJobs :: m Bool,

       -- | Try to cancel any outstanding but not-yet-started jobs.
       -- Call 'remainingJobs' after this to find out if any jobs are left
       -- (ie could not be cancelled).
       cancelJobs  :: m ()
     }


-- | Make a 'JobControl' that executes all jobs serially and in order.
-- It only executes jobs on demand when they are collected, not eagerly.
--
-- Cancelling will cancel /all/ jobs that have not been collected yet.
--
newSerialJobControl :: IO (JobControl IO a)
newSerialJobControl = do
    qVar <- newTChanIO
    return JobControl {
      spawnJob      = spawn     qVar,
      collectJob    = collect   qVar,
      remainingJobs = remaining qVar,
      cancelJobs    = cancel    qVar
    }
  where
    spawn :: TChan (IO a) -> IO a -> IO ()
    spawn qVar job = atomically $ writeTChan qVar job

    collect :: TChan (IO a) -> IO a
    collect qVar =
      join $ atomically $ readTChan qVar

    remaining :: TChan (IO a) -> IO Bool
    remaining qVar  = fmap not $ atomically $ isEmptyTChan qVar

    cancel :: TChan (IO a) -> IO ()
    cancel qVar = do
      _ <- atomically $ readAllTChan qVar
      return ()

-- | Make a 'JobControl' that eagerly executes jobs in parallel, with a given
-- maximum degree of parallelism.
--
-- Cancelling will cancel jobs that have not yet begun executing, but jobs
-- that have already been executed or are currently executing cannot be
-- cancelled.
--
newParallelJobControl :: Int -> IO (JobControl IO a)
newParallelJobControl n | n < 1 || n > 1000 =
  error $ "newParallelJobControl: not a sensible number of jobs: " ++ show n
newParallelJobControl maxJobLimit = do
    inqVar   <- newTChanIO
    outqVar  <- newTChanIO
    countVar <- newTVarIO 0
    replicateM_ maxJobLimit $
      forkIO $
        worker inqVar outqVar
    return JobControl {
      spawnJob      = spawn   inqVar  countVar,
      collectJob    = collect outqVar countVar,
      remainingJobs = remaining       countVar,
      cancelJobs    = cancel  inqVar  countVar
    }
  where
    worker ::  TChan (IO a) -> TChan (Either SomeException a) -> IO ()
    worker inqVar outqVar =
      forever $ do
        job <- atomically $ readTChan inqVar
        res <- try job
        atomically $ writeTChan outqVar res

    spawn :: TChan (IO a) -> TVar Int -> IO a -> IO ()
    spawn inqVar countVar job =
      atomically $ do
        modifyTVar' countVar (+1)
        writeTChan inqVar job

    collect :: TChan (Either SomeException a) -> TVar Int -> IO a
    collect outqVar countVar = do
      res <- atomically $ do
        modifyTVar' countVar (subtract 1)
        readTChan outqVar
      either throwIO return res

    remaining :: TVar Int -> IO Bool
    remaining countVar = fmap (/=0) $ atomically $ readTVar countVar

    cancel :: TChan (IO a) -> TVar Int -> IO ()
    cancel inqVar countVar =
      atomically $ do
        xs <- readAllTChan inqVar
        modifyTVar' countVar (subtract (length xs))

readAllTChan :: TChan a -> STM [a]
readAllTChan qvar = go []
  where
    go xs = do
      mx <- tryReadTChan qvar
      case mx of
        Nothing -> return (reverse xs)
        Just x  -> go (x:xs)

-------------------------
-- Job limits and locks
--

data JobLimit = JobLimit QSem

newJobLimit :: Int -> IO JobLimit
newJobLimit n =
  fmap JobLimit (newQSem n)

withJobLimit :: JobLimit -> IO a -> IO a
withJobLimit (JobLimit sem) =
  bracket_ (waitQSem sem) (signalQSem sem)

newtype Lock = Lock (MVar ())

newLock :: IO Lock
newLock = fmap Lock $ newMVar ()

criticalSection :: Lock -> IO a -> IO a
criticalSection (Lock lck) act = bracket_ (takeMVar lck) (putMVar lck ()) act