76 lines
3.1 KiB
Haskell
76 lines
3.1 KiB
Haskell
-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>
|
|
--
|
|
-- SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
module UnliftIO.Async.Utils
|
|
( allocateAsync, allocateLinkedAsync
|
|
, allocateAsyncWithUnmask, allocateLinkedAsyncWithUnmask
|
|
, allocateAsyncMasked, allocateLinkedAsyncMasked
|
|
) where
|
|
|
|
import ClassyPrelude hiding (cancel, async, link, finally, mask)
|
|
import Control.Lens
|
|
import Control.Lens.Extras (is)
|
|
|
|
import qualified UnliftIO.Async as UnliftIO
|
|
import qualified Control.Concurrent.Async as A
|
|
|
|
import Control.Monad.Trans.Resource
|
|
import qualified Control.Monad.Trans.Resource.Internal as ResourceT.Internal
|
|
import Data.Acquire
|
|
|
|
import Control.Monad.Catch
|
|
|
|
|
|
withReference :: forall m a. (MonadUnliftIO m, MonadResource m) => ((IO (), IO ()) -> m a) -> m a
|
|
withReference act = do
|
|
releaseAct <- newEmptyTMVarIO
|
|
|
|
let doAlloc = do
|
|
iSt <- liftResourceT getInternalState
|
|
liftIO $ mask $ \_ -> do
|
|
ResourceT.Internal.stateAlloc iSt
|
|
atomically $ putTMVar releaseAct ()
|
|
return iSt
|
|
doRelease iSt eCase = liftIO . whenM (atomically $ is _Just <$> tryTakeTMVar releaseAct) $ do
|
|
flip ResourceT.Internal.stateCleanup iSt $ case eCase of
|
|
ExitCaseSuccess _ -> ReleaseNormal
|
|
ExitCaseException _ -> ReleaseException
|
|
ExitCaseAbort -> ReleaseEarly
|
|
|
|
withRunInIO $ \run ->
|
|
fmap fst . generalBracket (run doAlloc) doRelease $ \iSt -> do
|
|
res <- run $ act
|
|
( atomically $ takeTMVar releaseAct
|
|
, ResourceT.Internal.stateCleanup ReleaseNormal iSt
|
|
)
|
|
atomically $ guard =<< isEmptyTMVar releaseAct
|
|
return res
|
|
|
|
|
|
allocateAsync :: forall m a.
|
|
( MonadUnliftIO m, MonadResource m )
|
|
=> m a -> m (Async a)
|
|
allocateAsync act = withReference $ \(signalReady, releaseRef) -> withRunInIO $ \run -> run . fmap (view _2) . flip allocate A.cancel . A.async . flip finally releaseRef $ signalReady >> run act
|
|
|
|
allocateLinkedAsync :: forall m a. (MonadUnliftIO m, MonadResource m) => m a -> m (Async a)
|
|
allocateLinkedAsync = uncurry (<$) . (id &&& UnliftIO.link) <=< allocateAsync
|
|
|
|
|
|
allocateAsyncWithUnmask :: forall m a.
|
|
( MonadUnliftIO m, MonadResource m)
|
|
=> ((forall b. m b -> m b) -> m a) -> m (Async a)
|
|
allocateAsyncWithUnmask act = withReference $ \(signalReady, releaseRef) -> withRunInIO $ \run -> run . fmap (view _2) . flip allocate A.cancel $ A.asyncWithUnmask $ \unmask -> flip finally releaseRef $ signalReady >> run (act $ liftIO . unmask . run)
|
|
|
|
allocateLinkedAsyncWithUnmask :: forall m a. (MonadUnliftIO m, MonadResource m) => ((forall b. m b -> m b) -> m a) -> m (Async a)
|
|
allocateLinkedAsyncWithUnmask act = uncurry (<$) . (id &&& UnliftIO.link) =<< allocateAsyncWithUnmask act
|
|
|
|
|
|
allocateAsyncMasked :: forall m a.
|
|
( MonadUnliftIO m, MonadResource m)
|
|
=> m a -> m (Async a)
|
|
allocateAsyncMasked act = allocateAsyncWithUnmask (const act)
|
|
|
|
allocateLinkedAsyncMasked :: forall m a. (MonadUnliftIO m, MonadResource m) => m a -> m (Async a)
|
|
allocateLinkedAsyncMasked act = allocateLinkedAsyncWithUnmask (const act)
|