{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE BlockArguments      #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE NamedFieldPuns      #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeData            #-}


-- | This module allows the management of a multiple Async jobs which
-- are grouped by an 'Ord group => group' type.
--
module Control.Concurrent.JobPool
  ( JobPool
  , HasQueue (..)
  , Job (..)
  , withJobPool
  , withJobPool_
  , forkJob
  , forkJobOn
  , readSize
  , readGroupSize
  , waitForJob
  , cancelGroup
  ) where

import Data.Functor (($>))
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map

import Control.Concurrent.Class.MonadSTM
import Control.Exception (SomeAsyncException (..))
import Control.Monad (void, when)
import Control.Monad.Class.MonadAsync
import Control.Monad.Class.MonadFork (MonadThread (..))
import Control.Monad.Class.MonadThrow


-- | Whether a 'JobPool' was created with a completion queue.  Only a
-- 'WithQueue' pool can be passed to 'waitForJob'.
--
type data HasQueue = WithQueue | WithoutQueue

-- | A completion queue, indexed by whether it exists at all.  Matching on
-- this GADT is what makes 'waitForJob' total: at index 'WithQueue' the only
-- constructor is 'CompletionQueue'.
--
data CompletionQueue (q :: HasQueue) m a where
  CompletionQueue   :: !(TQueue m a) -> CompletionQueue WithQueue    m a
  NoCompletionQueue ::                  CompletionQueue WithoutQueue m a

-- | JobPool allows to submit asynchronous jobs, wait for their completion or
-- cancel.  Jobs are grouped, each group can be cancelled separately.
--
data JobPool (q :: HasQueue) group m a = JobPool {
       forall (q :: HasQueue) group (m :: * -> *) a.
JobPool q group m a
-> TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar         :: !(TVar m (Map (group, ThreadId m) (Async m ()))),
       forall (q :: HasQueue) group (m :: * -> *) a.
JobPool q group m a -> CompletionQueue q m a
completionQueue :: !(CompletionQueue q m a)
     }

-- | An asynchronous job which belongs to some group and its exception handler.
--
data Job group m a =
    Job (m a)                  -- ^ job
        (SomeException -> m a) -- ^ error handler
        group                  -- ^ job group
        String                 -- ^ thread label

withJobPool :: forall group m a b.
               (MonadAsync m, MonadThrow m, MonadLabelledSTM m)
            => (JobPool WithQueue group m a -> m b) -> m b
withJobPool :: forall group (m :: * -> *) a b.
(MonadAsync m, MonadThrow m, MonadLabelledSTM m) =>
(JobPool WithQueue group m a -> m b) -> m b
withJobPool =
    m (JobPool WithQueue group m a)
-> (JobPool WithQueue group m a -> m ())
-> (JobPool WithQueue group m a -> m b)
-> m b
forall a b c. m a -> (a -> m b) -> (a -> m c) -> m c
forall (m :: * -> *) a b c.
MonadThrow m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket m (JobPool WithQueue group m a)
create JobPool WithQueue group m a -> m ()
forall (m :: * -> *) (q :: HasQueue) group a.
(MonadAsync m, MonadThrow m) =>
JobPool q group m a -> m ()
closeJobPool
  where
    create :: m (JobPool WithQueue group m a)
    create :: m (JobPool WithQueue group m a)
create =
      STM m (JobPool WithQueue group m a)
-> m (JobPool WithQueue group m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (JobPool WithQueue group m a)
 -> m (JobPool WithQueue group m a))
-> STM m (JobPool WithQueue group m a)
-> m (JobPool WithQueue group m a)
forall a b. (a -> b) -> a -> b
$
        TVar m (Map (group, ThreadId m) (Async m ()))
-> CompletionQueue WithQueue m a -> JobPool WithQueue group m a
forall (q :: HasQueue) group (m :: * -> *) a.
TVar m (Map (group, ThreadId m) (Async m ()))
-> CompletionQueue q m a -> JobPool q group m a
JobPool (TVar m (Map (group, ThreadId m) (Async m ()))
 -> CompletionQueue WithQueue m a -> JobPool WithQueue group m a)
-> STM m (TVar m (Map (group, ThreadId m) (Async m ())))
-> STM
     m (CompletionQueue WithQueue m a -> JobPool WithQueue group m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Map (group, ThreadId m) (Async m ())
-> STM m (TVar m (Map (group, ThreadId m) (Async m ())))
forall a. a -> STM m (TVar m a)
forall (m :: * -> *) a. MonadSTM m => a -> STM m (TVar m a)
newTVar Map (group, ThreadId m) (Async m ())
forall k a. Map k a
Map.empty STM m (TVar m (Map (group, ThreadId m) (Async m ())))
-> (TVar m (Map (group, ThreadId m) (Async m ()))
    -> STM m (TVar m (Map (group, ThreadId m) (Async m ()))))
-> STM m (TVar m (Map (group, ThreadId m) (Async m ())))
forall a b. STM m a -> (a -> STM m b) -> STM m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \TVar m (Map (group, ThreadId m) (Async m ()))
v -> TVar m (Map (group, ThreadId m) (Async m ())) -> String -> STM m ()
forall a. TVar m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TVar m a -> String -> STM m ()
labelTVar TVar m (Map (group, ThreadId m) (Async m ()))
v String
"job-pool" STM m ()
-> TVar m (Map (group, ThreadId m) (Async m ()))
-> STM m (TVar m (Map (group, ThreadId m) (Async m ())))
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> TVar m (Map (group, ThreadId m) (Async m ()))
v)
                STM
  m (CompletionQueue WithQueue m a -> JobPool WithQueue group m a)
-> STM m (CompletionQueue WithQueue m a)
-> STM m (JobPool WithQueue group m a)
forall a b. STM m (a -> b) -> STM m a -> STM m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TQueue m a -> CompletionQueue WithQueue m a
forall (m :: * -> *) a. TQueue m a -> CompletionQueue WithQueue m a
CompletionQueue (TQueue m a -> CompletionQueue WithQueue m a)
-> STM m (TQueue m a) -> STM m (CompletionQueue WithQueue m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM m (TQueue m a)
forall a. STM m (TQueue m a)
forall (m :: * -> *) a. MonadSTM m => STM m (TQueue m a)
newTQueue)

-- | Like 'withJobPool', but for a pool whose jobs' results nobody ever
-- inspects: no 'TQueue' is created, so there is nothing for 'forkJob'\/
-- 'forkJobOn' to write a finished job's result into, and nothing to
-- remember to drain.  'waitForJob' does not typecheck against a pool
-- created this way.
--
withJobPool_ :: forall group m a b.
                (MonadAsync m, MonadThrow m, MonadLabelledSTM m)
             => (JobPool WithoutQueue group m a -> m b) -> m b
withJobPool_ :: forall group (m :: * -> *) a b.
(MonadAsync m, MonadThrow m, MonadLabelledSTM m) =>
(JobPool WithoutQueue group m a -> m b) -> m b
withJobPool_ =
    m (JobPool WithoutQueue group m a)
-> (JobPool WithoutQueue group m a -> m ())
-> (JobPool WithoutQueue group m a -> m b)
-> m b
forall a b c. m a -> (a -> m b) -> (a -> m c) -> m c
forall (m :: * -> *) a b c.
MonadThrow m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket m (JobPool WithoutQueue group m a)
create JobPool WithoutQueue group m a -> m ()
forall (m :: * -> *) (q :: HasQueue) group a.
(MonadAsync m, MonadThrow m) =>
JobPool q group m a -> m ()
closeJobPool
  where
    create :: m (JobPool WithoutQueue group m a)
    create :: m (JobPool WithoutQueue group m a)
create =
      STM m (JobPool WithoutQueue group m a)
-> m (JobPool WithoutQueue group m a)
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m (JobPool WithoutQueue group m a)
 -> m (JobPool WithoutQueue group m a))
