{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE RecordWildCards #-}

module Control.Concurrent.TokenLimiter.Concurrent
  ( -- * Create
    Count,
    TokenLimitConfig (..),
    MonotonicTime,
    TokenLimiter (..),
    makeTokenLimiter,

    -- * Use
    tryDebit,
    waitDebit,
    MonotonicDiffNanos (..),

    -- * Helper functions
    computeCurrentCount,
  )
where

import Control.Concurrent
import Control.Exception
import Control.Monad
import Data.Maybe (fromMaybe)
import Data.Word
import GHC.Clock
import GHC.Generics (Generic)
import Numeric.Natural

-- | An amount of tokens
type Count = Word64

-- | A configuration for 'TokenLimiter'
data TokenLimitConfig = TokenLimitConfig
  { -- | How many tokens should be in the bucket when it's created
    TokenLimitConfig -> Word64
tokenLimitConfigInitialTokens :: !Count,
    -- | Maximum number of tokens the bucket can hold at any one time
    TokenLimitConfig -> Word64
tokenLimitConfigMaxTokens :: !Count,
    -- | How many tokens are added to the bucket per second
    TokenLimitConfig -> Word64
tokenLimitConfigTokensPerSecond :: !Count
  }
  deriving (Int -> TokenLimitConfig -> ShowS
[TokenLimitConfig] -> ShowS
TokenLimitConfig -> String
(Int -> TokenLimitConfig -> ShowS)
-> (TokenLimitConfig -> String)
-> ([TokenLimitConfig] -> ShowS)
-> Show TokenLimitConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TokenLimitConfig -> ShowS
showsPrec :: Int -> TokenLimitConfig -> ShowS
$cshow :: TokenLimitConfig -> String
show :: TokenLimitConfig -> String
$cshowList :: [TokenLimitConfig] -> ShowS
showList :: [TokenLimitConfig] -> ShowS
Show, TokenLimitConfig -> TokenLimitConfig -> Bool
(TokenLimitConfig -> TokenLimitConfig -> Bool)
-> (TokenLimitConfig -> TokenLimitConfig -> Bool)
-> Eq TokenLimitConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TokenLimitConfig -> TokenLimitConfig -> Bool
== :: TokenLimitConfig -> TokenLimitConfig -> Bool
$c/= :: TokenLimitConfig -> TokenLimitConfig -> Bool
/= :: TokenLimitConfig -> TokenLimitConfig -> Bool
Eq, (forall x. TokenLimitConfig -> Rep TokenLimitConfig x)
-> (forall x. Rep TokenLimitConfig x -> TokenLimitConfig)
-> Generic TokenLimitConfig
forall x. Rep TokenLimitConfig x -> TokenLimitConfig
forall x. TokenLimitConfig -> Rep TokenLimitConfig x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. TokenLimitConfig -> Rep TokenLimitConfig x
from :: forall x. TokenLimitConfig -> Rep TokenLimitConfig x
$cto :: forall x. Rep TokenLimitConfig x -> TokenLimitConfig
to :: forall x. Rep TokenLimitConfig x -> TokenLimitConfig
Generic)

-- | A type synonym for a number of "monotonic time" nanoseconds.
--
-- This only exists because it is also a 'Word64' and would be too easy to confuse with a 'Count'.
type MonotonicTime = Word64

newtype MonotonicDiffNanos = MonotonicDiffNanos {MonotonicDiffNanos -> Word64
unMonotonicDiffNanos :: Word64}
  deriving (Int -> MonotonicDiffNanos -> ShowS
[MonotonicDiffNanos] -> ShowS
MonotonicDiffNanos -> String
(Int -> MonotonicDiffNanos -> ShowS)
-> (MonotonicDiffNanos -> String)
-> ([MonotonicDiffNanos] -> ShowS)
-> Show MonotonicDiffNanos
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MonotonicDiffNanos -> ShowS
showsPrec :: Int -> MonotonicDiffNanos -> ShowS
$cshow :: MonotonicDiffNanos -> String
show :: MonotonicDiffNanos -> String
$cshowList :: [MonotonicDiffNanos] -> ShowS
showList :: [MonotonicDiffNanos] -> ShowS
Show, MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
(MonotonicDiffNanos -> MonotonicDiffNanos -> Bool)
-> (MonotonicDiffNanos -> MonotonicDiffNanos -> Bool)
-> Eq MonotonicDiffNanos
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
== :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
$c/= :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
/= :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
Eq, Eq MonotonicDiffNanos
Eq MonotonicDiffNanos =>
(MonotonicDiffNanos -> MonotonicDiffNanos -> Ordering)
-> (MonotonicDiffNanos -> MonotonicDiffNanos -> Bool)
-> (MonotonicDiffNanos -> MonotonicDiffNanos -> Bool)
-> (MonotonicDiffNanos -> MonotonicDiffNanos -> Bool)
-> (MonotonicDiffNanos -> MonotonicDiffNanos -> Bool)
-> (MonotonicDiffNanos -> MonotonicDiffNanos -> MonotonicDiffNanos)
-> (MonotonicDiffNanos -> MonotonicDiffNanos -> MonotonicDiffNanos)
-> Ord MonotonicDiffNanos
MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
MonotonicDiffNanos -> MonotonicDiffNanos -> Ordering
MonotonicDiffNanos -> MonotonicDiffNanos -> MonotonicDiffNanos
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: MonotonicDiffNanos -> MonotonicDiffNanos -> Ordering
compare :: MonotonicDiffNanos -> MonotonicDiffNanos -> Ordering
$c< :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
< :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
$c<= :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
<= :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
$c> :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
> :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
$c>= :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
>= :: MonotonicDiffNanos -> MonotonicDiffNanos -> Bool
$cmax :: MonotonicDiffNanos -> MonotonicDiffNanos -> MonotonicDiffNanos
max :: MonotonicDiffNanos -> MonotonicDiffNanos -> MonotonicDiffNanos
$cmin :: MonotonicDiffNanos -> MonotonicDiffNanos -> MonotonicDiffNanos
min :: MonotonicDiffNanos -> MonotonicDiffNanos -> MonotonicDiffNanos
Ord, (forall x. MonotonicDiffNanos -> Rep MonotonicDiffNanos x)
-> (forall x. Rep MonotonicDiffNanos x -> MonotonicDiffNanos)
-> Generic MonotonicDiffNanos
forall x. Rep MonotonicDiffNanos x -> MonotonicDiffNanos
forall x. MonotonicDiffNanos -> Rep MonotonicDiffNanos x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. MonotonicDiffNanos -> Rep MonotonicDiffNanos x
from :: forall x. MonotonicDiffNanos -> Rep MonotonicDiffNanos x
$cto :: forall x. Rep MonotonicDiffNanos x -> MonotonicDiffNanos
to :: forall x. Rep MonotonicDiffNanos x -> MonotonicDiffNanos
Generic)

-- | A token bucket-based rate limiter
--
-- This token limiter is thread-safe and guarantees that:
--
-- * <https://en.wikipedia.org/wiki/Thundering_herd_problem There will be no thundering herd problem>
-- * <https://hackage.haskell.org/package/base-4.14.1.0/docs/Control-Concurrent-MVar.html#v:modifyMVar Fairness: Waiting processes will be serviced in a first-come first-service order.>
data TokenLimiter = TokenLimiter
  { TokenLimiter -> TokenLimitConfig
tokenLimiterConfig :: !TokenLimitConfig,
    -- | The last time the limiter was used, and what the token count was at that time.
    --
    -- Not that this library assumes that you never put anything into this mvar
    -- yourself and only use the functions in this library to interact with it.
    TokenLimiter -> MVar (Word64, Word64)
tokenLimiterLastServiced :: !(MVar (MonotonicTime, Count))
  }
  deriving (TokenLimiter -> TokenLimiter -> Bool
(TokenLimiter -> TokenLimiter -> Bool)
-> (TokenLimiter -> TokenLimiter -> Bool) -> Eq TokenLimiter
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TokenLimiter -> TokenLimiter -> Bool
== :: TokenLimiter -> TokenLimiter -> Bool
$c/= :: TokenLimiter -> TokenLimiter -> Bool
/= :: TokenLimiter -> TokenLimiter -> Bool
Eq, (forall x. TokenLimiter -> Rep TokenLimiter x)
-> (forall x. Rep TokenLimiter x -> TokenLimiter)
-> Generic TokenLimiter
forall x. Rep TokenLimiter x -> TokenLimiter
forall x. TokenLimiter -> Rep TokenLimiter x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. TokenLimiter -> Rep TokenLimiter x
from :: forall x. TokenLimiter -> Rep TokenLimiter x
$cto :: forall x. Rep TokenLimiter x -> TokenLimiter
to :: forall x. Rep TokenLimiter x -> TokenLimiter
Generic)

-- | Make a token limiter
--
-- The initial number of tokens will be the minimum of the 'tokenLimitConfigInitialTokens' and the 'tokenLimitConfigMaxTokens',
makeTokenLimiter :: TokenLimitConfig -> IO TokenLimiter
makeTokenLimiter :: TokenLimitConfig -> IO TokenLimiter
makeTokenLimiter TokenLimitConfig
tokenLimiterConfig = do
  Word64
now <- IO Word64
getMonotonicTimeNSec
  MVar (Word64, Word64)
tokenLimiterLastServiced <- (Word64, Word64) -> IO (MVar (Word64, Word64))
forall a. a -> IO (MVar a)
newMVar (Word64
now, Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
min (TokenLimitConfig -> Word64
tokenLimitConfigInitialTokens TokenLimitConfig
tokenLimiterConfig) (TokenLimitConfig -> Word64
tokenLimitConfigMaxTokens TokenLimitConfig
tokenLimiterConfig))
  TokenLimiter -> IO TokenLimiter
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TokenLimiter {MVar (Word64, Word64)
TokenLimitConfig
tokenLimiterConfig :: TokenLimitConfig
tokenLimiterLastServiced :: MVar (Word64, Word64)
tokenLimiterConfig :: TokenLimitConfig
tokenLimiterLastServiced :: MVar (Word64, Word64)
..}

-- | Check if we can debit a number of tokens, and do it if possible.
--
-- The returned boolean represents whether the tokens were debited.
--
-- Note that there is a small race-condition in which `tryDebit` sometimes
-- returns `False` eventhough it could (maybe) have debited because another
-- thread was currently `waitDebit`-ing without actually waiting (because it
-- didn't need to wait).
tryDebit :: TokenLimiter -> Word64 -> IO Bool
tryDebit :: TokenLimiter -> Word64 -> IO Bool
tryDebit TokenLimiter {MVar (Word64, Word64)
TokenLimitConfig
tokenLimiterConfig :: TokenLimiter -> TokenLimitConfig
tokenLimiterLastServiced :: TokenLimiter -> MVar (Word64, Word64)
tokenLimiterConfig :: TokenLimitConfig
tokenLimiterLastServiced :: MVar (Word64, Word64)
..} Word64
debit =
  (Maybe Bool -> Bool) -> IO (Maybe Bool) -> IO Bool
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
False) (IO (Maybe Bool) -> IO Bool) -> IO (Maybe Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$
    MVar (Word64, Word64)
-> ((Word64, Word64) -> IO ((Word64, Word64), Bool))
-> IO (Maybe Bool)
forall a b. MVar a -> (a -> IO (a, b)) -> IO (Maybe b)
tryModifyMVar MVar (Word64, Word64)
tokenLimiterLastServiced (((Word64, Word64) -> IO ((Word64, Word64), Bool))
 -> IO (Maybe Bool))
-> ((Word64, Word64) -> IO ((Word64, Word64), Bool))
-> IO (Maybe Bool)
forall a b. (a -> b) -> a -> b
$ \(Word64
lastServiced, Word64
countThen) -> do
      Word64
now <- IO Word64
getMonotonicTimeNSec
      let currentCount :: Word64
currentCount = TokenLimitConfig -> Word64 -> Word64 -> Word64 -> Word64
computeCurrentCount TokenLimitConfig
tokenLimiterConfig Word64
lastServiced Word64
countThen Word64
now
      let enoughAvailable :: Bool
enoughAvailable = Word64
currentCount Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word64
debit
      ((Word64, Word64), Bool) -> IO ((Word64, Word64), Bool)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (((Word64, Word64), Bool) -> IO ((Word64, Word64), Bool))
-> ((Word64, Word64), Bool) -> IO ((Word64, Word64), Bool)
forall a b. (a -> b) -> a -> b
$
        if Bool
enoughAvailable
          then
            let newCount :: Word64
newCount = Word64
currentCount Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
debit
             in ((Word64
now, Word64
newCount), Bool
True)
          else ((Word64
lastServiced, Word64
countThen), Bool
False)

-- | Wait until the given number of tokens can be debited.
--
-- Returns the time waited, for stats recording purposes.
--
-- Note: only reports the time waited due to rate-limiting this specific action,
-- not the wall-clock time waited (that might include waiting for previous
-- actions limited by this rate limiter to finish).
--
-- Note: debitor threads are serviced in FIFO order, so a request for a small
-- (and currently satisfiable) number of tokens can still be delayed by a debit
-- request for a larger amount of tokens.
--
-- Note: the wait time reported can be inflated due to scheduling inaccuracy.
-- See https://gitlab.haskell.org/ghc/ghc/-/issues/16601.
waitDebit :: TokenLimiter -> Word64 -> IO (Maybe MonotonicDiffNanos)
waitDebit :: TokenLimiter -> Word64 -> IO (Maybe MonotonicDiffNanos)
waitDebit TokenLimiter {MVar (Word64, Word64)
TokenLimitConfig
tokenLimiterConfig :: TokenLimiter -> TokenLimitConfig
tokenLimiterLastServiced :: TokenLimiter -> MVar (Word64, Word64)
tokenLimiterConfig :: TokenLimitConfig
tokenLimiterLastServiced :: MVar (Word64, Word64)
..} Word64
debit = MVar (Word64, Word64)
-> ((Word64, Word64)
    -> IO ((Word64, Word64), Maybe MonotonicDiffNanos))
-> IO (Maybe MonotonicDiffNanos)
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar (Word64, Word64)
tokenLimiterLastServiced (((Word64, Word64)
  -> IO ((Word64, Word64), Maybe MonotonicDiffNanos))
 -> IO (Maybe MonotonicDiffNanos))
-> ((Word64, Word64)
    -> IO ((Word64, Word64), Maybe MonotonicDiffNanos))
-> IO (Maybe MonotonicDiffNanos)
forall a b. (a -> b) -> a -> b
$ \(Word64
lastServiced, Word64
countThen) -> do
  Word64
now <- IO Word64
getMonotonicTimeNSec
  let currentCount :: Word64
currentCount = TokenLimitConfig -> Word64 -> Word64 -> Word64 -> Word64
computeCurrentCount TokenLimitConfig
tokenLimiterConfig Word64
lastServiced Word64
countThen Word64
now
  let enoughAvailable :: Bool
enoughAvailable = Word64
currentCount Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word64
debit
  if Bool
enoughAvailable
    then do
      let newCount :: Word64
newCount = Word64
currentCount Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
debit
      ((Word64, Word64), Maybe MonotonicDiffNanos)
-> IO ((Word64, Word64), Maybe MonotonicDiffNanos)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Word64
now, Word64
newCount), Maybe MonotonicDiffNanos
forall a. Maybe a
Nothing)
    else do
      let extraTokensNeeded :: Word64
