{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE NamedFieldPuns      #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}


-- | This module allows the management of a multiple Async jobs which
-- are grouped by an 'Ord group => group' type.
--
module Control.Concurrent.JobPool
  ( JobPool
  , Job (..)
  , withJobPool
  , forkJob
  , 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

-- | JobPool allows to submit asynchronous jobs, wait for their completion or
-- cancel.  Jobs are grouped, each group can be cancelled separately.
--
data JobPool group m a = JobPool {
       forall group (m :: * -> *) a.
JobPool group m a -> TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar         :: !(TVar m (Map (group, ThreadId m) (Async m ()))),
       forall group (m :: * -> *) a. JobPool group m a -> TQueue m a
completionQueue :: !(TQueue 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 group m a -> m b) -> m b
withJobPool :: forall group (m :: * -> *) a b.
(MonadAsync m, MonadThrow m, MonadLabelledSTM m) =>
(JobPool group m a -> m b) -> m b
withJobPool =
    m (JobPool group m a)
-> (JobPool group m a -> m ()) -> (JobPool 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 group m a)
create JobPool group m a -> m ()
close
  where
    create :: m (JobPool group m a)
    create :: m (JobPool group m a)
create =
      STM m (JobPool group m a) -> m (JobPool 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 group m a) -> m (JobPool group m a))
-> STM m (JobPool group m a) -> m (JobPool group m a)
forall a b. (a -> b) -> a -> b
$
        TVar m (Map (group, ThreadId m) (Async m ()))
-> TQueue m a -> JobPool group m a
forall group (m :: * -> *) a.
TVar m (Map (group, ThreadId m) (Async m ()))
-> TQueue m a -> JobPool group m a
JobPool (TVar m (Map (group, ThreadId m) (Async m ()))
 -> TQueue m a -> JobPool group m a)
-> STM m (TVar m (Map (group, ThreadId m) (Async m ())))
-> STM m (TQueue m a -> JobPool 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 (TQueue m a -> JobPool group m a)
-> STM m (TQueue m a) -> STM m (JobPool 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
<*> STM m (TQueue m a)
forall a. STM m (TQueue m a)
forall (m :: * -> *) a. MonadSTM m => STM m (TQueue m a)
newTQueue

    -- 'bracket' requires that the 'close' callback is uninterruptible.  Note
    -- also that 'async' library is using 'uninterruptibleClose' 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).
    close :: JobPool group m a -> m ()
    close :: JobPool group m a -> m ()
close JobPool{TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall group (m :: * -> *) a.
JobPool 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 group m a.
           ( MonadAsync m, MonadMask m
           , Ord group
           )
        => JobPool group m a
        -> Job     group m a
        -> m ()
forkJob :: forall group (m :: * -> *) a.
(MonadAsync m, MonadMask m, Ord group) =>
JobPool group m a -> Job group m a -> m ()
forkJob JobPool{TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall group (m :: * -> *) a.
JobPool group m a -> TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar, TQueue m a
completionQueue :: forall group (m :: * -> *) a. JobPool group m a -> TQueue m a
completionQueue :: TQueue m a
completionQueue} (Job m a
action SomeException -> m a
handler group
group String
label) =
    ((forall a. m a -> m a) -> m ()) -> m ()
forall b. ((forall a. m a -> m a) -> m b) -> m b
forall (m :: * -> *) b.
MonadMask m =>
((forall a. m a -> m a) -> m b) -> m b
mask (((forall a. m a -> m a) -> m ()) -> m ())
-> ((forall a. m a -> m a) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> m a
restore -> do
      jobAsync <- m () -> m (Async m ())
forall a. m a -> m (Async m a)
forall (m :: * -> *) a. MonadAsync m => m a -> m (Async m a)
async (m () -> m (Async m ())) -> m () -> m (Async m ())
forall a b. (a -> b) -> a -> b
$ do
        tid <- m (ThreadId m)
forall (m :: * -> *). MonadThread m => m (ThreadId m)
myThreadId
        io tid restore
          `onException`
          atomically (modifyTVar' jobsVar (Map.delete (group, tid)))
        atomically (modifyTVar' jobsVar (Map.delete (group, 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
    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 a. m a -> m a) -> m ()
io ThreadId m
tid forall a. m a -> m a
restore = do
      ThreadId m -> String -> m ()
forall (m :: * -> *). MonadThread m => ThreadId m -> String -> m ()
labelThread ThreadId m
tid String
label
      !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 a. m a -> m a
restore m a
action
      atomically $ writeTQueue completionQueue res

readSize :: MonadSTM m => JobPool group m a -> STM m Int
readSize :: forall (m :: * -> *) group a.
MonadSTM m =>
JobPool group m a -> STM m Int
readSize JobPool{TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall group (m :: * -> *) a.
JobPool 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 group m a -> group -> STM m Int
readGroupSize :: forall (m :: * -> *) group a.
(MonadSTM m, Eq group) =>
JobPool group m a -> group -> STM m Int
readGroupSize JobPool{TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall group (m :: * -> *) a.
JobPool 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 group m a -> STM m a
waitForJob :: forall (m :: * -> *) group a.
MonadSTM m =>
JobPool group m a -> STM m a
waitForJob JobPool{TQueue m a
completionQueue :: forall group (m :: * -> *) a. JobPool group m a -> TQueue m a
completionQueue :: TQueue m a
completionQueue} = 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
completionQueue

-- | Cancel all threads in a given group.  Blocks until all threads terminated.
--
cancelGroup :: ( MonadAsync m
               , Eq group
               )
            => JobPool group m a -> group -> m ()
cancelGroup :: forall (m :: * -> *) group a.
(MonadAsync m, Eq group) =>
JobPool group m a -> group -> m ()
cancelGroup JobPool { TVar m (Map (group, ThreadId m) (Async m ()))
jobsVar :: forall group (m :: * -> *) a.
JobPool 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