-> STM m (JobPool WithoutQueue group m a)
-> m (JobPool WithoutQueue group m a)
forall a b. (a -> b) -> a -> b
$
        TVar m (Map (group, ThreadId m) (Async m ()))
-> CompletionQueue WithoutQueue m a
-> JobPool WithoutQueue group m a
forall (q :: HasQueue) group (m :: * -> *) a.
TVar m (Map (group, ThreadId m) (Async m ()))
-> CompletionQueue q m a -> JobPool q group m a
JobPool (TVar m (Map (group, ThreadId m) (Async m ()))
 -> CompletionQueue WithoutQueue m a
 -> JobPool WithoutQueue group m a)
-> STM m (TVar m (Map (group, ThreadId m) (Async m ())))
-> STM
     m
     (CompletionQueue WithoutQueue m a
      -> JobPool WithoutQueue group m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Map (group, ThreadId m) (Async m ())
-> STM m (TVar m (Map (group, ThreadId m) (Async m ())))
forall a. a -> STM m (TVar m a)
forall (m :: * -> *) a. MonadSTM m => a -> STM m (TVar m a)
newTVar Map (group, ThreadId m) (Async m ())
forall k a. Map k a
Map.empty STM m (TVar m (Map (group, ThreadId m) (Async m ())))
-> (TVar m (Map (group, ThreadId m) (Async m ()))
    -> STM m (TVar m (Map (group, ThreadId m) (Async m ()))))
-> STM m (TVar m (Map (group, ThreadId m) (Async m ())))
forall a b. STM m a -> (a -> STM m b) -> STM m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \TVar m (Map (group, ThreadId m) (Async m ()))
v -> TVar m (Map (group, ThreadId m) (Async m ())) -> String -> STM m ()
forall a. TVar m a -> String -> STM m ()
forall (m :: * -> *) a.
MonadLabelledSTM m =>
TVar m a -> String -> STM m ()
labelTVar TVar m (Map (group, ThreadId m) (Async m ()))
v String
"job-pool" STM m ()
-> TVar m (Map (group, ThreadId m) (Async m ()))
-> STM m (TVar m (Map (group, ThreadId m) (Async m ())))
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> TVar m (Map (group, ThreadId m) (Async m ()))
v)
                STM
  m
  (CompletionQueue WithoutQueue m a
   -> JobPool WithoutQueue group m a)
