60cee48 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add exception handling combinators with guaranteed cleanup

60cee48Harendra Kumar7 years ago

benchmark/FileIO.hs+18

compare with master

unifiedsplit

@@ -91,9 +91,15 @@ main = do

91 , mkBench "catBracket" href $ do

92 Handles inh _ <- readIORef href

93 BFA.catBracket devNull inh

94+ , mkBench "catBracketIO" href $ do

95+ Handles inh _ <- readIORef href

96+ BFA.catBracketIO devNull inh

94 , mkBench "catBracketStream" href $ do

97 , mkBench "catBracketStream" href $ do

95 Handles inh _ <- readIORef href

98 Handles inh _ <- readIORef href

96 BFA.catBracketStream devNull inh

99 BFA.catBracketStream devNull inh

100+ , mkBench "catBracketStreamIO" href $ do

101+ Handles inh _ <- readIORef href

102+ BFA.catBracketStreamIO devNull inh

97 , mkBench "catOnException" href $ do

103 , mkBench "catOnException" href $ do

98 Handles inh _ <- readIORef href

104 Handles inh _ <- readIORef href

99 BFA.catOnException devNull inh

105 BFA.catOnException devNull inh

@@ -142,15 +148,27 @@ main = do

142 , mkBench "catFinally" href $ do

148 , mkBench "catFinally" href $ do

143 Handles inh _ <- readIORef href

149 Handles inh _ <- readIORef href

144 BFS.catFinally devNull inh

150 BFS.catFinally devNull inh

151+ , mkBench "catFinallyIO" href $ do

152+ Handles inh _ <- readIORef href

153+ BFS.catFinallyIO devNull inh

145 , mkBench "catFinallyStream" href $ do

154 , mkBench "catFinallyStream" href $ do

146 Handles inh _ <- readIORef href

155 Handles inh _ <- readIORef href

147 BFS.catFinallyStream devNull inh

156 BFS.catFinallyStream devNull inh

157+ , mkBench "catFinallyStreamIO" href $ do

158+ Handles inh _ <- readIORef href

159+ BFS.catFinallyStreamIO devNull inh

148 , mkBench "catBracketStream" href $ do

160 , mkBench "catBracketStream" href $ do

149 Handles inh _ <- readIORef href

161 Handles inh _ <- readIORef href

150 BFS.catBracketStream devNull inh

162 BFS.catBracketStream devNull inh

163+ , mkBench "catBracketStreamIO" href $ do

164+ Handles inh _ <- readIORef href

165+ BFS.catBracketStreamIO devNull inh

151 , mkBench "catBracket" href $ do

166 , mkBench "catBracket" href $ do

152 Handles inh _ <- readIORef href

167 Handles inh _ <- readIORef href

153 BFS.catBracket devNull inh

168 BFS.catBracket devNull inh

169+ , mkBench "catBracketIO" href $ do

170+ Handles inh _ <- readIORef href

171+ BFS.catBracketIO devNull inh

154#endif

172#endif

155 , mkBench "read-word8" href $ do

173 , mkBench "read-word8" href $ do

156 Handles inh _ <- readIORef href

174 Handles inh _ <- readIORef href

benchmark/Streamly/Benchmark/FileIO/Array.hs+17

@@ -28,7 +28,9 @@ module Streamly.Benchmark.FileIO.Array

28 , cat

29 , catOnException

30 , catBracket

31+ , catBracketIO

31 , catBracketStream

32 , catBracketStream

33+ , catBracketStreamIO

32 , copy

34 , copy

33 , linesUnlinesCopy

35 , linesUnlinesCopy

34 , wordsUnwordsCopy

36 , wordsUnwordsCopy

@@ -52,6 +54,7 @@ import qualified Streamly.Internal.FileSystem.Handle as IFH

52import qualified Streamly.Internal.Memory.Array as IA

54import qualified Streamly.Internal.Memory.Array as IA

53import qualified Streamly.Internal.Memory.ArrayStream as AS

55import qualified Streamly.Internal.Memory.ArrayStream as AS

54import qualified Streamly.Internal.Data.Unfold as IUF

56import qualified Streamly.Internal.Data.Unfold as IUF

57+import qualified Streamly.Internal.Prelude as IP

56#ifdef INSPECTION

59#ifdef INSPECTION

57import Foreign.Storable (Storable)

60import Foreign.Storable (Storable)

@@ -144,6 +147,13 @@ inspect $ hasNoTypeClasses 'catBracket

144-- inspect $ 'catBracket `hasNoType` ''Step

147-- inspect $ 'catBracket `hasNoType` ''Step

145#endif

148#endif

146

149

