diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,20 @@ +# 1.5.0 + +Despite that the major part of the version was bumped up, this release does not include +any breaking changes, only improvements and additions. + +* Addition of a batch concept with `BacthId` which: + + * makes it possible for a `Scheduler` to be "resumable" with the help of `waitForBatch`, + `waitForBatch_` and `waitForBatchR` + * allows to cancel batches prematurely while keeping the results and not terminating the + scheduler itself. This can be done with new functions: `cancelBatch` and + `cancelBatchWith`. In order to check for batch status `getCurrentBatchId` and + `hasBatchFinished` functions can be used. + +* Addition of `GlobalScheduler` that can be reused throughout the codebase, thus reducing + initialization overhead. + # 1.4.2 * Add `withTrivialScheduler` diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -5,9 +5,13 @@ Whenever you have many actions you'd like to perform in parallel, but would only like to use a few threads to do the actual computation, this package is for you. -| Language | Travis | AppVeyor | Hackage | Nightly | LTS | -|:--------:|:------:|:--------:|:-------:|:-------:|:---:| -| ![GitHub top language](https://img.shields.io/github/languages/top/lehins/haskell-scheduler.svg) | [![Travis](https://img.shields.io/travis/lehins/haskell-scheduler/master.svg?label=Linux%20%26%20OS%20X)](https://travis-ci.org/lehins/haskell-scheduler) | [![AppVeyor](https://img.shields.io/appveyor/ci/lehins/haskell-scheduler/master.svg?label=Windows)](https://ci.appveyor.com/project/lehins/haskell-scheduler) | [![Hackage](https://img.shields.io/hackage/v/scheduler.svg)](https://hackage.haskell.org/package/scheduler)| [![Nightly](https://www.stackage.org/package/scheduler/badge/nightly)](https://www.stackage.org/nightly/package/scheduler) | [![Nightly](https://www.stackage.org/package/scheduler/badge/lts)](https://www.stackage.org/lts/package/scheduler) | +| Language | Travis | Azure | Coveralls | +|:--------:|:------:|:-----:|:---------:| +| ![GitHub top language](https://img.shields.io/github/languages/top/lehins/haskell-scheduler.svg) | [![Travis](https://img.shields.io/travis/lehins/haskell-scheduler/master.svg?label=Linux%20%26%20OS%20X)](https://travis-ci.org/lehins/haskell-scheduler) | [![Build Status](https://dev.azure.com/kuleshevich/haskell-scheduler/_apis/build/status/lehins.haskell-scheduler?branchName=master)](https://dev.azure.com/kuleshevich/haskell-scheduler/_build?branchName=master) | [![Coverage Status](https://coveralls.io/repos/github/lehins/haskell-scheduler/badge.svg?branch=master)](https://coveralls.io/github/lehins/haskell-scheduler?branch=master) | + +| Gihub | Hackage | Nightly | LTS | +|:------|:-------:|:-------:|:---:| +| [`scheduler`](https://github.com/lehins/haskell-scheduler) | [![Hackage](https://img.shields.io/hackage/v/scheduler.svg)](https://hackage.haskell.org/package/scheduler)| [![Nightly](https://www.stackage.org/package/scheduler/badge/nightly)](https://www.stackage.org/nightly/package/scheduler) | [![Nightly](https://www.stackage.org/package/scheduler/badge/lts)](https://www.stackage.org/lts/package/scheduler) | ## QuickStart diff --git a/scheduler.cabal b/scheduler.cabal --- a/scheduler.cabal +++ b/scheduler.cabal @@ -1,13 +1,13 @@ name: scheduler -version: 1.4.2.3 +version: 1.5.0 synopsis: Work stealing scheduler. -description: A work stealing scheduler that is primarily developed for [massiv](https://github.com/lehins/massiv) array library, but it is general enough to be useful for any computation that fits the model of few workers and many jobs. +description: A work stealing scheduler that is designed for parallelization of heavy work loads. It was primarily developed for [massiv](https://github.com/lehins/massiv) array library, but it is general enough to be useful for any computation that fits the model of few workers and many jobs. homepage: https://github.com/lehins/haskell-scheduler license: BSD3 license-file: LICENSE author: Alexey Kuleshevich maintainer: alexey@kuleshevi.ch -copyright: 2018-2019 Alexey Kuleshevich +copyright: 2018-2020 Alexey Kuleshevich category: Parallelism, Concurrency build-type: Simple extra-source-files: README.md @@ -17,16 +17,18 @@ library hs-source-dirs: src exposed-modules: Control.Scheduler + , Control.Scheduler.Global , Control.Scheduler.Internal , Control.Scheduler.Computation , Control.Scheduler.Queue + , Control.Scheduler.Types build-depends: base >= 4.9 && < 5 , atomic-primops , deepseq , exceptions , unliftio-core , primitive >= 0.6.4 - + , pvar < 2.0 default-language: Haskell2010 ghc-options: -Wall diff --git a/src/Control/Scheduler.hs b/src/Control/Scheduler.hs --- a/src/Control/Scheduler.hs +++ b/src/Control/Scheduler.hs @@ -1,14 +1,9 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE CPP #-} -{-# LANGUAGE LambdaCase #-} {-# LANGUAGE NamedFieldPuns #-} -{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE RankNTypes #-} -{-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE ScopedTypeVariables #-} -- | -- Module : Control.Scheduler --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2020 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental @@ -40,6 +35,17 @@ , scheduleWorkState , scheduleWorkState_ , replicateWork + -- * Batches + , Batch + , runBatch + , runBatch_ + , runBatchR + , cancelBatch + , cancelBatch_ + , cancelBatchWith + , hasBatchFinished + , getCurrentBatch + -- * Early termination , terminate , terminate_ , terminateWith @@ -63,22 +69,18 @@ , MutexException(..) ) where -import Control.Concurrent -import Control.Exception import Control.Monad import Control.Monad.IO.Unlift import Control.Monad.Primitive (PrimMonad) import Control.Scheduler.Computation import Control.Scheduler.Internal +import Control.Scheduler.Types import Control.Scheduler.Queue -import Data.Atomics (atomicModifyIORefCAS, atomicModifyIORefCAS_) -import qualified Data.Foldable as F (foldl', traverse_, toList) -import Data.IORef -import Data.Maybe (catMaybes) +import qualified Data.Foldable as F (traverse_, toList) import Data.Primitive.SmallArray -import Data.Primitive.MutVar import Data.Traversable + -- | Get the underlying `Scheduler`, which cannot access `WorkerStates`. -- -- @since 1.4.0 @@ -86,51 +88,13 @@ unwrapSchedulerWS = _getScheduler --- | Initialize a separate state for each worker. --- --- @since 1.4.0 -initWorkerStates :: MonadIO m => Comp -> (WorkerId -> m s) -> m (WorkerStates s) -initWorkerStates comp initState = do - nWorkers <- getCompWorkers comp - arr <- liftIO $ newSmallArray nWorkers (error "Uninitialized") - let go i = when (i < nWorkers) $ do - state <- initState (WorkerId i) - liftIO $ writeSmallArray arr i state - go (i + 1) - go 0 - workerStates <- liftIO $ unsafeFreezeSmallArray arr - mutex <- liftIO $ newIORef False - pure - WorkerStates - {_workerStatesComp = comp, _workerStatesArray = workerStates, _workerStatesMutex = mutex} - -- | Get the computation strategy the states where initialized with. -- -- @since 1.4.0 workerStatesComp :: WorkerStates s -> Comp workerStatesComp = _workerStatesComp -withSchedulerWSInternal :: - MonadUnliftIO m - => (Comp -> (Scheduler m a -> t) -> m b) - -> WorkerStates s - -> (SchedulerWS s m a -> t) - -> m b -withSchedulerWSInternal withScheduler' states action = - withRunInIO $ \run -> bracket lockState unlockState (run . runSchedulerWS) - where - mutex = _workerStatesMutex states - lockState = atomicModifyIORef' mutex ((,) True) - unlockState wasLocked - | wasLocked = pure () - | otherwise = writeIORef mutex False - runSchedulerWS isLocked - | isLocked = liftIO $ throwIO MutexException - | otherwise = - withScheduler' (_workerStatesComp states) $ \scheduler -> - action (SchedulerWS states scheduler) - -- | Run a scheduler with stateful workers. Throws `MutexException` if an attempt is made -- to concurrently use the same `WorkerStates` with another `SchedulerWS`. -- @@ -211,7 +175,10 @@ -- | Schedule an action to be picked up and computed by a worker from a pool of --- jobs. Argument supplied to the job will be the id of the worker doing the job. +-- jobs. Argument supplied to the job will be the id of the worker doing the job. This is +-- useful for identification of a thread that will be doing the work, since there is +-- one-to-one mapping from `Control.Concurrent.ThreadId` to `WorkerId` for a particular +-- scheduler. -- -- @since 1.2.0 scheduleWorkId :: Scheduler m a -> (WorkerId -> m a) -> m () @@ -219,15 +186,15 @@ -- | As soon as possible try to terminate any computation that is being performed by all -- workers managed by this scheduler and collect whatever results have been computed, with --- supplied element guaranteed to being the last one. In case when `Results` is the return --- type this function will cause the scheduler to produce `FinishedEarly` +-- supplied element guaranteed to being the last one. In case when `Results` type is +-- returned this function will cause the scheduler to produce `FinishedEarly` -- -- /Important/ - With `Seq` strategy this will not stop other scheduled tasks from being computed, -- although it will make sure their results are discarded. -- -- @since 1.1.0 terminate :: Scheduler m a -> a -> m a -terminate = _terminate +terminate scheduler a = _terminate scheduler (Early a) -- | Same as `terminate`, but returning a single element list containing the supplied -- argument. This can be very useful for parallel search algorithms. In case when @@ -240,7 +207,7 @@ -- -- @since 1.1.0 terminateWith :: Scheduler m a -> a -> m a -terminateWith = _terminateWith +terminateWith scheduler a = _terminate scheduler (EarlyWith a) -- | Schedule an action to be picked up and computed by a worker from a pool of -- jobs. Similar to `scheduleWorkId`, except the job doesn't get the worker id. @@ -249,6 +216,9 @@ scheduleWork :: Scheduler m a -> m a -> m () scheduleWork scheduler f = _scheduleWorkId scheduler (const f) + +-- FIXME: get rid of scheduleJob and decide at `scheduleWork` level if we should use Job or Job_ +-- Type here should be `scheduleWork_ :: Scheduler m a -> m () -> m () -- | Same as `scheduleWork`, but only for a `Scheduler` that doesn't keep the results. -- -- @since 1.1.0 @@ -280,19 +250,7 @@ -- -- @since 1.1.0 terminate_ :: Scheduler m () -> m () -terminate_ = (`_terminateWith` ()) - --- | The most basic scheduler that simply runs the task instead of scheduling it. Early termination --- requests are bluntly ignored. --- --- @since 1.1.0 -trivialScheduler_ :: Applicative f => Scheduler f () -trivialScheduler_ = Scheduler - { _numWorkers = 1 - , _scheduleWorkId = \f -> f (WorkerId 0) - , _terminate = const $ pure () - , _terminateWith = const $ pure () - } +terminate_ = (`_terminate` Early ()) -- | This trivial scheduler will behave in the same way as `withScheduler` with `Seq` @@ -302,39 +260,8 @@ withTrivialScheduler :: PrimMonad m => (Scheduler m a -> m b) -> m [a] withTrivialScheduler action = F.toList <$> withTrivialSchedulerR action --- | This trivial scheduler will behave in a similar way as `withSchedulerR` with `Seq` --- computation strategy, except it is restricted to `PrimMonad`, instead of --- `MonadUnliftIO` and the work isn't scheduled, but rather computed immediately. --- --- @since 1.4.2 -withTrivialSchedulerR :: PrimMonad m => (Scheduler m a -> m b) -> m (Results a) -withTrivialSchedulerR action = do - resVar <- newMutVar [] - finResVar <- newMutVar Nothing - _ <- action $ Scheduler - { _numWorkers = 1 - , _scheduleWorkId = \f -> do - r <- f (WorkerId 0) - modifyMutVar' resVar (r:) - , _terminate = \r -> do - rs <- readMutVar resVar - writeMutVar finResVar (Just (FinishedEarly rs r)) - pure r - , _terminateWith = \r -> do - writeMutVar finResVar (Just (FinishedEarlyWith r)) - pure r - } - readMutVar finResVar >>= \case - Just rs -> pure $ reverseResults rs - Nothing -> Finished . Prelude.reverse <$> readMutVar resVar --- | This is generally a faster way to traverse while ignoring the result rather than using `mapM_`. --- --- @since 1.0.0 -traverse_ :: (Applicative f, Foldable t) => (a -> f ()) -> t a -> f () -traverse_ f = F.foldl' (\c a -> c *> f a) (pure ()) - -- | Map an action over each element of the `Traversable` @t@ acccording to the supplied computation -- strategy. -- @@ -374,60 +301,8 @@ withScheduler_ comp $ \s -> scheduleWork s $ replicateM_ n (scheduleWork s $ void f) -scheduleJobs :: MonadIO m => Jobs m a -> (WorkerId -> m a) -> m () -scheduleJobs = scheduleJobsWith mkJob --- | Similarly to `scheduleWork`, but ignores the result of computation, thus having less overhead. --- --- @since 1.0.0 -scheduleJobs_ :: MonadIO m => Jobs m a -> (WorkerId -> m b) -> m () -scheduleJobs_ = scheduleJobsWith (\job -> pure (Job_ (void . job))) -scheduleJobsWith :: - MonadIO m => ((WorkerId -> m b) -> m (Job m a)) -> Jobs m a -> (WorkerId -> m b) -> m () -scheduleJobsWith mkJob' jobs action = do - liftIO $ atomicModifyIORefCAS_ (jobsCountRef jobs) (+ 1) - job <- - mkJob' $ \ i -> do - res <- action i - res `seq` - dropCounterOnZero (jobsCountRef jobs) $ - retireWorkersN (jobsQueue jobs) (jobsNumWorkers jobs) - return res - pushJQueue (jobsQueue jobs) job - --- | Helper function to place required number of @Retire@ instructions on the job queue. -retireWorkersN :: MonadIO m => JQueue m a -> Int -> m () -retireWorkersN jobsQueue n = traverse_ (pushJQueue jobsQueue) $ replicate n Retire - --- | Decrease a counter by one and perform an action when it drops down to zero. -dropCounterOnZero :: MonadIO m => IORef Int -> m () -> m () -dropCounterOnZero counterRef onZero = do - jc <- - liftIO $ - atomicModifyIORefCAS - counterRef - (\ !i' -> - let !i = i' - 1 - in (i, i)) - when (jc == 0) onZero - - --- | Runs the worker until the job queue is exhausted, at which point it will execute the final task --- of retirement and return -runWorker :: MonadIO m => - WorkerId - -> JQueue m a - -> m () -- ^ Action to run upon retirement - -> m () -runWorker wId jQueue onRetire = go - where - go = - popJQueue jQueue >>= \case - Just job -> job wId >> go - Nothing -> onRetire - - -- | Initialize a scheduler and submit jobs that will be computed sequentially or in parallelel, -- which is determined by the `Comp`utation strategy. -- @@ -462,7 +337,10 @@ -> (Scheduler m a -> m b) -- ^ Action that will be scheduling all the work. -> m [a] -withScheduler comp = withSchedulerInternal comp scheduleJobs readResults (reverse . resultsToList) +withScheduler Seq = fmap (reverse . resultsToList) . withTrivialSchedulerRIO +withScheduler comp = + fmap (reverse . resultsToList) . withSchedulerInternal comp scheduleJobs readResults +{-# INLINE withScheduler #-} -- | Same as `withScheduler`, except instead of a list it produces `Results`, which allows -- for distinguishing between the ways computation was terminated. @@ -474,7 +352,9 @@ -> (Scheduler m a -> m b) -- ^ Action that will be scheduling all the work. -> m (Results a) -withSchedulerR comp = withSchedulerInternal comp scheduleJobs readResults reverseResults +withSchedulerR Seq = fmap reverseResults . withTrivialSchedulerRIO +withSchedulerR comp = fmap reverseResults . withSchedulerInternal comp scheduleJobs readResults +{-# INLINE withSchedulerR #-} -- | Same as `withScheduler`, but discards results of submitted jobs. @@ -486,125 +366,118 @@ -> (Scheduler m a -> m b) -- ^ Action that will be scheduling all the work. -> m () -withScheduler_ comp = void . withSchedulerInternal comp scheduleJobs_ (const (pure [])) (const ()) +withScheduler_ Seq = void . withTrivialSchedulerRIO +withScheduler_ comp = void . withSchedulerInternal comp scheduleJobs_ (const (pure [])) +{-# INLINE withScheduler_ #-} -withSchedulerInternal :: - MonadUnliftIO m - => Comp -- ^ Computation strategy - -> (Jobs m a -> (WorkerId -> m a) -> m ()) -- ^ How to schedule work - -> (JQueue m a -> m [Maybe a]) -- ^ How to collect results - -> (Results a -> c) -- ^ Adjust results in some way - -> (Scheduler m a -> m b) - -- ^ Action that will be scheduling all the work. - -> m c -withSchedulerInternal comp submitWork collect adjust onScheduler = do - jobsNumWorkers <- getCompWorkers comp - sWorkersCounterRef <- liftIO $ newIORef jobsNumWorkers - jobsQueue <- newJQueue - jobsCountRef <- liftIO $ newIORef 0 - workDoneMVar <- liftIO newEmptyMVar - let jobs = Jobs {..} - scheduler = - Scheduler - { _numWorkers = jobsNumWorkers - , _scheduleWorkId = submitWork jobs - , _terminate = - \a -> do - mas <- collect jobsQueue - let as = catMaybes mas - liftIO $ - void $ tryPutMVar workDoneMVar $ SchedulerTerminatedEarly (FinishedEarly as a) - pure a - , _terminateWith = - \a -> do - liftIO $ - void $ tryPutMVar workDoneMVar $ SchedulerTerminatedEarly (FinishedEarlyWith a) - pure a - } - onRetire = - dropCounterOnZero sWorkersCounterRef $ - void $ liftIO (tryPutMVar workDoneMVar SchedulerFinished) - -- / Wait for the initial jobs to get scheduled before spawining off the workers, otherwise it - -- would be trickier to identify the beginning and the end of a job pool. - _ <- onScheduler scheduler - -- / Ensure at least something gets scheduled, so retirement can be triggered - jc <- liftIO $ readIORef jobsCountRef - when (jc == 0) $ scheduleJobs_ jobs (\_ -> pure ()) - let spawnWorkersWith fork ws = - withRunInIO $ \run -> - forM (zip [0 ..] ws) $ \(wId, on) -> - fork on $ \unmask -> - catch - (unmask $ run $ runWorker wId jobsQueue onRetire) - (run . handleWorkerException jobsQueue workDoneMVar jobsNumWorkers) - spawnWorkers = - case comp of - Seq -> return [] - -- \ no need to fork threads for a sequential computation - Par -> spawnWorkersWith forkOnWithUnmask [1 .. jobsNumWorkers] - ParOn ws -> spawnWorkersWith forkOnWithUnmask ws - ParN _ -> spawnWorkersWith (\_ -> forkIOWithUnmask) [1 .. jobsNumWorkers] - terminateWorkers = liftIO . traverse_ (`throwTo` SomeAsyncException WorkerTerminateException) - doWork tids = do - when (comp == Seq) $ runWorker 0 jobsQueue onRetire - mExc <- liftIO (readMVar workDoneMVar) - -- \ wait for all worker to finish. If any one of the workers had a problem, then - -- this MVar will contain an exception - case mExc of - SchedulerFinished -> adjust . Finished . catMaybes <$> collect jobsQueue - -- \ Now we are sure all workers have done their job we can safely read all of the - -- IORefs with results - SchedulerTerminatedEarly as -> terminateWorkers tids >> pure (adjust as) - SchedulerWorkerException (WorkerException exc) -> liftIO $ throwIO exc - -- \ Here we need to unwrap the legit worker exception and rethrow it, so the main thread - -- will think like it's his own - safeBracketOnError spawnWorkers terminateWorkers doWork -resultsToList :: Results a -> [a] -resultsToList = \case - Finished rs -> rs - FinishedEarly rs r -> r:rs - FinishedEarlyWith r -> [r] +-- | Check if the supplied batch has already finished. +-- +-- @since 1.5.0 +hasBatchFinished :: Functor m => Batch m a -> m Bool +hasBatchFinished = batchHasFinished +{-# INLINE hasBatchFinished #-} -reverseResults :: Results a -> Results a -reverseResults = \case - Finished rs -> Finished (reverse rs) - FinishedEarly rs r -> FinishedEarly (reverse rs) r - res -> res +-- | Cancel batch with supplied identifier, which will lead to scheduler to return +-- `FinishedEarly` result. This is an idempotent operation and has no affect if currently +-- running batch does not match supplied identifier. Returns `False` when cancelling did +-- not succeed due to mismatched identifier or does not return at all since all jobs get +-- cancelled immediately. For trivial schedulers however there is no way to perform +-- concurrent cancelation and it will return `True`. +-- +-- @since 1.5.0 +cancelBatch :: Batch m a -> a -> m Bool +cancelBatch = batchCancel +{-# INLINE cancelBatch #-} +-- | Same as `cancelBatch`, but only works with schedulers that don't care about results +-- +-- @since 1.5.0 +cancelBatch_ :: Batch m () -> m Bool +cancelBatch_ b = batchCancel b () +{-# INLINE cancelBatch_ #-} --- | Specialized exception handler for the work scheduler. -handleWorkerException :: - MonadIO m => JQueue m a -> MVar (SchedulerOutcome a) -> Int -> SomeException -> m () -handleWorkerException jQueue workDoneMVar nWorkers exc = - case asyncExceptionFromException exc of - Just WorkerTerminateException -> return () - -- \ some co-worker died, we can just move on with our death. - _ -> do - _ <- liftIO $ tryPutMVar workDoneMVar $ SchedulerWorkerException $ WorkerException exc - -- \ Main thread must know how we died - -- / Do the co-worker cleanup - retireWorkersN jQueue (nWorkers - 1) +-- | Same as `cancelBatch_`, but the result of computation will be set to `FinishedEarlyWith` +-- +-- @since 1.5.0 +cancelBatchWith :: Batch m a -> a -> m Bool +cancelBatchWith = batchCancelWith +{-# INLINE cancelBatchWith #-} --- Copy from unliftio: -safeBracketOnError :: MonadUnliftIO m => m a -> (a -> m b) -> (a -> m c) -> m c -safeBracketOnError before after thing = withRunInIO $ \run -> mask $ \restore -> do - x <- run before - res1 <- try $ restore $ run $ thing x - case res1 of - Left (e1 :: SomeException) -> do - _ :: Either SomeException b <- - try $ uninterruptibleMask_ $ run $ after x - throwIO e1 - Right y -> return y +-- | This function gives a way to get access to the main batch that started implicitely. +-- +-- @since 1.5.0 +getCurrentBatch :: + Monad m => Scheduler m a -> m (Batch m a) +getCurrentBatch scheduler = do + batchId <- _currentBatchId scheduler + pure $ Batch + { batchCancel = _cancelBatch scheduler batchId . Early + , batchCancelWith = _cancelBatch scheduler batchId . EarlyWith + , batchHasFinished = (batchId /=) <$> _currentBatchId scheduler + } +{-# INLINE getCurrentBatch #-} + +-- | Run a single batch of jobs. Supplied action will not return until all jobs placed on +-- the queue are done or the whole batch is cancelled with one of these `cancelBatch`, +-- `cancelBatch_` or `cancelBatchWith`. +-- +-- It waits for all scheduled jobs to finish and collects the computed results into a +-- list. It is a blocking operation, but if there are no jobs in progress it will return +-- immediately. It is safe to continue using the supplied scheduler after this function +-- returns. However, if any of the jobs resulted in an exception it will be rethrown by this +-- function, which, unless caught, will further put the scheduler in a terminated state. +-- +-- It is important to note that any job that hasn't had its results collected from the +-- scheduler prior to starting the batch it will end up on the batch result list. +-- +-- @since 1.5.0 +runBatch :: + Monad m => Scheduler m a -> (Batch m a -> m c) -> m [a] +runBatch scheduler f = do + _ <- f =<< getCurrentBatch scheduler + reverse . resultsToList <$> _waitForCurrentBatch scheduler +{-# INLINE runBatch #-} + +-- | Same as `runBatch`, except it ignores results of computation +-- +-- @since 1.5.0 +runBatch_ :: + Monad m => Scheduler m () -> (Batch m () -> m c) -> m () +runBatch_ scheduler f = do + _ <- f =<< getCurrentBatch scheduler + void (_waitForCurrentBatch scheduler) +{-# INLINE runBatch_ #-} + + +-- | Same as `runBatch`, except it produces `Results` instead of a list. +-- +-- @since 1.5.0 +runBatchR :: + Monad m => Scheduler m a -> (Batch m a -> m c) -> m (Results a) +runBatchR scheduler f = do + _ <- f =<< getCurrentBatch scheduler + reverseResults <$> _waitForCurrentBatch scheduler +{-# INLINE runBatchR #-} + +{- $setup + +>>> import Control.Exception +>>> import Control.Concurrent +>>> import Control.Concurrent.MVar + +-} + + {- $exceptions If any one of the workers dies with an exception, even if that exceptions is asynchronous, it will be re-thrown in the scheduling thread. + >>> let didAWorkerDie = handleJust asyncExceptionFromException (return . (== ThreadKilled)) . fmap or >>> :t didAWorkerDie diff --git a/src/Control/Scheduler/Global.hs b/src/Control/Scheduler/Global.hs new file mode 100644 --- /dev/null +++ b/src/Control/Scheduler/Global.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RecordWildCards #-} +-- | +-- Module : Control.Scheduler.Global +-- Copyright : (c) Alexey Kuleshevich 2020 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Control.Scheduler.Global + ( GlobalScheduler + , globalScheduler + , newGlobalScheduler + , withGlobalScheduler_ + ) where + +import Data.Maybe +import Control.Concurrent (ThreadId) +import Control.Concurrent.MVar +import Control.Exception +import Control.Monad +import Control.Monad.IO.Unlift +import Control.Scheduler +import Control.Scheduler.Internal +import Control.Scheduler.Types +import Data.IORef +import System.IO.Unsafe (unsafePerformIO) + +-- | Global scheduler with `Par` computation strategy that can be used anytime using +-- `withGlobalScheduler_` +globalScheduler :: GlobalScheduler IO +globalScheduler = unsafePerformIO (newGlobalScheduler Par) +{-# NOINLINE globalScheduler #-} + + +initGlobalScheduler :: + MonadUnliftIO m => Comp -> (Scheduler m a -> [ThreadId] -> m b) -> m b +initGlobalScheduler comp action = do + (jobs, mkScheduler) <- initScheduler comp scheduleJobs_ (const (pure [])) + safeBracketOnError (spawnWorkers jobs comp) (liftIO . terminateWorkers) $ \tids -> + action (mkScheduler tids) tids + + +-- | Create a new global scheduler, in case a single one `globalScheduler` is not +-- sufficient. It is important to note that too much parallelization can significantly +-- degrate performance, therefore it is best not to use more than one scheduler at a time. +-- +-- @since 1.5.0 +newGlobalScheduler :: MonadUnliftIO m => Comp -> m (GlobalScheduler m) +newGlobalScheduler comp = + initGlobalScheduler comp $ \scheduler tids -> + liftIO $ do + mvar <- newMVar scheduler + tidsRef <- newIORef tids + _ <- mkWeakMVar mvar (readIORef tidsRef >>= terminateWorkers) + pure $ + GlobalScheduler + { globalSchedulerComp = comp + , globalSchedulerMVar = mvar + , globalSchedulerThreadIdsRef = tidsRef + } + +-- | Use the global scheduler if it is not busy, otherwise initialize a temporary one. It +-- means that this function by itself will not block, but if the same global scheduler +-- used concurrently other schedulers might get created. +-- +-- @since 1.5.0 +withGlobalScheduler_ :: MonadUnliftIO m => GlobalScheduler m -> (Scheduler m () -> m a) -> m () +withGlobalScheduler_ GlobalScheduler {..} action = + withRunInIO $ \run -> do + let initializeNewScheduler = do + initGlobalScheduler globalSchedulerComp $ \scheduler tids -> + liftIO $ do + oldTids <- atomicModifyIORef' globalSchedulerThreadIdsRef $ (,) tids + terminateWorkers oldTids + putMVar globalSchedulerMVar scheduler + mask $ \restore -> + tryTakeMVar globalSchedulerMVar >>= \case + Nothing -> restore $ run $ withScheduler_ globalSchedulerComp action + Just scheduler -> do + let runScheduler = + run $ do + _ <- action scheduler + mEarly <- _earlyResults scheduler + mEarly <$ when (isNothing mEarly) (void (_waitForCurrentBatch scheduler)) + mEarly <- restore runScheduler `onException` run initializeNewScheduler + -- Whenever a scheduler is terminated it is no longer usable, need to re-initialize + case mEarly of + Nothing -> putMVar globalSchedulerMVar scheduler + Just _ -> run initializeNewScheduler diff --git a/src/Control/Scheduler/Internal.hs b/src/Control/Scheduler/Internal.hs --- a/src/Control/Scheduler/Internal.hs +++ b/src/Control/Scheduler/Internal.hs @@ -1,141 +1,486 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE LambdaCase #-} -{-# OPTIONS_HADDOCK hide, not-home #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE Unsafe #-} +{-# OPTIONS_HADDOCK hide, not-home #-} -- | -- Module : Control.Scheduler.Internal --- Copyright : (c) Alexey Kuleshevich 2018-2019 +-- Copyright : (c) Alexey Kuleshevich 2018-2020 -- License : BSD3 -- Maintainer : Alexey Kuleshevich -- Stability : experimental -- Portability : non-portable -- module Control.Scheduler.Internal - ( Scheduler(..) - , WorkerStates(..) - , SchedulerWS(..) - , Jobs(..) - , Results(..) - , SchedulerOutcome(..) - , WorkerException(..) - , WorkerTerminateException(..) - , MutexException(..) + ( withSchedulerInternal + , initWorkerStates + , withSchedulerWSInternal + , trivialScheduler_ + , withTrivialSchedulerR + , withTrivialSchedulerRIO + , initScheduler + , spawnWorkers + , terminateWorkers + , scheduleJobs + , scheduleJobs_ + , scheduleJobsWith + , reverseResults + , resultsToList + , traverse_ + , safeBracketOnError ) where +import Data.Coerce +import Control.Concurrent import Control.Exception +import Control.Monad +import Control.Monad.IO.Unlift import Control.Scheduler.Computation +import Control.Scheduler.Types import Control.Scheduler.Queue +import Data.Atomics (atomicModifyIORefCAS, atomicModifyIORefCAS_) +import qualified Data.Foldable as F (foldl') import Data.IORef import Data.Primitive.SmallArray +import Data.Primitive.MutVar +import Data.Primitive.PVar --- | Computed outcome of scheduled jobs. --- --- @since 1.4.2 -data Results a - = Finished ![a] - -- ^ Finished normally with all scheduled jobs completed - | FinishedEarly ![a] !a - -- ^ Finished early by the means of `Control.Scheduler.terminate`. - | FinishedEarlyWith !a - -- ^ Finished early by the means of `Control.Scheduler.terminateWith`. - deriving (Show, Read, Eq) -instance Functor Results where - fmap f = - \case - Finished xs -> Finished (fmap f xs) - FinishedEarly xs x -> FinishedEarly (fmap f xs) (f x) - FinishedEarlyWith x -> FinishedEarlyWith (f x) -instance Foldable Results where - foldr f acc = - \case - Finished xs -> foldr f acc xs - FinishedEarly xs x -> foldr f (f x acc) xs - FinishedEarlyWith x -> f x acc - foldr1 f = - \case - Finished xs -> foldr1 f xs - FinishedEarly xs x -> foldr f x xs - FinishedEarlyWith x -> x +-- | Initialize a separate state for each worker. +-- +-- @since 1.4.0 +initWorkerStates :: MonadIO m => Comp -> (WorkerId -> m s) -> m (WorkerStates s) +initWorkerStates comp initState = do + nWorkers <- getCompWorkers comp + arr <- liftIO $ newSmallArray nWorkers (error "Uninitialized") + let go i = + when (i < nWorkers) $ do + state <- initState (WorkerId i) + liftIO $ writeSmallArray arr i state + go (i + 1) + go 0 + workerStates <- liftIO $ unsafeFreezeSmallArray arr + mutex <- liftIO $ newIORef False + pure + WorkerStates + {_workerStatesComp = comp, _workerStatesArray = workerStates, _workerStatesMutex = mutex} -instance Traversable Results where - traverse f = - \case - Finished xs -> Finished <$> traverse f xs - FinishedEarly xs x -> FinishedEarly <$> traverse f xs <*> f x - FinishedEarlyWith x -> FinishedEarlyWith <$> f x +withSchedulerWSInternal :: + MonadUnliftIO m + => (Comp -> (Scheduler m a -> t) -> m b) + -> WorkerStates s + -> (SchedulerWS s m a -> t) + -> m b +withSchedulerWSInternal withScheduler' states action = + withRunInIO $ \run -> bracket lockState unlockState (run . runSchedulerWS) + where + mutex = _workerStatesMutex states + lockState = atomicModifyIORefCAS mutex $ (,) True + unlockState wasLocked + | wasLocked = pure () + | otherwise = atomicWriteIORef mutex False + runSchedulerWS isLocked + | isLocked = liftIO $ throwIO MutexException + | otherwise = + withScheduler' (_workerStatesComp states) $ \scheduler -> + action (SchedulerWS states scheduler) -data Jobs m a = Jobs - { jobsNumWorkers :: {-# UNPACK #-} !Int - , jobsQueue :: !(JQueue m a) - , jobsCountRef :: !(IORef Int) - } +-- | The most basic scheduler that simply runs the task instead of scheduling it. Early termination +-- requests are bluntly ignored. +-- +-- @since 1.1.0 +trivialScheduler_ :: Applicative f => Scheduler f () +trivialScheduler_ = + Scheduler + { _numWorkers = 1 + , _scheduleWorkId = \f -> f (WorkerId 0) + , _terminate = const $ pure () + , _waitForCurrentBatch = pure $ Finished [] + , _earlyResults = pure Nothing + , _currentBatchId = pure $ BatchId 0 + , _cancelBatch = \_ _ -> pure False + , _batchEarly = pure Nothing + } --- | Main type for scheduling work. See `Control.Scheduler.withScheduler` or --- `Control.Scheduler.withScheduler_` for ways to construct and use this data type. + +-- | This trivial scheduler will behave in a similar way as +-- `Control.Scheduler.withSchedulerR` with `Seq` computation strategy, except it is +-- restricted to `PrimMonad`, instead of `MonadUnliftIO` and the work isn't scheduled, but +-- rather computed immediately. -- --- @since 1.0.0 -data Scheduler m a = Scheduler - { _numWorkers :: {-# UNPACK #-} !Int - , _scheduleWorkId :: (WorkerId -> m a) -> m () - , _terminate :: a -> m a - , _terminateWith :: a -> m a - } +-- @since 1.4.2 +withTrivialSchedulerR :: PrimMonad m => (Scheduler m a -> m b) -> m (Results a) +withTrivialSchedulerR action = do + resVar <- newMutVar [] + batchVar <- newMutVar $ BatchId 0 + finResVar <- newMutVar Nothing + batchEarlyVar <- newMutVar Nothing + let bumpCurrentBatchId = atomicModifyMutVar' batchVar (\(BatchId x) -> (BatchId (x + 1), ())) + bumpBatchId (BatchId c) = + atomicModifyMutVar' batchVar $ \b@(BatchId x) -> + if x == c + then (BatchId (x + 1), True) + else (b, False) + takeBatchEarly = atomicModifyMutVar' batchEarlyVar $ \mEarly -> (Nothing, mEarly) + takeResults = atomicModifyMutVar' resVar $ \res -> ([], res) + _ <- + action $ + Scheduler + { _numWorkers = 1 + , _scheduleWorkId = + \f -> do + r <- f (WorkerId 0) + r `seq` atomicModifyMutVar' resVar (\rs -> (r : rs, ())) + , _terminate = + \early -> do + bumpCurrentBatchId + finishEarly <- collectResults (Just early) takeResults + unEarly early <$ writeMutVar finResVar (Just finishEarly) + , _waitForCurrentBatch = + do mEarly <- takeBatchEarly + bumpCurrentBatchId + collectResults mEarly . pure =<< takeResults + , _earlyResults = readMutVar finResVar + , _currentBatchId = readMutVar batchVar + , _batchEarly = takeBatchEarly + , _cancelBatch = + \batchId early -> do + b <- bumpBatchId batchId + when b $ writeMutVar batchEarlyVar (Just early) + pure b + } + readMutVar finResVar >>= \case + Just rs -> pure $ reverseResults rs + Nothing -> do + mEarly <- takeBatchEarly + reverseResults <$> collectResults mEarly takeResults --- | This is a wrapper around `Scheduler`, but it also keeps a separate state for each --- individual worker. See `Control.Scheduler.withSchedulerWS` or --- `Control.Scheduler.withSchedulerWS_` for ways to construct and use this data type. + + +-- | Same as `Control.Scheduler.withTrivialScheduler`, but works in `MonadUnliftIO` and +-- returns results in an original LIFO order. -- --- @since 1.4.0 -data SchedulerWS s m a = SchedulerWS - { _workerStates :: !(WorkerStates s) - , _getScheduler :: !(Scheduler m a) - } +-- @since 1.4.2 +withTrivialSchedulerRIO :: MonadUnliftIO m => (Scheduler m a -> m b) -> m (Results a) +withTrivialSchedulerRIO action = do + resRef <- liftIO $ newIORef [] + batchRef <- liftIO $ newIORef $ BatchId 0 + finResRef <- liftIO $ newIORef Nothing + batchEarlyRef <- liftIO $ newIORef Nothing + let bumpCurrentBatchId = atomicModifyIORefCAS_ (coerce batchRef) (+ (1 :: Int)) + bumpBatchId (BatchId c) = + atomicModifyIORefCAS batchRef $ \b@(BatchId x) -> + if x == c + then (BatchId (x + 1), True) + else (b, False) + takeBatchEarly = atomicModifyIORefCAS batchEarlyRef $ \mEarly -> (Nothing, mEarly) + takeResults = atomicModifyIORefCAS resRef $ \res -> ([], res) + scheduler = + Scheduler + { _numWorkers = 1 + , _scheduleWorkId = + \f -> do + r <- f (WorkerId 0) + r `seq` liftIO (atomicModifyIORefCAS_ resRef (r :)) + , _terminate = + \ !early -> + liftIO $ do + bumpCurrentBatchId + finishEarly <- collectResults (Just early) takeResults + atomicWriteIORef finResRef (Just finishEarly) + throwIO TerminateEarlyException + , _waitForCurrentBatch = + liftIO $ do + bumpCurrentBatchId + mEarly <- takeBatchEarly + collectResults mEarly . pure =<< takeResults + , _earlyResults = liftIO (readIORef finResRef) + , _currentBatchId = liftIO (readIORef batchRef) + , _batchEarly = liftIO takeBatchEarly + , _cancelBatch = + \batchId early -> liftIO $ do + b <- bumpBatchId batchId + when b $ atomicWriteIORef batchEarlyRef (Just early) + pure b + } + _ :: Either TerminateEarlyException b <- withRunInIO $ \run -> try $ run $ action scheduler + liftIO (readIORef finResRef) >>= \case + Just rs -> pure rs + Nothing -> + liftIO $ do + mEarly <- takeBatchEarly + collectResults mEarly takeResults +{-# INLINEABLE withTrivialSchedulerRIO #-} --- | Each worker is capable of keeping it's own state, that can be share for different --- schedulers, but not at the same time. In other words using the same `WorkerStates` on --- `Control.Scheduler.withSchedulerS` concurrently will result in an error. Can be initialized with --- `Control.Scheduler.initWorkerStates` + +-- | This is generally a faster way to traverse while ignoring the result rather than using `mapM_`. -- --- @since 1.4.0 -data WorkerStates s = WorkerStates - { _workerStatesComp :: !Comp - , _workerStatesArray :: !(SmallArray s) - , _workerStatesMutex :: !(IORef Bool) - } +-- @since 1.0.0 +traverse_ :: (Applicative f, Foldable t) => (a -> f ()) -> t a -> f () +traverse_ f = F.foldl' (\c a -> c *> f a) (pure ()) +{-# INLINE traverse_ #-} +scheduleJobs :: MonadIO m => Jobs m a -> (WorkerId -> m a) -> m () +scheduleJobs = scheduleJobsWith mkJob +{-# INLINEABLE scheduleJobs #-} -data SchedulerOutcome a - = SchedulerFinished - | SchedulerTerminatedEarly !(Results a) - | SchedulerWorkerException WorkerException +-- | Ignores the result of computation, thus avoiding some overhead. +scheduleJobs_ :: MonadIO m => Jobs m a -> (WorkerId -> m b) -> m () +scheduleJobs_ = scheduleJobsWith (\job -> pure (Job_ (void . job (\_ -> pure ())))) +{-# INLINEABLE scheduleJobs_ #-} +scheduleJobsWith :: + MonadIO m + => (((b -> m ()) -> WorkerId -> m ()) -> m (Job m a)) + -> Jobs m a + -> (WorkerId -> m b) + -> m () +scheduleJobsWith mkJob' Jobs {..} action = do + job <- + mkJob' $ \storeResult wid -> do + res <- action wid + res `seq` storeResult res + liftIO $ void $ atomicAddIntPVar jobsQueueCount 1 + pushJQueue jobsQueue job +{-# INLINEABLE scheduleJobsWith #-} --- | This exception should normally be never seen in the wild and is for internal use only. -newtype WorkerException = - WorkerException SomeException - -- ^ One of workers experienced an exception, main thread will receive the same `SomeException`. - deriving (Show) -instance Exception WorkerException +-- | Runs the worker until it is terminated with a `WorkerTerminateException` or is killed +-- by some other asynchronous exception, which will propagate to the user calling thread. +runWorker :: + MonadUnliftIO m + => (forall b. m b -> IO b) + -> (forall c. IO c -> IO c) + -> WorkerId + -> Jobs m a + -> IO () +runWorker run unmask wId Jobs {jobsQueue, jobsQueueCount, jobsSchedulerStatus} = go + where + onBlockedMVar eUnblocked = + case eUnblocked of + Right () -> go + Left uExc + | Just WorkerTerminateException <- asyncExceptionFromException uExc -> return () + Left uExc + | Just CancelBatchException <- asyncExceptionFromException uExc -> go + Left uExc -> throwIO uExc + go = do + eRes <- try $ do + job <- run (popJQueue jobsQueue) + unmask (run (job wId) >> atomicSubIntPVar jobsQueueCount 1) + -- \ popJQueue can block, but it is still interruptable + case eRes of + Right 1 -> try (putMVar jobsSchedulerStatus SchedulerIdle) >>= onBlockedMVar + Right _ -> go + Left exc + | Just WorkerTerminateException <- asyncExceptionFromException exc -> return () + Left exc + | Just CancelBatchException <- asyncExceptionFromException exc -> go + Left exc -> do + eUnblocked <- + try $ putMVar jobsSchedulerStatus (SchedulerWorkerException (WorkerException exc)) + -- \ without blocking with putMVar here we would not be able to report an + -- exception in the main thread, especially if `exc` is asynchronous. + unless (isSyncException exc) $ throwIO exc + onBlockedMVar eUnblocked +{-# INLINEABLE runWorker #-} -data WorkerTerminateException = - WorkerTerminateException - -- ^ When a co-worker dies of some exception, all the other ones will be terminated - -- asynchronously with this one. - deriving (Show) +initScheduler :: + MonadIO m + => Comp + -> (Jobs m a -> (WorkerId -> m a) -> m ()) + -> (JQueue m a -> m [a]) + -> m (Jobs m a, [ThreadId] -> Scheduler m a) +initScheduler comp submitWork collect = do + jobsNumWorkers <- getCompWorkers comp + jobsQueue <- newJQueue + jobsQueueCount <- liftIO $ newPVar 1 + jobsSchedulerStatus <- liftIO newEmptyMVar + earlyTerminationResultRef <- liftIO $ newIORef Nothing + batchIdRef <- liftIO $ newIORef $ BatchId 0 + batchEarlyRef <- liftIO $ newIORef Nothing + let jobs = + Jobs + { jobsNumWorkers = jobsNumWorkers + , jobsQueue = jobsQueue + , jobsQueueCount = jobsQueueCount + , jobsSchedulerStatus = jobsSchedulerStatus + } + bumpCurrentBatchId = atomicModifyIORefCAS_ (coerce batchIdRef) (+ (1 :: Int)) + bumpBatchId (BatchId c) = + atomicModifyIORefCAS batchIdRef $ \b@(BatchId x) -> + if x == c + then (BatchId (x + 1), True) + else (b, False) + mkScheduler tids = + Scheduler + { _numWorkers = jobsNumWorkers + , _scheduleWorkId = submitWork jobs + , _terminate = + \early -> do + finishEarly <- + case early of + Early r -> FinishedEarly <$> collect jobsQueue <*> pure r + EarlyWith r -> pure $ FinishedEarlyWith r + liftIO $ do + bumpCurrentBatchId + atomicWriteIORef earlyTerminationResultRef $ Just finishEarly + throwIO TerminateEarlyException + , _waitForCurrentBatch = + do scheduleJobs_ jobs (\_ -> liftIO $ void $ atomicSubIntPVar jobsQueueCount 1) + unblockPopJQueue jobsQueue + status <- liftIO $ takeMVar jobsSchedulerStatus + mEarly <- liftIO $ atomicModifyIORefCAS batchEarlyRef $ \mEarly -> (Nothing, mEarly) + rs <- + case status of + SchedulerWorkerException (WorkerException exc) -> + case fromException exc of + Just CancelBatchException -> do + _ <- clearPendingJQueue jobsQueue + liftIO $ + traverse_ (`throwTo` SomeAsyncException CancelBatchException) tids + collectResults mEarly . pure =<< collect jobsQueue + Nothing -> liftIO $ throwIO exc + SchedulerIdle -> do + blockPopJQueue jobsQueue + liftIO bumpCurrentBatchId + res <- collect jobsQueue + res `seq` collectResults mEarly (pure res) + rs <$ liftIO (atomicWriteIntPVar jobsQueueCount 1) + , _earlyResults = liftIO (readIORef earlyTerminationResultRef) + , _currentBatchId = liftIO (readIORef batchIdRef) + , _batchEarly = liftIO (readIORef batchEarlyRef) + , _cancelBatch = + \batchId early -> do + b <- liftIO $ bumpBatchId batchId + when b $ do + blockPopJQueue jobsQueue + liftIO $ do + atomicWriteIORef batchEarlyRef $ Just early + throwIO CancelBatchException + pure b + } + pure (jobs, mkScheduler) +{-# INLINEABLE initScheduler #-} -instance Exception WorkerTerminateException +withSchedulerInternal :: + MonadUnliftIO m + => Comp -- ^ Computation strategy + -> (Jobs m a -> (WorkerId -> m a) -> m ()) -- ^ How to schedule work + -> (JQueue m a -> m [a]) -- ^ How to collect results + -> (Scheduler m a -> m b) + -- ^ Action that will be scheduling all the work. + -> m (Results a) +withSchedulerInternal comp submitWork collect onScheduler = do + (jobs@Jobs {..}, mkScheduler) <- initScheduler comp submitWork collect + -- / Wait for the initial jobs to get scheduled before spawining off the workers, otherwise it + -- would be trickier to identify the beginning and the end of a job pool. + withRunInIO $ \run -> do + bracket (run (spawnWorkers jobs comp)) terminateWorkers $ \tids -> + let scheduler = mkScheduler tids + readEarlyTermination = + _earlyResults scheduler >>= \case + Nothing -> error "Impossible: uninitialized early termination value" + Just rs -> pure rs + in try (run (onScheduler scheduler)) >>= \case + Left TerminateEarlyException -> run readEarlyTermination + Right _ -> do + run $ scheduleJobs_ jobs (\_ -> liftIO $ void $ atomicSubIntPVar jobsQueueCount 1) + run $ unblockPopJQueue jobsQueue + status <- takeMVar jobsSchedulerStatus + -- \ wait for all worker to finish. If any one of the workers had a problem, then + -- this MVar will contain an exception + case status of + SchedulerWorkerException (WorkerException exc) + | Just TerminateEarlyException <- fromException exc -> run readEarlyTermination + | Just CancelBatchException <- fromException exc -> + run $ do + mEarly <- _batchEarly scheduler + collectResults mEarly (collect jobsQueue) + | otherwise -> throwIO exc + -- \ Here we need to unwrap the legit worker exception and rethrow it, so + -- the main thread will think like it's his own + SchedulerIdle -> + run $ do + mEarly <- _batchEarly scheduler + collectResults mEarly (collect jobsQueue) + -- \ Now we are sure all workers have done their job we can safely read + -- all of the IORefs with results +{-# INLINEABLE withSchedulerInternal #-} --- | Exception that gets thrown whenever concurrent access is attempted to the `WorkerStates` --- --- @since 1.4.0 -data MutexException = - MutexException - deriving (Eq, Show) -instance Exception MutexException where - displayException MutexException = - "MutexException: WorkerStates cannot be used at the same time by different schedulers" +collectResults :: Applicative f => Maybe (Early a) -> f [a] -> f (Results a) +collectResults mEarly collect = + case mEarly of + Nothing -> Finished <$> collect + Just (Early r) -> FinishedEarly <$> collect <*> pure r + Just (EarlyWith r) -> pure $ FinishedEarlyWith r +{-# INLINEABLE collectResults #-} + + +spawnWorkers :: forall m a. MonadUnliftIO m => Jobs m a -> Comp -> m [ThreadId] +spawnWorkers jobs@Jobs {jobsNumWorkers} = + \case + Par -> spawnWorkersWith forkOnWithUnmask [1 .. jobsNumWorkers] + ParOn ws -> spawnWorkersWith forkOnWithUnmask ws + ParN _ -> spawnWorkersWith (\_ -> forkIOWithUnmask) [1 .. jobsNumWorkers] + Seq -> spawnWorkersWith (\_ -> forkIOWithUnmask) [1 :: Int] + -- \ sequential computation is suboptimal when used in this way. + where + spawnWorkersWith :: + MonadUnliftIO m + => (Int -> ((forall c. IO c -> IO c) -> IO ()) -> IO ThreadId) + -> [Int] + -> m [ThreadId] + spawnWorkersWith fork ws = + withRunInIO $ \run -> + forM (zip [0 ..] ws) $ \(wId, on) -> + fork on $ \unmask -> runWorker run unmask wId jobs +{-# INLINEABLE spawnWorkers #-} + +terminateWorkers :: [ThreadId] -> IO () +terminateWorkers = traverse_ (`throwTo` SomeAsyncException WorkerTerminateException) + +-- | Conversion to a list. Elements are expected to be in the orignal LIFO order, so +-- calling `reverse` is still necessary for getting the results in FIFO order. +resultsToList :: Results a -> [a] +resultsToList = \case + Finished rs -> rs + FinishedEarly rs r -> r:rs + FinishedEarlyWith r -> [r] +{-# INLINEABLE resultsToList #-} + + +reverseResults :: Results a -> Results a +reverseResults = \case + Finished rs -> Finished (reverse rs) + FinishedEarly rs r -> FinishedEarly (reverse rs) r + res -> res +{-# INLINEABLE reverseResults #-} + + + +-- Copies from unliftio + +isSyncException :: Exception e => e -> Bool +isSyncException exc = + case fromException (toException exc) of + Just (SomeAsyncException _) -> False + Nothing -> True + +safeBracketOnError :: MonadUnliftIO m => m a -> (a -> m b) -> (a -> m c) -> m c +safeBracketOnError before after thing = withRunInIO $ \run -> mask $ \restore -> do + x <- run before + res1 <- try $ restore $ run $ thing x + case res1 of + Left (e1 :: SomeException) -> do + _ :: Either SomeException b <- + try $ uninterruptibleMask_ $ run $ after x + throwIO e1 + Right y -> return y diff --git a/src/Control/Scheduler/Queue.hs b/src/Control/Scheduler/Queue.hs --- a/src/Control/Scheduler/Queue.hs +++ b/src/Control/Scheduler/Queue.hs @@ -1,6 +1,8 @@ {-# OPTIONS_HADDOCK hide, not-home #-} -{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} -- | -- Module : Control.Scheduler.Queue -- Copyright : (c) Alexey Kuleshevich 2018-2019 @@ -11,30 +13,34 @@ -- module Control.Scheduler.Queue ( -- * Job queue - Job(Retire, Job_) + Job(Job_) , mkJob - , JQueue + , Queue(..) + , JQueue(..) , WorkerId(..) , newJQueue , pushJQueue , popJQueue + , clearPendingJQueue , readResults + , blockPopJQueue + , unblockPopJQueue ) where import Control.Concurrent.MVar -import Control.Monad (join, void) +import Control.Monad (join) import Control.Monad.IO.Unlift -import Data.Atomics (atomicModifyIORefCAS) +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 :: !(MVar ()) + , qBaton :: {-# UNPACK #-}!(MVar ()) } @@ -45,7 +51,7 @@ -- @since 1.4.0 newtype WorkerId = WorkerId { getWorkerId :: Int - } deriving (Show, Read, Eq, Ord, Enum, Num) + } deriving (Show, Read, Eq, Ord, Enum, Bounded, Num, Real, Integral) popQueue :: Queue m a -> Maybe (Job m a, Queue m a) @@ -56,69 +62,94 @@ case reverse (qStack queue) of [] -> Nothing y:ys -> Just (y, queue {qQueue = ys, qStack = []}) +{-# INLINEABLE popQueue #-} data Job m a - = Job !(IORef (Maybe a)) (WorkerId -> m a) + = Job {-# UNPACK #-} !(IORef (Maybe a)) (WorkerId -> m ()) | Job_ (WorkerId -> m ()) - | Retire -mkJob :: MonadIO m => (WorkerId -> m a) -> m (Job m a) +mkJob :: MonadIO m => ((a -> m ()) -> WorkerId -> m ()) -> m (Job m a) mkJob action = do resRef <- liftIO $ newIORef Nothing - return $! - Job resRef $ \ i -> do - res <- action i - liftIO $ writeIORef resRef $ Just res - return $! res + return $ Job resRef (action (liftIO . writeIORef resRef . Just)) +{-# INLINEABLE mkJob #-} -newtype JQueue m a = JQueue (IORef (Queue m a)) +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 - + 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 = +pushJQueue (JQueue jQueueRef _) job = liftIO $ do newBaton <- newEmptyMVar join $ - atomicModifyIORefCAS - jQueueRef - (\Queue {qQueue, qStack, qResults, qBaton} -> - ( Queue - qQueue - (job : qStack) - (case job of + atomicModifyIORefCAS jQueueRef $ \queue@Queue {qStack, qResults, qBaton} -> + ( queue + { qResults = + case job of Job resRef _ -> resRef : qResults - _ -> qResults) - newBaton - , liftIO $ putMVar qBaton ())) - + _ -> qResults + , qStack = job : qStack + , qBaton = newBaton + } + , putMVar qBaton ()) +{-# INLINEABLE pushJQueue #-} -popJQueue :: MonadIO m => JQueue m a -> m (Maybe (WorkerId -> m ())) -popJQueue (JQueue jQueueRef) = liftIO inner +-- | 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 = + inner = do + readMVar lock join $ - atomicModifyIORefCAS jQueueRef $ \queue -> - case popQueue queue of - Nothing -> (queue, readMVar (qBaton queue) >> inner) - Just (job, newQueue) -> - ( newQueue - , case job of - Job _ action -> return $ Just (void . action) - Job_ action_ -> return $ Just action_ - Retire -> return Nothing) + 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 #-} -readResults :: MonadIO m => JQueue m a -> m [Maybe a] -readResults (JQueue jQueueRef) = +-- | 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 - resRefs <- qResults <$> readIORef jQueueRef - mapM readIORef resRefs + results <- + atomicModifyIORefCAS jQueueRef $ \queue -> + (queue {qQueue = [], qStack = [], qResults = []}, qResults queue) + catMaybes <$> mapM readIORef results +{-# INLINEABLE readResults #-} + + diff --git a/src/Control/Scheduler/Types.hs b/src/Control/Scheduler/Types.hs new file mode 100644 --- /dev/null +++ b/src/Control/Scheduler/Types.hs @@ -0,0 +1,217 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE Unsafe #-} +{-# OPTIONS_HADDOCK hide, not-home #-} +-- | +-- Module : Control.Scheduler.Types +-- Copyright : (c) Alexey Kuleshevich 2018-2020 +-- License : BSD3 +-- Maintainer : Alexey Kuleshevich +-- Stability : experimental +-- Portability : non-portable +-- +module Control.Scheduler.Types + ( Scheduler(..) + , WorkerStates(..) + , SchedulerWS(..) + , GlobalScheduler(..) + , Batch(..) + , BatchId(..) + , Jobs(..) + , Early(..) + , unEarly + , Results(..) + , SchedulerStatus(..) + , WorkerException(..) + , CancelBatchException(..) + , TerminateEarlyException(..) + , WorkerTerminateException(..) + , MutexException(..) + ) where + +import Control.Concurrent (ThreadId) +import Control.Concurrent.MVar +import Control.Exception +import Control.Scheduler.Computation +import Control.Scheduler.Queue +import Data.IORef +import Data.Primitive.SmallArray +import Data.Primitive.PVar + +-- | Computed results of scheduled jobs. +-- +-- @since 1.4.2 +data Results a + = Finished [a] + -- ^ Finished normally with all scheduled jobs completed + | FinishedEarly [a] !a + -- ^ Finished early by the means of `Control.Scheduler.cancelBatch` or + -- `Control.Scheduler.terminate`. + | FinishedEarlyWith !a + -- ^ Finished early by the means of `Control.Scheduler.cancelBatchWith` or + -- `Control.Scheduler.terminateWith`. + deriving (Show, Read, Eq) + +instance Functor Results where + fmap f = + \case + Finished xs -> Finished (fmap f xs) + FinishedEarly xs x -> FinishedEarly (fmap f xs) (f x) + FinishedEarlyWith x -> FinishedEarlyWith (f x) + +instance Foldable Results where + foldr f acc = + \case + Finished xs -> foldr f acc xs + FinishedEarly xs x -> foldr f (f x acc) xs + FinishedEarlyWith x -> f x acc + foldr1 f = + \case + Finished xs -> foldr1 f xs + FinishedEarly xs x -> foldr f x xs + FinishedEarlyWith x -> x + +instance Traversable Results where + traverse f = + \case + Finished xs -> Finished <$> traverse f xs + FinishedEarly xs x -> FinishedEarly <$> traverse f xs <*> f x + FinishedEarlyWith x -> FinishedEarlyWith <$> f x + +data Jobs m a = Jobs + { jobsNumWorkers :: {-# UNPACK #-} !Int + , jobsQueue :: !(JQueue m a) +#if MIN_VERSION_pvar(1,0,0) + , jobsQueueCount :: !(PVar Int RealWorld) +#else + , jobsQueueCount :: !(PVar IO Int) +#endif + , jobsSchedulerStatus :: !(MVar SchedulerStatus) + } + + +-- | This is a result for premature ending of computation. +data Early a + = Early a + -- ^ This value along with all results computed up to the moment when computation was + -- cancelled or termianted will be returned + | EarlyWith a + -- ^ Only this value will be returned all other results will get discarded + +unEarly :: Early a -> a +unEarly (Early r) = r +unEarly (EarlyWith r) = r + +-- | Main type for scheduling work. See `Control.Scheduler.withScheduler` or +-- `Control.Scheduler.withScheduler_` for ways to construct and use this data type. +-- +-- @since 1.0.0 +data Scheduler m a = Scheduler + { _numWorkers :: {-# UNPACK #-} !Int + , _scheduleWorkId :: (WorkerId -> m a) -> m () + , _terminate :: Early a -> m a + , _waitForCurrentBatch :: m (Results a) + , _earlyResults :: m (Maybe (Results a)) + , _currentBatchId :: m BatchId + -- ^ Returns an opaque identifier for current batch of jobs, which can be used to either + -- cancel the batch early or simply check if the batch has finished or not. + , _cancelBatch :: BatchId -> Early a -> m Bool + -- ^ Stops current batch and cancells all the outstanding jobs and the ones that are + -- currently in progress. + , _batchEarly :: m (Maybe (Early a)) + } + +-- | This is a wrapper around `Scheduler`, but it also keeps a separate state for each +-- individual worker. See `Control.Scheduler.withSchedulerWS` or +-- `Control.Scheduler.withSchedulerWS_` for ways to construct and use this data type. +-- +-- @since 1.4.0 +data SchedulerWS s m a = SchedulerWS + { _workerStates :: !(WorkerStates s) + , _getScheduler :: !(Scheduler m a) + } + +-- | Each worker is capable of keeping it's own state, that can be share for different +-- schedulers, but not at the same time. In other words using the same `WorkerStates` on +-- `Control.Scheduler.withSchedulerS` concurrently will result in an error. Can be initialized with +-- `Control.Scheduler.initWorkerStates` +-- +-- @since 1.4.0 +data WorkerStates s = WorkerStates + { _workerStatesComp :: !Comp + , _workerStatesArray :: !(SmallArray s) + , _workerStatesMutex :: !(IORef Bool) + } + +-- | This identifier is needed for tracking batches. +newtype BatchId = BatchId { getBatchId :: Int } + deriving (Show, Eq, Ord) + + +-- | Batch is an artifical checkpoint that can be controlled by the user throughout the +-- lifetime of a scheduler. +-- +-- @since 1.5.0 +data Batch m a = Batch + { batchCancel :: a -> m Bool + , batchCancelWith :: a -> m Bool + , batchHasFinished :: m Bool + } + + +-- | A thread safe wrapper around `Scheduler`, which allows it to be reused indefinitely +-- and globally if need be. There is one already created in this library: +-- `Control.Scheduler.Global.globalSchdeuler` +-- +-- @since 1.5.0 +data GlobalScheduler m = + GlobalScheduler + { globalSchedulerComp :: !Comp + , globalSchedulerMVar :: !(MVar (Scheduler m ())) + , globalSchedulerThreadIdsRef :: !(IORef [ThreadId]) + } + + +data SchedulerStatus + = SchedulerIdle + | SchedulerWorkerException WorkerException + +data TerminateEarlyException = + TerminateEarlyException + deriving (Show) + +instance Exception TerminateEarlyException + +data CancelBatchException = + CancelBatchException + deriving (Show) + +instance Exception CancelBatchException + +-- | This exception should normally be never seen in the wild and is for internal use only. +newtype WorkerException = + WorkerException SomeException + -- ^ One of workers experienced an exception, main thread will receive the same `SomeException`. + deriving (Show) + +instance Exception WorkerException + +data WorkerTerminateException = + WorkerTerminateException + -- ^ When a co-worker dies of some exception, all the other ones will be terminated + -- asynchronously with this one. + deriving (Show) + + +instance Exception WorkerTerminateException + +-- | Exception that gets thrown whenever concurrent access is attempted to the `WorkerStates` +-- +-- @since 1.4.0 +data MutexException = + MutexException + deriving (Eq, Show) + +instance Exception MutexException where + displayException MutexException = + "MutexException: WorkerStates cannot be used at the same time by different schedulers" diff --git a/tests/Control/SchedulerSpec.hs b/tests/Control/SchedulerSpec.hs --- a/tests/Control/SchedulerSpec.hs +++ b/tests/Control/SchedulerSpec.hs @@ -5,6 +5,7 @@ ( spec ) where +import Data.Int import Control.Concurrent (killThread, myThreadId, threadDelay, yield) import Control.Concurrent.MVar import Control.DeepSeq @@ -18,6 +19,7 @@ import Data.IORef import Data.List (groupBy, sort, sortOn) import Test.Hspec +import Test.Hspec.QuickCheck import Test.QuickCheck import Test.QuickCheck.Function import Test.QuickCheck.Monadic @@ -32,27 +34,36 @@ import Data.Semigroup #endif -concurrentProperty :: IO Property -> Property -concurrentProperty = within 1000000 . monadicIO . run +concurrentProperty :: Testable prop => prop -> Property +concurrentProperty = within 2000000 + +concurrentExpectation :: Expectation -> Property +concurrentExpectation = concurrentProperty + +concurrentPropertyIO :: IO Property -> Property +concurrentPropertyIO = concurrentProperty . monadicIO . run + instance Arbitrary Comp where + arbitrary = frequency [(20, pure Seq), (80, getNonSeq <$> arbitrary)] + +newtype NonSeq = NonSeq {getNonSeq :: Comp } + deriving (Show, Eq) + +instance Arbitrary NonSeq where arbitrary = - frequency - [ (20, pure Seq) - , (10, pure Par) - , (35, ParOn <$> arbitrary) - , (35, ParN . getSmall <$> arbitrary) - ] + NonSeq <$> + frequency [(10, pure Par), (35, ParOn <$> arbitrary), (35, ParN . getSmall <$> arbitrary)] prop_SameList :: Comp -> [Int] -> Property prop_SameList comp xs = - concurrentProperty $ do + concurrentPropertyIO $ do xs' <- withScheduler comp $ \scheduler -> mapM_ (scheduleWork scheduler . return) xs return (xs === xs') prop_Recursive :: Comp -> [Int] -> Property prop_Recursive comp xs = - concurrentProperty $ do + concurrentPropertyIO $ do xs' <- withScheduler comp (schedule xs) return (sort xs === sort xs') where @@ -62,7 +73,7 @@ prop_Serially :: Comp -> [Int] -> Property prop_Serially comp xs = - concurrentProperty $ do + concurrentPropertyIO $ do xs' <- schedule xs return (xs === concat xs') where @@ -74,7 +85,7 @@ prop_Nested :: Comp -> [Int] -> Property prop_Nested comp xs = - concurrentProperty $ do + concurrentPropertyIO $ do xs' <- schedule xs return (sort xs === sort (concat xs')) where @@ -85,34 +96,45 @@ -- | Check whether all jobs have been completed (similar roprop_Traverse) prop_AllJobsProcessed :: Comp -> [Int] -> Property prop_AllJobsProcessed comp jobs = + concurrentProperty $ monadicIO ((=== jobs) <$> run (withScheduler comp $ \scheduler -> mapM_ (scheduleWork scheduler . pure) jobs)) prop_Traverse :: Comp -> [Int] -> Fun Int Int -> Property prop_Traverse comp xs f = - concurrentProperty $ (===) <$> traverse f' xs <*> traverseConcurrently comp f' xs + concurrentPropertyIO $ (===) <$> traverse f' xs <*> traverseConcurrently comp f' xs where f' = pure . apply f replicateSeq :: (Int -> IO Int -> IO [Int]) -> Int -> Fun Int Int -> Property replicateSeq justAs n f = - concurrentProperty $ do + concurrentPropertyIO $ do iRef <- newIORef 0 jRef <- newIORef 0 let g ref = atomicModifyIORef' ref (\i -> (apply f i, i + 1)) (===) <$> S.replicateConcurrently Seq n (g jRef) <*> justAs n (g iRef) prop_ReplicateM :: Int -> Fun Int Int -> Property -prop_ReplicateM = replicateSeq replicateM +prop_ReplicateM i = concurrentProperty . replicateSeq replicateM i prop_ReplicateWorkSeq :: Int -> Fun Int Int -> Property -prop_ReplicateWorkSeq = replicateSeq (\ n g -> withScheduler Seq (\s -> replicateWork n s g)) +prop_ReplicateWorkSeq i = + concurrentProperty . replicateSeq (\ n g -> withScheduler Seq (\s -> replicateWork n s g)) i +prop_ManyJobsInChunks :: Property +prop_ManyJobsInChunks = noShrinking $ \ comp (jss :: [[Int]]) -> + concurrentExpectation $ do + xs <- withScheduler comp $ \s -> + forM_ jss $ \js -> do + rs <- runBatch s $ \ _ -> mapM_ (scheduleWork s . pure) js + rs `shouldBe` js + xs `shouldBe` [] + prop_ArbitraryCompNested :: [(Comp, Int)] -> Property prop_ArbitraryCompNested xs = - concurrentProperty $ do + concurrentPropertyIO $ do xs' <- schedule xs return (sort (map snd xs) === sort (concat xs')) where @@ -123,6 +145,7 @@ -- | Ensure proper exception handling. prop_CatchDivideByZero :: Comp -> Int -> [Positive Int] -> Property prop_CatchDivideByZero comp k xs = + concurrentProperty $ assertExceptionIO (== DivideByZero) (traverseConcurrently @@ -132,7 +155,8 @@ -- | Ensure proper exception handling. prop_CatchDivideByZeroNested :: Comp -> Int -> Positive Int -> Property -prop_CatchDivideByZeroNested comp a (Positive k) = assertExceptionIO (== DivideByZero) (schedule k) +prop_CatchDivideByZeroNested comp a (Positive k) = + concurrentProperty $ assertExceptionIO (== DivideByZero) (schedule k) where schedule i | i < 0 = return [] @@ -143,6 +167,7 @@ -- | Make sure one co-worker can kill another one, of course when there are at least two of. prop_KillBlockedCoworker :: Comp -> Property prop_KillBlockedCoworker comp = + concurrentProperty $ assertExceptionIO (== DivideByZero) (withScheduler_ comp $ \scheduler -> @@ -156,6 +181,7 @@ -- | Make sure one co-worker can kill another one, of course when there are at least two of. prop_KillSleepingCoworker :: Comp -> Property prop_KillSleepingCoworker comp = + concurrentProperty $ assertExceptionIO (== DivideByZero) (withScheduler_ comp $ \scheduler -> do @@ -167,6 +193,7 @@ prop_ExpectAsyncException :: Comp -> Property prop_ExpectAsyncException comp = + concurrentProperty $ let didAWorkerDie = EUnsafe.handleJust EUnsafe.asyncExceptionFromException (return . (== EUnsafe.ThreadKilled)) . fmap or @@ -177,6 +204,7 @@ prop_WorkerCaughtAsyncException :: Positive Int -> Property prop_WorkerCaughtAsyncException (Positive n) = + concurrentProperty $ assertExceptionIO (== DivideByZero) $ do lock <- newEmptyMVar result <- @@ -201,6 +229,7 @@ -- | Make sure there is no problems if sub-schedules worker get killed prop_AllWorkersDied :: Comp -> Comp -> Positive Int -> Property prop_AllWorkersDied comp1 comp (Positive n) = + concurrentProperty $ assertAsyncExceptionIO (== ThreadKilled) (withScheduler_ comp1 $ \scheduler1 -> @@ -209,9 +238,9 @@ (withScheduler_ comp $ \scheduler -> replicateM_ n (scheduleWork scheduler (myThreadId >>= killThread)))) -prop_FinishEarly_ :: Comp -> Property -prop_FinishEarly_ comp = - comp /= Seq ==> concurrentProperty $ do +prop_FinishEarly_ :: NonSeq -> Property +prop_FinishEarly_ (NonSeq comp) = + concurrentPropertyIO $ do ref <- newIORef True withScheduler_ comp $ \scheduler -> scheduleWork_ @@ -221,7 +250,7 @@ prop_FinishEarly :: Comp -> Property prop_FinishEarly comp = - concurrentProperty $ do + concurrentPropertyIO $ do let scheduleJobs scheduler = do scheduleWork scheduler (pure (2 :: Int)) scheduleWork scheduler (threadDelay 10000 >> terminate scheduler 3 >> pure 1) @@ -231,7 +260,7 @@ prop_FinishEarlyWith :: Comp -> Int -> Property prop_FinishEarlyWith comp n = - concurrentProperty $ do + concurrentPropertyIO $ do let scheduleJobs scheduler = do scheduleWork scheduler $ pure (complement (n + 1)) scheduleWork scheduler $ terminateWith scheduler n >> pure (complement n) @@ -241,7 +270,7 @@ prop_FinishBeforeStarting :: Comp -> Property prop_FinishBeforeStarting comp = - concurrentProperty $ do + concurrentPropertyIO $ do res <- withScheduler comp $ \scheduler -> do void $ terminate scheduler 1 @@ -250,7 +279,7 @@ prop_FinishWithBeforeStarting :: Comp -> Int -> Property prop_FinishWithBeforeStarting comp n = - concurrentProperty $ do + concurrentPropertyIO $ do res <- withScheduler comp $ \scheduler -> do void $ terminateWith scheduler n @@ -259,7 +288,7 @@ prop_TrivialSchedulerSameAsSeq_ :: [Int] -> Property prop_TrivialSchedulerSameAsSeq_ zs = - concurrentProperty $ do + concurrentPropertyIO $ do let consRef xsRef x = atomicModifyIORef' xsRef $ \ xs -> (x:xs, ()) trivial = trivialScheduler_ nRef <- newIORef False @@ -276,7 +305,7 @@ prop_SameAsTrivialScheduler :: Comp -> [Int] -> Fun Int Int -> Property prop_SameAsTrivialScheduler comp zs f = - concurrentProperty $ do + concurrentPropertyIO $ do let schedule scheduler = forM_ zs (scheduleWork scheduler . pure . apply f) xs <- withScheduler comp schedule ys <- withTrivialScheduler schedule @@ -290,13 +319,15 @@ -> [Int] -> Int -> [Int] - -> Expectation -prop_Terminate withSchedulerR' term expected xs x ys = do - rs <- withSchedulerR' $ \ scheduler -> do - forM_ xs (scheduleWork scheduler . pure) - _ <- scheduleWork scheduler $ term scheduler x - forM_ ys (scheduleWork scheduler . pure) - rs `shouldBe` expected xs x + -> Property +prop_Terminate withSchedulerR' term expected xs x ys = + concurrentExpectation $ do + rs <- + withSchedulerR' $ \scheduler -> do + forM_ xs (scheduleWork scheduler . pure) + _ <- scheduleWork scheduler $ term scheduler x + forM_ ys (scheduleWork scheduler . pure) + rs `shouldBe` expected xs x -- prop_TerminateSeq :: -- ((Scheduler IO Int -> IO ()) -> IO (Results Int)) -> [Int] -> Int -> [Int] -> Expectation @@ -325,7 +356,7 @@ -- | Check if an element is in the list with an exception prop_TraverseConcurrently_ :: Comp -> [Int] -> Int -> Property prop_TraverseConcurrently_ comp xs x = - concurrentProperty $ do + concurrentPropertyIO $ do let f i | i == x = throwIO $ Elem x | otherwise = pure () @@ -336,9 +367,9 @@ -- TODO: fix the infinite property for single worker schedulers -- | Check if an element is in the list with an exception, where we know that list is infinite and -- element is part of that list. -prop_TraverseConcurrentlyInfinite_ :: Comp -> [Int] -> Int -> Property -prop_TraverseConcurrentlyInfinite_ comp xs x = - comp /= Seq ==> concurrentProperty $ do +prop_TraverseConcurrentlyInfinite_ :: NonSeq -> [Int] -> Int -> Property +prop_TraverseConcurrentlyInfinite_ (NonSeq comp) xs x = + concurrentPropertyIO $ do let f i | i == x = throwIO $ Elem x | otherwise = pure () @@ -348,44 +379,119 @@ return (eRes === eRes') -prop_WorkerStateExclusive :: Comp -> NonNegative Int -> Expectation -prop_WorkerStateExclusive comp (NonNegative n) = do - state <- initWorkerStates comp (\wid -> (,) wid <$> newIORef (0 :: Int)) - workerStatesComp state `shouldBe` comp - nWorkers <- getCompWorkers comp - let scheduleJobs schedulerWS = do - replicateM n $ - scheduleWorkState schedulerWS $ \(wid, ref) -> do - counter <- readIORef ref - writeIORef ref (counter + 1) - pure (wid, counter) - gather = map (sortOn snd) . groupBy (\x y -> fst x == fst y) . sortOn fst - isMonotonicStartingAt _ [] = True - isMonotonicStartingAt k (k':ks) = k == k' && isMonotonicStartingAt (k + 1) ks - baseIds = [(wid, -1) | wid <- [0 .. WorkerId nWorkers - 1]] - ids <- withSchedulerWS state scheduleJobs - length ids `shouldBe` n - let gathered = gather (ids ++ baseIds) - map (map snd) gathered `shouldSatisfy` all (isMonotonicStartingAt (-1)) - ids' <- withSchedulerWSR state scheduleJobs - length ids' `shouldBe` n - let gathered' = gather (baseIds ++ ids ++ F.toList ids') - map (map snd) gathered' `shouldSatisfy` all (isMonotonicStartingAt (-1)) - withSchedulerWS_ state $ \schedulerWS -> do - numWorkers (unwrapSchedulerWS schedulerWS) `shouldBe` nWorkers - replicateM (10 * n) $ - scheduleWorkState_ schedulerWS $ \(wid, ref) -> do - counter <- readIORef ref - when (counter > 0) $ snd (last (gathered' !! getWorkerId wid)) `shouldBe` pred counter +prop_WorkerStateExclusive :: Comp -> NonNegative Int -> Property +prop_WorkerStateExclusive comp (NonNegative n) = + concurrentExpectation $ do + state <- initWorkerStates comp (\wid -> (,) wid <$> newIORef (0 :: Int)) + workerStatesComp state `shouldBe` comp + nWorkers <- getCompWorkers comp + let scheduleJobs schedulerWS = do + replicateM n $ + scheduleWorkState schedulerWS $ \(wid, ref) -> do + counter <- readIORef ref + writeIORef ref (counter + 1) + pure (wid, counter) + gather = map (sortOn snd) . groupBy (\x y -> fst x == fst y) . sortOn fst + isMonotonicStartingAt _ [] = True + isMonotonicStartingAt k (k':ks) = k == k' && isMonotonicStartingAt (k + 1) ks + baseIds = [(wid, -1) | wid <- [0 .. WorkerId nWorkers - 1]] + ids <- withSchedulerWS state scheduleJobs + length ids `shouldBe` n + let gathered = gather (ids ++ baseIds) + map (map snd) gathered `shouldSatisfy` all (isMonotonicStartingAt (-1)) + ids' <- withSchedulerWSR state scheduleJobs + length ids' `shouldBe` n + let gathered' = gather (baseIds ++ ids ++ F.toList ids') + map (map snd) gathered' `shouldSatisfy` all (isMonotonicStartingAt (-1)) + withSchedulerWS_ state $ \schedulerWS -> do + numWorkers (unwrapSchedulerWS schedulerWS) `shouldBe` nWorkers + replicateM (10 * n) $ + scheduleWorkState_ schedulerWS $ \(wid, ref) -> do + counter <- readIORef ref + when (counter > 0) $ snd (last (gathered' !! getWorkerId wid)) `shouldBe` pred counter prop_MutexException :: Comp -> Property prop_MutexException comp = + concurrentProperty $ assertExceptionIO (== MutexException) $ do state <- initWorkerStates comp (pure . getWorkerId) withSchedulerWS_ state $ \schedulerWS -> scheduleWorkState_ schedulerWS $ \_s -> withSchedulerWS_ state $ \_s' -> pure () +prop_FindCancelResume :: Comp -> Int64 -> ([Int64], [Int64]) -> [Int64] -> Property +prop_FindCancelResume comp x' (xs1', xs2') ys = + concurrentExpectation $ do + let f = (10 *) + g = (100 *) + xs1 = filter (/= x') xs1' + xs2 = filter (/= x') xs2' + xs = concat [xs1, [x'], xs2] + res <- + withSchedulerR comp $ \s -> do + r <- runBatchR s $ \batch -> do + forM_ xs $ \x -> + scheduleWork s $ do + if x == x' + then Just x <$ cancelBatchWith batch (Just (f x)) + else pure Nothing + r `shouldBe` FinishedEarlyWith (Just (f x')) + r' <- runBatchR s $ \_batch -> forM_ ys (scheduleWork s . pure . Just) + r' `shouldBe` Finished (map Just ys) + batch <- getCurrentBatch s + forM_ xs $ \x -> + scheduleWork s $ do + if x == x' + then Just x <$ cancelBatch batch (Just (g x)) + else pure $ Just (f x) + case res of + FinishedEarly rs r -> do + r `shouldBe` Just (g x') + rs `satisfyOrderedPartialPrefix` concat [map (Just . f) xs1, [Just x'], map (Just . f) xs2] + fr -> expectationFailure $ "Unexpected result: " ++ show fr + where + satisfyOrderedPartialPrefix as bs = + unless (orderedPartialPrefixOf as bs) $ + expectationFailure $ + "Expected " ++ + show as ++ " to be prefix of " ++ show bs ++ " possibly with some elements skipped" + -- Make sure the first list is the prefix of the second + orderedPartialPrefixOf [] _ = True + orderedPartialPrefixOf (_:_) [] = False + orderedPartialPrefixOf (a:as) (b:bs) + | a == b = orderedPartialPrefixOf as bs + | otherwise = orderedPartialPrefixOf (a : as) bs +-- prop_CancelBatchEarly_ :: NonSeq -> Property +-- prop_CancelBatchEarly_ (NonSeq comp) = +-- concurrentPropertyIO $ do +-- ref <- newIORef True +-- withScheduler_ comp $ \scheduler -> +-- scheduleWork_ +-- scheduler +-- (cancelBatch_ scheduler >> yield >> threadDelay 10000 >> writeIORef ref False) +-- counterexample "Scheduler did not terminate early" <$> readIORef ref + +-- prop_FinishEarly :: Comp -> Property +-- prop_FinishEarly comp = +-- concurrentPropertyIO $ do +-- let scheduleJobs scheduler = do +-- scheduleWork scheduler (pure (2 :: Int)) +-- scheduleWork scheduler (threadDelay 10000 >> terminate scheduler 3 >> pure 1) +-- res <- withScheduler comp scheduleJobs +-- res' <- withSchedulerR comp scheduleJobs +-- pure (res === [2, 3] .&&. res' === FinishedEarly [2] 3) + +-- prop_FinishEarlyWith :: Comp -> Int -> Property +-- prop_FinishEarlyWith comp n = +-- concurrentPropertyIO $ do +-- let scheduleJobs scheduler = do +-- scheduleWork scheduler $ pure (complement (n + 1)) +-- scheduleWork scheduler $ terminateWith scheduler n >> pure (complement n) +-- res <- withScheduler comp scheduleJobs +-- res' <- withSchedulerR comp scheduleJobs +-- pure (res === [n] .&&. res' === FinishedEarlyWith n) + + spec :: Spec spec = do describe "Comp" $ do @@ -431,55 +537,58 @@ terminate_ trivialScheduler_ `shouldReturn` () terminate trivialScheduler_ () `shouldReturn` () terminateWith trivialScheduler_ () `shouldReturn` () - it "TerminateSeq" $ timed $ prop_Terminate withTrivialScheduler terminate (\xs x -> xs ++ [x]) - it "TerminateWithSeq" $ timed $ prop_Terminate withTrivialScheduler terminateWith (\_ x -> [x]) - it "TerminateSeqR" $ timed $ prop_Terminate withTrivialSchedulerR terminate FinishedEarly - it "TerminateWithSeqR" $ - timed $ prop_Terminate withTrivialSchedulerR terminateWith (const FinishedEarlyWith) + prop "TerminateSeq" $ prop_Terminate withTrivialScheduler terminate (\xs x -> xs ++ [x]) + prop "TerminateWithSeq" $ prop_Terminate withTrivialScheduler terminateWith (\_ x -> [x]) + prop "TerminateSeqR" $ prop_Terminate withTrivialSchedulerR terminate FinishedEarly + prop "TerminateWithSeqR" $ + prop_Terminate withTrivialSchedulerR terminateWith (const FinishedEarlyWith) describe "Seq" $ do - it "SameList" $ timed $ prop_SameList Seq - it "Recursive" $ timed $ prop_Recursive Seq - it "Nested" $ timed $ prop_Nested Seq - it "Serially" $ timed $ prop_Serially Seq - it "TrivialAsSeq_" $ timed prop_TrivialSchedulerSameAsSeq_ - it "replicateConcurrently == replicateM" $ timed prop_ReplicateM - it "replicateConcurrently == replicateWork" $ timed prop_ReplicateWorkSeq + prop "SameList" $ prop_SameList Seq + prop "Recursive" $ prop_Recursive Seq + prop "Nested" $ prop_Nested Seq + prop "Serially" $ prop_Serially Seq + prop "TrivialAsSeq_" prop_TrivialSchedulerSameAsSeq_ + prop "replicateConcurrently == replicateM" prop_ReplicateM + prop "replicateConcurrently == replicateWork" prop_ReplicateWorkSeq it "WorkerIdIsZero" $ withScheduler Seq (`scheduleWorkId` pure) `shouldReturn` [0] - it "TerminateSeq" $ timed $ prop_Terminate (withScheduler Seq) terminate (\xs x -> xs ++ [x]) - it "TerminateWithSeq" $ timed $ prop_Terminate (withScheduler Seq) terminateWith (\_ x -> [x]) - it "TerminateSeqR" $ timed $ prop_Terminate (withSchedulerR Seq) terminate FinishedEarly - it "TerminateWithSeqR" $ - timed $ prop_Terminate (withSchedulerR Seq) terminateWith (const FinishedEarlyWith) + prop "TerminateSeq" $ prop_Terminate (withScheduler Seq) terminate (\xs x -> xs ++ [x]) + prop "TerminateWithSeq" $ prop_Terminate (withScheduler Seq) terminateWith (\_ x -> [x]) + prop "TerminateSeqR" $ prop_Terminate (withSchedulerR Seq) terminate FinishedEarly + prop "TerminateWithSeqR" $ + prop_Terminate (withSchedulerR Seq) terminateWith (const FinishedEarlyWith) describe "ParOn" $ do - it "SameList" $ timed $ \cs -> prop_SameList (ParOn cs) - it "Recursive" $ timed $ \cs -> prop_Recursive (ParOn cs) - it "Nested" $ timed $ \cs -> prop_Nested (ParOn cs) - it "Serially" $ timed $ \cs -> prop_Serially (ParOn cs) + prop "SameList" $ \cs -> prop_SameList (ParOn cs) + prop "Recursive" $ \cs -> prop_Recursive (ParOn cs) + prop "Nested" $ \cs -> prop_Nested (ParOn cs) + prop "Serially" $ \cs -> prop_Serially (ParOn cs) describe "Arbitrary Comp" $ do - it "Trivial" $ timed prop_SameAsTrivialScheduler - it "ArbitraryNested" $ timed prop_ArbitraryCompNested - it "AllJobsProcessed" $ timed prop_AllJobsProcessed - it "traverseConcurrently == traverse" $ timed prop_Traverse + prop "Trivial" prop_SameAsTrivialScheduler + prop "ArbitraryCompNested" prop_ArbitraryCompNested + prop "AllJobsProcessed" prop_AllJobsProcessed + prop "traverseConcurrently == traverse" prop_Traverse describe "Exceptions" $ do - it "CatchDivideByZero" $ timed prop_CatchDivideByZero - it "CatchDivideByZeroNested" $ timed prop_CatchDivideByZeroNested - it "KillBlockedCoworker" $ timed prop_KillBlockedCoworker - it "KillSleepingCoworker" $ timed prop_KillSleepingCoworker - it "ExpectAsyncException" $ timed prop_ExpectAsyncException - it "WorkerCaughtAsyncException" $ timed prop_WorkerCaughtAsyncException - it "AllWorkersDied" $ timed prop_AllWorkersDied - it "traverseConcurrently_" $ timed prop_TraverseConcurrently_ - it "traverseConcurrentlyInfinite_" $ property prop_TraverseConcurrentlyInfinite_ + prop "CatchDivideByZero" prop_CatchDivideByZero + prop "CatchDivideByZeroNested" prop_CatchDivideByZeroNested + prop "KillBlockedCoworker" prop_KillBlockedCoworker + prop "KillSleepingCoworker" prop_KillSleepingCoworker + prop "ExpectAsyncException" prop_ExpectAsyncException + prop "WorkerCaughtAsyncException" prop_WorkerCaughtAsyncException + prop "AllWorkersDied" prop_AllWorkersDied + prop "traverseConcurrently_" prop_TraverseConcurrently_ + prop "traverseConcurrentlyInfinite_" prop_TraverseConcurrentlyInfinite_ describe "Premature" $ do - it "FinishEarly" $ timed prop_FinishEarly - it "FinishEarly_" $ timed prop_FinishEarly_ - it "FinishEarlyWith" $ timed prop_FinishEarlyWith - it "FinishBeforeStarting" $ timed prop_FinishBeforeStarting - it "FinishWithBeforeStarting" $ timed prop_FinishWithBeforeStarting + prop "FinishEarly" prop_FinishEarly + prop "FinishEarly_" prop_FinishEarly_ + prop "FinishEarlyWith" prop_FinishEarlyWith + prop "FinishBeforeStarting" prop_FinishBeforeStarting + prop "FinishWithBeforeStarting" prop_FinishWithBeforeStarting describe "WorkerState" $ do - it "MutexException" $ timed prop_MutexException - it "WorkerStateExclusive" $ timed prop_WorkerStateExclusive + prop "MutexException" prop_MutexException + prop "WorkerStateExclusive" prop_WorkerStateExclusive + describe "Restartable" $ do + prop "ManyJobsInChunks" prop_ManyJobsInChunks + prop "FindCancelResume" prop_FindCancelResume instance Arbitrary WorkerId where arbitrary = WorkerId <$> arbitrary @@ -491,9 +600,6 @@ , FinishedEarly <$> arbitrary <*> arbitrary , FinishedEarlyWith <$> arbitrary ] - -timed :: Testable prop => prop -> Property -timed = within 2000000 -- | Assert a synchronous exception assertExceptionIO :: (NFData a, Exception exc) =>