-> STM m (CompletionQueue WithoutQueue m a)
-> STM m (JobPool WithoutQueue group m a)
forall a b. STM m (a -> b) -> STM m a -> STM m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> CompletionQueue WithoutQueue m a
-> STM m (CompletionQueue WithoutQueue m a)
forall a. a -> STM m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CompletionQueue WithoutQueue m a
forall (m :: * -> *) a. CompletionQueue WithoutQueue m a
NoCompletionQueue

-- 'bracket' requires that this callback is uninterruptible.  Note also that
-- 'async' library is using 'uninterruptibleCancel' in 'withAsync' combinator.
-- This can only deadlock if the threads in 'JobPool' got deadlocked so that
-- the asynchronous exception cannot be delivered, e.g. deadlock in an ffi
-- call or a tight loop which does not allocate (which is not a deadlock per
-- se, but rather a rare unfortunate condition).
closeJobPool :: (MonadAsync m, MonadThrow m)
             => JobPool q group m a -> m ()
closeJobPool :: forall (m :: * -> *) (q :: HasQueue) group a.
(MonadAsync m, MonadThrow m) =>
JobPool q group m a -> m ()
closeJobPool JobPool{TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall (q :: HasQueue) group (m :: * -> *) a.
JobPool q group m a
-> TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar} = do
  jobs <- TVar m (Map (group, ThreadId m) (Async m ()))
