@@ -281,7 +281,7 @@ module Streamly.Internal.Data.Stream.StreamD
281281 -- * Exceptions
282282 , newFinalizedIORef
283283 , runIORefFinalizer
284− , clearIORefFinalizer
284+ , withIORefFinalizer
285285 , gbracket_
286286 , gbracket
287287 , before
@@ -313,7 +313,7 @@ import Control.Monad.IO.Class (MonadIO(..))
313313import Control.Monad.Trans.Reader (ReaderT)
314314import Control.Monad.Trans.State.Strict (StateT)
315315import Control.Monad.Trans.Class (MonadTrans(lift))
316−import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
316+import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)
317317import Data.Bits (shiftR, shiftL, (.|.), (.&.))
318318import Data.Functor.Identity (Identity(..))
319319import Data.Int (Int64)
@@ -3224,10 +3224,15 @@ runIORefFinalizer ref = liftIO $ do
32243224 writeIORef ref Nothing
32253225 action
32263226
3227−-- | Deactivate the finalizer stored in an IORef without running it.
3228−--
3229−clearIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3230−clearIORefFinalizer ref = liftIO $ writeIORef ref Nothing
3227+-- | Run an action clearing the finalizer IORef atomically wrt async
3228+-- exceptions. The action is run with async exceptions masked.
3229+withIORefFinalizer :: MonadBaseControl IO m
3230+ => IORef (Maybe (IO ())) -> m a -> m a
3231+withIORefFinalizer ref action = do
3232+ control $ \runinio ->
3233+ mask_ $ do
3234+ writeIORef ref Nothing
3235+ runinio action
32313236
32323237------------------------------------------------------------------------------
32333238
@@ -3251,7 +3256,7 @@ gbracket
32513256 => m c -- ^ before
32523257 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
32533258 -> (c -> m d) -- ^ after, on normal stop or GC
3254− -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
3259+ -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
32553260 -> (c -> Stream m b) -- ^ stream generator
32563261 -> Stream m b
32573262gbracket bef exc aft fexc fnormal =
@@ -3286,8 +3291,12 @@ gbracket bef exc aft fexc fnormal =
32863291 runIORefFinalizer ref
32873292 return Stop
32883293 Left e -> do
3289− clearIORefFinalizer ref
3290− return $ Skip (GBracketIOException (fexc v e (UnStream step1 st)))
3294+ -- Clearing of finalizer and running of exception handler must
3295+ -- be atomic wrt async exceptions. Otherwise if we have cleared
3296+ -- the finalizer and have not run the exception handler then we
3297+ -- may leak the resource.
3298+ stream <- withIORefFinalizer ref (fexc v e (UnStream step1 st))
3299+ return $ Skip (GBracketIOException stream)
32913300 step gst (GBracketIOException (UnStream step1 st)) = do
32923301 res <- step1 gst st
32933302 case res of
@@ -3404,7 +3413,7 @@ bracket :: (MonadAsync m, MonadCatch m)
34043413 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
34053414bracket bef aft bet =
34063415 gbracket bef MC.try aft
3407− (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e)) bet
3416+ (\a (e :: SomeException) _ -> aft a >> return (nilM (MC.throwM e))) bet
34083417
34093418data BracketState s v = BracketInit | BracketRun s v
34103419