git.iko.soy/streamly
on 209e990 help sign in
streamly/src/Streamly/Internal/Data/Stream/Exception/Lifted.hs
335 lines 11 kB history raw
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
9module Streamly.Internal.Data.Stream.Exception.Lifted
10 (
11 after
12 , bracket
13 , bracket3
14 , finally
15 , retry
17 -- For IsStream module
18 , afterD
19 , bracket3D
20 , retryD
21 )
22where
24#include "inline.hs"
26import Control.Exception (Exception, SomeException, mask_)
27import Control.Monad.Catch (MonadCatch)
28#ifdef USE_UNLIFTIO
29import Control.Monad.IO.Unlift (MonadUnliftIO)
30#else
31import Control.Monad.IO.Class (MonadIO)
32import Control.Monad.Trans.Control (MonadBaseControl)
33#endif
34import Data.Map.Strict (Map)
35import GHC.Exts (inline)
36import Streamly.Internal.Control.Concurrent
37 (MonadRunInIO, MonadAsync, withRunInIO)
38import Streamly.Internal.Data.Stream (Stream)
39import Streamly.Internal.Data.IOFinalizer.Lifted
40 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
41import Streamly.Internal.Data.Stream (Step(..))
43import qualified Control.Monad.Catch as MC
44import qualified Data.Map.Strict as Map
45import qualified Streamly.Internal.Data.Stream as D
47-- $setup
48-- >>> :m
49-- >>> import qualified Streamly.Internal.Data.Stream.Exception.Lifted as Stream
51-- XXX Implement in terms of the corresponding IO operation (gbracketIO).
53data GbracketIOState s1 s2 v wref
54 = GBracketIOInit
55 | GBracketIONormal s1 v wref
56 | GBracketIOException s2
58{-# INLINE_NORMAL gbracket #-}
59gbracket
60 :: MonadRunInIO m
61 => m c -- ^ before
62 -> (c -> m d1) -- ^ on normal stop
63 -> (c -> e -> D.Stream m b -> m (D.Stream m b)) -- ^ on exception
64 -> (c -> m d2) -- ^ on GC without normal stop or exception
65 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
66 -> (c -> D.Stream m b) -- ^ stream generator
67 -> D.Stream m b
68gbracket bef aft onExc onGC ftry action =
69 D.Stream step GBracketIOInit
71 where
73 -- If the stream is never evaluated the "aft" action will never be
74 -- called. For that to occur we will need the user of this API to pass a
75 -- weak pointer to us.
76 {-# INLINE_LATE step #-}
77 step _ GBracketIOInit = do
78 -- We mask asynchronous exceptions to make the execution
79 -- of 'bef' and the registration of 'aft' atomic.
80 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
81 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
82 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
83 r <- bef
84 ref <- newIOFinalizer (onGC r)
85 return (r, ref)
86 return $ Skip $ GBracketIONormal (action r) r ref
88 step gst (GBracketIONormal (D.UnStream step1 st) v ref) = do
89 res <- ftry $ step1 gst st
90 case res of
91 Right r -> case r of
92 Yield x s ->
93 return $ Yield x (GBracketIONormal (D.Stream step1 s) v ref)
94 Skip s ->
95 return $ Skip (GBracketIONormal (D.Stream step1 s) v ref)
96 Stop ->
97 clearingIOFinalizer ref (aft v) >> return Stop
98 -- XXX Do not handle async exceptions, just rethrow them.
99 Left e -> do
100 -- Clearing of finalizer and running of exception handler must
101 -- be atomic wrt async exceptions. Otherwise if we have cleared
102 -- the finalizer and have not run the exception handler then we
103 -- may leak the resource.
104 stream <-
105 clearingIOFinalizer ref (onExc v e (D.UnStream step1 st))
106 return $ Skip (GBracketIOException stream)
107 step gst (GBracketIOException (D.UnStream step1 st)) = do
108 res <- step1 gst st
109 case res of
110 Yield x s ->
111 return $ Yield x (GBracketIOException (D.Stream step1 s))
112 Skip s -> return $ Skip (GBracketIOException (D.Stream step1 s))
113 Stop -> return Stop
115{-# INLINE_NORMAL bracket3D #-}
116bracket3D :: (MonadAsync m, MonadCatch m) =>
117 m b
118 -> (b -> m c)
119 -> (b -> m d)
120 -> (b -> m e)
121 -> (b -> D.Stream m a)
122 -> D.Stream m a
123bracket3D bef aft onExc onGC =
124 gbracket
125 bef
126 aft
127 (\a (e :: SomeException) _ -> onExc a >> MC.throwM e)
128 onGC
129 (inline MC.try)
131-- For a use case of this see the "streamly-process" package. It needs to kill
132-- the process in case of exception or garbage collection, but waits for the
133-- process to terminate in normal cases.
135-- | Like 'bracket' but can use 3 separate cleanup actions depending on the
136-- mode of termination:
137--
138-- 1. When the stream stops normally
139-- 2. When the stream is garbage collected
140-- 3. When the stream encounters an exception
141--
142-- @bracket3 before onStop onGC onException action@ runs @action@ using the
143-- result of @before@. If the stream stops, @onStop@ action is executed, if the
144-- stream is abandoned @onGC@ is executed, if the stream encounters an
145-- exception @onException@ is executed.
146--
147-- The exception is not caught, it is rethrown.
148--
149-- /Pre-release/
150{-# INLINE bracket3 #-}
151bracket3 :: (MonadAsync m, MonadCatch m)
152 => m b
153 -> (b -> m c)
154 -> (b -> m d)
155 -> (b -> m e)
156 -> (b -> Stream m a)
157 -> Stream m a
158bracket3 = bracket3D
160-- | Run the alloc action @IO b@ with async exceptions disabled but keeping
161-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
162-- output @b@ of the IO action as input to the function @b -> Stream m a@ to
163-- generate an output stream.
164--
165-- @b@ is usually a resource under the IO monad, e.g. a file handle, that
166-- requires a cleanup after use. The cleanup action @b -> m c@, runs whenever
167-- (1) the stream ends normally, (2) due to a sync or async exception or, (3)
168-- if it gets garbage collected after a partial lazy evaluation. The exception
169-- is not caught, it is rethrown.
170--
171-- 'bracket' only guarantees that the cleanup action runs, and it runs with
172-- async exceptions enabled. The action must ensure that it can successfully
173-- cleanup the resource in the face of sync or async exceptions.
174--
175-- When the stream ends normally or on a sync exception, cleanup action runs
176-- immediately in the current thread context, whereas in other cases it runs in
177-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
178--
179-- /See also: 'bracket_'/
180--
181-- /Inhibits stream fusion/
182--
183{-# INLINE bracket #-}
184bracket :: (MonadAsync m, MonadCatch m)
185 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
186bracket bef aft = bracket3 bef aft aft aft
188-- | Run the action @m b@ whenever the stream @Stream m a@ stops normally,
189-- aborts due to an exception or if it is garbage collected after a partial
190-- lazy evaluation.
191--
192-- The semantics of running the action @m b@ are similar to the cleanup action
193-- semantics described in 'bracket'.
194--
195-- >>> finally action xs = Stream.bracket (return ()) (const action) (const xs)
196--
197-- /See also 'finally_'/
198--
199-- /Inhibits stream fusion/
200--
201{-# INLINE finally #-}
202finally :: (MonadAsync m, MonadCatch m) =>
203 m b -> Stream m a -> Stream m a
204finally action xs = bracket (return ()) (const action) (const xs)
206{-# INLINE_NORMAL afterD #-}
207afterD :: MonadRunInIO m
208 => m b -> D.Stream m a -> D.Stream m a
209afterD action (D.Stream step state) = D.Stream step' Nothing
211 where
213 {-# INLINE_LATE step' #-}
214 step' _ Nothing = do
215 ref <- newIOFinalizer action
216 return $ Skip $ Just (state, ref)
217 step' gst (Just (st, ref)) = do
218 res <- step gst st
219 case res of
220 Yield x s -> return $ Yield x (Just (s, ref))
221 Skip s -> return $ Skip (Just (s, ref))
222 Stop -> do
223 runIOFinalizer ref
224 return Stop
226-- | Run the action @m b@ whenever the stream @Stream m a@ stops normally, or
227-- if it is garbage collected after a partial lazy evaluation.
228--
229-- The semantics of the action @m b@ are similar to the semantics of cleanup
230-- action in 'bracket'.
231--
232-- /See also 'after_'/
233--
234{-# INLINE after #-}
235after ::
236#ifdef USE_UNLIFTIO
237 MonadUnliftIO m
238#else
239 (MonadIO m, MonadBaseControl IO m)
240#endif
241 => m b -> Stream m a -> Stream m a
242after = afterD
244data RetryState emap s1 s2
245 = RetryWithMap emap s1
246 | RetryDefault s2
248-- | See 'Streamly.Internal.Data.Stream.retry'
249--
250{-# INLINE_NORMAL retryD #-}
251retryD
252 :: forall e m a. (Exception e, Ord e, MonadCatch m)
253 => Map e Int
254 -- ^ map from exception to retry count
255 -> (e -> D.Stream m a)
256 -- ^ default handler for those exceptions that are not in the map
257 -> D.Stream m a
258 -> D.Stream m a
259retryD emap0 defaultHandler (D.Stream step0 state0) = D.Stream step state
261 where
263 state = RetryWithMap emap0 state0
265 {-# INLINE_LATE step #-}
266 step gst (RetryWithMap emap st) = do
267 eres <- MC.try $ step0 gst st
268 case eres of
269 Left e -> handler e emap st
270 Right res ->
271 return
272 $ case res of
273 Yield x st1 -> Yield x $ RetryWithMap emap st1
274 Skip st1 -> Skip $ RetryWithMap emap st1
275 Stop -> Stop
276 step gst (RetryDefault (D.UnStream step1 state1)) = do
277 res <- step1 gst state1
278 return
279 $ case res of
280 Yield x st1 -> Yield x $ RetryDefault (D.Stream step1 st1)
281 Skip st1 -> Skip $ RetryDefault (D.Stream step1 st1)
282 Stop -> Stop
284 {-# INLINE handler #-}
285 handler e emap st =
286 return
287 $ Skip
288 $ case Map.lookup e emap of
289 Just i
290 | i > 0 ->
291 let emap1 = Map.insert e (i - 1) emap
292 in RetryWithMap emap1 st
293 | otherwise -> RetryDefault $ defaultHandler e
294 Nothing -> RetryDefault $ defaultHandler e
296-- | @retry@ takes 3 arguments
297--
298-- 1. A map @m@ whose keys are exceptions and values are the number of times to
299-- retry the action given that the exception occurs.
300--
301-- 2. A handler @han@ that decides how to handle an exception when the exception
302-- cannot be retried.
303--
304-- 3. The stream itself that we want to run this mechanism on.
305--
306-- When evaluating a stream if an exception occurs,
307--
308-- 1. The stream evaluation aborts
309--
310-- 2. The exception is looked up in @m@
311--
312-- a. If the exception exists and the mapped value is > 0 then,
313--
314-- i. The value is decreased by 1.
315--
316-- ii. The stream is resumed from where the exception was called, retrying
317-- the action.
318--
319-- b. If the exception exists and the mapped value is == 0 then the stream
320-- evaluation stops.
321--
322-- c. If the exception does not exist then we handle the exception using
323-- @han@.
324--
325-- /Internal/
326--
327{-# INLINE retry #-}
328retry :: (MonadCatch m, Exception e, Ord e)
329 => Map e Int
330 -- ^ map from exception to retry count
331 -> (e -> Stream m a)
332 -- ^ default handler for those exceptions that are not in the map
333 -> Stream m a
334 -> Stream m a
335retry = retryD