150+{-# INLINE catBracketIO #-}

151+catBracketIO :: Handle -> Handle -> IO ()

152+catBracketIO devNull inh =

153+ let readEx = IUF.bracketIO return (\_ -> hClose inh)

154+ (IUF.supplyFirst FH.readChunksWithBufferOf (256*1024))

155+ in IUF.fold readEx (IFH.writeChunks devNull) inh

156+

147-- | Send the file contents to /dev/null with exception handling

157-- | Send the file contents to /dev/null with exception handling

148{-# INLINE catBracketStream #-}

158{-# INLINE catBracketStream #-}

149catBracketStream :: Handle -> Handle -> IO ()

159catBracketStream :: Handle -> Handle -> IO ()

@@ -157,6 +167,13 @@ inspect $ hasNoTypeClasses 'catBracketStream

157-- inspect $ 'catBracketStream `hasNoType` ''Step

167-- inspect $ 'catBracketStream `hasNoType` ''Step

158#endif

168#endif

159

169

170+{-# INLINE catBracketStreamIO #-}

171+catBracketStreamIO :: Handle -> Handle -> IO ()

172+catBracketStreamIO devNull inh =

173+ let readEx = IP.bracketIO (return ()) (\_ -> hClose inh)

174+ (\_ -> IFH.toChunksWithBufferOf (256*1024) inh)

175+ in S.fold (IFH.writeChunks devNull) $ readEx

176+

160-- | Send the file contents to /dev/null with exception handling

177-- | Send the file contents to /dev/null with exception handling

161{-# INLINE catOnException #-}

178{-# INLINE catOnException #-}

162catOnException :: Handle -> Handle -> IO ()

179catOnException :: Handle -> Handle -> IO ()

benchmark/Streamly/Benchmark/FileIO/Stream.hs+29

@@ -31,13 +31,17 @@ module Streamly.Benchmark.FileIO.Stream

31 , cat

32 , catStreamWrite

33 , catBracket

34+ , catBracketIO

34 , catBracketStream

35 , catBracketStream

36+ , catBracketStreamIO

35 , catOnException

37 , catOnException

36 , catOnExceptionStream

38 , catOnExceptionStream

37 , catHandle

39 , catHandle

38 , catHandleStream

40 , catHandleStream

39 , catFinally

41 , catFinally

42+ , catFinallyIO

40 , catFinallyStream

43 , catFinallyStream

44+ , catFinallyStreamIO

41 , copy

45 , copy

42 , linesUnlinesCopy

46 , linesUnlinesCopy

43 , linesUnlinesArrayWord8Copy

47 , linesUnlinesArrayWord8Copy

@@ -211,6 +215,12 @@ inspect $ hasNoTypeClasses 'catBracket

211-- inspect $ 'catBracket `hasNoType` ''D.ConcatMapUState

215-- inspect $ 'catBracket `hasNoType` ''D.ConcatMapUState

212#endif

216#endif

213

217

218+{-# INLINE catBracketIO #-}

219+catBracketIO :: Handle -> Handle -> IO ()

220+catBracketIO devNull inh =

221+ let readEx = IUF.bracketIO return (\_ -> hClose inh) FH.read

222+ in S.fold (FH.write devNull) $ S.unfold readEx inh

223+

214-- | Send the file contents to /dev/null with exception handling

224-- | Send the file contents to /dev/null with exception handling

215{-# INLINE catBracketStream #-}

225{-# INLINE catBracketStream #-}

216catBracketStream :: Handle -> Handle -> IO ()

226catBracketStream :: Handle -> Handle -> IO ()

@@ -224,6 +234,13 @@ inspect $ hasNoTypeClasses 'catBracketStream

224-- inspect $ 'catBracketStream `hasNoType` ''Step

234-- inspect $ 'catBracketStream `hasNoType` ''Step

225#endif

235#endif

226

236

237+{-# INLINE catBracketStreamIO #-}

238+catBracketStreamIO :: Handle -> Handle -> IO ()

239+catBracketStreamIO devNull inh =

240+ let readEx = IP.bracketIO (return ()) (\_ -> hClose inh)

241+ (\_ -> IFH.toBytes inh)

242+ in IFH.fromBytes devNull $ readEx

243+

227-- | Send the file contents to /dev/null with exception handling

244-- | Send the file contents to /dev/null with exception handling

228{-# INLINE catOnException #-}

245{-# INLINE catOnException #-}

229catOnException :: Handle -> Handle -> IO ()

246catOnException :: Handle -> Handle -> IO ()

@@ -259,6 +276,12 @@ inspect $ hasNoTypeClasses 'catFinally

259-- inspect $ 'catFinally `hasNoType` ''D.ConcatMapUState

276-- inspect $ 'catFinally `hasNoType` ''D.ConcatMapUState

260#endif

277#endif

261

278

279+{-# INLINE catFinallyIO #-}

280+catFinallyIO :: Handle -> Handle -> IO ()

281+catFinallyIO devNull inh =

282+ let readEx = IUF.finallyIO (\_ -> hClose inh) FH.read

283+ in S.fold (FH.write devNull) $ S.unfold readEx inh

284+

262-- | Send the file contents to /dev/null with exception handling

285-- | Send the file contents to /dev/null with exception handling

263{-# INLINE catFinallyStream #-}

286{-# INLINE catFinallyStream #-}

264catFinallyStream :: Handle -> Handle -> IO ()

287catFinallyStream :: Handle -> Handle -> IO ()

@@ -266,6 +289,12 @@ catFinallyStream devNull inh =

266 let readEx = S.finally (hClose inh) (S.unfold FH.read inh)

289 let readEx = S.finally (hClose inh) (S.unfold FH.read inh)

267 in S.fold (FH.write devNull) readEx

290 in S.fold (FH.write devNull) readEx

268

291

292+{-# INLINE catFinallyStreamIO #-}

293+catFinallyStreamIO :: Handle -> Handle -> IO ()

294+catFinallyStreamIO devNull inh =

295+ let readEx = IP.finallyIO (hClose inh) (S.unfold FH.read inh)

296+ in S.fold (FH.write devNull) readEx

297+

269-- | Send the file contents to /dev/null with exception handling

298-- | Send the file contents to /dev/null with exception handling

270{-# INLINE catHandle #-}

299{-# INLINE catHandle #-}

271catHandle :: Handle -> Handle -> IO ()

300catHandle :: Handle -> Handle -> IO ()

src/Streamly/Internal/Data/Stream/StreamD.hs+116−1

@@ -293,12 +293,16 @@ module Streamly.Internal.Data.Stream.StreamD

293 , the

294

295 -- * Exceptions

296+ , newFinalizedIORef

296 , gbracket

297 , gbracket

297 , before

298 , before

298 , after

299 , after

300+ , afterIO

299 , bracket

301 , bracket

302+ , bracketIO

300 , onException

303 , onException

301 , finally

304 , finally

305+ , finallyIO

302 , handle

306 , handle

303

307

304 -- * Concurrent Application

308 -- * Concurrent Application

@@ -323,9 +327,10 @@ import Control.Monad.IO.Class (MonadIO(..))

323import Control.Monad.Reader (ReaderT)

327import Control.Monad.Reader (ReaderT)

324import Control.Monad.State.Strict (StateT)

328import Control.Monad.State.Strict (StateT)

325import Control.Monad.Trans (MonadTrans(lift))

329import Control.Monad.Trans (MonadTrans(lift))

330+import Control.Monad.Trans.Control (MonadBaseControl)

326import Data.Bits (shiftR, shiftL, (.|.), (.&.))

331import Data.Bits (shiftR, shiftL, (.|.), (.&.))

327import Data.Functor.Identity (Identity(..))

332import Data.Functor.Identity (Identity(..))

328−import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

333+import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)

329import Data.Maybe (fromJust, isJust, isNothing)

334import Data.Maybe (fromJust, isJust, isNothing)

330import Data.Word (Word32)

335import Data.Word (Word32)

331import Foreign.Ptr (Ptr)

336import Foreign.Ptr (Ptr)

@@ -2876,6 +2881,84 @@ gbracket bef exc aft fexc fnormal =

2876 Skip s -> return $ Skip (GBracketException (Stream step1 s))

2881 Skip s -> return $ Skip (GBracketException (Stream step1 s))

2877 Stop -> return Stop

2882 Stop -> return Stop

2878

2883

2884+-- | Create an IORef holding a finalizer that is called automatically when the

2885+-- IORef is garbage collected. The IORef can be written to with a 'Nothing'

2886+-- value to deactivate the finalizer.

2887+newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)

2888+ => m a -> m (IORef (Maybe (IO ())))

2889+newFinalizedIORef finalizer = do

2890+ mrun <- captureMonadState

2891+ ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do

2892+ _ <- runInIO mrun finalizer

2893+ return ()

2894+ let finalizer1 = do

2895+ res <- readIORef ref

2896+ case res of

2897+ Nothing -> return ()

2898+ Just f -> f

2899+ _ <- liftIO $ mkWeakIORef ref finalizer1

2900+ return ref

2901+

2902+data GbracketIOState s1 s2 v wref

2903+ = GBracketIOInit

2904+ | GBracketIONormal s1 v wref

2905+ | GBracketIOException s2

2906+

2907+-- | Like gbracket but also uses a finalizer to make sure when the stream is

2908+-- garbage collected we run the finalizing action. This requires a MonadIO and

2909+-- MonadBaseControl IO constraint.

2910+--

2911+-- | The most general bracketing and exception combinator. All other

2912+-- combinators can be expressed in terms of this combinator. This can also be

2913+-- used for cases which are not covered by the standard combinators.

2914+--

2915+-- /Internal/

2916+--

2917+{-# INLINE_NORMAL gbracketIO #-}

2918+gbracketIO

2919+ :: (MonadIO m, MonadBaseControl IO m)

2920+ => m c -- ^ before

2921+ -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)

2922+ -> (c -> m d) -- ^ after, on normal stop or GC

2923+ -> (c -> e -> Stream m b) -- ^ on exception

2924+ -> (c -> Stream m b) -- ^ stream generator

2925+ -> Stream m b

2926+gbracketIO bef exc aft fexc fnormal =

2927+ Stream step GBracketIOInit

2928+

2929+ where

2930+

2931+ -- If the stream is never evaluated the "aft" action will never be

2932+ -- called. For that to occur we will need the user of this API to pass a

2933+ -- weak pointer to us.

2934+ {-# INLINE_LATE step #-}

2935+ step _ GBracketIOInit = do

2936+ r <- bef

2937+ ref <- newFinalizedIORef (aft r)

2938+ return $ Skip $ GBracketIONormal (fnormal r) r ref

2939+

2940+ step gst (GBracketIONormal (UnStream step1 st) v ref) = do

2941+ res <- exc $ step1 gst st

2942+ case res of

2943+ Right r -> case r of

2944+ Yield x s ->

2945+ return $ Yield x (GBracketIONormal (Stream step1 s) v ref)

2946+ Skip s ->

2947+ return $ Skip (GBracketIONormal (Stream step1 s) v ref)

2948+ Stop -> do

2949+ liftIO (writeIORef ref Nothing)

2950+ aft v >> return Stop

2951+ Left e -> do

2952+ liftIO (writeIORef ref Nothing)

2953+ return $ Skip (GBracketIOException (fexc v e))

2954+ step gst (GBracketIOException (UnStream step1 st)) = do

2955+ res <- step1 gst st

2956+ case res of

2957+ Yield x s ->

2958+ return $ Yield x (GBracketIOException (Stream step1 s))

2959+ Skip s -> return $ Skip (GBracketIOException (Stream step1 s))

2960+ Stop -> return Stop

2961+

2879-- | Run a side effect before the stream yields its first element.

2962-- | Run a side effect before the stream yields its first element.

2880{-# INLINE_NORMAL before #-}

2963{-# INLINE_NORMAL before #-}

2881before :: Monad m => m b -> Stream m a -> Stream m a

2964before :: Monad m => m b -> Stream m a -> Stream m a

@@ -2908,6 +2991,27 @@ after action (Stream step state) = Stream step' state

2908 Skip s -> return $ Skip s

2991 Skip s -> return $ Skip s

2909 Stop -> action >> return Stop

2992 Stop -> action >> return Stop

2910

2993

2994+{-# INLINE_NORMAL afterIO #-}

2995+afterIO :: (MonadIO m, MonadBaseControl IO m)

2996+ => m b -> Stream m a -> Stream m a

2997+afterIO action (Stream step state) = Stream step' Nothing

2998+

2999+ where

3000+

3001+ {-# INLINE_LATE step' #-}

3002+ step' _ Nothing = do

3003+ ref <- newFinalizedIORef action

3004+ return $ Skip $ Just (state, ref)

3005+ step' gst (Just (st, ref)) = do

3006+ res <- step gst st

3007+ case res of

3008+ Yield x s -> return $ Yield x (Just (s, ref))

3009+ Skip s -> return $ Skip (Just (s, ref))

3010+ Stop -> liftIO $ do

3011+ Just f <- readIORef ref

3012+ writeIORef ref Nothing

3013+ f >> return Stop

3014+

2911-- XXX These combinators are expensive due to the call to

3015-- XXX These combinators are expensive due to the call to

2912-- onException/handle/try on each step. Therefore, when possible, they should

3016-- onException/handle/try on each step. Therefore, when possible, they should

2913-- be called in an outer loop where we perform less iterations. For example, we

3017-- be called in an outer loop where we perform less iterations. For example, we

@@ -2954,6 +3058,13 @@ bracket bef aft bet =

2954 gbracket bef MC.try aft

3058 gbracket bef MC.try aft

2955 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet

3059 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet

2956

3060

3061+{-# INLINE_NORMAL bracketIO #-}

3062+bracketIO :: (MonadAsync m, MonadCatch m)

3063+ => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a

3064+bracketIO bef aft bet =

3065+ gbracketIO bef MC.try aft

3066+ (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet

3067+

2957data BracketState s v = BracketInit | BracketRun s v

3068data BracketState s v = BracketInit | BracketRun s v

2958

3069

2959{-# INLINE_NORMAL _bracket #-}

3070{-# INLINE_NORMAL _bracket #-}

@@ -2984,6 +3095,10 @@ finally :: MonadCatch m => m b -> Stream m a -> Stream m a

2984-- finally action xs = after action $ onException action xs

3095-- finally action xs = after action $ onException action xs

2985finally action xs = bracket (return ()) (\_ -> action) (const xs)

3096finally action xs = bracket (return ()) (\_ -> action) (const xs)

2986

3097

3098+{-# INLINE finallyIO #-}

3099+finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a

3100+finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)

3101+

2987-- | When evaluating a stream if an exception occurs, stream evaluation aborts

3102-- | When evaluating a stream if an exception occurs, stream evaluation aborts

2988-- and the specified exception handler is run with the exception as argument.

3103-- and the specified exception handler is run with the exception as argument.

2989{-# INLINE_NORMAL handle #-}

3104{-# INLINE_NORMAL handle #-}

src/Streamly/Internal/Data/Unfold.hs+150−1

@@ -112,16 +112,23 @@ module Streamly.Internal.Data.Unfold

112

113 -- * Exceptions

114 , gbracket

115+ , gbracketIO

115 , before

116 , before

116 , after

117 , after

118+ , afterIO

117 , onException

119 , onException

118 , finally

120 , finally

121+ , finallyIO

119 , bracket

122 , bracket

123+ , bracketIO

120 , handle

124 , handle

121 )

125 )

122where

126where

123

127

124import Control.Exception (Exception)

128import Control.Exception (Exception)

129+import Control.Monad.IO.Class (MonadIO(..))

130+import Control.Monad.Trans.Control (MonadBaseControl)

131+import Data.IORef (readIORef, writeIORef)

125import Data.Void (Void)

132import Data.Void (Void)

126import GHC.Types (SPEC(..))

133import GHC.Types (SPEC(..))

127import Prelude hiding (concat, map, mapM, takeWhile, take, filter, const)

134import Prelude hiding (concat, map, mapM, takeWhile, take, filter, const)

@@ -132,7 +139,7 @@ import Streamly.Internal.Data.Stream.StreamD.Type (pattern Stream)

132#endif

139#endif

133import Streamly.Internal.Data.Unfold.Types (Unfold(..))

140import Streamly.Internal.Data.Unfold.Types (Unfold(..))

134import Streamly.Internal.Data.Fold.Types (Fold(..))

141import Streamly.Internal.Data.Fold.Types (Fold(..))

135−import Streamly.Internal.Data.SVar (defState)

142+import Streamly.Internal.Data.SVar (defState, MonadAsync)

136import Control.Monad.Catch (MonadCatch)

143import Control.Monad.Catch (MonadCatch)

137

144

138import qualified Prelude

145import qualified Prelude

@@ -653,6 +660,53 @@ gbracket bef exc aft (Unfold estep einject) (Unfold step1 inject1) =

653 Skip s -> return $ Skip (Left s)

660 Skip s -> return $ Skip (Left s)

654 Stop -> return Stop

661 Stop -> return Stop

655

662

663+-- | The most general bracketing and exception combinator. All other

664+-- combinators can be expressed in terms of this combinator. This can also be

665+-- used for cases which are not covered by the standard combinators.

666+--

667+-- /Internal/

668+--

669+{-# INLINE_NORMAL gbracketIO #-}

670+gbracketIO

671+ :: (MonadIO m, MonadBaseControl IO m)

672+ => (a -> m c) -- ^ before

673+ -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)

674+ -> (c -> m d) -- ^ after, on normal stop, or GC

675+ -> Unfold m (c, e) b -- ^ on exception

676+ -> Unfold m c b -- ^ unfold to run

677+ -> Unfold m a b

678+gbracketIO bef exc aft (Unfold estep einject) (Unfold step1 inject1) =

679+ Unfold step inject

680+

681+ where

682+

683+ inject x = do

684+ r <- bef x

685+ ref <- D.newFinalizedIORef (aft r)

686+ s <- inject1 r

687+ return $ Right (s, r, ref)

688+

689+ {-# INLINE_LATE step #-}

690+ step (Right (st, v, ref)) = do

691+ res <- exc $ step1 st

692+ case res of

693+ Right r -> case r of

694+ Yield x s -> return $ Yield x (Right (s, v, ref))

695+ Skip s -> return $ Skip (Right (s, v, ref))

696+ Stop -> do

697+ liftIO $ writeIORef ref Nothing

698+ aft v >> return Stop

699+ Left e -> do

700+ liftIO (writeIORef ref Nothing)

701+ r <- einject (v, e)

702+ return $ Skip (Left r)

703+ step (Left st) = do

704+ res <- estep st

705+ case res of

706+ Yield x s -> return $ Yield x (Left s)

707+ Skip s -> return $ Skip (Left s)

708+ Stop -> return Stop

709+

656-- The custom implementation of "before" is slightly faster (5-7%) than

710-- The custom implementation of "before" is slightly faster (5-7%) than

657-- "_before". This is just to document and make sure that we can always use

711-- "_before". This is just to document and make sure that we can always use

658-- gbracket to implement before. The same applies to other combinators as well.

712-- gbracket to implement before. The same applies to other combinators as well.

@@ -690,6 +744,10 @@ _after aft = gbracket return (fmap Right) aft undefined

690

744

691-- | Run a side effect whenever the unfold stops normally.

745-- | Run a side effect whenever the unfold stops normally.

692--

746--

747+-- Prefer afterIO over this as the @after@ action in this combinator is not

748+-- executed if the unfold is partially evaluated lazily and then garbage

749+-- collected.

750+--

693-- /Internal/

751-- /Internal/

694{-# INLINE_NORMAL after #-}

752{-# INLINE_NORMAL after #-}

695after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b

753after :: Monad m => (a -> m c) -> Unfold m a b -> Unfold m a b

@@ -709,6 +767,33 @@ after action (Unfold step1 inject1) = Unfold step inject

709 Skip s -> return $ Skip (s, v)

767 Skip s -> return $ Skip (s, v)

710 Stop -> action v >> return Stop

768 Stop -> action v >> return Stop

711

769

770+-- | Run a side effect whenever the unfold stops normally

771+-- or is garbage collected after a partial lazy evaluation.

772+--

773+-- /Internal/

774+{-# INLINE_NORMAL afterIO #-}

775+afterIO :: (MonadIO m, MonadBaseControl IO m)

776+ => (a -> m c) -> Unfold m a b -> Unfold m a b

777+afterIO action (Unfold step1 inject1) = Unfold step inject

778+

779+ where

780+

781+ inject x = do

782+ s <- inject1 x

783+ ref <- D.newFinalizedIORef (action x)

784+ return (s, ref)

785+

786+ {-# INLINE_LATE step #-}

787+ step (st, ref) = do

788+ res <- step1 st

789+ case res of

790+ Yield x s -> return $ Yield x (s, ref)

791+ Skip s -> return $ Skip (s, ref)

792+ Stop -> liftIO $ do

793+ Just f <- readIORef ref

794+ writeIORef ref Nothing

795+ f >> return Stop

796+

712{-# INLINE_NORMAL _onException #-}

797{-# INLINE_NORMAL _onException #-}

713_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b

798_onException :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b

714_onException action unf =

799_onException action unf =

@@ -746,6 +831,10 @@ _finally action unf =

746-- | Run a side effect whenever the unfold stops normally or aborts due to an

831-- | Run a side effect whenever the unfold stops normally or aborts due to an

747-- exception.

832-- exception.

748--

833--

834+-- Prefer finallyIO over this as the @after@ action in this combinator is not

835+-- executed if the unfold is partially evaluated lazily and then garbage

836+-- collected.

837+--

749-- /Internal/

838-- /Internal/

750{-# INLINE_NORMAL finally #-}

839{-# INLINE_NORMAL finally #-}

751finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b

840finally :: MonadCatch m => (a -> m c) -> Unfold m a b -> Unfold m a b

@@ -765,6 +854,32 @@ finally action (Unfold step1 inject1) = Unfold step inject

765 Skip s -> return $ Skip (s, v)

854 Skip s -> return $ Skip (s, v)

766 Stop -> action v >> return Stop

855 Stop -> action v >> return Stop

767

856

857+-- | Run a side effect whenever the unfold stops normally, aborts due to an

858+-- exception or if it is garbage collected after a partial lazy evaluation.

859+--

860+-- /Internal/

861+{-# INLINE_NORMAL finallyIO #-}

862+finallyIO :: (MonadAsync m, MonadCatch m)

863+ => (a -> m c) -> Unfold m a b -> Unfold m a b

864+finallyIO action (Unfold step1 inject1) = Unfold step inject

865+

866+ where

867+

868+ inject x = do

869+ s <- inject1 x

870+ ref <- D.newFinalizedIORef (action x)

871+ return (s, x, ref)

872+

873+ {-# INLINE_LATE step #-}

874+ step (st, v, ref) = do

875+ res <- step1 st `MC.onException` action v

876+ case res of

877+ Yield x s -> return $ Yield x (s, v, ref)

878+ Skip s -> return $ Skip (s, v, ref)

879+ Stop -> do

880+ liftIO (writeIORef ref Nothing)

881+ action v >> return Stop

882+

768{-# INLINE_NORMAL _bracket #-}

883{-# INLINE_NORMAL _bracket #-}

769_bracket :: MonadCatch m

884_bracket :: MonadCatch m

770 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b

885 => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b

@@ -777,6 +892,10 @@ _bracket bef aft unf =

777-- if an exception occurs then the @after@ action is run with the output of

892-- if an exception occurs then the @after@ action is run with the output of

778-- @before@ as argument.

893-- @before@ as argument.

779--

894--

895+-- Prefer bracketIO over this as the @after@ action in this combinator is not

896+-- executed if the unfold is partially evaluated lazily and then garbage

897+-- collected.

898+--

780-- /Internal/

899-- /Internal/

781{-# INLINE_NORMAL bracket #-}

900{-# INLINE_NORMAL bracket #-}

782bracket :: MonadCatch m

901bracket :: MonadCatch m

@@ -798,6 +917,36 @@ bracket bef aft (Unfold step1 inject1) = Unfold step inject

798 Skip s -> return $ Skip (s, v)

917 Skip s -> return $ Skip (s, v)

799 Stop -> aft v >> return Stop

918 Stop -> aft v >> return Stop

800

919

920+-- | @bracket before after between@ runs the @before@ action and then unfolds

921+-- its output using the @between@ unfold. When the @between@ unfold is done or

922+-- if an exception occurs then the @after@ action is run with the output of

923+-- @before@ as argument. The after action is also executed if the unfold is

924+-- paritally evaluated and then garbage collected.

925+--

926+-- /Internal/

927+{-# INLINE_NORMAL bracketIO #-}

928+bracketIO :: (MonadAsync m, MonadCatch m)

929+ => (a -> m c) -> (c -> m d) -> Unfold m c b -> Unfold m a b

930+bracketIO bef aft (Unfold step1 inject1) = Unfold step inject

931+

932+ where

933+

934+ inject x = do

935+ r <- bef x

936+ s <- inject1 r

937+ ref <- D.newFinalizedIORef (aft r)

938+ return (s, r, ref)

939+

940+ {-# INLINE_LATE step #-}

941+ step (st, v, ref) = do

942+ res <- step1 st `MC.onException` aft v

943+ case res of

944+ Yield x s -> return $ Yield x (s, v, ref)

945+ Skip s -> return $ Skip (s, v, ref)

946+ Stop -> do

947+ liftIO (writeIORef ref Nothing)

948+ aft v >> return Stop

949+

801-- | When unfolding if an exception occurs, unfold the exception using the

950-- | When unfolding if an exception occurs, unfold the exception using the

802-- exception unfold supplied as the first argument to 'handle'.

951-- exception unfold supplied as the first argument to 'handle'.

803--

952--

src/Streamly/Internal/Prelude.hs+47

@@ -390,9 +390,12 @@ module Streamly.Internal.Prelude

390 -- * Exceptions

391 , before

392 , after

393+ , afterIO

393 , bracket

394 , bracket

395+ , bracketIO

394 , onException

396 , onException

395 , finally

397 , finally

398+ , finallyIO

396 , handle

399 , handle

397

400

398 -- * Generalize Inner Monad

401 -- * Generalize Inner Monad

@@ -438,6 +441,7 @@ import Control.Monad.IO.Class (MonadIO(..))

438import Control.Monad.Reader (ReaderT)

441import Control.Monad.Reader (ReaderT)

439import Control.Monad.State.Strict (StateT)

442import Control.Monad.State.Strict (StateT)

440import Control.Monad.Trans (MonadTrans(..))

443import Control.Monad.Trans (MonadTrans(..))

444+import Control.Monad.Trans.Control (MonadBaseControl)

441import Data.Functor.Identity (Identity (..))

445import Data.Functor.Identity (Identity (..))

442import Data.Heap (Entry(..))

446import Data.Heap (Entry(..))

443import Data.Maybe (isJust, fromJust, isNothing)

447import Data.Maybe (isJust, fromJust, isNothing)

@@ -4109,11 +4113,24 @@ before action xs = D.fromStreamD $ D.before action $ D.toStreamD xs

4109

4113

4110-- | Run a side effect whenever the stream stops normally.

4114-- | Run a side effect whenever the stream stops normally.

4111--

4115--

4116+-- Prefer 'afterIO' over this as the @after@ action in this combinator is not

4117+-- executed if the unfold is partially evaluated lazily and then garbage

4118+-- collected.

4119+--

4112-- @since 0.7.0

4120-- @since 0.7.0

4113{-# INLINE after #-}

4121{-# INLINE after #-}

4114after :: (IsStream t, Monad m) => m b -> t m a -> t m a

4122after :: (IsStream t, Monad m) => m b -> t m a -> t m a

4115after action xs = D.fromStreamD $ D.after action $ D.toStreamD xs

4123after action xs = D.fromStreamD $ D.after action $ D.toStreamD xs

4116

4124

4125+-- | Run a side effect whenever the stream stops normally

4126+-- or is garbage collected after a partial lazy evaluation.

4127+--

4128+-- /Internal/

4129+--

4130+{-# INLINE afterIO #-}

4131+afterIO :: (IsStream t, MonadIO m, MonadBaseControl IO m) => m b -> t m a -> t m a

4132+afterIO action xs = D.fromStreamD $ D.afterIO action $ D.toStreamD xs

4133+

4117-- | Run a side effect whenever the stream aborts due to an exception.

4134-- | Run a side effect whenever the stream aborts due to an exception.

4118--

4135--

4119-- @since 0.7.0

4136-- @since 0.7.0

@@ -4124,16 +4141,33 @@ onException action xs = D.fromStreamD $ D.onException action $ D.toStreamD xs

4124-- | Run a side effect whenever the stream stops normally or aborts due to an

4141-- | Run a side effect whenever the stream stops normally or aborts due to an

4125-- exception.

4142-- exception.

4126--

4143--

4144+-- Prefer 'finallyIO' over this as the @after@ action in this combinator is not

4145+-- executed if the unfold is partially evaluated lazily and then garbage

4146+-- collected.

4147+--

4127-- @since 0.7.0

4148-- @since 0.7.0

4128{-# INLINE finally #-}

4149{-# INLINE finally #-}

4129finally :: (IsStream t, MonadCatch m) => m b -> t m a -> t m a

4150finally :: (IsStream t, MonadCatch m) => m b -> t m a -> t m a

4130finally action xs = D.fromStreamD $ D.finally action $ D.toStreamD xs

4151finally action xs = D.fromStreamD $ D.finally action $ D.toStreamD xs

4131

4152

4153+-- | Run a side effect whenever the stream stops normally, aborts due to an

4154+-- exception or if it is garbage collected after a partial lazy evaluation.

4155+--

4156+-- /Internal/

4157+--

4158+{-# INLINE finallyIO #-}

4159+finallyIO :: (IsStream t, MonadAsync m, MonadCatch m) => m b -> t m a -> t m a

4160+finallyIO action xs = D.fromStreamD $ D.finallyIO action $ D.toStreamD xs

4161+

4132-- | Run the first action before the stream starts and remember its output,

4162-- | Run the first action before the stream starts and remember its output,

4133-- generate a stream using the output, run the second action using the

4163-- generate a stream using the output, run the second action using the

4134-- remembered value as an argument whenever the stream ends normally or due to

4164-- remembered value as an argument whenever the stream ends normally or due to

4135-- an exception.

4165-- an exception.

4136--

4166--

4167+-- Prefer 'bracketIO' over this as the @after@ action in this combinator is not

4168+-- executed if the unfold is partially evaluated lazily and then garbage

4169+-- collected.

4170+--

4137-- @since 0.7.0

4171-- @since 0.7.0

4138{-# INLINE bracket #-}

4172{-# INLINE bracket #-}

4139bracket :: (IsStream t, MonadCatch m)

4173bracket :: (IsStream t, MonadCatch m)

@@ -4141,6 +4175,19 @@ bracket :: (IsStream t, MonadCatch m)

4141bracket bef aft bet = D.fromStreamD $

4175bracket bef aft bet = D.fromStreamD $

4142 D.bracket bef aft (\x -> toStreamD $ bet x)

4176 D.bracket bef aft (\x -> toStreamD $ bet x)

4143

4177

4178+-- | Run the first action before the stream starts and remember its output,

4179+-- generate a stream using the output, run the second action using the

4180+-- remembered value as an argument whenever the stream ends normally, due to

4181+-- an exception or if it is garbage collected after a partial lazy evaluation.

4182+--

4183+-- /Internal/

4184+--

4185+{-# INLINE bracketIO #-}

4186+bracketIO :: (IsStream t, MonadAsync m, MonadCatch m)

4187+ => m b -> (b -> m c) -> (b -> t m a) -> t m a

4188+bracketIO bef aft bet = D.fromStreamD $

4189+ D.bracketIO bef aft (\x -> toStreamD $ bet x)

4190+

4144-- | When evaluating a stream if an exception occurs, stream evaluation aborts

4191-- | When evaluating a stream if an exception occurs, stream evaluation aborts

4145-- and the specified exception handler is run with the exception as argument.

4192-- and the specified exception handler is run with the exception as argument.

4146--

4193--