extraTokensNeeded = Word64
debit Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
currentCount
      let microsecondsToWaitDouble :: Double
          microsecondsToWaitDouble :: Double
microsecondsToWaitDouble =
            Double
1_000_000
              Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Word64 -> Double) Word64
extraTokensNeeded
              Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Word64 -> Double) (TokenLimitConfig -> Word64
tokenLimitConfigTokensPerSecond TokenLimitConfig
tokenLimiterConfig)

      let microsecondsToWait :: Int
microsecondsToWait = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling Double
microsecondsToWaitDouble
      -- threadDelay guarantees that _at least_ the given number of microseconds will have passed.
      Int -> IO ()
threadDelay Int
microsecondsToWait
      -- However, it could be MUCH longer than that, so we will recalculate the time instead of
      -- adding that number of microseconds to the old time.
      Word64
nowAfterWaiting <- IO Word64
getMonotonicTimeNSec
      let delta :: MonotonicDiffNanos
delta = Word64 -> MonotonicDiffNanos
MonotonicDiffNanos (Word64
nowAfterWaiting Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
now)
      -- We do assume here that we will now have enough tokens and do not need to recalculate whether there will be enough.
      -- (We would not know what to do if there weren't, anyway.)
      -- BUT this assumption _should_ hold because _modifyMVar_ guarantees
      -- atomicity if there are no other producers for this MVar, which there
      -- aren't.
      let currentCountAfterWaiting :: Word64
currentCountAfterWaiting = TokenLimitConfig -> Word64 -> Word64 -> Word64 -> Word64
computeCurrentCount TokenLimitConfig
tokenLimiterConfig Word64
lastServiced Word64
countThen Word64
nowAfterWaiting
      let newCount :: Word64
newCount = Word64
currentCountAfterWaiting Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
debit
      ((Word64, Word64), Maybe MonotonicDiffNanos)
-> IO ((Word64, Word64), Maybe MonotonicDiffNanos)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Word64
nowAfterWaiting, Word64
newCount), MonotonicDiffNanos -> Maybe MonotonicDiffNanos
forall a. a -> Maybe a
Just MonotonicDiffNanos
delta)