-> m (Map (group, ThreadId m) (Async m ()))
forall a. TVar m a -> m a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> m a
readTVarIO TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar
  mapM_ uninterruptibleCancel jobs


forkJob' :: forall q group m a.
            ( MonadAsync m, MonadMask m
            , Ord group
            )
         => (((forall x. m x -> m x) -> m ()) -> m (Async m ()))
         -- ^ how to fork a thread, e.g. `async`, `asyncOn`.
         -> JobPool q group m a
         -> Job     group m a
         -> m ()
forkJob' :: forall (q :: HasQueue) group (m :: * -> *) a.
(MonadAsync m, MonadMask m, Ord group) =>
(((forall x. m x -> m x) -> m ()) -> m (Async m ()))
-> JobPool q group m a -> Job group m a -> m ()
forkJob' ((forall x. m x -> m x) -> m ()) -> m (Async m ())
doFork JobPool{TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall (q :: HasQueue) group (m :: * -> *) a.
JobPool q group m a
-> TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar, CompletionQueue q m a
completionQueue :: forall (q :: HasQueue) group (m :: * -> *) a.
JobPool q group m a -> CompletionQueue q m a
completionQueue :: CompletionQueue q m a
completionQueue} (Job m a
action SomeException -> m a
handler group
group String
label) =
    m () -> m ()
forall x. m x -> m x
forall (m :: * -> *) a. MonadMask m => m a -> m a
mask_ do
      jobAsync <- ((forall x. m x -> m x) -> m ()) -> m (Async m ())
doFork \forall x. m x -> m x
restore -> do
        tid <- m (ThreadId m)
forall (m :: * -> *). MonadThread m => m (ThreadId m)
myThreadId
        io tid restore
          `onException`
          deregister tid
        deregister tid

      let !tid = Async m () -> ThreadId m
forall a. Async m a -> ThreadId m
forall (m :: * -> *) a. MonadAsync m => Async m a -> ThreadId m
asyncThreadId Async m ()
jobAsync
      atomically $ modifyTVar' jobsVar (Map.insert (group, tid) $! jobAsync)
      return ()
  where
    -- | Remove this job's own @(group, tid)@ entry from 'jobsVar', once it is
    -- actually there.  This can only block before the parent's insert, which
    -- is guaranteed to follow, since the parent is masked.
    --
    deregister :: ThreadId m -> m ()
    deregister :: ThreadId m -> m ()
deregister ThreadId m
tid =
      STM m () -> m ()
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m () -> m ()) -> STM m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        registered <- (group, ThreadId m) -> Map (group, ThreadId m) (Async m ()) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member (group
group, ThreadId m
tid) (Map (group, ThreadId m) (Async m ()) -> Bool)
-> STM m (Map (group, ThreadId m) (Async m ())) -> STM m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar m (Map (group, ThreadId m) (Async m ()))
-> STM m (Map (group, ThreadId m) (Async m ()))
forall a. TVar m a -> STM m a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> STM m a
readTVar TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar
        if registered
          then modifyTVar' jobsVar (Map.delete (group, tid))
          else retry

    notAsyncExceptions :: SomeException -> Maybe SomeException
    notAsyncExceptions :: SomeException -> Maybe SomeException
notAsyncExceptions SomeException
e
      | Just (SomeAsyncException e
_) <- SomeException -> Maybe SomeAsyncException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e
                  = Maybe SomeException
forall a. Maybe a
Nothing
      | Bool
otherwise = SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just SomeException
e

    io :: ThreadId m
       -> (forall x. m x -> m x)
       -> m ()
    io :: ThreadId m -> (forall x. m x -> m x) -> m ()
io ThreadId m
tid forall x. m x -> m x
restore = do
      ThreadId m -> String -> m ()
forall (m :: * -> *). MonadThread m => ThreadId m -> String -> m ()
labelThread ThreadId m
tid String
label
      -- NOTE: `network-mux` relies on not passing asynchronous exceptions to
      -- the exception handler, see `Network.Mux.miniProtocolJob`.
      !res <- (SomeException -> Maybe SomeException)
