cabal-install/src/Distribution/Client/JobControl.hs

   1 {-# LANGUAGE FlexibleContexts #-}
   2 -----------------------------------------------------------------------------
   3 -- |
   4 -- Module      :  Distribution.Client.JobControl
   5 -- Copyright   :  (c) Duncan Coutts 2012
   6 -- License     :  BSD-like
   7 --
   8 -- Maintainer  :  cabal-devel@haskell.org
   9 -- Stability   :  provisional
  10 -- Portability :  portable
  11 --
  12 -- A job control concurrency abstraction
  13 -----------------------------------------------------------------------------
  14 module Distribution.Client.JobControl (
  15     JobControl,
  16     newSerialJobControl,
  17     newParallelJobControl,
  18     spawnJob,
  19     collectJob,
  20     remainingJobs,
  21     cancelJobs,
  22
  23     JobLimit,
  24     newJobLimit,
  25     withJobLimit,
  26
  27     Lock,
  28     newLock,
  29     criticalSection
  30   ) where
  31
  32 import Distribution.Client.Compat.Prelude
  33 import Prelude ()
  34
  35 import Control.Monad (forever, replicateM_)
  36 import Control.Concurrent (forkIO)
  37 import Control.Concurrent.MVar
  38 import Control.Concurrent.STM (STM, atomically)
  39 import Control.Concurrent.STM.TVar
  40 import Control.Concurrent.STM.TChan
  41 import Control.Exception (bracket_, try)
  42 import Distribution.Compat.Stack
  43 import Distribution.Client.Compat.Semaphore
  44
  45
  46 -- | A simple concurrency abstraction. Jobs can be spawned and can complete
  47 -- in any order. This allows both serial and parallel implementations.
  48 --
  49 data JobControl m a = JobControl {
  50        -- | Add a new job to the pool of jobs
  51        spawnJob    :: m a -> m (),
  52
  53        -- | Wait until one job is complete
  54        collectJob  :: m a,
  55
  56        -- | Returns True if there are any outstanding jobs
  57        -- (ie spawned but yet to be collected)
  58        remainingJobs :: m Bool,
  59
  60        -- | Try to cancel any outstanding but not-yet-started jobs.
  61        -- Call 'remainingJobs' after this to find out if any jobs are left
  62        -- (ie could not be cancelled).
  63        cancelJobs  :: m ()
  64      }
  65
  66
  67 -- | Make a 'JobControl' that executes all jobs serially and in order.
  68 -- It only executes jobs on demand when they are collected, not eagerly.
  69 --
  70 -- Cancelling will cancel /all/ jobs that have not been collected yet.
  71 --
  72 newSerialJobControl :: IO (JobControl IO a)
  73 newSerialJobControl = do
  74     qVar <- newTChanIO
  75     return JobControl {
  76       spawnJob      = spawn     qVar,
  77       collectJob    = collect   qVar,
  78       remainingJobs = remaining qVar,
  79       cancelJobs    = cancel    qVar
  80     }
  81   where
  82     spawn :: TChan (IO a) -> IO a -> IO ()
  83     spawn qVar job = atomically $ writeTChan qVar job
  84
  85     collect :: TChan (IO a) -> IO a
  86     collect qVar =
  87       join $ atomically $ readTChan qVar
  88
  89     remaining :: TChan (IO a) -> IO Bool
  90     remaining qVar  = fmap not $ atomically $ isEmptyTChan qVar
  91
  92     cancel :: TChan (IO a) -> IO ()
  93     cancel qVar = do
  94       _ <- atomically $ readAllTChan qVar
  95       return ()
  96
  97 -- | Make a 'JobControl' that eagerly executes jobs in parallel, with a given
  98 -- maximum degree of parallelism.
  99 --
 100 -- Cancelling will cancel jobs that have not yet begun executing, but jobs
 101 -- that have already been executed or are currently executing cannot be
 102 -- cancelled.
 103 --
 104 newParallelJobControl :: WithCallStack (Int -> IO (JobControl IO a))
 105 newParallelJobControl n | n < 1 || n > 1000 =
 106   error $ "newParallelJobControl: not a sensible number of jobs: " ++ show n
 107 newParallelJobControl maxJobLimit = do
 108     inqVar   <- newTChanIO
 109     outqVar  <- newTChanIO
 110     countVar <- newTVarIO 0
 111     replicateM_ maxJobLimit $
 112       forkIO $
 113         worker inqVar outqVar
 114     return JobControl {
 115       spawnJob      = spawn   inqVar  countVar,
 116       collectJob    = collect outqVar countVar,
 117       remainingJobs = remaining       countVar,
 118       cancelJobs    = cancel  inqVar  countVar
 119     }
 120   where
 121     worker ::  TChan (IO a) -> TChan (Either SomeException a) -> IO ()
 122     worker inqVar outqVar =
 123       forever $ do
 124         job <- atomically $ readTChan inqVar
 125         res <- try job
 126         atomically $ writeTChan outqVar res
 127
 128     spawn :: TChan (IO a) -> TVar Int -> IO a -> IO ()
 129     spawn inqVar countVar job =
 130       atomically $ do
 131         modifyTVar' countVar (+1)
 132         writeTChan inqVar job
 133
 134     collect :: TChan (Either SomeException a) -> TVar Int -> IO a
 135     collect outqVar countVar = do
 136       res <- atomically $ do
 137         modifyTVar' countVar (subtract 1)
 138         readTChan outqVar
 139       either throwIO return res
 140
 141     remaining :: TVar Int -> IO Bool
 142     remaining countVar = fmap (/=0) $ atomically $ readTVar countVar
 143
 144     cancel :: TChan (IO a) -> TVar Int -> IO ()
 145     cancel inqVar countVar =
 146       atomically $ do
 147         xs <- readAllTChan inqVar
 148         modifyTVar' countVar (subtract (length xs))
 149
 150 readAllTChan :: TChan a -> STM [a]
 151 readAllTChan qvar = go []
 152   where
 153     go xs = do
 154       mx <- tryReadTChan qvar
 155       case mx of
 156         Nothing -> return (reverse xs)
 157         Just x  -> go (x:xs)
 158
 159 -------------------------
 160 -- Job limits and locks
 161 --
 162
 163 data JobLimit = JobLimit QSem
 164
 165 newJobLimit :: Int -> IO JobLimit
 166 newJobLimit n =
 167   fmap JobLimit (newQSem n)
 168
 169 withJobLimit :: JobLimit -> IO a -> IO a
 170 withJobLimit (JobLimit sem) =
 171   bracket_ (waitQSem sem) (signalQSem sem)
 172
 173 newtype Lock = Lock (MVar ())
 174
 175 newLock :: IO Lock
 176 newLock = fmap Lock $ newMVar ()
 177
 178 criticalSection :: Lock -> IO a -> IO a
 179 criticalSection (Lock lck) act = bracket_ (takeMVar lck) (putMVar lck ()) act