scheduler 1.4.2.3 → 1.5.0

raw patch · 9 files changed

+1212/−526 lines, 9 filesdep +pvarPVP ok

version bump matches the API change (PVP)

Dependencies added: pvar

API changes (from Hackage documentation)

Files

• CHANGELOG.md +17/−0
• README.md +7/−3
• scheduler.cabal +6/−4
• src/Control/Scheduler.hs +132/−259
• src/Control/Scheduler/Global.hs +91/−0
• src/Control/Scheduler/Internal.hs +447/−102
• src/Control/Scheduler/Queue.hs +78/−47
• src/Control/Scheduler/Types.hs +217/−0
• tests/Control/SchedulerSpec.hs +217/−111

@@ -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`

@@ -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

@@ -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 

@@ -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 <lehins@yandex.ru> -- 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

@@ -0,0 +1,91 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE RecordWildCards #-}+-- |+-- Module      : Control.Scheduler.Global+-- Copyright   : (c) Alexey Kuleshevich 2020+-- License     : BSD3+-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>+-- 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

@@ -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 <lehins@yandex.ru> -- 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

@@ -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 #-}++

@@ -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 <lehins@yandex.ru>+-- 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"

@@ -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) =>