scheduler 1.5.0 → 2.0.0
raw patch · 8 files changed
+258/−179 lines, 8 filesdep ~QuickCheckdep ~primitive
Dependency ranges changed: QuickCheck, primitive
Files
- CHANGELOG.md +9/−0
- README.md +3/−3
- scheduler.cabal +4/−3
- src/Control/Scheduler.hs +99/−65
- src/Control/Scheduler/Global.hs +24/−24
- src/Control/Scheduler/Internal.hs +46/−43
- src/Control/Scheduler/Types.hs +26/−21
- tests/Control/SchedulerSpec.hs +47/−20
CHANGELOG.md view
@@ -1,3 +1,12 @@+# 2.0.0++* Switch type parameter of `Scheduler` from monad `m` to state token `s`. Which+ also means that constraints on many functions got tighter (i.e. `MonadPrim`, `MonadPrimBase`)+* Remove `m` type parameter for `SchedulerWS`, since it can only work in `IO` like monad anyways.+* Switch type parameter of `Batch` from monad `m` to state token `s`.+* Swap order of arguments for `replicateWork` for consistency+* Add `replicateWork_` that is slightly more efficient than `replicateWork`+ # 1.5.0 Despite that the major part of the version was bumped up, this release does not include
README.md view
@@ -5,9 +5,9 @@ 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 | Azure | Coveralls |-|:--------:|:------:|:-----:|:---------:|-|  | [](https://travis-ci.org/lehins/haskell-scheduler) | [](https://dev.azure.com/kuleshevich/haskell-scheduler/_build?branchName=master) | [](https://coveralls.io/github/lehins/haskell-scheduler?branch=master) |+| Language | Github Actions | Coveralls |Gitter.im |+|:--------:|:--------------:|:---------:|:--------:|+|  | [](https://github.com/lehins/haskell-scheduler/actions) | [](https://coveralls.io/github/lehins/haskell-scheduler?branch=master) | [](https://gitter.im/haskell-massiv/Lobby?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge) | Gihub | Hackage | Nightly | LTS | |:------|:-------:|:-------:|:---:|
scheduler.cabal view
@@ -1,5 +1,5 @@ name: scheduler-version: 1.5.0+version: 2.0.0 synopsis: Work stealing scheduler. 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@@ -27,10 +27,10 @@ , deepseq , exceptions , unliftio-core- , primitive >= 0.6.4+ , primitive >= 0.7.1 , pvar < 2.0 default-language: Haskell2010- ghc-options: -Wall+ ghc-options: -Wall -Wno-simplifiable-class-constraints test-suite tests@@ -39,6 +39,7 @@ main-is: Main.hs other-modules: Spec , Control.SchedulerSpec+ build-tool-depends: hspec-discover:hspec-discover build-depends: base , deepseq , genvalidity-hspec
src/Control/Scheduler.hs view
@@ -1,9 +1,11 @@ {-# LANGUAGE BangPatterns #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE RankNTypes #-} -- | -- Module : Control.Scheduler--- Copyright : (c) Alexey Kuleshevich 2018-2020+-- Copyright : (c) Alexey Kuleshevich 2018-2021 -- License : BSD3 -- Maintainer : Alexey Kuleshevich <lehins@yandex.ru> -- Stability : experimental@@ -35,6 +37,7 @@ , scheduleWorkState , scheduleWorkState_ , replicateWork+ , replicateWork_ -- * Batches , Batch , runBatch@@ -70,8 +73,9 @@ ) where import Control.Monad+import Control.Monad.ST import Control.Monad.IO.Unlift-import Control.Monad.Primitive (PrimMonad)+import Control.Monad.Primitive import Control.Scheduler.Computation import Control.Scheduler.Internal import Control.Scheduler.Types@@ -84,14 +88,14 @@ -- | Get the underlying `Scheduler`, which cannot access `WorkerStates`. -- -- @since 1.4.0-unwrapSchedulerWS :: SchedulerWS s m a -> Scheduler m a+unwrapSchedulerWS :: SchedulerWS ws a -> Scheduler RealWorld a unwrapSchedulerWS = _getScheduler -- | Get the computation strategy the states where initialized with. -- -- @since 1.4.0-workerStatesComp :: WorkerStates s -> Comp+workerStatesComp :: WorkerStates ws -> Comp workerStatesComp = _workerStatesComp @@ -132,27 +136,33 @@ -- except for the doctests to pass. -- -- @since 1.4.0-withSchedulerWS :: MonadUnliftIO m => WorkerStates s -> (SchedulerWS s m a -> m b) -> m [a]+withSchedulerWS ::+ MonadUnliftIO m => WorkerStates ws -> (SchedulerWS ws a -> m b) -> m [a] withSchedulerWS = withSchedulerWSInternal withScheduler -- | Run a scheduler with stateful workers, while discarding computation results. -- -- @since 1.4.0-withSchedulerWS_ :: MonadUnliftIO m => WorkerStates s -> (SchedulerWS s m () -> m b) -> m ()+withSchedulerWS_ ::+ MonadUnliftIO m => WorkerStates ws -> (SchedulerWS ws () -> m b) -> m () withSchedulerWS_ = withSchedulerWSInternal withScheduler_ -- | Same as `withSchedulerWS`, except instead of a list it produces `Results`, which -- allows for distinguishing between the ways computation was terminated. -- -- @since 1.4.2-withSchedulerWSR :: MonadUnliftIO m => WorkerStates s -> (SchedulerWS s m a -> m b) -> m (Results a)+withSchedulerWSR ::+ MonadUnliftIO m+ => WorkerStates ws+ -> (SchedulerWS ws a -> m b)+ -> m (Results a) withSchedulerWSR = withSchedulerWSInternal withSchedulerR -- | Schedule a job that will get a worker state passed as an argument -- -- @since 1.4.0-scheduleWorkState :: SchedulerWS s m a -> (s -> m a) -> m ()+scheduleWorkState :: MonadPrimBase RealWorld m => SchedulerWS ws a -> (ws -> m a) -> m () scheduleWorkState schedulerS withState = scheduleWorkId (_getScheduler schedulerS) $ \(WorkerId i) -> withState (indexSmallArray (_workerStatesArray (_workerStates schedulerS)) i)@@ -160,7 +170,7 @@ -- | Same as `scheduleWorkState`, but dont' keep the result of computation. -- -- @since 1.4.0-scheduleWorkState_ :: SchedulerWS s m () -> (s -> m ()) -> m ()+scheduleWorkState_ :: MonadPrimBase RealWorld m => SchedulerWS ws () -> (ws -> m ()) -> m () scheduleWorkState_ schedulerS withState = scheduleWorkId_ (_getScheduler schedulerS) $ \(WorkerId i) -> withState (indexSmallArray (_workerStatesArray (_workerStates schedulerS)) i)@@ -170,7 +180,7 @@ -- capabilities you have. Related function is `getCompWorkers`. -- -- @since 1.0.0-numWorkers :: Scheduler m a -> Int+numWorkers :: Scheduler s a -> Int numWorkers = _numWorkers @@ -181,8 +191,8 @@ -- scheduler. -- -- @since 1.2.0-scheduleWorkId :: Scheduler m a -> (WorkerId -> m a) -> m ()-scheduleWorkId =_scheduleWorkId+scheduleWorkId :: MonadPrimBase s m => Scheduler s a -> (WorkerId -> m a) -> m ()+scheduleWorkId s f = stToPrim (_scheduleWorkId s (primToPrim . f)) -- | 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@@ -193,8 +203,8 @@ -- although it will make sure their results are discarded. -- -- @since 1.1.0-terminate :: Scheduler m a -> a -> m a-terminate scheduler a = _terminate scheduler (Early a)+terminate :: MonadPrim s m => Scheduler s a -> a -> m a+terminate scheduler a = stToPrim (_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@@ -206,58 +216,69 @@ -- function. -- -- @since 1.1.0-terminateWith :: Scheduler m a -> a -> m a-terminateWith scheduler a = _terminate scheduler (EarlyWith a)+terminateWith :: MonadPrim s m => Scheduler s a -> a -> m a+terminateWith scheduler a = stToPrim $ _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. -- -- @since 1.0.0-scheduleWork :: Scheduler m a -> m a -> m ()-scheduleWork scheduler f = _scheduleWorkId scheduler (const f)+scheduleWork :: MonadPrimBase s m => Scheduler s a -> m a -> m ()+scheduleWork scheduler f = stToPrim $ _scheduleWorkId scheduler (const (primToPrim 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 ()+-- Type here should be `scheduleWork_ :: Scheduler s a -> m () -> m () -- | Same as `scheduleWork`, but only for a `Scheduler` that doesn't keep the results. -- -- @since 1.1.0-scheduleWork_ :: Scheduler m () -> m () -> m ()-scheduleWork_ = scheduleWork+scheduleWork_ :: MonadPrimBase s m => Scheduler s () -> m () -> m ()+scheduleWork_ s = stToPrim . scheduleWork s . primToPrim -- | Same as `scheduleWorkId`, but only for a `Scheduler` that doesn't keep the results. -- -- @since 1.2.0-scheduleWorkId_ :: Scheduler m () -> (WorkerId -> m ()) -> m ()-scheduleWorkId_ = _scheduleWorkId+scheduleWorkId_ :: MonadPrimBase s m => Scheduler s () -> (WorkerId -> m ()) -> m ()+scheduleWorkId_ scheduler f = stToPrim $ _scheduleWorkId scheduler (primToPrim . f) -- | Schedule the same action to run @n@ times concurrently. This differs from -- `replicateConcurrently` by allowing the caller to use the `Scheduler` freely, -- or to allow early termination via `terminate` across all (identical) threads. -- To be called within a `withScheduler` block. ----- @since 1.4.1-replicateWork :: Applicative m => Int -> Scheduler m a -> m a -> m ()-replicateWork !n scheduler f = go n+-- @since 2.0.0+replicateWork :: MonadPrimBase s m => Scheduler s a -> Int -> m a -> m ()+replicateWork scheduler n f = go n where go !k | k <= 0 = pure ()- | otherwise = scheduleWork scheduler f *> go (k - 1)+ | otherwise = stToPrim (scheduleWork scheduler (primToPrim f)) *> go (k - 1) ++-- | Same as `replicateWork`, but it does not retain the results of scheduled jobs+--+-- @since 2.0.0+replicateWork_ :: MonadPrimBase s m => Scheduler s () -> Int -> m a -> m ()+replicateWork_ scheduler n f = go n+ where+ go !k+ | k <= 0 = pure ()+ | otherwise = stToPrim (scheduleWork_ scheduler (primToPrim (void f))) *> go (k - 1)+ -- | Similar to `terminate`, but for a `Scheduler` that does not keep any results of computation. -- -- /Important/ - In case of `Seq` computation strategy this function has no affect. -- -- @since 1.1.0-terminate_ :: Scheduler m () -> m ()-terminate_ = (`_terminate` Early ())+terminate_ :: MonadPrim s m => Scheduler s () -> m ()+terminate_ s = stToPrim $ _terminate s (Early ()) -- | This trivial scheduler will behave in the same way as `withScheduler` with `Seq` -- computation strategy, except it is restricted to `PrimMonad`, instead of `MonadUnliftIO`. -- -- @since 1.4.2-withTrivialScheduler :: PrimMonad m => (Scheduler m a -> m b) -> m [a]+withTrivialScheduler :: MonadPrim s m => (Scheduler s a -> m b) -> m [a] withTrivialScheduler action = F.toList <$> withTrivialSchedulerR action @@ -267,8 +288,8 @@ -- -- @since 1.0.0 traverseConcurrently :: (MonadUnliftIO m, Traversable t) => Comp -> (a -> m b) -> t a -> m (t b)-traverseConcurrently comp f xs = do- ys <- withScheduler comp $ \s -> traverse_ (scheduleWork s . f) xs+traverseConcurrently comp f xs = withRunInIO $ \run -> do+ ys <- withScheduler comp $ \s -> traverse_ (scheduleWork s . run . f) xs pure $ transList ys xs transList :: Traversable t => [a] -> t b -> t a@@ -283,7 +304,8 @@ -- @since 1.0.0 traverseConcurrently_ :: (MonadUnliftIO m, Foldable t) => Comp -> (a -> m b) -> t a -> m () traverseConcurrently_ comp f xs =- withScheduler_ comp $ \s -> scheduleWork s $ F.traverse_ (scheduleWork s . void . f) xs+ withRunInIO $ \run ->+ withScheduler_ comp $ \s -> scheduleWork s $ F.traverse_ (scheduleWork s . void . run . f) xs -- | Replicate an action @n@ times and schedule them acccording to the supplied computation -- strategy.@@ -291,14 +313,16 @@ -- @since 1.1.0 replicateConcurrently :: MonadUnliftIO m => Comp -> Int -> m a -> m [a] replicateConcurrently comp n f =- withScheduler comp $ \s -> replicateM_ n $ scheduleWork s f+ withRunInIO $ \run ->+ withScheduler comp $ \s -> replicateM_ n $ scheduleWork s (run f) -- | Just like `replicateConcurrently`, but discards the results of computation. -- -- @since 1.1.0 replicateConcurrently_ :: MonadUnliftIO m => Comp -> Int -> m a -> m () replicateConcurrently_ comp n f =- withScheduler_ comp $ \s -> scheduleWork s $ replicateM_ n (scheduleWork s $ void f)+ withRunInIO $ \run -> do+ withScheduler_ comp $ \s -> scheduleWork s $ replicateM_ n (scheduleWork s $ void $ run f) @@ -334,12 +358,15 @@ withScheduler :: MonadUnliftIO m => Comp -- ^ Computation strategy- -> (Scheduler m a -> m b)+ -> (Scheduler RealWorld a -> m b) -- ^ Action that will be scheduling all the work. -> m [a]-withScheduler Seq = fmap (reverse . resultsToList) . withTrivialSchedulerRIO-withScheduler comp =- fmap (reverse . resultsToList) . withSchedulerInternal comp scheduleJobs readResults+withScheduler Seq f =+ withRunInIO $ \run -> do+ reverse . resultsToList <$> withTrivialSchedulerRIO (run . f)+withScheduler comp f =+ withRunInIO $ \run -> do+ reverse . resultsToList <$> withSchedulerInternal comp scheduleJobs readResults (run . f) {-# INLINE withScheduler #-} -- | Same as `withScheduler`, except instead of a list it produces `Results`, which allows@@ -349,11 +376,15 @@ withSchedulerR :: MonadUnliftIO m => Comp -- ^ Computation strategy- -> (Scheduler m a -> m b)+ -> (Scheduler RealWorld a -> m b) -- ^ Action that will be scheduling all the work. -> m (Results a)-withSchedulerR Seq = fmap reverseResults . withTrivialSchedulerRIO-withSchedulerR comp = fmap reverseResults . withSchedulerInternal comp scheduleJobs readResults+withSchedulerR Seq f =+ withRunInIO $ \run -> do+ reverseResults <$> withTrivialSchedulerRIO (run . f)+withSchedulerR comp f =+ withRunInIO $ \run -> do+ reverseResults <$> withSchedulerInternal comp scheduleJobs readResults (run .f) {-# INLINE withSchedulerR #-} @@ -363,11 +394,15 @@ withScheduler_ :: MonadUnliftIO m => Comp -- ^ Computation strategy- -> (Scheduler m a -> m b)+ -> (Scheduler RealWorld a -> m b) -- ^ Action that will be scheduling all the work. -> m ()-withScheduler_ Seq = void . withTrivialSchedulerRIO-withScheduler_ comp = void . withSchedulerInternal comp scheduleJobs_ (const (pure []))+withScheduler_ Seq f =+ withRunInIO $ \run -> do+ void $ withTrivialSchedulerRIO (run . f)+withScheduler_ comp f =+ withRunInIO $ \run -> do+ void $ withSchedulerInternal comp scheduleJobs_ (const (pure [])) (run . f) {-# INLINE withScheduler_ #-} @@ -375,8 +410,8 @@ -- | Check if the supplied batch has already finished. -- -- @since 1.5.0-hasBatchFinished :: Functor m => Batch m a -> m Bool-hasBatchFinished = batchHasFinished+hasBatchFinished :: MonadPrim s m => Batch s a -> m Bool+hasBatchFinished = stToPrim . batchHasFinished {-# INLINE hasBatchFinished #-} @@ -388,22 +423,22 @@ -- concurrent cancelation and it will return `True`. -- -- @since 1.5.0-cancelBatch :: Batch m a -> a -> m Bool-cancelBatch = batchCancel+cancelBatch :: MonadPrim s m => Batch s a -> a -> m Bool+cancelBatch b = stToPrim . batchCancel b {-# 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 ()+cancelBatch_ :: MonadPrim s m => Batch s () -> m Bool+cancelBatch_ b = stToPrim $ batchCancel b () {-# INLINE cancelBatch_ #-} -- | 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+cancelBatchWith :: MonadPrim s m => Batch s a -> a -> m Bool+cancelBatchWith b = stToPrim . batchCancelWith b {-# INLINE cancelBatchWith #-} @@ -411,8 +446,8 @@ -- -- @since 1.5.0 getCurrentBatch ::- Monad m => Scheduler m a -> m (Batch m a)-getCurrentBatch scheduler = do+ MonadPrim s m => Scheduler s a -> m (Batch s a)+getCurrentBatch scheduler = stToPrim $ do batchId <- _currentBatchId scheduler pure $ Batch { batchCancel = _cancelBatch scheduler batchId . Early@@ -436,10 +471,9 @@ -- 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+runBatch :: MonadPrimBase s m => Scheduler s a -> (Batch s a -> m c) -> m [a]+runBatch scheduler f = stToPrim $ do+ _ <- primToPrim . f =<< getCurrentBatch scheduler reverse . resultsToList <$> _waitForCurrentBatch scheduler {-# INLINE runBatch #-} @@ -447,9 +481,9 @@ -- -- @since 1.5.0 runBatch_ ::- Monad m => Scheduler m () -> (Batch m () -> m c) -> m ()-runBatch_ scheduler f = do- _ <- f =<< getCurrentBatch scheduler+ MonadPrimBase s m => Scheduler s () -> (Batch s () -> m c) -> m ()+runBatch_ scheduler f = stToPrim $ do+ _ <- primToPrim . f =<< getCurrentBatch scheduler void (_waitForCurrentBatch scheduler) {-# INLINE runBatch_ #-} @@ -458,9 +492,9 @@ -- -- @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+ MonadPrimBase s m => Scheduler s a -> (Batch s a -> m c) -> m (Results a)+runBatchR scheduler f = stToPrim $ do+ _ <- primToPrim . f =<< getCurrentBatch scheduler reverseResults <$> _waitForCurrentBatch scheduler {-# INLINE runBatchR #-}
src/Control/Scheduler/Global.hs view
@@ -15,31 +15,33 @@ , 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.Monad.ST+import Control.Monad.Primitive import Control.Scheduler import Control.Scheduler.Internal import Control.Scheduler.Types import Data.IORef+import Data.Maybe import System.IO.Unsafe (unsafePerformIO) -- | Global scheduler with `Par` computation strategy that can be used anytime using -- `withGlobalScheduler_`-globalScheduler :: GlobalScheduler IO+globalScheduler :: GlobalScheduler globalScheduler = unsafePerformIO (newGlobalScheduler Par) {-# NOINLINE globalScheduler #-} initGlobalScheduler ::- MonadUnliftIO m => Comp -> (Scheduler m a -> [ThreadId] -> m b) -> m b-initGlobalScheduler comp action = do+ MonadUnliftIO m => Comp -> (Scheduler RealWorld a -> [ThreadId] -> m b) -> m b+initGlobalScheduler comp action = withRunInIO $ \run -> do (jobs, mkScheduler) <- initScheduler comp scheduleJobs_ (const (pure []))- safeBracketOnError (spawnWorkers jobs comp) (liftIO . terminateWorkers) $ \tids ->- action (mkScheduler tids) tids+ safeBracketOnError (spawnWorkers jobs comp) terminateWorkers $ \tids ->+ run (action (mkScheduler tids) tids) -- | Create a new global scheduler, in case a single one `globalScheduler` is not@@ -47,26 +49,25 @@ -- 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 :: MonadIO m => Comp -> m GlobalScheduler 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- }+ liftIO $ initGlobalScheduler comp $ \scheduler tids -> 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_ :: MonadUnliftIO m => GlobalScheduler -> (Scheduler RealWorld () -> m a) -> m () withGlobalScheduler_ GlobalScheduler {..} action = withRunInIO $ \run -> do let initializeNewScheduler = do@@ -79,11 +80,10 @@ 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))+ let runScheduler = do+ _ <- run $ action scheduler+ mEarly <- stToPrim (_earlyResults scheduler)+ mEarly <$ when (isNothing mEarly) (void (stToPrim (_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
src/Control/Scheduler/Internal.hs view
@@ -1,10 +1,10 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE Unsafe #-} {-# OPTIONS_HADDOCK hide, not-home #-} -- | -- Module : Control.Scheduler.Internal@@ -38,6 +38,7 @@ import Control.Exception import Control.Monad import Control.Monad.IO.Unlift+import Control.Monad.Primitive import Control.Scheduler.Computation import Control.Scheduler.Types import Control.Scheduler.Queue@@ -53,7 +54,7 @@ -- | Initialize a separate state for each worker. -- -- @since 1.4.0-initWorkerStates :: MonadIO m => Comp -> (WorkerId -> m s) -> m (WorkerStates s)+initWorkerStates :: MonadIO m => Comp -> (WorkerId -> m ws) -> m (WorkerStates ws) initWorkerStates comp initState = do nWorkers <- getCompWorkers comp arr <- liftIO $ newSmallArray nWorkers (error "Uninitialized")@@ -64,27 +65,27 @@ go (i + 1) go 0 workerStates <- liftIO $ unsafeFreezeSmallArray arr- mutex <- liftIO $ newIORef False+ mutex <- liftIO $ newPVar 0 pure WorkerStates {_workerStatesComp = comp, _workerStatesArray = workerStates, _workerStatesMutex = mutex} withSchedulerWSInternal :: MonadUnliftIO m- => (Comp -> (Scheduler m a -> t) -> m b)- -> WorkerStates s- -> (SchedulerWS s m a -> t)+ => (Comp -> (Scheduler RealWorld a -> t) -> m b)+ -> WorkerStates ws+ -> (SchedulerWS ws a -> t) -> m b withSchedulerWSInternal withScheduler' states action = withRunInIO $ \run -> bracket lockState unlockState (run . runSchedulerWS) where mutex = _workerStatesMutex states- lockState = atomicModifyIORefCAS mutex $ (,) True+ lockState = atomicOrIntPVar mutex 1 unlockState wasLocked- | wasLocked = pure ()- | otherwise = atomicWriteIORef mutex False+ | wasLocked == 1 = pure ()+ | otherwise = void $ liftIO $ atomicAndIntPVar mutex 0 runSchedulerWS isLocked- | isLocked = liftIO $ throwIO MutexException+ | isLocked == 1 = liftIO $ throwIO MutexException | otherwise = withScheduler' (_workerStatesComp states) $ \scheduler -> action (SchedulerWS states scheduler)@@ -94,7 +95,7 @@ -- requests are bluntly ignored. -- -- @since 1.1.0-trivialScheduler_ :: Applicative f => Scheduler f ()+trivialScheduler_ :: Scheduler s () trivialScheduler_ = Scheduler { _numWorkers = 1@@ -114,19 +115,23 @@ -- rather computed immediately. -- -- @since 1.4.2-withTrivialSchedulerR :: PrimMonad m => (Scheduler m a -> m b) -> m (Results a)+withTrivialSchedulerR :: forall a b m s. MonadPrim s m => (Scheduler s 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), ()))+ let bumpCurrentBatchId :: MonadPrim s m' => m' ()+ bumpCurrentBatchId = atomicModifyMutVar' batchVar (\(BatchId x) -> (BatchId (x + 1), ()))+ bumpBatchId :: MonadPrim s m' => BatchId -> m' Bool bumpBatchId (BatchId c) = atomicModifyMutVar' batchVar $ \b@(BatchId x) -> if x == c then (BatchId (x + 1), True) else (b, False)+ takeBatchEarly :: MonadPrim s m' => m' (Maybe (Early a)) takeBatchEarly = atomicModifyMutVar' batchEarlyVar $ \mEarly -> (Nothing, mEarly)+ takeResults :: MonadPrim s m' => m' [a] takeResults = atomicModifyMutVar' resVar $ \res -> ([], res) _ <- action $@@ -166,7 +171,7 @@ -- returns results in an original LIFO order. -- -- @since 1.4.2-withTrivialSchedulerRIO :: MonadUnliftIO m => (Scheduler m a -> m b) -> m (Results a)+withTrivialSchedulerRIO :: MonadUnliftIO m => (Scheduler RealWorld a -> m b) -> m (Results a) withTrivialSchedulerRIO action = do resRef <- liftIO $ newIORef [] batchRef <- liftIO $ newIORef $ BatchId 0@@ -186,24 +191,24 @@ , _scheduleWorkId = \f -> do r <- f (WorkerId 0)- r `seq` liftIO (atomicModifyIORefCAS_ resRef (r :))+ r `seq` ioToPrim (atomicModifyIORefCAS_ resRef (r :)) , _terminate = \ !early ->- liftIO $ do+ ioToPrim $ do bumpCurrentBatchId finishEarly <- collectResults (Just early) takeResults atomicWriteIORef finResRef (Just finishEarly) throwIO TerminateEarlyException , _waitForCurrentBatch =- liftIO $ do+ ioToPrim $ do bumpCurrentBatchId mEarly <- takeBatchEarly collectResults mEarly . pure =<< takeResults- , _earlyResults = liftIO (readIORef finResRef)- , _currentBatchId = liftIO (readIORef batchRef)- , _batchEarly = liftIO takeBatchEarly+ , _earlyResults = ioToPrim (readIORef finResRef)+ , _currentBatchId = ioToPrim (readIORef batchRef)+ , _batchEarly = ioToPrim takeBatchEarly , _cancelBatch =- \batchId early -> liftIO $ do+ \batchId early -> ioToPrim $ do b <- bumpBatchId batchId when b $ atomicWriteIORef batchEarlyRef (Just early) pure b@@ -292,11 +297,10 @@ initScheduler ::- MonadIO m- => Comp- -> (Jobs m a -> (WorkerId -> m a) -> m ())- -> (JQueue m a -> m [a])- -> m (Jobs m a, [ThreadId] -> Scheduler m a)+ Comp+ -> (Jobs IO a -> (WorkerId -> IO a) -> IO ())+ -> (JQueue IO a -> IO [a])+ -> IO (Jobs IO a, [ThreadId] -> Scheduler RealWorld a) initScheduler comp submitWork collect = do jobsNumWorkers <- getCompWorkers comp jobsQueue <- newJQueue@@ -321,18 +325,18 @@ mkScheduler tids = Scheduler { _numWorkers = jobsNumWorkers- , _scheduleWorkId = submitWork jobs+ , _scheduleWorkId = \f -> ioToPrim $ submitWork jobs (stToPrim . f) , _terminate =- \early -> do+ \early -> ioToPrim $ do finishEarly <- case early of Early r -> FinishedEarly <$> collect jobsQueue <*> pure r EarlyWith r -> pure $ FinishedEarlyWith r- liftIO $ do+ ioToPrim $ do bumpCurrentBatchId atomicWriteIORef earlyTerminationResultRef $ Just finishEarly throwIO TerminateEarlyException- , _waitForCurrentBatch =+ , _waitForCurrentBatch = ioToPrim $ do scheduleJobs_ jobs (\_ -> liftIO $ void $ atomicSubIntPVar jobsQueueCount 1) unblockPopJQueue jobsQueue status <- liftIO $ takeMVar jobsSchedulerStatus@@ -353,11 +357,11 @@ 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)+ , _earlyResults = ioToPrim (readIORef earlyTerminationResultRef)+ , _currentBatchId = ioToPrim (readIORef batchIdRef)+ , _batchEarly = ioToPrim (readIORef batchEarlyRef) , _cancelBatch =- \batchId early -> do+ \batchId early -> ioToPrim $ do b <- liftIO $ bumpBatchId batchId when b $ do blockPopJQueue jobsQueue@@ -370,13 +374,12 @@ {-# INLINEABLE initScheduler #-} 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)+ Comp -- ^ Computation strategy+ -> (Jobs IO a -> (WorkerId -> IO a) -> IO ()) -- ^ How to schedule work+ -> (JQueue IO a -> IO [a]) -- ^ How to collect results+ -> (Scheduler RealWorld a -> IO b) -- ^ Action that will be scheduling all the work.- -> m (Results a)+ -> IO (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@@ -385,7 +388,7 @@ bracket (run (spawnWorkers jobs comp)) terminateWorkers $ \tids -> let scheduler = mkScheduler tids readEarlyTermination =- _earlyResults scheduler >>= \case+ stToPrim (_earlyResults scheduler) >>= \case Nothing -> error "Impossible: uninitialized early termination value" Just rs -> pure rs in try (run (onScheduler scheduler)) >>= \case@@ -401,14 +404,14 @@ | Just TerminateEarlyException <- fromException exc -> run readEarlyTermination | Just CancelBatchException <- fromException exc -> run $ do- mEarly <- _batchEarly scheduler+ mEarly <- stToPrim $ _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+ mEarly <- stToPrim $ _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
src/Control/Scheduler/Types.hs view
@@ -106,29 +106,30 @@ -- `Control.Scheduler.withScheduler_` for ways to construct and use this data type. -- -- @since 1.0.0-data Scheduler m a = Scheduler+data Scheduler s 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+ , _scheduleWorkId :: (WorkerId -> ST s a) -> ST s ()+ , _terminate :: Early a -> ST s a+ , _waitForCurrentBatch :: ST s (Results a)+ , _earlyResults :: ST s (Maybe (Results a))+ , _currentBatchId :: ST s 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+ , _cancelBatch :: BatchId -> Early a -> ST s Bool -- ^ Stops current batch and cancells all the outstanding jobs and the ones that are -- currently in progress.- , _batchEarly :: m (Maybe (Early a))+ , _batchEarly :: ST s (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)+data SchedulerWS ws a = SchedulerWS+ { _workerStates :: !(WorkerStates ws)+ , _getScheduler :: !(Scheduler RealWorld a) } -- | Each worker is capable of keeping it's own state, that can be share for different@@ -137,10 +138,14 @@ -- `Control.Scheduler.initWorkerStates` -- -- @since 1.4.0-data WorkerStates s = WorkerStates+data WorkerStates ws = WorkerStates { _workerStatesComp :: !Comp- , _workerStatesArray :: !(SmallArray s)- , _workerStatesMutex :: !(IORef Bool)+ , _workerStatesArray :: !(SmallArray ws)+#if MIN_VERSION_pvar(1,0,0)+ , _workerStatesMutex :: !(PVar Int RealWorld)+#else+ , _workerStatesMutex :: !(PVar IO Int)+#endif } -- | This identifier is needed for tracking batches.@@ -152,22 +157,22 @@ -- lifetime of a scheduler. -- -- @since 1.5.0-data Batch m a = Batch- { batchCancel :: a -> m Bool- , batchCancelWith :: a -> m Bool- , batchHasFinished :: m Bool+data Batch s a = Batch+ { batchCancel :: a -> ST s Bool+ , batchCancelWith :: a -> ST s Bool+ , batchHasFinished :: ST s 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`+-- `Control.Scheduler.Global.globalScheduler` -- -- @since 1.5.0-data GlobalScheduler m =+data GlobalScheduler = GlobalScheduler { globalSchedulerComp :: !Comp- , globalSchedulerMVar :: !(MVar (Scheduler m ()))+ , globalSchedulerMVar :: !(MVar (Scheduler RealWorld ())) , globalSchedulerThreadIdsRef :: !(IORef [ThreadId]) }
tests/Control/SchedulerSpec.hs view
@@ -1,6 +1,7 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-} module Control.SchedulerSpec ( spec ) where@@ -13,6 +14,7 @@ import Control.Exception.Base (ArithException(DivideByZero), AsyncException(ThreadKilled)) import Control.Monad+import Control.Monad.ST import Control.Scheduler as S import Data.Bits (complement) import qualified Data.Foldable as F (toList, traverse_)@@ -41,7 +43,7 @@ concurrentExpectation :: Expectation -> Property concurrentExpectation = concurrentProperty -concurrentPropertyIO :: IO Property -> Property+concurrentPropertyIO :: Testable prop => IO prop -> Property concurrentPropertyIO = concurrentProperty . monadicIO . run instance Arbitrary Comp where@@ -55,6 +57,12 @@ NonSeq <$> frequency [(10, pure Par), (35, ParOn <$> arbitrary), (35, ParN . getSmall <$> arbitrary)] +newtype SeqLike = SeqLike {getSeqLike :: Comp }+ deriving (Show, Eq)++instance Arbitrary SeqLike where+ arbitrary = SeqLike <$> oneof [pure Seq, ParOn . pure <$> arbitrary, pure $ ParN 1]+ prop_SameList :: Comp -> [Int] -> Property prop_SameList comp xs = concurrentPropertyIO $ do@@ -107,22 +115,37 @@ where f' = pure . apply f -replicateSeq :: (Int -> IO Int -> IO [Int]) -> Int -> Fun Int Int -> Property-replicateSeq justAs n f =- 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)+replicateLike :: ([Word] -> [Word]) -> (Int -> IO Word -> IO [Word]) -> Int -> Fun Word Word -> IO ()+replicateLike adjust justAs n f = do+ iRef <- newIORef 0+ jRef <- newIORef 0+ let g ref = atomicModifyIORef' ref (\i -> (apply f i, i + 1))+ xs <- replicateM n (g jRef)+ ys <- justAs n (g iRef)+ adjust ys `shouldBe` adjust xs -prop_ReplicateM :: Int -> Fun Int Int -> Property-prop_ReplicateM i = concurrentProperty . replicateSeq replicateM i+prop_ReplicateM :: ([Word] -> [Word]) -> Comp -> Int -> Fun Word Word -> Property+prop_ReplicateM adjust comp i =+ concurrentPropertyIO . replicateLike adjust (S.replicateConcurrently comp) i -prop_ReplicateWorkSeq :: Int -> Fun Int Int -> Property-prop_ReplicateWorkSeq i =- concurrentProperty . replicateSeq (\ n g -> withScheduler Seq (\s -> replicateWork n s g)) i+prop_ReplicateWork :: ([Word] -> [Word]) -> Comp -> Int -> Fun Word Word -> Property+prop_ReplicateWork adjust comp i =+ concurrentPropertyIO .+ replicateLike adjust (\n g -> withScheduler comp (\s -> replicateWork s n g)) i +prop_ReplicateWork_ :: ([Word] -> [Word]) -> Comp -> Int -> Fun Word Word -> Property+prop_ReplicateWork_ adjust comp i =+ concurrentPropertyIO . replicateLike adjust scheduleAndCollect i+ where+ scheduleAndCollect n g = do+ ref <- newIORef []+ withScheduler_ comp $ \s ->+ replicateWork_ s n $ do+ x <- g+ atomicModifyIORef' ref (\xs -> (x : xs, ()))+ reverse <$> readIORef ref + prop_ManyJobsInChunks :: Property prop_ManyJobsInChunks = noShrinking $ \ comp (jss :: [[Int]]) -> concurrentExpectation $ do@@ -286,15 +309,15 @@ scheduleWork scheduler $ pure (complement n) pure (res === [n]) -prop_TrivialSchedulerSameAsSeq_ :: [Int] -> Property-prop_TrivialSchedulerSameAsSeq_ zs =+prop_TrivialSchedulerSameAsSeq_ :: SeqLike -> [Int] -> Property+prop_TrivialSchedulerSameAsSeq_ (SeqLike comp) zs = concurrentPropertyIO $ do let consRef xsRef x = atomicModifyIORef' xsRef $ \ xs -> (x:xs, ()) trivial = trivialScheduler_ nRef <- newIORef False xRefs <- newIORef [] yRefs <- newIORef []- withScheduler_ Seq $ \scheduler -> do+ withScheduler_ comp $ \scheduler -> do writeIORef nRef (numWorkers scheduler == numWorkers trivial) mapM_ (scheduleWork_ scheduler . consRef xRefs) zs mapM_ (scheduleWork_ trivial . consRef yRefs) zs@@ -313,8 +336,8 @@ prop_Terminate :: (Show a, Eq a)- => ((Scheduler IO Int -> IO ()) -> IO a)- -> (Scheduler IO Int -> Int -> IO Int)+ => ((Scheduler RealWorld Int -> IO ()) -> IO a)+ -> (Scheduler RealWorld Int -> Int -> IO Int) -> ([Int] -> Int -> a) -> [Int] -> Int@@ -548,8 +571,9 @@ prop "Nested" $ prop_Nested Seq prop "Serially" $ prop_Serially Seq prop "TrivialAsSeq_" prop_TrivialSchedulerSameAsSeq_- prop "replicateConcurrently == replicateM" prop_ReplicateM- prop "replicateConcurrently == replicateWork" prop_ReplicateWorkSeq+ prop "replicateConcurrently == replicateM" $ prop_ReplicateM id . getSeqLike+ prop "replicateWork == replicateM" $ prop_ReplicateWork id . getSeqLike+ prop "replicateWork_ == replicateM" $ prop_ReplicateWork_ id . getSeqLike it "WorkerIdIsZero" $ withScheduler Seq (`scheduleWorkId` pure) `shouldReturn` [0] prop "TerminateSeq" $ prop_Terminate (withScheduler Seq) terminate (\xs x -> xs ++ [x])@@ -567,6 +591,9 @@ prop "ArbitraryCompNested" prop_ArbitraryCompNested prop "AllJobsProcessed" prop_AllJobsProcessed prop "traverseConcurrently == traverse" prop_Traverse+ prop "replicateConcurrently == replicateM" $ prop_ReplicateM sort+ prop "replicateWork == replicateM" $ prop_ReplicateWork sort+ prop "replicateWork_ == replicateM" $ prop_ReplicateWork_ sort describe "Exceptions" $ do prop "CatchDivideByZero" prop_CatchDivideByZero prop "CatchDivideByZeroNested" prop_CatchDivideByZeroNested