git.iko.soy/streamly
on d08d73a help sign in
streamly/core/src/Streamly/Internal/Data/Stream/StreamD/Exception.hs
424 lines 14 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.StreamD.Exception
3-- Copyright : (c) 2020 Composewell Technologies and Contributors
4-- License : BSD-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
9module Streamly.Internal.Data.Stream.StreamD.Exception
10 (
11 gbracket_
12 , gbracket
13 , before
14 , after_
15 , after
16 , bracket_
17 , bracket'
18 , onException
19 , finally_
20 , finally
21 , ghandle
22 , handle
23 , retry
24 )
25where
27#include "inline.hs"
29import Control.Exception (Exception, SomeException, mask_)
30import Control.Monad.Catch (MonadCatch)
31import Data.Map.Strict (Map)
32import GHC.Exts (inline)
33import Streamly.Internal.Control.Concurrent (MonadRunInIO, MonadAsync, withRunInIO)
34import Streamly.Internal.Data.IOFinalizer
35 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
37import qualified Control.Monad.Catch as MC
38import qualified Data.Map.Strict as Map
40import Streamly.Internal.Data.Stream.StreamD.Type
42data GbracketState s1 s2 v
43 = GBracketInit
44 | GBracketNormal s1 v
45 | GBracketException s2
47-- | Like 'gbracket' but with following differences:
48--
49-- * alloc action @m c@ runs with async exceptions enabled
50-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
51-- after partial evaluation.
52-- * does not require a 'MonadAsync' constraint.
53--
54-- /Inhibits stream fusion/
55--
56-- /Pre-release/
57--
58{-# INLINE_NORMAL gbracket_ #-}
59gbracket_
60 :: Monad m
61 => m c -- ^ before
62 -> (c -> m d) -- ^ after, on normal stop
63 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
64 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
65 -> (c -> Stream m b) -- ^ stream generator
66 -> Stream m b
67gbracket_ bef aft onExc ftry action =
68 Stream step GBracketInit
70 where
72 {-# INLINE_LATE step #-}
73 step _ GBracketInit = do
74 r <- bef
75 return $ Skip $ GBracketNormal (action r) r
77 step gst (GBracketNormal (UnStream step1 st) v) = do
78 res <- ftry $ step1 gst st
79 case res of
80 Right r -> case r of
81 Yield x s ->
82 return $ Yield x (GBracketNormal (Stream step1 s) v)
83 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
84 Stop -> aft v >> return Stop
85 -- XXX Do not handle async exceptions, just rethrow them.
86 Left e ->
87 return
88 $ Skip (GBracketException (onExc v e (UnStream step1 st)))
89 step gst (GBracketException (UnStream step1 st)) = do
90 res <- step1 gst st
91 case res of
92 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
93 Skip s -> return $ Skip (GBracketException (Stream step1 s))
94 Stop -> return Stop
96data GbracketIOState s1 s2 v wref
97 = GBracketIOInit
98 | GBracketIONormal s1 v wref
99 | GBracketIOException s2
101-- | Run the alloc action @m c@ with async exceptions disabled but keeping
102-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
103-- output @c@ as input to @c -> Stream m b@ to generate an output stream. When
104-- generating the stream use the supplied @try@ operation @forall s. m s -> m
105-- (Either e s)@ to catch synchronous exceptions. If an exception occurs run
106-- the exception handler @c -> e -> Stream m b -> m (Stream m b)@. Note that
107-- 'gbracket' does not rethrow the exception, it has to be done by the
108-- exception handler if desired.
109--
110-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
111-- to a sync or async exception or if it gets garbage collected after a partial
112-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
113--
114-- 'gbracket' can express all other exception handling combinators.
115--
116-- /Inhibits stream fusion/
117--
118-- /Pre-release/
119{-# INLINE_NORMAL gbracket #-}
120gbracket
121 :: MonadRunInIO m
122 => m c -- ^ before
123 -> (c -> m d1) -- ^ on normal stop
124 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
125 -> (c -> m d2) -- ^ on GC without normal stop or exception
126 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
127 -> (c -> Stream m b) -- ^ stream generator
128 -> Stream m b
129gbracket bef aft onExc onGC ftry action =
130 Stream step GBracketIOInit
132 where
134 -- If the stream is never evaluated the "aft" action will never be
135 -- called. For that to occur we will need the user of this API to pass a
136 -- weak pointer to us.
137 {-# INLINE_LATE step #-}
138 step _ GBracketIOInit = do
139 -- We mask asynchronous exceptions to make the execution
140 -- of 'bef' and the registration of 'aft' atomic.
141 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
142 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
143 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
144 r <- bef
145 ref <- newIOFinalizer (onGC r)
146 return (r, ref)
147 return $ Skip $ GBracketIONormal (action r) r ref
149 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
150 res <- ftry $ step1 gst st
151 case res of
152 Right r -> case r of
153 Yield x s ->
154 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
155 Skip s ->
156 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
157 Stop ->
158 clearingIOFinalizer ref (aft v) >> return Stop
159 -- XXX Do not handle async exceptions, just rethrow them.
160 Left e -> do
161 -- Clearing of finalizer and running of exception handler must
162 -- be atomic wrt async exceptions. Otherwise if we have cleared
163 -- the finalizer and have not run the exception handler then we
164 -- may leak the resource.
165 stream <-
166 clearingIOFinalizer ref (onExc v e (UnStream step1 st))
167 return $ Skip (GBracketIOException stream)
168 step gst (GBracketIOException (UnStream step1 st)) = do
169 res <- step1 gst st
170 case res of
171 Yield x s ->
172 return $ Yield x (GBracketIOException (Stream step1 s))
173 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
174 Stop -> return Stop
176-- | See 'Streamly.Internal.Data.Stream.IsStream.before'.
177--
178{-# INLINE_NORMAL before #-}
179before :: Monad m => m b -> Stream m a -> Stream m a
180before action (Stream step state) = Stream step' Nothing
182 where
184 {-# INLINE_LATE step' #-}
185 step' _ Nothing = action >> return (Skip (Just state))
187 step' gst (Just st) = do
188 res <- step gst st
189 case res of
190 Yield x s -> return $ Yield x (Just s)
191 Skip s -> return $ Skip (Just s)
192 Stop -> return Stop
194-- | See 'Streamly.Internal.Data.Stream.IsStream.after_'.
195--
196{-# INLINE_NORMAL after_ #-}
197after_ :: Monad m => m b -> Stream m a -> Stream m a
198after_ action (Stream step state) = Stream step' state
200 where
202 {-# INLINE_LATE step' #-}
203 step' gst st = do
204 res <- step gst st
205 case res of
206 Yield x s -> return $ Yield x s
207 Skip s -> return $ Skip s
208 Stop -> action >> return Stop
210-- | See 'Streamly.Internal.Data.Stream.IsStream.after'.
211--
212{-# INLINE_NORMAL after #-}
213after :: MonadRunInIO m
214 => m b -> Stream m a -> Stream m a
215after action (Stream step state) = Stream step' Nothing
217 where
219 {-# INLINE_LATE step' #-}
220 step' _ Nothing = do
221 ref <- newIOFinalizer action
222 return $ Skip $ Just (state, ref)
223 step' gst (Just (st, ref)) = do
224 res <- step gst st
225 case res of
226 Yield x s -> return $ Yield x (Just (s, ref))
227 Skip s -> return $ Skip (Just (s, ref))
228 Stop -> do
229 runIOFinalizer ref
230 return Stop
232-- XXX For high performance error checks in busy streams we may need another
233-- Error constructor in step.
234--
235-- | See 'Streamly.Internal.Data.Stream.IsStream.onException'.
236--
237{-# INLINE_NORMAL onException #-}
238onException :: MonadCatch m => m b -> Stream m a -> Stream m a
239onException action stream =
240 gbracket_
241 (return ()) -- before
242 return -- after
243 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
244 (inline MC.try)
245 (const stream)
247{-# INLINE_NORMAL _onException #-}
248_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
249_onException action (Stream step state) = Stream step' state
251 where
253 {-# INLINE_LATE step' #-}
254 step' gst st = do
255 res <- step gst st `MC.onException` action
256 case res of
257 Yield x s -> return $ Yield x s
258 Skip s -> return $ Skip s
259 Stop -> return Stop
261-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket_'.
262--
263{-# INLINE_NORMAL bracket_ #-}
264bracket_ :: MonadCatch m
265 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
266bracket_ bef aft =
267 gbracket_
268 bef
269 aft
270 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e))
271 (inline MC.try)
273-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket'.
274--
275{-# INLINE_NORMAL bracket' #-}
276bracket' :: (MonadAsync m, MonadCatch m) =>
277 m b
278 -> (b -> m c)
279 -> (b -> m d)
280 -> (b -> m e)
281 -> (b -> Stream m a)
282 -> Stream m a
283bracket' bef aft onExc onGC =
284 gbracket
285 bef
286 aft
287 (\a (e :: SomeException) _ -> onExc a >> return (nilM (MC.throwM e)))
288 onGC
289 (inline MC.try)
291data BracketState s v = BracketInit | BracketRun s v
293-- | Alternate (custom) implementation of 'bracket'.
294--
295{-# INLINE_NORMAL _bracket #-}
296_bracket :: MonadCatch m
297 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
298_bracket bef aft bet = Stream step' BracketInit
300 where
302 {-# INLINE_LATE step' #-}
303 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
305 -- NOTE: It is important to use UnStream instead of the Stream pattern
306 -- here, otherwise we get huge perf degradation, see note in concatMap.
307 step' gst (BracketRun (UnStream step state) v) = do
308 -- res <- step gst state `MC.onException` aft v
309 res <- inline MC.try $ step gst state
310 case res of
311 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
312 Right r -> case r of
313 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
314 Skip s -> return $ Skip (BracketRun (Stream step s) v)
315 Stop -> aft v >> return Stop
317-- | See 'Streamly.Internal.Data.Stream.IsStream.finally_'.
318--
319{-# INLINE finally_ #-}
320finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
321finally_ action xs = bracket_ (return ()) (const action) (const xs)
323-- | See 'Streamly.Internal.Data.Stream.IsStream.finally'.
324--
325-- finally action xs = after action $ onException action xs
326--
327{-# INLINE finally #-}
328finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
329finally action xs = bracket' (return ()) act act act (const xs)
330 where act _ = action
332-- | See 'Streamly.Internal.Data.Stream.IsStream.ghandle'.
333--
334{-# INLINE_NORMAL ghandle #-}
335ghandle :: (MonadCatch m, Exception e)
336 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
337ghandle f stream =
338 gbracket_ (return ()) return (const f) (inline MC.try) (const stream)
340-- | See 'Streamly.Internal.Data.Stream.IsStream.handle'.
341--
342{-# INLINE_NORMAL handle #-}
343handle :: (MonadCatch m, Exception e)
344 => (e -> Stream m a) -> Stream m a -> Stream m a
345handle f stream =
346 gbracket_ (return ()) return (\_ e _ -> f e) (inline MC.try) (const stream)
348-- | Alternate (custom) implementation of 'handle'.
349--
350{-# INLINE_NORMAL _handle #-}
351_handle :: (MonadCatch m, Exception e)
352 => (e -> Stream m a) -> Stream m a -> Stream m a
353_handle f (Stream step state) = Stream step' (Left state)
355 where
357 {-# INLINE_LATE step' #-}
358 step' gst (Left st) = do
359 res <- inline MC.try $ step gst st
360 case res of
361 Left e -> return $ Skip $ Right (f e)
362 Right r -> case r of
363 Yield x s -> return $ Yield x (Left s)
364 Skip s -> return $ Skip (Left s)
365 Stop -> return Stop
367 step' gst (Right (UnStream step1 st)) = do
368 res <- step1 gst st
369 case res of
370 Yield x s -> return $ Yield x (Right (Stream step1 s))
371 Skip s -> return $ Skip (Right (Stream step1 s))
372 Stop -> return Stop
374data RetryState emap s1 s2
375 = RetryWithMap emap s1
376 | RetryDefault s2
378-- | See 'Streamly.Internal.Data.Stream.IsStream.retry'
379--
380{-# INLINE_NORMAL retry #-}
381retry
382 :: forall e m a. (Exception e, Ord e, MonadCatch m)
383 => Map e Int
384 -- ^ map from exception to retry count
385 -> (e -> Stream m a)
386 -- ^ default handler for those exceptions that are not in the map
387 -> Stream m a
388 -> Stream m a
389retry emap0 defaultHandler (Stream step0 state0) = Stream step state
391 where
393 state = RetryWithMap emap0 state0
395 {-# INLINE_LATE step #-}
396 step gst (RetryWithMap emap st) = do
397 eres <- MC.try $ step0 gst st
398 case eres of
399 Left e -> handler e emap st
400 Right res ->
401 return
402 $ case res of
403 Yield x st1 -> Yield x $ RetryWithMap emap st1
404 Skip st1 -> Skip $ RetryWithMap emap st1
405 Stop -> Stop
406 step gst (RetryDefault (UnStream step1 state1)) = do
407 res <- step1 gst state1
408 return
409 $ case res of
410 Yield x st1 -> Yield x $ RetryDefault (Stream step1 st1)
411 Skip st1 -> Skip $ RetryDefault (Stream step1 st1)
412 Stop -> Stop
414 {-# INLINE handler #-}
415 handler e emap st =
416 return
417 $ Skip
418 $ case Map.lookup e emap of
419 Just i
420 | i > 0 ->
421 let emap1 = Map.insert e (i - 1) emap
422 in RetryWithMap emap1 st
423 | otherwise -> RetryDefault $ defaultHandler e
424 Nothing -> RetryDefault $ defaultHandler e