-- | Compute the current number of tokens in a bucket purely.
--
-- You should not need this function.
computeCurrentCount :: TokenLimitConfig -> MonotonicTime -> Count -> MonotonicTime -> Count
computeCurrentCount :: TokenLimitConfig -> Word64 -> Word64 -> Word64 -> Word64
computeCurrentCount TokenLimitConfig {Word64
tokenLimitConfigInitialTokens :: TokenLimitConfig -> Word64
tokenLimitConfigMaxTokens :: TokenLimitConfig -> Word64
tokenLimitConfigTokensPerSecond :: TokenLimitConfig -> Word64
tokenLimitConfigInitialTokens :: Word64
tokenLimitConfigMaxTokens :: Word64
tokenLimitConfigTokensPerSecond :: Word64
..} Word64
lastServiced Word64
countThen Word64
now =
  let nanoDiff :: Word64
      nanoDiff :: Word64
nanoDiff = Word64
now Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
lastServiced
      countToAddDouble :: Double
      countToAddDouble :: Double
countToAddDouble =
        (Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Word64 -> Double) Word64
nanoDiff
          Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Word64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Word64 -> Double) Word64
tokenLimitConfigTokensPerSecond
          Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
1_000_000_000
      countToAdd :: Word64
      countToAdd :: Word64
countToAdd = Double -> Word64
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
countToAddDouble
      totalPrecise :: Natural
      totalPrecise :: Natural
