diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -89,12 +89,6 @@ aaaaaaaaabcdd*** Exception: divide by zero ``` -A special case is when a thread is killed by an async exception. Whenever that happens than the -exception will be re-thrown in a scheduling thread, but it will be wrapped in a custom -`WorkerAsyncException` exception. If for some reason you need to recover the original async -exception you can use `fromWorkerAsyncException`. See function documentation for an example. - - ### Nested jobs Scheduling actions can themselves schedule actions indefinitely. That of course means that order of diff --git a/massiv-scheduler.cabal b/massiv-scheduler.cabal --- a/massiv-scheduler.cabal +++ b/massiv-scheduler.cabal @@ -1,5 +1,5 @@ name: massiv-scheduler -version: 0.1.0.0 +version: 0.1.1.0 synopsis: Work stealing scheduler for Massiv (Массив) and other parallel applications. description: A work stealing scheduler that is used by [massiv](https://www.stackage.org/package/massiv) array librarry, but can be useful for any other library or application that fits such model of computation. homepage: https://github.com/lehins/massiv @@ -25,7 +25,7 @@ other-modules: Control.Massiv.Scheduler.Computation , Control.Massiv.Scheduler.Queue - build-depends: base >= 4.8 && < 5 + build-depends: base >= 4.9 && < 5 , atomic-primops , deepseq , exceptions @@ -41,11 +41,12 @@ main-is: Main.hs other-modules: Spec , Control.Massiv.SchedulerSpec - build-depends: base >= 4.8 && < 5 + build-depends: base , deepseq , massiv-scheduler , hspec , QuickCheck + , unliftio default-language: Haskell2010 ghc-options: -Wall -fno-warn-orphans -threaded -with-rtsopts=-N diff --git a/src/Control/Massiv/Scheduler.hs b/src/Control/Massiv/Scheduler.hs --- a/src/Control/Massiv/Scheduler.hs +++ b/src/Control/Massiv/Scheduler.hs @@ -35,6 +35,7 @@ import Data.Atomics (atomicModifyIORefCAS, atomicModifyIORefCAS_) import Data.Foldable as F (foldl') import Data.IORef +import Data.Traversable data Jobs m a = Jobs { jobsNumWorkers :: {-# UNPACK #-} !Int @@ -51,7 +52,7 @@ -- -- @since 0.1.0 data Scheduler m a = Scheduler - { numWorkers :: {-# UNPACK #-} !Int + { numWorkers :: {-# UNPACK #-} !Int -- ^ Get the number of workers. , scheduleWork :: m a -> m () -- ^ Schedule an action to be picked up and computed by a worker from a pool. @@ -60,7 +61,7 @@ -- | This is generally a faster way to traverse while ignoring the result rather than using `mapM_`. -- -- @since 0.1.0 -traverse_ :: (Foldable t, Applicative f) => (a -> f ()) -> t a -> f () +traverse_ :: (Applicative f, Foldable t) => (a -> f ()) -> t a -> f () traverse_ f = F.foldl' (\c a -> c *> f a) (pure ()) @@ -68,9 +69,17 @@ -- strategy. -- -- @since 0.1.0 -traverseConcurrently :: (MonadUnliftIO m, Foldable t) => Comp -> (a -> m b) -> t a -> m [b] -traverseConcurrently comp f xs = withScheduler comp $ \s -> traverse_ (scheduleWork s . f) xs +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 + pure $ transList ys xs +transList :: Traversable t => [a] -> t b -> t a +transList xs' = snd . mapAccumL withR xs' + where + withR (x:xs) _ = (xs, x) + withR _ _ = error "Impossible - Mismatched sizes" + -- | Just like `traverseConcurrently`, but discard the results of computation. -- -- @since 0.1.0 @@ -96,7 +105,7 @@ return res pushJQueue jobsQueue job --- | Helper function to place `Retire` instructions on the job queue. +-- | 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 @@ -186,11 +195,11 @@ withSchedulerInternal comp submitWork collect onScheduler = do jobsNumWorkers <- case comp of - Seq -> return 1 - Par -> liftIO getNumCapabilities + Seq -> return 1 + Par -> liftIO getNumCapabilities ParOn ws -> return $ length ws - ParN 0 -> liftIO getNumCapabilities - ParN n -> return $ fromIntegral n + ParN 0 -> liftIO getNumCapabilities + ParN n -> return $ fromIntegral n sWorkersCounterRef <- liftIO $ newIORef jobsNumWorkers jobsQueue <- newJQueue jobsCountRef <- liftIO $ newIORef 0 @@ -205,13 +214,12 @@ jc <- liftIO $ readIORef jobsCountRef when (jc == 0) $ scheduleJobs_ jobs (pure ()) let spawnWorkersWith fork ws = - forM ws $ \w -> - withRunInIO $ \run -> - mask_ $ + withRunInIO $ \run -> + forM ws $ \w -> fork w $ \unmask -> catch (unmask $ run $ runWorker jobsQueue onRetire) - (unmask . run . handleWorkerException jobsQueue workDoneMVar jobsNumWorkers) + (run . handleWorkerException jobsQueue workDoneMVar jobsNumWorkers) {-# INLINE spawnWorkersWith #-} spawnWorkers = case comp of @@ -227,76 +235,61 @@ -- \ wait for all worker to finish. If any one of them had a problem this MVar will -- contain an exception case mExc of - Nothing - -- / Now we are sure all workers have done their job we can safely read all of the - -- IORefs with results - -> collect jobsQueue - Just exc -- / Here we need to unwrap the legit worker exception and rethrow it, so the - -- main thread will think like it's his own - | Just (WorkerException wexc) <- fromException exc -> liftIO $ throwIO wexc - Just exc -> liftIO $ throwIO exc -- Somethig funky is happening, propagate it. + Nothing -> collect jobsQueue + -- \ Now we are sure all workers have done their job we can safely read all of the + -- IORefs with results + Just (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 {-# INLINE doWork #-} safeBracketOnError spawnWorkers - (liftIO . traverse_ (`throwTo` WorkerTerminateException)) + (liftIO . traverse_ (`throwTo` SomeAsyncException WorkerTerminateException)) (const doWork) -- | Specialized exception handler for the work scheduler. handleWorkerException :: - MonadIO m => JQueue m a -> MVar (Maybe SomeException) -> Int -> SomeException -> m () + MonadIO m => JQueue m a -> MVar (Maybe WorkerException) -> Int -> SomeException -> m () handleWorkerException jQueue workDoneMVar nWorkers exc = - case fromException exc of + case asyncExceptionFromException exc of Just WorkerTerminateException -> return () -- \ some co-worker died, we can just move on with our death. - _ -> - case fromException exc of - Just asyncExc -> do - liftIO $ putMVar workDoneMVar $ Just $ toException $ WorkerAsyncException asyncExc - -- \ Let the main thread know about this terrible async exception. - -- / Do the co-worker cleanup - retireWorkersN jQueue (nWorkers - 1) - liftIO $ throwIO asyncExc - -- \ do not recover from an async exception - Nothing -> do - liftIO $ putMVar workDoneMVar $ Just $ toException $ WorkerException exc - -- \ Main thread must know how we died - -- / Do the co-worker cleanup - retireWorkersN jQueue (nWorkers - 1) - -- / As to one self, gracefully leave off into the outer world + _ -> do + _ <- liftIO $ tryPutMVar workDoneMVar $ Just $ WorkerException exc + -- \ Main thread must know how we died + -- / Do the co-worker cleanup + retireWorkersN jQueue (nWorkers - 1) -- | This exception should normally be not seen in the wild. The only one that could possibly pop up -- is the `WorkerAsyncException`. -data WorkerException - = WorkerException !SomeException +newtype WorkerException = + WorkerException SomeException -- ^ One of workers experienced an exception, main thread will receive the same `SomeException`. - | WorkerAsyncException !SomeAsyncException - -- ^ If a worker recieves an async exception, something is utterly wrong. This exception in - -- particular would mean that one of the workers have died of an asyncronous exception. We don't - -- want to raise that async exception synchronously in the main thread, therefore it will receive - -- this wrapper exception instead. - | WorkerTerminateException - -- ^ When a brother worker dies of some exception, all the other ones will be terminated - -- asynchronously with this one. - deriving Show + deriving (Show) instance Exception WorkerException where displayException workerExc = case workerExc of WorkerException exc -> "A worker handled a job that ended with exception: " ++ displayException exc - WorkerAsyncException exc -> - "A worker was killed with an async exception: " ++ displayException exc - WorkerTerminateException -> "A worker was terminated by the scheduler" +data WorkerTerminateException = + WorkerTerminateException + -- ^ When a brother worker dies of some exception, all the other ones will be terminated + -- asynchronously with this one. + deriving (Show) --- | If any one of the workers dies with an async exception, it is possible to recover that --- exception in the main thread with this function. This does not apply to `Seq`uential computation, --- since no workers are created in such case and async exception will be received by the main thread --- directly. + +instance Exception WorkerTerminateException where + displayException WorkerTerminateException = "A worker was terminated by the scheduler" + + +-- | If any one of the workers dies with any exception, even the async exception, it will be +-- re-thrown in the main thread. -- --- >>> let didAWorkerDie = handleJust fromWorkerAsyncException (return . (== ThreadKilled)) . fmap or +-- >>> let didAWorkerDie = handleJust asyncExceptionFromException (return . (== ThreadKilled)) . fmap or -- >>> :t didAWorkerDie -- didAWorkerDie :: Foldable t => IO (t Bool) -> IO Bool -- >>> didAWorkerDie $ withScheduler Par $ \ s -> scheduleWork s $ pure False @@ -304,14 +297,12 @@ -- >>> didAWorkerDie $ withScheduler Par $ \ s -> scheduleWork s $ myThreadId >>= killThread >> pure False -- True -- >>> withScheduler Par $ \ s -> scheduleWork s $ myThreadId >>= killThread >> pure False --- *** Exception: WorkerAsyncException thread killed +-- *** Exception: thread killed -- -- @since 0.1.0 fromWorkerAsyncException :: Exception e => SomeException -> Maybe e -fromWorkerAsyncException exc = - case fromException exc of - Just (WorkerAsyncException asyncExc) -> asyncExceptionFromException (toException asyncExc) - _ -> Nothing +fromWorkerAsyncException = asyncExceptionFromException +{-# DEPRECATED fromWorkerAsyncException "In favor of `asyncExceptionFromException`" #-} -- Copy from unliftio: @@ -321,7 +312,6 @@ res1 <- try $ restore $ run $ thing x case res1 of Left (e1 :: SomeException) -> do - -- ignore the exception, see bracket for explanation _ :: Either SomeException b <- try $ uninterruptibleMask_ $ run $ after x throwIO e1 diff --git a/tests/Control/Massiv/SchedulerSpec.hs b/tests/Control/Massiv/SchedulerSpec.hs --- a/tests/Control/Massiv/SchedulerSpec.hs +++ b/tests/Control/Massiv/SchedulerSpec.hs @@ -3,13 +3,18 @@ import Control.Concurrent (killThread, myThreadId, threadDelay) import Control.Concurrent.MVar import Control.DeepSeq -import Control.Exception hiding (assert) -import Control.Exception.Base (ArithException(DivideByZero)) +import qualified Control.Exception as EUnsafe +import Control.Exception.Base (ArithException(DivideByZero), + AsyncException(ThreadKilled)) import Control.Massiv.Scheduler +import Control.Monad import Data.List (sort) import Test.Hspec import Test.QuickCheck +import Test.QuickCheck.Function import Test.QuickCheck.Monadic +import UnliftIO.Async +import UnliftIO.Exception hiding (assert) instance Arbitrary Comp where @@ -47,7 +52,7 @@ where schedule [] = return [] schedule (y:ys) = do - y' <- withScheduler comp (\s -> scheduleWork s (return y)) + y' <- withScheduler comp (`scheduleWork` pure y) ys' <- schedule ys return (y':ys') @@ -62,6 +67,12 @@ schedule (y:ys) = withScheduler comp (\s -> scheduleWork s (schedule ys >>= \ys' -> return (y : concat ys'))) +prop_Traversable :: Comp -> [Int] -> Fun Int Int -> Property +prop_Traversable comp xs f = + monadicIO $ run $ (===) <$> traverse f' xs <*> traverseConcurrently comp f' xs + where + f' = pure . apply f + prop_ArbitraryCompNested :: [(Comp, Int)] -> Property prop_ArbitraryCompNested xs = monadicIO $ @@ -120,18 +131,50 @@ prop_ExpectAsyncException :: Comp -> Property prop_ExpectAsyncException comp = - let didAWorkerDie = handleJust fromWorkerAsyncException' (return . (== ThreadKilled)) . fmap or - fromWorkerAsyncException' = - case comp of - Seq -> asyncExceptionFromException - _ -> fromWorkerAsyncException - in (monadicIO . - run . - didAWorkerDie . - withScheduler comp $ \s -> scheduleWork s (myThreadId >>= killThread >> pure False)) .&&. - (monadicIO . - run . fmap not . didAWorkerDie . withScheduler Par $ \s -> scheduleWork s $ pure False) + let didAWorkerDie = + EUnsafe.handleJust EUnsafe.asyncExceptionFromException (return . (== EUnsafe.ThreadKilled)) . + fmap or + in (monadicIO . run . didAWorkerDie . withScheduler comp $ \s -> + scheduleWork s (myThreadId >>= killThread >> pure False)) .&&. + (monadicIO . run . fmap not . didAWorkerDie . withScheduler Par $ \s -> + scheduleWork s $ pure False) +prop_WorkerCaughtAsyncException :: Positive Int -> Property +prop_WorkerCaughtAsyncException (Positive n) = + assertExceptionIO (== DivideByZero) $ do + lock <- newEmptyMVar + result <- + race (readMVar lock) $ + withScheduler_ (ParN 2) $ \scheduler -> do + scheduleWork scheduler $ do + threadDelay (n `mod` 1000000) + EUnsafe.throwIO DivideByZero + scheduleWork scheduler $ do + e <- tryAny $ replicateM_ 5 $ threadDelay 1000000 + case e of + Right _ -> throwString "Impossible, shouldn't have waited for so long" + Left exc -> do + putMVar lock exc + throwString $ + "I should not have survived: " ++ displayException (exc :: SomeException) + void $ throwString $ + case result of + Left innerError -> "Scheduled job cought async exception: " ++ displayException innerError + Right () -> "Scheduler terminated properly. Should not have happened" + +-- | 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) = + assertAsyncExceptionIO + (== ThreadKilled) + (withScheduler_ comp1 $ \scheduler1 -> + scheduleWork + scheduler1 + (withScheduler_ comp $ \scheduler -> + replicateM_ n (scheduleWork scheduler (myThreadId >>= killThread)))) + + + spec :: Spec spec = do describe "Seq" $ do @@ -144,15 +187,20 @@ it "Recursive" $ property $ \ cs -> prop_Recursive (ParOn cs) it "Nested" $ property $ \ cs -> prop_Nested (ParOn cs) it "Serially" $ property $ \ cs -> prop_Serially (ParOn cs) - describe "Arbitrary Comp" $ + describe "Arbitrary Comp" $ do it "ArbitraryNested" $ property prop_ArbitraryCompNested + it "traverseConcurrently == traverse" $ property prop_Traversable describe "Exceptions" $ do it "CatchDivideByZero" $ property prop_CatchDivideByZero it "CatchDivideByZeroNested" $ property prop_CatchDivideByZeroNested it "KillBlockedCoworker" $ property prop_KillBlockedCoworker it "KillSleepingCoworker" $ property prop_KillSleepingCoworker it "ExpectAsyncException" $ property prop_ExpectAsyncException + it "WorkerCaughtAsyncException" $ property prop_WorkerCaughtAsyncException + it "AllWorkersDied" $ property prop_AllWorkersDied + +-- | Assert a synchronous exception assertExceptionIO :: (NFData a, Exception exc) => (exc -> Bool) -- ^ Return True if that is the exception that was expected -> IO a -- ^ IO Action that should throw an exception @@ -160,17 +208,23 @@ assertExceptionIO isExc action = monadicIO $ do hasFailed <- - run - (catch - (do res <- action - res `deepseq` return False) $ \exc -> - displayException exc `deepseq` return (isExc exc)) + run $ + catch + (do res <- action + res `deepseq` return False) $ \exc -> displayException exc `deepseq` return (isExc exc) assert hasFailed - - - - - - - +assertAsyncExceptionIO :: (Exception e, NFData a) => (e -> Bool) -> IO a -> Property +assertAsyncExceptionIO isAsyncExc action = + monadicIO $ do + hasFailed <- + run $ + EUnsafe.catch + (do res <- action + res `deepseq` return False) + (\exc -> + case EUnsafe.asyncExceptionFromException exc of + Just asyncExc + | isAsyncExc asyncExc -> displayException asyncExc `deepseq` pure True + _ -> EUnsafe.throwIO exc) + assert hasFailed