-> (SomeException -> m a) -> m a -> m a
forall e b a.
Exception e =>
(e -> Maybe b) -> (b -> m a) -> m a -> m a
forall (m :: * -> *) e b a.
(MonadCatch m, Exception e) =>
(e -> Maybe b) -> (b -> m a) -> m a -> m a
handleJust SomeException -> Maybe SomeException
notAsyncExceptions SomeException -> m a
handler (m a -> m a) -> m a -> m a
forall a b. (a -> b) -> a -> b
$
              m a -> m a
forall x. m x -> m x
restore m a
action
      case completionQueue of
        CompletionQueue q m a
NoCompletionQueue  -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        CompletionQueue TQueue m a
cq -> STM m () -> m ()
forall a. HasCallStack => STM m a -> m a
forall (m :: * -> *) a.
(MonadSTM m, HasCallStack) =>
STM m a -> m a
atomically (STM m () -> m ()) -> STM m () -> m ()
forall a b. (a -> b) -> a -> b
$ TQueue m a -> a -> STM m ()
forall a. TQueue m a -> a -> STM m ()
forall (m :: * -> *) a. MonadSTM m => TQueue m a -> a -> STM m ()
writeTQueue TQueue m a
cq a
res



-- | Fork a `Job` using `async`.
--
forkJob :: forall q group m a.
           ( MonadAsync m, MonadMask m
           , Ord group
           )
        => JobPool q group m a
        -> Job     group m a
        -> m ()
forkJob :: forall (q :: HasQueue) group (m :: * -> *) a.
(MonadAsync m, MonadMask m, Ord group) =>
JobPool q group m a -> Job group m a -> m ()
forkJob = (((forall x. m x -> m x) -> m ()) -> m (Async m ()))
-> JobPool q group m a -> Job group m a -> m ()
forall (q :: HasQueue) group (m :: * -> *) a.
(MonadAsync m, MonadMask m, Ord group) =>
(((forall x. m x -> m x) -> m ()) -> m (Async m ()))
-> JobPool q group m a -> Job group m a -> m ()
forkJob' ((forall x. m x -> m x) -> m ()) -> m (Async m ())
forall a. ((forall x. m x -> m x) -> m a) -> m (Async m a)
forall (m :: * -> *) a.
MonadAsync m =>
((forall b. m b -> m b) -> m a) -> m (Async m a)
asyncWithUnmask


-- | Fork a `Job` using `asyncOn`.
--
forkJobOn :: forall q group m a.
             ( MonadAsync m, MonadMask m
             , Ord group
             )
          => Int
          -> JobPool q group m a
          -> Job     group m a
          -> m ()
forkJobOn :: forall (q :: HasQueue) group (m :: * -> *) a.
(MonadAsync m, MonadMask m, Ord group) =>
Int -> JobPool q group m a -> Job group m a -> m ()
forkJobOn Int
cap = (((forall x. m x -> m x) -> m ()) -> m (Async m ()))
-> JobPool q group m a -> Job group m a -> m ()
forall (q :: HasQueue) group (m :: * -> *) a.
(MonadAsync m, MonadMask m, Ord group) =>
(((forall x. m x -> m x) -> m ()) -> m (Async m ()))
-> JobPool q group m a -> Job group m a -> m ()
forkJob' (Int -> ((forall x. m x -> m x) -> m ()) -> m (Async m ())
forall a. Int -> ((forall x. m x -> m x) -> m a) -> m (Async m a)
forall (m :: * -> *) a.
MonadAsync m =>
Int -> ((forall b. m b -> m b) -> m a) -> m (Async m a)
asyncOnWithUnmask Int
cap)