totalPrecise = Word64 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
countThen Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Word64 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
countToAdd
      willOverflow :: Bool
willOverflow = Natural
totalPrecise Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Word64 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
forall a. Bounded a => a
maxBound :: Word64)
      totalCount :: Word64
totalCount = Word64
countThen Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
countToAdd
   in if Bool
willOverflow
        then Word64
tokenLimitConfigMaxTokens
        else Word64 -> Word64 -> Word64
forall a. Ord a => a -> a -> a
min Word64
tokenLimitConfigMaxTokens Word64
totalCount

tryModifyMVar :: MVar a -> (a -> IO (a, b)) -> IO (Maybe b)
tryModifyMVar :: forall a b. MVar a -> (a -> IO (a, b)) -> IO (Maybe b)
tryModifyMVar MVar a
m a -> IO (a, b)
io =
  ((forall a. IO a -> IO a) -> IO (Maybe b)) -> IO (Maybe b)
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO (Maybe b)) -> IO (Maybe b))
-> ((forall a. IO a -> IO a) -> IO (Maybe b)) -> IO (Maybe b)
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
    Maybe a
mA <- MVar a -> IO (Maybe a)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar a
m
    Maybe a -> (a -> IO b) -> IO (Maybe b)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Maybe a
mA ((a -> IO b) -> IO (Maybe b)) -> (a -> IO b) -> IO (Maybe b)
forall a b. (a -> b) -> a -> b
$ \a
a -> do
      (a
a', b
b) <-
        IO (a, b) -> IO (a, b)
forall a. IO a -> IO a
restore (a -> IO (a, b)
io a
a)
          IO (a, b) -> IO () -> IO (a, b)
forall a b. IO a -> IO b -> IO a
`onException` MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
m a
a
      MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
m a
a'
      b -> IO b
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
b