Control.Concurrent.Async.Extra

Extra.hs

{-# LANGUAGE ImplicitPrelude #-}

module Control.Concurrent.Async.Extra where

import           Control.Concurrent.Async
import           Control.Monad.Trans.Except
import           Data.Foldable

concurrentlyExcept :: IO (Except e a) -> IO (Except e b) -> IO (Except e (a, b))
concurrentlyExcept actA actB = do
    asyncA <- async actA
    asyncB <- async actB

    res <- waitEither asyncA asyncB
    case res of
      Left  ma -> waitPiece (\a b -> (a, b)) ma asyncB
      Right mb -> waitPiece (\b a -> (a, b)) mb asyncA
  where
    waitPiece f m = waitPiece' f $ runExcept m

    waitPiece' _ (Left x) ay = do
      cancel ay
      return $ throwE x
    waitPiece' f (Right x) ay = do
      y <- wait ay
      return $ f x <$> y

concatConcurrentExcepts :: (Foldable t, Monoid a) => t (IO (Except e a)) -> IO (Except e a)
concatConcurrentExcepts = foldl' go $ pure (pure mempty)
  where
    go r act = do
      acts <- concurrentlyExcept act r
      return $ uncurry mappend <$> acts