readSize :: MonadSTM m => JobPool q group m a -> STM m Int
readSize :: forall (m :: * -> *) (q :: HasQueue) group a.
MonadSTM m =>
JobPool q group m a -> STM m Int
readSize JobPool{TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall (q :: HasQueue) group (m :: * -> *) a.
JobPool q group m a
-> TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar} = Map (group, ThreadId m) (Async m ()) -> Int
forall k a. Map k a -> Int
Map.size (Map (group, ThreadId m) (Async m ()) -> Int)
-> STM m (Map (group, ThreadId m) (Async m ())) -> STM m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar m (Map (group, ThreadId m) (Async m ()))
-> STM m (Map (group, ThreadId m) (Async m ()))
forall a. TVar m a -> STM m a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> STM m a
readTVar TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar

readGroupSize :: ( MonadSTM m
                 , Eq group
                 )
              => JobPool q group m a -> group -> STM m Int
readGroupSize :: forall (m :: * -> *) group (q :: HasQueue) a.
(MonadSTM m, Eq group) =>
JobPool q group m a -> group -> STM m Int
readGroupSize JobPool{TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall (q :: HasQueue) group (m :: * -> *) a.
JobPool q group m a
-> TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar} group
group =
      Map (group, ThreadId m) (Async m ()) -> Int
forall k a. Map k a -> Int
Map.size
    (Map (group, ThreadId m) (Async m ()) -> Int)
-> (Map (group, ThreadId m) (Async m ())
    -> Map (group, ThreadId m) (Async m ()))
-> Map (group, ThreadId m) (Async m ())
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((group, ThreadId m) -> Async m () -> Bool)
-> Map (group, ThreadId m) (Async m ())
-> Map (group, ThreadId m) (Async m ())
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (\(group
group', ThreadId m
_) Async m ()
_ -> group
group' group -> group -> Bool
forall a. Eq a => a -> a -> Bool
== group
group)
  (Map (group, ThreadId m) (Async m ()) -> Int)
-> STM m (Map (group, ThreadId m) (Async m ())) -> STM m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar m (Map (group, ThreadId m) (Async m ()))
-> STM m (Map (group, ThreadId m) (Async m ()))
forall a. TVar m a -> STM m a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> STM m a
readTVar TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar

-- | Wait for next successfully completed job.  Unlike 'wait' it will not throw
-- if a job errors.
--
waitForJob :: MonadSTM m => JobPool WithQueue group m a -> STM m a
waitForJob :: forall (m :: * -> *) group a.
MonadSTM m =>
JobPool WithQueue group m a -> STM m a
waitForJob JobPool{completionQueue :: forall (q :: HasQueue) group (m :: * -> *) a.
JobPool q group m a -> CompletionQueue q m a
completionQueue = CompletionQueue TQueue m a
cq} = TQueue m a -> STM m a
forall a. TQueue m a -> STM m a
forall (m :: * -> *) a. MonadSTM m => TQueue m a -> STM m a
readTQueue TQueue m a
cq

-- | Cancel all threads in a given group.  Blocks until all threads terminated.
--
cancelGroup :: ( MonadAsync m
               , Eq group
               )
            => JobPool q group m a -> group -> m ()
cancelGroup :: forall (m :: * -> *) group (q :: HasQueue) a.
(MonadAsync m, Eq group) =>
JobPool q group m a -> group -> m ()
cancelGroup JobPool { TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall (q :: HasQueue) group (m :: * -> *) a.
JobPool q group m a
-> TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar } group
group = do
    jobs <- TVar m (Map (group, ThreadId m) (Async m ()))
-> m (Map (group, ThreadId m) (Async m ()))
forall a. TVar m a -> m a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> m a
readTVarIO TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar
    void $ Map.traverseWithKey
             (\(group
group', ThreadId m
_) Async m ()
thread ->
                Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (group
group' group -> group -> Bool
forall a. Eq a => a -> a -> Bool
== group
group) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
                  Async m () -> m ()
forall a. Async m a -> m ()
forall (m :: * -> *) a. MonadAsync m => Async m a -> m ()
cancel Async m ()
thread
             )
             jobs