diff --git a/Control/Concurrent/Async/Pool.hs b/Control/Concurrent/Async/Pool.hs new file mode 100644 --- /dev/null +++ b/Control/Concurrent/Async/Pool.hs @@ -0,0 +1,130 @@ +----------------------------------------------------------------------------- +-- | +-- Module : Control.Concurrent.Async.Pool +-- Copyright : (c) Simon Marlow 2012, John Wiegley 2014 +-- License : BSD3 (see the file LICENSE) +-- +-- Maintainer : John Wiegley +-- Stability : provisional +-- Portability : non-portable (requires concurrency) +-- +-- This module provides a set of operations for running IO operations +-- asynchronously and waiting for their results. It is a thin layer over the +-- basic concurrency operations provided by "Control.Concurrent". The main +-- additional functionality it provides is the ability to wait for the return +-- value of a thread, plus functions for managing task pools, work groups, and +-- many-to-many dependencies between tasks. The interface also provides some +-- additional safety and robustness over using threads and @MVar@ directly. +-- +-- The basic type is @'Async' a@, which represents an asynchronous @IO@ action +-- that will return a value of type @a@, or die with an exception. An @Async@ +-- corresponds to either a thread, or a @Handle@ to an action waiting to be +-- spawned. This makes it possible to submit very large numbers of tasks, +-- with only N threads active at one time. +-- +-- For example, to fetch two web pages at the same time, we could do +-- this (assuming a suitable @getURL@ function): +-- +-- > withTaskGroup 4 $ \g -> do +-- > a1 <- async g (getURL url1) +-- > a2 <- async g (getURL url2) +-- > page1 <- wait a1 +-- > page2 <- wait a2 +-- > ... +-- +-- where 'async' submits the operation to the worker group (and from which it +-- is spawned in a separate thread), and 'wait' waits for and returns the +-- result. The number 4 indicates the maximum number of threads which may be +-- spawned at one time. If the operation throws an exception, then that +-- exception is re-thrown by 'wait'. This is one of the ways in which this +-- library provides some additional safety: it is harder to accidentally +-- forget about exceptions thrown in child threads. +-- +-- A slight improvement over the previous example is this: +-- +-- > withTaskGroup 4 $ \g -> do +-- > withAsync g (getURL url1) $ \a1 -> do +-- > withAsync g (getURL url2) $ \a2 -> do +-- > page1 <- wait a1 +-- > page2 <- wait a2 +-- > ... +-- +-- 'withAsync' is like 'async', except that the 'Async' is automatically +-- killed (or unscheduled, using 'cancel') if the enclosing IO operation +-- returns before it has completed. Consider the case when the first 'wait' +-- throws an exception; then the second 'Async' will be automatically killed +-- rather than being left to run in the background, possibly indefinitely. +-- This is the second way that the library provides additional safety: using +-- 'withAsync' means we can avoid accidentally leaving threads running. +-- Furthermore, 'withAsync' allows a tree of threads to be built, such that +-- children are automatically killed if their parents die for any reason. +-- +-- The pattern of performing two IO actions concurrently and waiting for their +-- results is packaged up in a combinator 'concurrently', so we can further +-- shorten the above example to: +-- +-- > withTaskGroup 4 $ \g -> do +-- > (page1, page2) <- concurrently g (getURL url1) (getURL url2) +-- > ... +-- +-- The 'Functor' instance can be used to change the result of an 'Async'. For +-- example: +-- +-- > ghci> a <- async g (return 3) +-- > ghci> wait a +-- > 3 +-- > ghci> wait (fmap (+1) a) +-- > 4 + +module Control.Concurrent.Async.Pool + ( + -- * Asynchronous actions + Async, + + -- * Task pools and groups + withTaskGroup, withTaskGroupIn, + Pool, createPool, + TaskGroup, createTaskGroup, runTaskGroup, + + -- ** Spawning tasks + async, asyncBound, asyncOn, asyncWithUnmask, asyncOnWithUnmask, + asyncSTM, + + -- ** Dependent tasks + taskHandle, asyncAfter, asyncAfterAll, + makeDependent, unsafeMakeDependent, + + -- ** Spawning with automatic 'cancel'ation + withAsync, withAsyncBound, withAsyncOn, withAsyncWithUnmask, + withAsyncOnWithUnmask, + + -- ** Quering 'Async's + wait, poll, waitCatch, cancel, cancelWith, + + -- ** STM operations + waitSTM, pollSTM, waitCatchSTM, + + -- ** Waiting for multiple 'Async's + waitAny, waitAnyCatch, waitAnyCancel, waitAnyCatchCancel, + waitEither, waitEitherCatch, waitEitherCancel, waitEitherCatchCancel, + waitEither_, + waitBoth, + + -- ** Linking + link, link2, + + -- ** Lists of actions + mapTasks, mapTasks_, mapTasksE, mapTasksE_, + mapRace, mapReduce, + scatterFoldMapM, + + -- ** The Task Monad and Applicative + Task, runTask, task, + + -- * Other utilities + race, race_, + concurrently, mapConcurrently, Concurrently(..) + ) where + +import Control.Concurrent.Async.Pool.Async +import Control.Concurrent.Async.Pool.Internal diff --git a/Control/Concurrent/Async/Pool/Async.hs b/Control/Concurrent/Async/Pool/Async.hs new file mode 100644 --- /dev/null +++ b/Control/Concurrent/Async/Pool/Async.hs @@ -0,0 +1,706 @@ +{-# LANGUAGE CPP, MagicHash, UnboxedTuples, RankNTypes #-} +#if __GLASGOW_HASKELL__ >= 701 +{-# LANGUAGE Trustworthy #-} +#endif +{-# OPTIONS -Wall #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Control.Concurrent.Async +-- Copyright : (c) Simon Marlow 2012 +-- License : BSD3 (see the file LICENSE) +-- +-- Maintainer : Simon Marlow +-- Stability : provisional +-- Portability : non-portable (requires concurrency) +-- +-- This module provides a set of operations for running IO operations +-- asynchronously and waiting for their results. It is a thin layer +-- over the basic concurrency operations provided by +-- "Control.Concurrent". The main additional functionality it +-- provides is the ability to wait for the return value of a thread, +-- but the interface also provides some additional safety and +-- robustness over using threads and @MVar@ directly. +-- +-- The basic type is @'Async' a@, which represents an asynchronous +-- @IO@ action that will return a value of type @a@, or die with an +-- exception. An @Async@ corresponds to a thread, and its 'ThreadId' +-- can be obtained with 'asyncThreadId', although that should rarely +-- be necessary. +-- +-- For example, to fetch two web pages at the same time, we could do +-- this (assuming a suitable @getURL@ function): +-- +-- > do a1 <- async (getURL url1) +-- > a2 <- async (getURL url2) +-- > page1 <- wait a1 +-- > page2 <- wait a2 +-- > ... +-- +-- where 'async' starts the operation in a separate thread, and +-- 'wait' waits for and returns the result. If the operation +-- throws an exception, then that exception is re-thrown by +-- 'wait'. This is one of the ways in which this library +-- provides some additional safety: it is harder to accidentally +-- forget about exceptions thrown in child threads. +-- +-- A slight improvement over the previous example is this: +-- +-- > withAsync (getURL url1) $ \a1 -> do +-- > withAsync (getURL url2) $ \a2 -> do +-- > page1 <- wait a1 +-- > page2 <- wait a2 +-- > ... +-- +-- 'withAsync' is like 'async', except that the 'Async' is +-- automatically killed (using 'cancel') if the enclosing IO operation +-- returns before it has completed. Consider the case when the first +-- 'wait' throws an exception; then the second 'Async' will be +-- automatically killed rather than being left to run in the +-- background, possibly indefinitely. This is the second way that the +-- library provides additional safety: using 'withAsync' means we can +-- avoid accidentally leaving threads running. Furthermore, +-- 'withAsync' allows a tree of threads to be built, such that +-- children are automatically killed if their parents die for any +-- reason. +-- +-- The pattern of performing two IO actions concurrently and waiting +-- for their results is packaged up in a combinator 'concurrently', so +-- we can further shorten the above example to: +-- +-- > (page1, page2) <- concurrently (getURL url1) (getURL url2) +-- > ... +-- +-- The 'Functor' instance can be used to change the result of an +-- 'Async'. For example: +-- +-- > ghci> a <- async (return 3) +-- > ghci> wait a +-- > 3 +-- > ghci> wait (fmap (+1) a) +-- > 4 + +----------------------------------------------------------------------------- + +module Control.Concurrent.Async.Pool.Async + ( module Control.Concurrent.Async.Pool.Async + , module Gr + ) where + +import Control.Concurrent.STM +import Control.Exception +import Control.Concurrent +import Control.Applicative +import Control.Monad hiding (forM, forM_, mapM, mapM_) +import Data.Foldable +import Data.Graph.Inductive.Graph as Gr hiding ((&)) +import Data.Graph.Inductive.PatriciaTree as Gr +import Data.Graph.Inductive.Query.BFS as Gr +import Data.IntMap (IntMap) +import qualified Data.IntMap as IntMap +import Data.Traversable +import Prelude hiding (mapM_, mapM, foldr, all, any, concatMap, foldl1) + +import GHC.Exts +import GHC.IO hiding (finally, onException) +import GHC.Conc + +-- | A 'Handle' is a unique identifier for a task submitted to a 'Pool'. +type Handle = Node +data State = Ready | Starting | Started ThreadId deriving (Eq, Show) +data Status = Pending | Completed deriving (Eq, Show) +type TaskGraph = Gr (TVar State) Status + +-- | A 'Pool' manages a collection of possibly interdependent tasks, such that +-- tasks await execution until the tasks they depend on have finished (and +-- tasks may depend on an arbitrary number of other tasks), while +-- independent tasks execute concurrently up to the number of available +-- resource slots in the pool. +-- +-- Results from each task are available until the status of the task is +-- polled or waited on. Further, the results are kept until that occurs, so +-- failing to ever wait will result in a memory leak. +-- +-- Tasks may be cancelled, in which case all dependent tasks are +-- unscheduled. +data Pool = Pool + { tasks :: TVar TaskGraph + -- ^ The task graph represents a partially ordered set P with subset S + -- such that for every x ∈ S and y ∈ P, either x ≤ y or x is unrelated + -- to y. Stated more simply, S is the set of least elements of all + -- maximal chains in P. In our case, ≤ relates two uncompleted tasks + -- by dependency. Therefore, S is equal to the set of tasks which may + -- execute concurrently, as none of them have incomplete dependencies. + -- + -- We use a graph representation to make determination of S more + -- efficient (where S is just the set of roots in P expressed as a + -- graph). Completion status is recorded on the edges, and nodes are + -- removed from the graph once no other incomplete node depends on + -- them. + , tokens :: TVar Int + -- ^ Tokens identify tasks, and are provisioned monotonically. + } + +data TaskGroup = TaskGroup + { pool :: Pool + , avail :: TVar Int + -- ^ The number of available execution slots in the pool. + , pending :: TVar (IntMap (IO ThreadId)) + -- ^ Nodes in the task graph that are waiting to start. + } + +-- ----------------------------------------------------------------------------- +-- STM Async API + + +-- | An asynchronous action spawned by 'async' or 'withAsync'. +-- Asynchronous actions are executed in a separate thread, and +-- operations are provided for waiting for asynchronous actions to +-- complete and obtaining their results (see e.g. 'wait'). +-- +data Async a = Async + { taskGroup :: TaskGroup + , taskHandle :: {-# UNPACK #-} !Handle + , _asyncWait :: STM (Either SomeException a) + } + +getTaskVar :: TaskGraph -> Handle -> TVar State +getTaskVar g h = let (_to, _, t, _from) = context g h in t + +getThreadId :: TaskGraph -> Node -> STM (Maybe ThreadId) +getThreadId g h = do + status <- readTVar (getTaskVar g h) + case status of + Ready -> return Nothing + Starting -> retry + Started x -> return $ Just x + +instance Eq (Async a) where + Async _ a _ == Async _ b _ = a == b + +instance Ord (Async a) where + Async _ a _ `compare` Async _ b _ = a `compare` b + +instance Functor Async where + fmap f (Async p a w) = Async p a (fmap (fmap f) w) + + +-- | Spawn an asynchronous action in a separate thread. +async :: TaskGroup -> IO a -> IO (Async a) +async p = atomically . inline asyncUsing p rawForkIO + +-- | Like 'async' but using 'forkOS' internally. +asyncBound :: TaskGroup -> IO a -> IO (Async a) +asyncBound p = atomically . asyncUsing p forkOS + +-- | Like 'async' but using 'forkOn' internally. +asyncOn :: TaskGroup -> Int -> IO a -> IO (Async a) +asyncOn p = (atomically .) . asyncUsing p . rawForkOn + +-- | Like 'async' but using 'forkIOWithUnmask' internally. +-- The child thread is passed a function that can be used to unmask asynchronous exceptions. +asyncWithUnmask :: TaskGroup -> ((forall b . IO b -> IO b) -> IO a) -> IO (Async a) +asyncWithUnmask p actionWith = + atomically $ asyncUsing p rawForkIO (actionWith unsafeUnmask) + +-- | Like 'asyncOn' but using 'forkOnWithUnmask' internally. +-- The child thread is passed a function that can be used to unmask asynchronous exceptions. +asyncOnWithUnmask :: TaskGroup -> Int -> ((forall b . IO b -> IO b) -> IO a) -> IO (Async a) +asyncOnWithUnmask p cpu actionWith = + atomically $ asyncUsing p (rawForkOn cpu) (actionWith unsafeUnmask) + +asyncUsing :: TaskGroup -> (IO () -> IO ThreadId) -> IO a -> STM (Async a) +asyncUsing p doFork action = do + h <- nextIdent (pool p) + + var <- newEmptyTMVar + let start = mask $ \restore -> + doFork $ try (restore (action `finally` cleanup h)) + >>= atomically . putTMVar var + + modifyTVar (pending p) (IntMap.insert h start) + tv <- newTVar Ready + modifyTVar (tasks (pool p)) (insNode (h, tv)) + + return $ Async p h (readTMVar var) + where + cleanup h = atomically $ do + modifyTVar (avail p) succ + cleanupTask (pool p) h + +-- | Return the next available thread identifier from the pool. These are +-- monotonically increasing integers. +nextIdent :: Pool -> STM Int +nextIdent p = do + tok <- readTVar (tokens p) + writeTVar (tokens p) (succ tok) + return tok + +cleanupTask :: Pool -> Handle -> STM () +cleanupTask p h = + -- Once the task is done executing, we must alter the graph so any + -- dependent children will know their parent has completed. + modifyTVar (tasks p) $ \g -> + case zip (repeat h) (Gr.suc g h) of + -- If nothing dependend on this task and if the final result value + -- has been observed, prune it from the graph, as well as any + -- parents which now have no dependents. Otherwise mark the edges + -- as Completed so dependent children can execute. + [] -> dropTask h g + es -> insEdges (completeEdges es) $ delEdges es g + where + completeEdges = map (\(f, t) -> (f, t, Completed)) + + dropTask k gr = foldl' f (delNode k gr) (Gr.pre gr k) + where + f g n = if outdeg g n == 0 then dropTask n g else g + +-- | Spawn an asynchronous action in a separate thread, and pass its +-- @Async@ handle to the supplied function. When the function returns +-- or throws an exception, 'cancel' is called on the @Async@. +-- +-- > withAsync action inner = bracket (async action) cancel inner +-- +-- This is a useful variant of 'async' that ensures an @Async@ is +-- never left running unintentionally. +-- +-- Since 'cancel' may block, 'withAsync' may also block; see 'cancel' +-- for details. +-- +withAsync :: TaskGroup -> IO a -> (Async a -> IO b) -> IO b +withAsync p = inline withAsyncUsing p rawForkIO + +-- | Like 'withAsync' but uses 'forkOS' internally. +withAsyncBound :: TaskGroup -> IO a -> (Async a -> IO b) -> IO b +withAsyncBound p = withAsyncUsing p forkOS + +-- | Like 'withAsync' but uses 'forkOn' internally. +withAsyncOn :: TaskGroup -> Int -> IO a -> (Async a -> IO b) -> IO b +withAsyncOn p = withAsyncUsing p . rawForkOn + +-- | Like 'withAsync' but uses 'forkIOWithUnmask' internally. +-- The child thread is passed a function that can be used to unmask asynchronous exceptions. +withAsyncWithUnmask :: TaskGroup -> ((forall c. IO c -> IO c) -> IO a) -> (Async a -> IO b) -> IO b +withAsyncWithUnmask p actionWith = + withAsyncUsing p rawForkIO (actionWith unsafeUnmask) + +-- | Like 'withAsyncOn' but uses 'forkOnWithUnmask' internally. +-- The child thread is passed a function that can be used to unmask asynchronous exceptions +withAsyncOnWithUnmask :: TaskGroup -> Int -> ((forall c. IO c -> IO c) -> IO a) -> (Async a -> IO b) -> IO b +withAsyncOnWithUnmask p cpu actionWith = withAsyncUsing p (rawForkOn cpu) (actionWith unsafeUnmask) + +withAsyncUsing :: TaskGroup -> (IO () -> IO ThreadId) -> IO a -> (Async a -> IO b) + -> IO b +-- The bracket version works, but is slow. We can do better by +-- hand-coding it: +withAsyncUsing p doFork = \action inner -> do + mask $ \restore -> do + a <- atomically $ asyncUsing p doFork $ restore action + r <- restore (inner a) `catchAll` \e -> do cancel a; throwIO e + cancel a + return r + +-- | Wait for an asynchronous action to complete, and return its +-- value. If the asynchronous action threw an exception, then the +-- exception is re-thrown by 'wait'. +-- +-- > wait = atomically . waitSTM +-- +{-# INLINE wait #-} +wait :: Async a -> IO a +wait = atomically . waitSTM + +-- | Wait for an asynchronous action to complete, and return either +-- @Left e@ if the action raised an exception @e@, or @Right a@ if it +-- returned a value @a@. +-- +-- > waitCatch = atomically . waitCatchSTM +-- +{-# INLINE waitCatch #-} +waitCatch :: Async a -> IO (Either SomeException a) +waitCatch = atomically . waitCatchSTM + +-- | Check whether an 'Async' has completed yet. If it has not +-- completed yet, then the result is @Nothing@, otherwise the result +-- is @Just e@ where @e@ is @Left x@ if the @Async@ raised an +-- exception @x@, or @Right a@ if it returned a value @a@. +-- +-- > poll = atomically . pollSTM +-- +{-# INLINE poll #-} +poll :: Async a -> IO (Maybe (Either SomeException a)) +poll = atomically . pollSTM + +-- | A version of 'wait' that can be used inside an STM transaction. +-- +waitSTM :: Async a -> STM a +waitSTM a = do + r <- waitCatchSTM a + either throwSTM return r + +-- | A version of 'waitCatch' that can be used inside an STM transaction. +-- +{-# INLINE waitCatchSTM #-} +waitCatchSTM :: Async a -> STM (Either SomeException a) +waitCatchSTM (Async _ _ w) = w + +-- | A version of 'poll' that can be used inside an STM transaction. +-- +{-# INLINE pollSTM #-} +pollSTM :: Async a -> STM (Maybe (Either SomeException a)) +pollSTM (Async _ _ w) = (Just <$> w) `orElse` return Nothing + +-- | Cancel an asynchronous action by throwing the @ThreadKilled@ +-- exception to it. Has no effect if the 'Async' has already +-- completed. +-- +-- > cancel a = throwTo (asyncThreadId a) ThreadKilled +-- +-- Note that 'cancel' is synchronous in the same sense as 'throwTo'. +-- It does not return until the exception has been thrown in the +-- target thread, or the target thread has completed. In particular, +-- if the target thread is making a foreign call, the exception will +-- not be thrown until the foreign call returns, and in this case +-- 'cancel' may block indefinitely. An asynchronous 'cancel' can +-- of course be obtained by wrapping 'cancel' itself in 'async'. +-- +{-# INLINE cancel #-} +cancel :: Async a -> IO () +cancel = flip cancelWith ThreadKilled + +-- | Cancel an asynchronous action by throwing the supplied exception +-- to it. +-- +-- > cancelWith a x = throwTo (asyncThreadId a) x +-- +-- The notes about the synchronous nature of 'cancel' also apply to +-- 'cancelWith'. +cancelWith' :: Exception e => Pool -> Handle -> e -> IO () +cancelWith' p h e = + (mapM_ (`throwTo` e) =<<) $ atomically $ do + g <- readTVar (tasks p) + let hs = if gelem h g then nodeList g h else [] + xs <- foldM (go g) [] hs + writeTVar (tasks p) $ foldl' (flip delNode) g hs + return xs + where + go g acc h' = maybe acc (:acc) <$> getThreadId g h' + + nodeList :: TaskGraph -> Node -> [Node] + nodeList g k = k : concatMap (nodeList g) (Gr.suc g k) + +cancelWith :: Exception e => Async a -> e -> IO () +cancelWith (Async p h _) = cancelWith' (pool p) h + +-- | Cancel an asynchronous action by throwing the @ThreadKilled@ exception to +-- it, or unregistering it from the task pool if it had not started yet. Has +-- no effect if the 'Async' has already completed. +-- +-- Note that 'cancel' is synchronous in the same sense as 'throwTo'. It does +-- not return until the exception has been thrown in the target thread, or the +-- target thread has completed. In particular, if the target thread is making +-- a foreign call, the exception will not be thrown until the foreign call +-- returns, and in this case 'cancel' may block indefinitely. An asynchronous +-- 'cancel' can of course be obtained by wrapping 'cancel' itself in 'async'. +cancelAll :: TaskGroup -> IO () +cancelAll p = do + hs <- atomically $ do + writeTVar (pending p) IntMap.empty + g <- readTVar (tasks (pool p)) + return $ nodes g + mapM_ (\h -> cancelWith' (pool p) h ThreadKilled) hs + +-- | Wait for any of the supplied asynchronous operations to complete. +-- The value returned is a pair of the 'Async' that completed, and the +-- result that would be returned by 'wait' on that 'Async'. +-- +-- If multiple 'Async's complete or have completed, then the value +-- returned corresponds to the first completed 'Async' in the list. +-- +waitAnyCatch :: [Async a] -> IO (Async a, Either SomeException a) +waitAnyCatch asyncs = + atomically $ + foldr orElse retry $ + map (\a -> do r <- waitCatchSTM a; return (a, r)) asyncs + +-- | Like 'waitAnyCatch', but also cancels the other asynchronous +-- operations as soon as one has completed. +-- +waitAnyCatchCancel :: [Async a] -> IO (Async a, Either SomeException a) +waitAnyCatchCancel asyncs = + waitAnyCatch asyncs `finally` mapM_ cancel asyncs + +-- | Wait for any of the supplied @Async@s to complete. If the first +-- to complete throws an exception, then that exception is re-thrown +-- by 'waitAny'. +-- +-- If multiple 'Async's complete or have completed, then the value +-- returned corresponds to the first completed 'Async' in the list. +-- +waitAny :: [Async a] -> IO (Async a, a) +waitAny asyncs = + atomically $ + foldr orElse retry $ + map (\a -> do r <- waitSTM a; return (a, r)) asyncs + +-- | Like 'waitAny', but also cancels the other asynchronous +-- operations as soon as one has completed. +-- +waitAnyCancel :: [Async a] -> IO (Async a, a) +waitAnyCancel asyncs = + waitAny asyncs `finally` mapM_ cancel asyncs + +-- | Wait for the first of two @Async@s to finish. +waitEitherCatch :: Async a -> Async b + -> IO (Either (Either SomeException a) + (Either SomeException b)) +waitEitherCatch left right = + atomically $ + (Left <$> waitCatchSTM left) + `orElse` + (Right <$> waitCatchSTM right) + +-- | Like 'waitEitherCatch', but also 'cancel's both @Async@s before +-- returning. +-- +waitEitherCatchCancel :: Async a -> Async b + -> IO (Either (Either SomeException a) + (Either SomeException b)) +waitEitherCatchCancel left right = + waitEitherCatch left right `finally` (cancel left >> cancel right) + +-- | Wait for the first of two @Async@s to finish. If the @Async@ +-- that finished first raised an exception, then the exception is +-- re-thrown by 'waitEither'. +-- +waitEither :: Async a -> Async b -> IO (Either a b) +waitEither left right = + atomically $ + (Left <$> waitSTM left) + `orElse` + (Right <$> waitSTM right) + +-- | Like 'waitEither', but the result is ignored. +-- +waitEither_ :: Async a -> Async b -> IO () +waitEither_ left right = + atomically $ + (void $ waitSTM left) + `orElse` + (void $ waitSTM right) + +-- | Like 'waitEither', but also 'cancel's both @Async@s before +-- returning. +-- +waitEitherCancel :: Async a -> Async b -> IO (Either a b) +waitEitherCancel left right = + waitEither left right `finally` (cancel left >> cancel right) + +-- | Waits for both @Async@s to finish, but if either of them throws +-- an exception before they have both finished, then the exception is +-- re-thrown by 'waitBoth'. +-- +waitBoth :: Async a -> Async b -> IO (a,b) +waitBoth left right = + atomically $ do + a <- waitSTM left + `orElse` + (waitSTM right >> retry) + b <- waitSTM right + return (a,b) + + +-- | Link the given @Async@ to the current thread, such that if the +-- @Async@ raises an exception, that exception will be re-thrown in +-- the current thread. +-- +link :: Async a -> IO () +link (Async _ _ w) = do + me <- myThreadId + void $ forkRepeat $ do + r <- atomically $ w + case r of + Left e -> throwTo me e + _ -> return () + +-- | Link two @Async@s together, such that if either raises an +-- exception, the same exception is re-thrown in the other @Async@. +-- +link2 :: Async a -> Async b -> IO () +link2 left right = + void $ forkRepeat $ do + r <- waitEitherCatch left right + case r of + Left (Left e) -> cancelWith right e + Right (Left e) -> cancelWith left e + _ -> return () + + +-- ----------------------------------------------------------------------------- + +-- | Run two @IO@ actions concurrently, and return the first to +-- finish. The loser of the race is 'cancel'led. +-- +-- > race left right = +-- > withAsync left $ \a -> +-- > withAsync right $ \b -> +-- > waitEither a b +-- +race :: TaskGroup -> IO a -> IO b -> IO (Either a b) + +-- | Like 'race', but the result is ignored. +-- +race_ :: TaskGroup -> IO a -> IO b -> IO () + +-- | Run two @IO@ actions concurrently, and return both results. If +-- either action throws an exception at any time, then the other +-- action is 'cancel'led, and the exception is re-thrown by +-- 'concurrently'. +-- +-- > concurrently left right = +-- > withAsync left $ \a -> +-- > withAsync right $ \b -> +-- > waitBoth a b +concurrently :: TaskGroup -> IO a -> IO b -> IO (a,b) + +#define USE_ASYNC_VERSIONS 1 + +#if USE_ASYNC_VERSIONS + +race p left right = + withAsync p left $ \a -> + withAsync p right $ \b -> + waitEither a b + +race_ p left right = + withAsync p left $ \a -> + withAsync p right $ \b -> + waitEither_ a b + +concurrently p left right = + withAsync p left $ \a -> + withAsync p right $ \b -> + waitBoth a b + +#else + +-- MVar versions of race/concurrently +-- More ugly than the Async versions, but quite a bit faster. + +-- race :: IO a -> IO b -> IO (Either a b) +race left right = concurrently' left right collect + where + collect m = do + e <- takeMVar m + case e of + Left ex -> throwIO ex + Right r -> return r + +-- race_ :: IO a -> IO b -> IO () +race_ left right = void $ race left right + +-- concurrently :: IO a -> IO b -> IO (a,b) +concurrently left right = concurrently' left right (collect []) + where + collect [Left a, Right b] _ = return (a,b) + collect [Right b, Left a] _ = return (a,b) + collect xs m = do + e <- takeMVar m + case e of + Left ex -> throwIO ex + Right r -> collect (r:xs) m + +concurrently' :: IO a -> IO b + -> (MVar (Either SomeException (Either a b)) -> IO r) + -> IO r +concurrently' left right collect = do + done <- newEmptyMVar + mask $ \restore -> do + lid <- forkIO $ restore (left >>= putMVar done . Right . Left) + `catchAll` (putMVar done . Left) + rid <- forkIO $ restore (right >>= putMVar done . Right . Right) + `catchAll` (putMVar done . Left) + let stop = killThread lid >> killThread rid + r <- restore (collect done) `onException` stop + stop + return r + +#endif + +-- | maps an @IO@-performing function over any @Traversable@ data +-- type, performing all the @IO@ actions concurrently, and returning +-- the original data structure with the arguments replaced by the +-- results. +-- +-- For example, @mapConcurrently@ works with lists: +-- +-- > pages <- mapConcurrently getURL ["url1", "url2", "url3"] +-- +mapConcurrently :: Traversable t => TaskGroup -> (a -> IO b) -> t a -> IO (t b) +mapConcurrently tg f = flip runConcurrently tg . traverse (\a -> Concurrently $ \_ -> f a) + +-- ----------------------------------------------------------------------------- + +-- | A value of type @Concurrently a@ is an @IO@ operation that can be +-- composed with other @Concurrently@ values, using the @Applicative@ +-- and @Alternative@ instances. +-- +-- Calling @runConcurrently@ on a value of type @Concurrently a@ will +-- execute the @IO@ operations it contains concurrently, before +-- delivering the result of type @a@. +-- +-- For example +-- +-- > (page1, page2, page3) +-- > <- runConcurrently $ (,,) +-- > <$> Concurrently (getURL "url1") +-- > <*> Concurrently (getURL "url2") +-- > <*> Concurrently (getURL "url3") +-- +newtype Concurrently a = Concurrently { runConcurrently :: TaskGroup -> IO a } + +instance Functor Concurrently where + fmap f (Concurrently a) = Concurrently $ fmap f <$> a + +instance Applicative Concurrently where + pure x = Concurrently $ \_ -> return x + Concurrently fs <*> Concurrently as = + Concurrently $ \tg -> (\(f, a) -> f a) <$> concurrently tg (fs tg) (as tg) + +instance Alternative Concurrently where + empty = Concurrently $ \_ -> forever (threadDelay maxBound) + Concurrently as <|> Concurrently bs = + Concurrently $ \tg -> either id id <$> race tg (as tg) (bs tg) + +-- ---------------------------------------------------------------------------- + +-- | Fork a thread that runs the supplied action, and if it raises an +-- exception, re-runs the action. The thread terminates only when the +-- action runs to completion without raising an exception. +forkRepeat :: IO a -> IO ThreadId +forkRepeat action = + mask $ \restore -> + let go = do r <- tryAll (restore action) + case r of + Left _ -> go + _ -> return () + in forkIO go + +catchAll :: IO a -> (SomeException -> IO a) -> IO a +catchAll = catch + +tryAll :: IO a -> IO (Either SomeException a) +tryAll = try + +-- A version of forkIO that does not include the outer exception +-- handler: saves a bit of time when we will be installing our own +-- exception handler. +{-# INLINE rawForkIO #-} +rawForkIO :: IO () -> IO ThreadId +rawForkIO action = IO $ \ s -> + case (fork# action s) of (# s1, tid #) -> (# s1, ThreadId tid #) + +{-# INLINE rawForkOn #-} +rawForkOn :: Int -> IO () -> IO ThreadId +rawForkOn (I# cpu) action = IO $ \ s -> + case (forkOn# cpu action s) of (# s1, tid #) -> (# s1, ThreadId tid #) diff --git a/Control/Concurrent/Async/Pool/Internal.hs b/Control/Concurrent/Async/Pool/Internal.hs new file mode 100644 --- /dev/null +++ b/Control/Concurrent/Async/Pool/Internal.hs @@ -0,0 +1,310 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} + +module Control.Concurrent.Async.Pool.Internal where + +import Control.Applicative (Applicative((<*>), pure), (<$>)) +import Control.Arrow (first) +import Control.Concurrent (ThreadId) +import qualified Control.Concurrent.Async as Async (withAsync) +import Control.Concurrent.Async.Pool.Async +import Control.Concurrent.STM +import Control.Exception (SomeException, throwIO, finally) +import Control.Monad hiding (forM, forM_) +import Control.Monad.Base +import Control.Monad.IO.Class (MonadIO(..)) +import Control.Monad.Trans.Control +import Data.Foldable (Foldable(foldMap), toList, forM_, all) +import Data.Graph.Inductive.Graph as Gr (Graph(empty)) +import Data.IntMap (IntMap) +import qualified Data.IntMap as IntMap +import Data.List (delete) +import Data.Monoid (Monoid(mempty), (<>)) +import Data.Traversable (Traversable(sequenceA), forM) +import Prelude hiding (mapM_, mapM, foldr, all, any, concatMap, foldl1) + +-- | Return a list of actions ready for execution, by checking the graph to +-- ensure that all dependencies have completed. +getReadyNodes :: TaskGroup -> TaskGraph -> STM (IntMap (IO ThreadId)) +getReadyNodes p g = do + availSlots <- readTVar (avail p) + check (availSlots > 0) + taskQueue <- readTVar (pending p) + check (not (IntMap.null taskQueue)) + let readyNodes = IntMap.fromList + . take availSlots + . IntMap.toAscList + . IntMap.filterWithKey (const . isReady) + $ taskQueue + check (not (IntMap.null readyNodes)) + writeTVar (avail p) (availSlots - IntMap.size readyNodes) + writeTVar (pending p) (taskQueue IntMap.\\ readyNodes) + return readyNodes + where + isReady = all isCompleted . inn g + + isCompleted (_, _, Completed) = True + isCompleted (_, _, _) = False + +-- | Return a list of tasks ready to execute, and their related state +-- variables from the dependency graph. +getReadyTasks :: TaskGroup -> STM [(TVar State, IO ThreadId)] +getReadyTasks p = do + g <- readTVar (tasks (pool p)) + map (first (getTaskVar g)) . IntMap.toList <$> getReadyNodes p g + +-- | Create a task pool for managing many-to-many acyclic dependencies among +-- tasks. +createPool :: IO Pool +createPool = Pool <$> newTVarIO Gr.empty + <*> newTVarIO 0 + +-- | Create a task group for executing interdependent tasks concurrently. The +-- number of available slots governs how many tasks may run at one time. +createTaskGroup :: Pool -> Int -> IO TaskGroup +createTaskGroup p cnt = TaskGroup <$> pure p + <*> newTVarIO cnt + <*> newTVarIO mempty + +-- | Execute tasks in a given task group. The number of slots determines how +-- many threads may execute concurrently. +runTaskGroup :: TaskGroup -> IO () +runTaskGroup p = forever $ do + ready <- atomically $ do + cnt <- readTVar (avail p) + check (cnt > 0) + ready <- getReadyTasks p + check (not (null ready)) + forM_ ready $ \(tv, _) -> writeTVar tv Starting + return ready + forM_ ready $ \(tv, go) -> do + t <- go + atomically $ swapTVar tv $ Started t + +-- | Create a task group within the given pool having a specified number of +-- execution slots, but with a bounded lifetime. Leaving the block cancels +-- every task still executing in the group. +withTaskGroupIn :: Pool -> Int -> (TaskGroup -> IO b) -> IO b +withTaskGroupIn p n f = do + g <- createTaskGroup p n + Async.withAsync (runTaskGroup g) $ const $ f g `finally` cancelAll g + +-- | Create both a pool and a task group with a given number of execution +-- slots. +withTaskGroup :: Int -> (TaskGroup -> IO b) -> IO b +withTaskGroup n f = do + p <- createPool + withTaskGroupIn p n f + +-- | Given parent and child tasks, link them so the child cannot execute until +-- the parent has finished. This function does not check for cycles being +-- introduced into the dependency graph, which would prevent a task from +-- ever running. +unsafeMakeDependent :: Pool + -> Handle -- ^ Handle of task doing the waiting + -> Handle -- ^ Handle of task we must wait on (the parent) + -> STM () +unsafeMakeDependent p child parent = do + g <- readTVar (tasks p) + -- If the parent is no longer in the graph, there is no need to establish + -- dependency. The child can begin executing in the next free slot. + when (gelem parent g) $ + modifyTVar (tasks p) (insEdge (parent, child, Pending)) + +-- | Given parent and child tasks, link them so the child cannot execute until +-- the parent has finished. +makeDependent :: Pool + -> Handle -- ^ Handle of task doing the waiting + -> Handle -- ^ Handle of task we must wait on (the parent) + -> STM () +makeDependent p child parent = do + g <- readTVar (tasks p) + -- Check whether the parent is in any way dependent on the child, which + -- would introduce a cycle. + when (gelem parent g) $ + case esp child parent g of + -- If the parent is no longer in the graph, there is no need to + -- establish a dependency. The child can begin executing in the + -- next free slot. + [] -> modifyTVar (tasks p) (insEdge (parent, child, Pending)) + _ -> error "makeDependent: Cycle in task graph" + +-- | Equivalent to async, but acts in STM so that 'makeDependent' may be +-- called after the task is created but before it may begin executing. +asyncSTM :: TaskGroup -> IO a -> STM (Async a) +asyncSTM p = asyncUsing p rawForkIO + +-- | Submit a task which begins executing after all its parents has completed. +-- This is equivalent to submitting a new task with 'asyncSTM' and linking +-- it to its parent using 'makeDependent'. +asyncAfterAll :: TaskGroup -> [Handle] -> IO a -> IO (Async a) +asyncAfterAll p parents t = atomically $ do + child <- asyncUsing p rawForkIO t + forM_ parents $ makeDependent (pool p) (taskHandle child) + return child + +-- | Submit a task which begins executing after its parent has completed. +-- This is equivalent to submitting a new task with 'asyncSTM' and linking +-- it to its parent using 'makeDependent'. +asyncAfter :: TaskGroup -> Async b -> IO a -> IO (Async a) +asyncAfter p parent = asyncAfterAll p [taskHandle parent] + +-- | Helper function used by several of the variants of 'mapTasks' below. +mapTasksWorker :: Traversable t + => TaskGroup + -> t (IO a) + -> (IO (t b) -> IO (t c)) + -> (Async a -> IO b) + -> IO (t c) +mapTasksWorker p fs f g = do + hs <- forM fs $ atomically . asyncUsing p rawForkIO + f $ forM hs g + +-- | Execute a group of tasks within the given pool, returning the results in +-- order. The order of execution is random, but the results are returned in +-- order. +mapTasks :: Traversable t => TaskGroup -> t (IO a) -> IO (t a) +mapTasks p fs = mapTasksWorker p fs id wait + +-- | Execute a group of tasks within the given pool, returning the results in +-- order as an Either type to represent exceptions from actions. The order +-- of execution is random, but the results are returned in order. +mapTasksE :: Traversable t => TaskGroup -> t (IO a) -> IO (t (Either SomeException a)) +mapTasksE p fs = mapTasksWorker p fs id waitCatch + +-- | Execute a group of tasks within the given pool, ignoring results. +mapTasks_ :: Foldable t => TaskGroup -> t (IO a) -> IO () +mapTasks_ p fs = forM_ fs $ atomically . asyncUsing p rawForkIO + +-- | Execute a group of tasks within the given pool, ignoring results, but +-- returning a list of all exceptions. +mapTasksE_ :: Traversable t => TaskGroup -> t (IO a) -> IO (t (Maybe SomeException)) +mapTasksE_ p fs = mapTasksWorker p fs (fmap (fmap leftToMaybe)) waitCatch + where + leftToMaybe :: Either a b -> Maybe a + leftToMaybe = either Just (const Nothing) + +-- | Execute a group of tasks, but return the first result or failure and +-- cancel the remaining tasks. +mapRace :: Foldable t + => TaskGroup -> t (IO a) -> IO (Async a, Either SomeException a) +mapRace p fs = do + hs <- atomically $ sequenceA $ foldMap ((:[]) <$> asyncUsing p rawForkIO) fs + waitAnyCatchCancel hs + +-- | Given a list of actions yielding 'Monoid' results, execute the actions +-- concurrently (up to N at time, based on available slots), and 'mappend' +-- each pair of results concurrently as they become ready. The immediate +-- result of this function is an 'Async' representing the final value. +-- +-- This is equivalent to the following: @mconcat <$> mapTasks n actions@, +-- except that intermediate results can be garbage collected as soon as +-- they've merged. Also, the value returned from this function is a 'Async' +-- which may be polled for the final result. +-- +-- Lastly, if an 'Exception' occurs in any subtask, the final result will +-- yield an exception, but not necessarily the first or last that was caught. +mapReduce :: (Foldable t, Monoid a) + => TaskGroup -- ^ Pool to execute the tasks within + -> t (IO a) -- ^ Set of Monoid-yielding IO actions + -> STM (Async a) -- ^ Returns the final result task +mapReduce p fs = do + -- Submit all the tasks right away, and jobs to combine all those results. + -- Since we're working with a Monoid, it doesn't matter what order they + -- complete in, or what order we combine the results in, just as long we + -- each combination waits on the results it intends to combine. + hs <- sequenceA $ foldMap ((:[]) <$> asyncUsing p rawForkIO) fs + loopM hs + where + loopM hs = do + hs' <- squeeze hs + case hs' of + [] -> error "mapReduce: impossible" + [x] -> return x + xs -> loopM xs + + squeeze [] = (:[]) <$> asyncUsing p rawForkIO (return mempty) + squeeze [x] = return [x] + squeeze (x:y:xs) = do + t <- asyncUsing p rawForkIO $ do + meres <- atomically $ do + -- These polls should by definition always succeed, since this + -- task should not start until results are available. + eres1 <- pollSTM x + eres2 <- pollSTM y + case liftM2 (<>) <$> eres1 <*> eres2 of + Nothing -> retry + Just a -> return a + case meres of + Left e -> throwIO e + Right a -> return a + forM_ [x, y] (unsafeMakeDependent (pool p) (taskHandle t) . taskHandle) + case xs of + [] -> return [t] + _ -> (t :) <$> squeeze xs + +-- | Execute a group of tasks concurrently (using up to N active threads, +-- depending on the pool), and feed results to a continuation as soon as +-- they become available, in random order. That function may return a +-- monoid value which is accumulated to yield a final result. +scatterFoldMapM :: (Foldable t, Monoid b, MonadBaseControl IO m) + => TaskGroup -> t (IO a) -> (Either SomeException a -> m b) -> m b +scatterFoldMapM p fs f = do + hs <- liftBase $ atomically + $ sequenceA + $ foldMap ((:[]) <$> asyncUsing p rawForkIO) fs + control $ \run -> loop run (run $ return mempty) (toList hs) + where + loop _ z [] = z + loop run z hs = do + (h, eres) <- atomically $ do + mres <- foldM go Nothing hs + maybe retry return mres + r' <- z + r <- run $ do + s <- restoreM r' + r <- f eres + return $ s <> r + loop run (return r) (delete h hs) + + go acc@(Just _) _ = return acc + go acc h = do + eres <- pollSTM h + return $ case eres of + Nothing -> acc + Just (Left e) -> Just (h, Left e) + Just (Right x) -> Just (h, Right x) + +-- | The 'Task' Applicative and Monad allow for task dependencies to be built +-- using applicative and do notation. Monadic evaluation is sequenced, +-- while applicative evaluation is done concurrently for each argument. In +-- this way, mixing the two allows you to build a dependency tree using +-- ordinary Haskell code. +newtype Task a = Task { runTask' :: TaskGroup -> IO (IO a) } + +-- | Run a value in the 'Task' monad and block until the final result is +-- computed. +runTask :: TaskGroup -> Task a -> IO a +runTask group ts = join $ runTask' ts group + +-- | Lift any 'IO' action into a 'Task'. This is a synonym for 'liftIO'. +task :: IO a -> Task a +task action = Task $ \_ -> return action + +instance Functor Task where + fmap f (Task k) = Task $ fmap (fmap (liftM f)) k + +instance Applicative Task where + pure x = Task $ \_ -> return (return x) + Task f <*> Task x = Task $ \tg -> do + xa <- x tg + x' <- wait <$> async tg xa + fa <- f tg + return $ fa <*> x' + +instance Monad Task where + return = pure + Task m >>= f = Task $ \tg -> join (m tg) >>= flip runTask' tg . f + +instance MonadIO Task where + liftIO = task diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,19 @@ +opyright (c) 2014 John Wiegley + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/async-pool.cabal b/async-pool.cabal new file mode 100644 --- /dev/null +++ b/async-pool.cabal @@ -0,0 +1,56 @@ +Name: async-pool +Version: 0.8.0 +Synopsis: A modified version of async that supports worker groups and many-to-many task dependencies +License-file: LICENSE +License: MIT +Author: Simon Marlow, John Wiegley +Maintainer: johnw@newartisans.com +Build-Type: Simple +Cabal-Version: >=1.10 +Category: System +Description: + This library modifies the @async@ package to allow for task pooling and + many-to-many dependencies between tasks. + +Source-repository head + type: git + location: git://github.com/jwiegley/async-pool.git + +Library + default-language: Haskell98 + ghc-options: -Wall + build-depends: + base >= 3 && < 5 + , fgl + , async + , stm + , containers + , transformers + , transformers-base + , monad-control + exposed-modules: + Control.Concurrent.Async.Pool + other-modules: + Control.Concurrent.Async.Pool.Async + Control.Concurrent.Async.Pool.Internal + +test-suite test + hs-source-dirs: . test + default-language: Haskell2010 + main-is: main.hs + type: exitcode-stdio-1.0 + ghc-options: -threaded -with-rtsopts "-N2" + other-modules: + Control.Concurrent.Async.Pool.Async + Control.Concurrent.Async.Pool.Internal + build-depends: + base + , async + , stm + , transformers + , transformers-base + , monad-control + , fgl + , containers + , hspec >= 1.4 + , time diff --git a/test/main.hs b/test/main.hs new file mode 100644 --- /dev/null +++ b/test/main.hs @@ -0,0 +1,193 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeSynonymInstances #-} + +import Control.Applicative +import Control.Concurrent +import qualified Control.Concurrent.Async as Async +import Control.Concurrent.STM +import Control.Exception +import Control.Monad +import Data.Functor +import Data.Graph.Inductive.Graph as Gr +import qualified Data.IntMap as M +import Data.Monoid +import Control.Concurrent.Async.Pool.Async +import Control.Concurrent.Async.Pool.Internal +import Data.Time +import Test.Hspec + +instance Show (TVar State) where + show _ = "Task" + +testAvail p x = do + a <- atomically $ readTVar (avail p) + a `shouldBe` x + +testGraph p f x = do + g <- atomically $ readTVar (tasks (pool p)) + (f g `shouldBe` x) `onException` prettyPrint g + +graphPict p x = do + g <- atomically $ readTVar (tasks (pool p)) + prettify g `shouldBe` x + +testProcs p f x = do + ps <- atomically $ do + g <- readTVar (tasks (pool p)) + foldM (go g) M.empty (nodes g) + (f ps `shouldBe` x) `onException` print (M.keys ps) + where + go g acc h' = do + mres <- getThreadId g h' + return $ case mres of + Nothing -> acc + Just x -> M.insert h' x acc + +main :: IO () +main = hspec $ do + describe "simple tasks" $ do + it "completes a task" $ do + p' <- createPool + p <- createTaskGroup p' 8 + + -- Upon creation of the pool, both the task graph and the process map + -- are empty. + testAvail p 8 + testGraph p isEmpty True + testProcs p M.null True + + -- We submit a task, so that the graph has an entry, but the process + -- map is still empty. + h <- async p $ return (42 :: Int) + testGraph p isEmpty False + testProcs p M.null True + + -- Start running the pool in another thread and wait 100ms. This is + -- time enough for the task to finish. + Async.withAsync (runTaskGroup p) $ \_ -> do + threadDelay 100000 + + -- Now the task graph should be empty. + testAvail p 8 + testProcs p M.null True + + -- Wait on the task and see the result value from the task. + res <- wait h + res `shouldBe` 42 + + -- Now the task graph should be empty, since observing the final + -- state removed the process entry from the map. + testGraph p isEmpty True + testProcs p M.null True + + it "completes two concurrent tasks" $ do + p' <- createPool + p <- createTaskGroup p' 8 + + testAvail p 8 + testGraph p isEmpty True + testProcs p M.null True + + h1 <- async p $ return (42 :: Int) + h2 <- async p $ return 43 + + testGraph p isEmpty False + testProcs p M.null True + + graphPict p "0:Task->[]\n1:Task->[]\n" + + Async.withAsync (runTaskGroup p) $ \_ -> do + threadDelay 100000 + + testAvail p 8 + testProcs p M.null True + + res <- wait h1 + res `shouldBe` 42 + res' <- wait h2 + res' `shouldBe` 43 + + testGraph p isEmpty True + testProcs p M.null True + + it "completes two linked tasks" $ do + p' <- createPool + p <- createTaskGroup p' 8 + + testAvail p 8 + testGraph p isEmpty True + testProcs p M.null True + + -- Start two interdependent tasks. The first task waits a bit and + -- then writes a value into a TVar. The second task does not wait, but + -- immediately reads the value from the TVar and adds to it. + -- Sequencing should cause these two to happen in series. + x <- atomically $ newTVar (0 :: Int) + h1 <- async p $ do + threadDelay 50000 + atomically $ writeTVar x 42 + return 42 + h2 <- asyncAfter p h1 $ do + y <- atomically $ readTVar x + return $ y + 100 + + testGraph p isEmpty False + testProcs p M.null True + + graphPict p "0:Task->[(Pending,1)]\n1:Task->[]\n" + + Async.withAsync (runTaskGroup p) $ \_ -> do + threadDelay 250000 + + testAvail p 8 + testProcs p M.null True + + res <- wait h1 + res `shouldBe` 42 + res' <- wait h2 + res' `shouldBe` 142 + + testGraph p isEmpty True + testProcs p M.null True + + describe "map reduce" $ do + it "sums a group of integers" $ do + p' <- createPool + p <- createTaskGroup p' 8 + h <- atomically $ mapReduce p $ map (return . Sum) [1..10] + g <- atomically $ readTVar (tasks (pool p)) + Async.withAsync (runTaskGroup p) $ const $ do + x <- wait h + x `shouldBe` Sum 55 + + describe "scatter fold" $ do + it "sums in random order" $ withTaskGroup 8 $ \p -> do + let go x = do + threadDelay (10000 * (x `mod` 3)) + return $ Sum x + res <- scatterFoldMapM p (map go [1..20]) $ \ex -> + case ex of + Left e -> mempty <$ print ("Hmmm... " ++ show e) + Right x -> return x + getSum res `shouldBe` 210 + + describe "applicative style" $ do + it "maps tasks" $ withTaskGroup 8 $ \p -> do + start <- getCurrentTime + x <- mapTasks p (replicate 8 (threadDelay 1000000 >> return (1 :: Int))) + sum x `shouldBe` 8 + end <- getCurrentTime + let diff = diffUTCTime end start + diff < 1.2 `shouldBe` True + + it "counts to ten in one second" $ withTaskGroup 8 $ \p -> do + start <- getCurrentTime + x <- runTask p $ + let k a b c d e f g h = a + b + c + d + e + f + g + h + h = task (threadDelay 1000000 >> return (1 :: Int)) + in k <$> h <*> h <*> h <*> h <*> h <*> h <*> h <*> h + x `shouldBe` 8 + end <- getCurrentTime + let diff = diffUTCTime end start + diff < 1.2 `shouldBe` True