@@ -0,0 +1,221 @@
1+-- |
2+-- Module : Streamly.Internal.Data.Stream.Exception.Lifted
3+-- Copyright : (c) 2019 Composewell Technologies
4+-- License : BSD-3-Clause
5+-- Maintainer : streamly@composewell.com
6+-- Stability : experimental
7+-- Portability : GHC
8+
9+module Streamly.Internal.Data.Stream.Exception.Lifted
10+ (
11+ after
12+ , bracket
13+ , bracket'
14+ , finally
15+
16+ -- For IsStream module
17+ , afterD
18+ , bracket3D
19+ )
20+where
21+
22+#include "inline.hs"
23+
24+import Control.Exception (SomeException, mask_)
25+import Control.Monad.Catch (MonadCatch)
26+import GHC.Exts (inline)
27+import Streamly.Internal.Control.Concurrent
28+ (MonadRunInIO, MonadAsync, withRunInIO)
29+import Streamly.Internal.Data.Stream.Type (Stream, fromStreamD, toStreamD)
30+import Streamly.Internal.Data.IOFinalizer.Lifted
31+ (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
32+
33+import qualified Control.Monad.Catch as MC
34+import qualified Streamly.Internal.Data.Stream.StreamD as D
35+
36+import Streamly.Internal.Data.Stream.StreamD.Type hiding (Stream)
37+
38+-- $setup
39+-- >>> :m
40+-- >>> import qualified Streamly.Internal.Data.Stream.Exception.Lifted as Stream
41+
42+-- XXX Implement in terms of the corresponding IO operation (gbracketIO).
43+
44+data GbracketIOState s1 s2 v wref
45+ = GBracketIOInit
46+ | GBracketIONormal s1 v wref
47+ | GBracketIOException s2
48+
49+{-# INLINE_NORMAL gbracket #-}
50+gbracket
51+ :: MonadRunInIO m
52+ => m c -- ^ before
53+ -> (c -> m d1) -- ^ on normal stop
54+ -> (c -> e -> D.Stream m b -> m (D.Stream m b)) -- ^ on exception
55+ -> (c -> m d2) -- ^ on GC without normal stop or exception
56+ -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
57+ -> (c -> D.Stream m b) -- ^ stream generator
58+ -> D.Stream m b
59+gbracket bef aft onExc onGC ftry action =
60+ D.Stream step GBracketIOInit
61+
62+ where
63+
64+ -- If the stream is never evaluated the "aft" action will never be
65+ -- called. For that to occur we will need the user of this API to pass a
66+ -- weak pointer to us.
67+ {-# INLINE_LATE step #-}
68+ step _ GBracketIOInit = do
69+ -- We mask asynchronous exceptions to make the execution
70+ -- of 'bef' and the registration of 'aft' atomic.
71+ -- A similar thing is done in the resourcet package: https://git.io/JvKV3
72+ -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
73+ (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
74+ r <- bef
75+ ref <- newIOFinalizer (onGC r)
76+ return (r, ref)
77+ return $ Skip $ GBracketIONormal (action r) r ref
78+
79+ step gst (GBracketIONormal (UnStream step1 st) v ref) = do
80+ res <- ftry $ step1 gst st
81+ case res of
82+ Right r -> case r of
83+ Yield x s ->
84+ return $ Yield x (GBracketIONormal (D.Stream step1 s) v ref)
85+ Skip s ->
86+ return $ Skip (GBracketIONormal (D.Stream step1 s) v ref)
87+ Stop ->
88+ clearingIOFinalizer ref (aft v) >> return Stop
89+ -- XXX Do not handle async exceptions, just rethrow them.
90+ Left e -> do
91+ -- Clearing of finalizer and running of exception handler must
92+ -- be atomic wrt async exceptions. Otherwise if we have cleared
93+ -- the finalizer and have not run the exception handler then we
94+ -- may leak the resource.
95+ stream <-
96+ clearingIOFinalizer ref (onExc v e (UnStream step1 st))
97+ return $ Skip (GBracketIOException stream)
98+ step gst (GBracketIOException (UnStream step1 st)) = do
99+ res <- step1 gst st
100+ case res of
101+ Yield x s ->
102+ return $ Yield x (GBracketIOException (D.Stream step1 s))
103+ Skip s -> return $ Skip (GBracketIOException (D.Stream step1 s))
104+ Stop -> return Stop
105+
106+{-# INLINE_NORMAL bracket3D #-}
107+bracket3D :: (MonadAsync m, MonadCatch m) =>
108+ m b
109+ -> (b -> m c)
110+ -> (b -> m d)
111+ -> (b -> m e)
112+ -> (b -> D.Stream m a)
113+ -> D.Stream m a
114+bracket3D bef aft onExc onGC =
115+ gbracket
116+ bef
117+ aft
118+ (\a (e :: SomeException) _ -> onExc a >> return (nilM (MC.throwM e)))
119+ onGC
120+ (inline MC.try)
121+
122+-- For a use case of this see the "streamly-process" package. It needs to kill
123+-- the process in case of exception or garbage collection, but waits for the
124+-- process to terminate in normal cases.
125+
126+-- | Like 'bracket' but can use separate cleanup actions depending on the mode
127+-- of termination. @bracket' before onStop onGC onException action@ runs
128+-- @action@ using the result of @before@. If the stream stops, @onStop@ action
129+-- is executed, if the stream is abandoned @onGC@ is executed, if the stream
130+-- encounters an exception @onException@ is executed.
131+--
132+-- /Pre-release/
133+{-# INLINE bracket' #-}
134+bracket' :: (MonadRunInIO m, MonadCatch m)
135+ => m b
136+ -> (b -> m c)
137+ -> (b -> m d)
138+ -> (b -> m e)
139+ -> (b -> Stream m a)
140+ -> Stream m a
141+bracket' bef aft gc exc bet = fromStreamD $
142+ bracket3D bef aft exc gc (toStreamD . bet)
143+
144+-- | Run the alloc action @m b@ with async exceptions disabled but keeping
145+-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
146+-- output @b@ as input to @b -> Stream m a@ to generate an output stream.
147+--
148+-- @b@ is usually a resource under the state of monad @m@, e.g. a file
149+-- handle, that requires a cleanup after use. The cleanup action @b -> m c@,
150+-- runs whenever the stream ends normally, due to a sync or async exception or
151+-- if it gets garbage collected after a partial lazy evaluation.
152+--
153+-- 'bracket' only guarantees that the cleanup action runs, and it runs with
154+-- async exceptions enabled. The action must ensure that it can successfully
155+-- cleanup the resource in the face of sync or async exceptions.
156+--
157+-- When the stream ends normally or on a sync exception, cleanup action runs
158+-- immediately in the current thread context, whereas in other cases it runs in
159+-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
160+--
161+-- /See also: 'bracket_'/
162+--
163+-- /Inhibits stream fusion/
164+--
165+-- @since 0.9.0
166+{-# INLINE bracket #-}
167+bracket :: (MonadRunInIO m, MonadCatch m)
168+ => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
169+bracket bef aft = bracket' bef aft aft aft
170+
171+-- | Run the action @m b@ whenever the stream @Stream m a@ stops normally,
172+-- aborts due to an exception or if it is garbage collected after a partial
173+-- lazy evaluation.
174+--
175+-- The semantics of running the action @m b@ are similar to the cleanup action
176+-- semantics described in 'bracket'.
177+--
178+-- >>> finally action xs = Stream.bracket (return ()) (const action) (const xs)
179+--
180+-- /See also 'finally_'/
181+--
182+-- /Inhibits stream fusion/
183+--
184+-- @since 0.9.0
185+{-# INLINE finally #-}
186+finally :: (MonadRunInIO m, MonadCatch m) =>
187+ m b -> Stream m a -> Stream m a
188+finally action xs = bracket (return ()) (const action) (const xs)
189+
190+{-# INLINE_NORMAL afterD #-}
191+afterD :: MonadRunInIO m
192+ => m b -> D.Stream m a -> D.Stream m a
193+afterD action (D.Stream step state) = D.Stream step' Nothing
194+
195+ where
196+
197+ {-# INLINE_LATE step' #-}
198+ step' _ Nothing = do
199+ ref <- newIOFinalizer action
200+ return $ Skip $ Just (state, ref)
201+ step' gst (Just (st, ref)) = do
202+ res <- step gst st
203+ case res of
204+ Yield x s -> return $ Yield x (Just (s, ref))
205+ Skip s -> return $ Skip (Just (s, ref))
206+ Stop -> do
207+ runIOFinalizer ref
208+ return Stop
209+
210+-- | Run the action @m b@ whenever the stream @Stream m a@ stops normally, or
211+-- if it is garbage collected after a partial lazy evaluation.
212+--
213+-- The semantics of the action @m b@ are similar to the semantics of cleanup
214+-- action in 'bracket'.
215+--
216+-- /See also 'after_'/
217+--
218+-- @since 0.9.0
219+{-# INLINE after #-}
220+after :: MonadRunInIO m => m b -> Stream m a -> Stream m a
221+after action xs = fromStreamD $ afterD action $ toStreamD xs