git.iko.soy/streamly
on 0717458 help sign in
streamly/core/src/Streamly/Internal/Data/Stream/StreamD/Exception.hs
423 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.Monad.IO.Class (MonadIO(..))
30import Control.Exception (Exception, SomeException, mask_)
31import Control.Monad.Catch (MonadCatch)
32import Data.Map.Strict (Map)
33import GHC.Exts (inline)
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--
53-- /Inhibits stream fusion/
54--
55-- /Pre-release/
56--
57{-# INLINE_NORMAL gbracket_ #-}
58gbracket_
59 :: Monad m
60 => m c -- ^ before
61 -> (c -> m d) -- ^ after, on normal stop
62 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
63 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
64 -> (c -> Stream m b) -- ^ stream generator
65 -> Stream m b
66gbracket_ bef aft onExc ftry action =
67 Stream step GBracketInit
69 where
71 {-# INLINE_LATE step #-}
72 step _ GBracketInit = do
73 r <- bef
74 return $ Skip $ GBracketNormal (action r) r
76 step gst (GBracketNormal (UnStream step1 st) v) = do
77 res <- ftry $ step1 gst st
78 case res of
79 Right r -> case r of
80 Yield x s ->
81 return $ Yield x (GBracketNormal (Stream step1 s) v)
82 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
83 Stop -> aft v >> return Stop
84 -- XXX Do not handle async exceptions, just rethrow them.
85 Left e ->
86 return
87 $ Skip (GBracketException (onExc 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 :: MonadIO m
121 => IO c -- ^ before
122 -> (c -> IO d1) -- ^ on normal stop
123 -> (c -> e -> Stream m b -> IO (Stream m b)) -- ^ on exception
124 -> (c -> IO d2) -- ^ on GC without normal stop or exception
125 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
126 -> (c -> Stream m b) -- ^ stream generator
127 -> Stream m b
128gbracket bef aft onExc onGC ftry action =
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) <- liftIO $ mask_ $ do
143 r <- bef
144 ref <- newIOFinalizer (onGC r)
145 return (r, ref)
146 return $ Skip $ GBracketIONormal (action r) r ref
148 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
149 res <- ftry $ 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 liftIO (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 <-
165 liftIO (clearingIOFinalizer ref (onExc v e (UnStream step1 st)))
166 return $ Skip (GBracketIOException stream)
167 step gst (GBracketIOException (UnStream step1 st)) = do
168 res <- step1 gst st
169 case res of
170 Yield x s ->
171 return $ Yield x (GBracketIOException (Stream step1 s))
172 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
173 Stop -> return Stop
175-- | See 'Streamly.Internal.Data.Stream.before'.
176--
177{-# INLINE_NORMAL before #-}
178before :: Monad m => m b -> Stream m a -> Stream m a
179before action (Stream step state) = Stream step' Nothing
181 where
183 {-# INLINE_LATE step' #-}
184 step' _ Nothing = action >> return (Skip (Just state))
186 step' gst (Just st) = do
187 res <- step gst st
188 case res of
189 Yield x s -> return $ Yield x (Just s)
190 Skip s -> return $ Skip (Just s)
191 Stop -> return Stop
193-- | See 'Streamly.Internal.Data.Stream.after_'.
194--
195{-# INLINE_NORMAL after_ #-}
196after_ :: Monad m => m b -> Stream m a -> Stream m a
197after_ action (Stream step state) = Stream step' state
199 where
201 {-# INLINE_LATE step' #-}
202 step' gst st = do
203 res <- step gst st
204 case res of
205 Yield x s -> return $ Yield x s
206 Skip s -> return $ Skip s
207 Stop -> action >> return Stop
209-- | See 'Streamly.Internal.Data.Stream.after'.
210--
211{-# INLINE_NORMAL after #-}
212after :: MonadIO m
213 => IO b -> Stream m a -> Stream m a
214after action (Stream step state) = Stream step' Nothing
216 where
218 {-# INLINE_LATE step' #-}
219 step' _ Nothing = do
220 ref <- liftIO $ newIOFinalizer action
221 return $ Skip $ Just (state, ref)
222 step' gst (Just (st, ref)) = do
223 res <- step gst st
224 case res of
225 Yield x s -> return $ Yield x (Just (s, ref))
226 Skip s -> return $ Skip (Just (s, ref))
227 Stop -> do
228 runIOFinalizer ref
229 return Stop
231-- XXX For high performance error checks in busy streams we may need another
232-- Error constructor in step.
233--
234-- | See 'Streamly.Internal.Data.Stream.onException'.
235--
236{-# INLINE_NORMAL onException #-}
237onException :: MonadCatch m => m b -> Stream m a -> Stream m a
238onException action stream =
239 gbracket_
240 (return ()) -- before
241 return -- after
242 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
243 (inline MC.try)
244 (const stream)
246{-# INLINE_NORMAL _onException #-}
247_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
248_onException action (Stream step state) = Stream step' state
250 where
252 {-# INLINE_LATE step' #-}
253 step' gst st = do
254 res <- step gst st `MC.onException` action
255 case res of
256 Yield x s -> return $ Yield x s
257 Skip s -> return $ Skip s
258 Stop -> return Stop
260-- | See 'Streamly.Internal.Data.Stream.bracket_'.
261--
262{-# INLINE_NORMAL bracket_ #-}
263bracket_ :: MonadCatch m
264 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
265bracket_ bef aft =
266 gbracket_
267 bef
268 aft
269 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e))
270 (inline MC.try)
272-- | See 'Streamly.Internal.Data.Stream.bracket'.
273--
274{-# INLINE_NORMAL bracket' #-}
275bracket' :: (MonadIO m, MonadCatch m) =>
276 IO b
277 -> (b -> IO c)
278 -> (b -> IO d)
279 -> (b -> IO e)
280 -> (b -> Stream m a)
281 -> Stream m a
282bracket' bef aft onExc onGC =
283 gbracket
284 bef
285 aft
286 (\a (e :: SomeException) _ -> onExc a >> return (nilM (MC.throwM e)))
287 onGC
288 (inline MC.try)
290data BracketState s v = BracketInit | BracketRun s v
292-- | Alternate (custom) implementation of 'bracket'.
293--
294{-# INLINE_NORMAL _bracket #-}
295_bracket :: MonadCatch m
296 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
297_bracket bef aft bet = Stream step' BracketInit
299 where
301 {-# INLINE_LATE step' #-}
302 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
304 -- NOTE: It is important to use UnStream instead of the Stream pattern
305 -- here, otherwise we get huge perf degradation, see note in concatMap.
306 step' gst (BracketRun (UnStream step state) v) = do
307 -- res <- step gst state `MC.onException` aft v
308 res <- inline MC.try $ step gst state
309 case res of
310 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
311 Right r -> case r of
312 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
313 Skip s -> return $ Skip (BracketRun (Stream step s) v)
314 Stop -> aft v >> return Stop
316-- | See 'Streamly.Internal.Data.Stream.finally_'.
317--
318{-# INLINE finally_ #-}
319finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
320finally_ action xs = bracket_ (return ()) (const action) (const xs)
322-- | See 'Streamly.Internal.Data.Stream.finally'.
323--
324-- finally action xs = after action $ onException action xs
325--
326{-# INLINE finally #-}
327finally :: (MonadIO m, MonadCatch m) => IO b -> Stream m a -> Stream m a
328finally action xs = bracket' (return ()) act act act (const xs)
329 where act _ = action
331-- | See 'Streamly.Internal.Data.Stream.ghandle'.
332--
333{-# INLINE_NORMAL ghandle #-}
334ghandle :: (MonadCatch m, Exception e)
335 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
336ghandle f stream =
337 gbracket_ (return ()) return (const f) (inline MC.try) (const stream)
339-- | See 'Streamly.Internal.Data.Stream.handle'.
340--
341{-# INLINE_NORMAL handle #-}
342handle :: (MonadCatch m, Exception e)
343 => (e -> Stream m a) -> Stream m a -> Stream m a
344handle f stream =
345 gbracket_ (return ()) return (\_ e _ -> f e) (inline MC.try) (const stream)
347-- | Alternate (custom) implementation of 'handle'.
348--
349{-# INLINE_NORMAL _handle #-}
350_handle :: (MonadCatch m, Exception e)
351 => (e -> Stream m a) -> Stream m a -> Stream m a
352_handle f (Stream step state) = Stream step' (Left state)
354 where
356 {-# INLINE_LATE step' #-}
357 step' gst (Left st) = do
358 res <- inline MC.try $ step gst st
359 case res of
360 Left e -> return $ Skip $ Right (f e)
361 Right r -> case r of
362 Yield x s -> return $ Yield x (Left s)
363 Skip s -> return $ Skip (Left s)
364 Stop -> return Stop
366 step' gst (Right (UnStream step1 st)) = do
367 res <- step1 gst st
368 case res of
369 Yield x s -> return $ Yield x (Right (Stream step1 s))
370 Skip s -> return $ Skip (Right (Stream step1 s))
371 Stop -> return Stop
373data RetryState emap s1 s2
374 = RetryWithMap emap s1
375 | RetryDefault s2
377-- | See 'Streamly.Internal.Data.Stream.retry'
378--
379{-# INLINE_NORMAL retry #-}
380retry
381 :: forall e m a. (Exception e, Ord e, MonadCatch m)
382 => Map e Int
383 -- ^ map from exception to retry count
384 -> (e -> Stream m a)
385 -- ^ default handler for those exceptions that are not in the map
386 -> Stream m a
387 -> Stream m a
388retry emap0 defaultHandler (Stream step0 state0) = Stream step state
390 where
392 state = RetryWithMap emap0 state0
394 {-# INLINE_LATE step #-}
395 step gst (RetryWithMap emap st) = do
396 eres <- MC.try $ step0 gst st
397 case eres of
398 Left e -> handler e emap st
399 Right res ->
400 return
401 $ case res of
402 Yield x st1 -> Yield x $ RetryWithMap emap st1
403 Skip st1 -> Skip $ RetryWithMap emap st1
404 Stop -> Stop
405 step gst (RetryDefault (UnStream step1 state1)) = do
406 res <- step1 gst state1
407 return
408 $ case res of
409 Yield x st1 -> Yield x $ RetryDefault (Stream step1 st1)
410 Skip st1 -> Skip $ RetryDefault (Stream step1 st1)
411 Stop -> Stop
413 {-# INLINE handler #-}
414 handler e emap st =
415 return
416 $ Skip
417 $ case Map.lookup e emap of
418 Just i
419 | i > 0 ->
420 let emap1 = Map.insert e (i - 1) emap
421 in RetryWithMap emap1 st
422 | otherwise -> RetryDefault $ defaultHandler e
423 Nothing -> RetryDefault $ defaultHandler e