git.iko.soy/streamly
on f8afdaf help sign in
streamly/core/src/Streamly/Internal/Data/Stream/StreamD/Exception.hs
412 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 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
63 -> (c -> m d) -- ^ after, on normal stop
64 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
65 -> (c -> Stream m b) -- ^ stream generator
66 -> Stream m b
67gbracket_ bef exc aft fexc fnormal =
68 Stream step GBracketInit
70 where
72 {-# INLINE_LATE step #-}
73 step _ GBracketInit = do
74 r <- bef
75 return $ Skip $ GBracketNormal (fnormal r) r
77 step gst (GBracketNormal (UnStream step1 st) v) = do
78 res <- exc $ 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 $ Skip (GBracketException (fexc v e (UnStream step1 st)))
88 step gst (GBracketException (UnStream step1 st)) = do
89 res <- step1 gst st
90 case res of
91 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
92 Skip s -> return $ Skip (GBracketException (Stream step1 s))
93 Stop -> return Stop
95data GbracketIOState s1 s2 v wref
96 = GBracketIOInit
97 | GBracketIONormal s1 v wref
98 | GBracketIOException s2
100-- | Run the alloc action @m c@ with async exceptions disabled but keeping
101-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
102-- output @c@ as input to @c -> Stream m b@ to generate an output stream. When
103-- generating the stream use the supplied @try@ operation @forall s. m s -> m
104-- (Either e s)@ to catch synchronous exceptions. If an exception occurs run
105-- the exception handler @c -> e -> Stream m b -> m (Stream m b)@. Note that
106-- 'gbracket' does not rethrow the exception, it has to be done by the
107-- exception handler if desired.
108--
109-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
110-- to a sync or async exception or if it gets garbage collected after a partial
111-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
112--
113-- 'gbracket' can express all other exception handling combinators.
114--
115-- /Inhibits stream fusion/
116--
117-- /Pre-release/
118{-# INLINE_NORMAL gbracket #-}
119gbracket
120 :: MonadRunInIO m
121 => m c -- ^ before
122 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
123 -> (c -> m d1) -- ^ on normal stop
124 -> (c -> m d2) -- ^ on GC without normal stop or exception
125 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
126 -> (c -> Stream m b) -- ^ stream generator
127 -> Stream m b
128gbracket bef exc aft gc fexc fnormal =
129 Stream step GBracketIOInit
131 where
133 -- If the stream is never evaluated the "aft" action will never be
134 -- called. For that to occur we will need the user of this API to pass a
135 -- weak pointer to us.
136 {-# INLINE_LATE step #-}
137 step _ GBracketIOInit = do
138 -- We mask asynchronous exceptions to make the execution
139 -- of 'bef' and the registration of 'aft' atomic.
140 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
141 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
142 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
143 r <- bef
144 ref <- newIOFinalizer (gc r)
145 return (r, ref)
146 return $ Skip $ GBracketIONormal (fnormal r) r ref
148 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
149 res <- exc $ step1 gst st
150 case res of
151 Right r -> case r of
152 Yield x s ->
153 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
154 Skip s ->
155 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
156 Stop ->
157 clearingIOFinalizer ref (aft v) >> return Stop
158 -- XXX Do not handle async exceptions, just rethrow them.
159 Left e -> do
160 -- Clearing of finalizer and running of exception handler must
161 -- be atomic wrt async exceptions. Otherwise if we have cleared
162 -- the finalizer and have not run the exception handler then we
163 -- may leak the resource.
164 stream <- clearingIOFinalizer ref (fexc v e (UnStream step1 st))
165 return $ Skip (GBracketIOException stream)
166 step gst (GBracketIOException (UnStream step1 st)) = do
167 res <- step1 gst st
168 case res of
169 Yield x s ->
170 return $ Yield x (GBracketIOException (Stream step1 s))
171 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
172 Stop -> return Stop
174-- | See 'Streamly.Internal.Data.Stream.IsStream.before'.
175--
176{-# INLINE_NORMAL before #-}
177before :: Monad m => m b -> Stream m a -> Stream m a
178before action (Stream step state) = Stream step' Nothing
180 where
182 {-# INLINE_LATE step' #-}
183 step' _ Nothing = action >> return (Skip (Just state))
185 step' gst (Just st) = do
186 res <- step gst st
187 case res of
188 Yield x s -> return $ Yield x (Just s)
189 Skip s -> return $ Skip (Just s)
190 Stop -> return Stop
192-- | See 'Streamly.Internal.Data.Stream.IsStream.after_'.
193--
194{-# INLINE_NORMAL after_ #-}
195after_ :: Monad m => m b -> Stream m a -> Stream m a
196after_ action (Stream step state) = Stream step' state
198 where
200 {-# INLINE_LATE step' #-}
201 step' gst st = do
202 res <- step gst st
203 case res of
204 Yield x s -> return $ Yield x s
205 Skip s -> return $ Skip s
206 Stop -> action >> return Stop
208-- | See 'Streamly.Internal.Data.Stream.IsStream.after'.
209--
210{-# INLINE_NORMAL after #-}
211after :: MonadRunInIO m
212 => m b -> Stream m a -> Stream m a
213after action (Stream step state) = Stream step' Nothing
215 where
217 {-# INLINE_LATE step' #-}
218 step' _ Nothing = do
219 ref <- newIOFinalizer action
220 return $ Skip $ Just (state, ref)
221 step' gst (Just (st, ref)) = do
222 res <- step gst st
223 case res of
224 Yield x s -> return $ Yield x (Just (s, ref))
225 Skip s -> return $ Skip (Just (s, ref))
226 Stop -> do
227 runIOFinalizer ref
228 return Stop
230-- XXX For high performance error checks in busy streams we may need another
231-- Error constructor in step.
232--
233-- | See 'Streamly.Internal.Data.Stream.IsStream.onException'.
234--
235{-# INLINE_NORMAL onException #-}
236onException :: MonadCatch m => m b -> Stream m a -> Stream m a
237onException action str =
238 gbracket_ (return ()) (inline MC.try) return
239 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
240 (const str)
242{-# INLINE_NORMAL _onException #-}
243_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
244_onException action (Stream step state) = Stream step' state
246 where
248 {-# INLINE_LATE step' #-}
249 step' gst st = do
250 res <- step gst st `MC.onException` action
251 case res of
252 Yield x s -> return $ Yield x s
253 Skip s -> return $ Skip s
254 Stop -> return Stop
256-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket_'.
257--
258{-# INLINE_NORMAL bracket_ #-}
259bracket_ :: MonadCatch m
260 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
261bracket_ bef aft =
262 gbracket_ bef (inline MC.try) aft
263 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e))
265-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket'.
266--
267{-# INLINE_NORMAL bracket' #-}
268bracket' :: (MonadAsync m, MonadCatch m) =>
269 m b
270 -> (b -> m c)
271 -> (b -> m d)
272 -> (b -> m e)
273 -> (b -> Stream m a)
274 -> Stream m a
275bracket' bef aft exc gc =
276 gbracket bef (inline MC.try) aft gc
277 (\a (e :: SomeException) _ -> exc a >> return (nilM (MC.throwM e)))
279data BracketState s v = BracketInit | BracketRun s v
281-- | Alternate (custom) implementation of 'bracket'.
282--
283{-# INLINE_NORMAL _bracket #-}
284_bracket :: MonadCatch m
285 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
286_bracket bef aft bet = Stream step' BracketInit
288 where
290 {-# INLINE_LATE step' #-}
291 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
293 -- NOTE: It is important to use UnStream instead of the Stream pattern
294 -- here, otherwise we get huge perf degradation, see note in concatMap.
295 step' gst (BracketRun (UnStream step state) v) = do
296 -- res <- step gst state `MC.onException` aft v
297 res <- inline MC.try $ step gst state
298 case res of
299 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
300 Right r -> case r of
301 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
302 Skip s -> return $ Skip (BracketRun (Stream step s) v)
303 Stop -> aft v >> return Stop
305-- | See 'Streamly.Internal.Data.Stream.IsStream.finally_'.
306--
307{-# INLINE finally_ #-}
308finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
309finally_ action xs = bracket_ (return ()) (const action) (const xs)
311-- | See 'Streamly.Internal.Data.Stream.IsStream.finally'.
312--
313-- finally action xs = after action $ onException action xs
314--
315{-# INLINE finally #-}
316finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
317finally action xs = bracket' (return ()) act act act (const xs)
318 where act _ = action
320-- | See 'Streamly.Internal.Data.Stream.IsStream.ghandle'.
321--
322{-# INLINE_NORMAL ghandle #-}
323ghandle :: (MonadCatch m, Exception e)
324 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
325ghandle f str =
326 gbracket_ (return ()) (inline MC.try) return (const f) (const str)
328-- | See 'Streamly.Internal.Data.Stream.IsStream.handle'.
329--
330{-# INLINE_NORMAL handle #-}
331handle :: (MonadCatch m, Exception e)
332 => (e -> Stream m a) -> Stream m a -> Stream m a
333handle f str =
334 gbracket_ (return ()) (inline MC.try) return (\_ e _ -> f e) (const str)
336-- | Alternate (custom) implementation of 'handle'.
337--
338{-# INLINE_NORMAL _handle #-}
339_handle :: (MonadCatch m, Exception e)
340 => (e -> Stream m a) -> Stream m a -> Stream m a
341_handle f (Stream step state) = Stream step' (Left state)
343 where
345 {-# INLINE_LATE step' #-}
346 step' gst (Left st) = do
347 res <- inline MC.try $ step gst st
348 case res of
349 Left e -> return $ Skip $ Right (f e)
350 Right r -> case r of
351 Yield x s -> return $ Yield x (Left s)
352 Skip s -> return $ Skip (Left s)
353 Stop -> return Stop
355 step' gst (Right (UnStream step1 st)) = do
356 res <- step1 gst st
357 case res of
358 Yield x s -> return $ Yield x (Right (Stream step1 s))
359 Skip s -> return $ Skip (Right (Stream step1 s))
360 Stop -> return Stop
362data RetryState emap s1 s2
363 = RetryWithMap emap s1
364 | RetryDefault s2
366-- | See 'Streamly.Internal.Data.Stream.IsStream.retry'
367--
368{-# INLINE_NORMAL retry #-}
369retry
370 :: forall e m a. (Exception e, Ord e, MonadCatch m)
371 => Map e Int
372 -- ^ map from exception to retry count
373 -> (e -> Stream m a)
374 -- ^ default handler for those exceptions that are not in the map
375 -> Stream m a
376 -> Stream m a
377retry emap0 defaultHandler (Stream step0 state0) = Stream step state
379 where
381 state = RetryWithMap emap0 state0
383 {-# INLINE_LATE step #-}
384 step gst (RetryWithMap emap st) = do
385 eres <- MC.try $ step0 gst st
386 case eres of
387 Left e -> handler e emap st
388 Right res ->
389 return
390 $ case res of
391 Yield x st1 -> Yield x $ RetryWithMap emap st1
392 Skip st1 -> Skip $ RetryWithMap emap st1
393 Stop -> Stop
394 step gst (RetryDefault (UnStream step1 state1)) = do
395 res <- step1 gst state1
396 return
397 $ case res of
398 Yield x st1 -> Yield x $ RetryDefault (Stream step1 st1)
399 Skip st1 -> Skip $ RetryDefault (Stream step1 st1)
400 Stop -> Stop
402 {-# INLINE handler #-}
403 handler e emap st =
404 return
405 $ Skip
406 $ case Map.lookup e emap of
407 Just i
408 | i > 0 ->
409 let emap1 = Map.insert e (i - 1) emap
410 in RetryWithMap emap1 st
411 | otherwise -> RetryDefault $ defaultHandler e
412 Nothing -> RetryDefault $ defaultHandler e