git.iko.soy/streamly
on a296fac help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamD/Exception.hs
414 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 Control.Monad.IO.Class (MonadIO(..))
32import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
33import Data.Map.Strict (Map)
34import GHC.Exts (inline)
35import Streamly.Internal.Control.Concurrent (MonadAsync)
36import Streamly.Internal.Data.IOFinalizer
37 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
39import qualified Control.Monad.Catch as MC
40import qualified Data.Map.Strict as Map
42import Streamly.Internal.Data.Stream.StreamD.Type
44data GbracketState s1 s2 v
45 = GBracketInit
46 | GBracketNormal s1 v
47 | GBracketException s2
49-- | Like 'gbracket' but with following differences:
50--
51-- * alloc action @m c@ runs with async exceptions enabled
52-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
53-- after partial evaluation.
54-- * does not require a 'MonadAsync' constraint.
55--
56-- /Inhibits stream fusion/
57--
58-- /Pre-release/
59--
60{-# INLINE_NORMAL gbracket_ #-}
61gbracket_
62 :: Monad m
63 => m c -- ^ before
64 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
65 -> (c -> m d) -- ^ after, on normal stop
66 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
67 -> (c -> Stream m b) -- ^ stream generator
68 -> Stream m b
69gbracket_ bef exc aft fexc fnormal =
70 Stream step GBracketInit
72 where
74 {-# INLINE_LATE step #-}
75 step _ GBracketInit = do
76 r <- bef
77 return $ Skip $ GBracketNormal (fnormal r) r
79 step gst (GBracketNormal (UnStream step1 st) v) = do
80 res <- exc $ step1 gst st
81 case res of
82 Right r -> case r of
83 Yield x s ->
84 return $ Yield x (GBracketNormal (Stream step1 s) v)
85 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
86 Stop -> aft v >> return Stop
87 -- XXX Do not handle async exceptions, just rethrow them.
88 Left e ->
89 return $ Skip (GBracketException (fexc v e (UnStream step1 st)))
90 step gst (GBracketException (UnStream step1 st)) = do
91 res <- step1 gst st
92 case res of
93 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
94 Skip s -> return $ Skip (GBracketException (Stream step1 s))
95 Stop -> return Stop
97data GbracketIOState s1 s2 v wref
98 = GBracketIOInit
99 | GBracketIONormal s1 v wref
100 | GBracketIOException s2
102-- | Run the alloc action @m c@ with async exceptions disabled but keeping
103-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
104-- output @c@ as input to @c -> Stream m b@ to generate an output stream. When
105-- generating the stream use the supplied @try@ operation @forall s. m s -> m
106-- (Either e s)@ to catch synchronous exceptions. If an exception occurs run
107-- the exception handler @c -> e -> Stream m b -> m (Stream m b)@. Note that
108-- 'gbracket' does not rethrow the exception, it has to be done by the
109-- exception handler if desired.
110--
111-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
112-- to a sync or async exception or if it gets garbage collected after a partial
113-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
114--
115-- 'gbracket' can express all other exception handling combinators.
116--
117-- /Inhibits stream fusion/
118--
119-- /Pre-release/
120{-# INLINE_NORMAL gbracket #-}
121gbracket
122 :: (MonadIO m, MonadBaseControl IO m)
123 => m c -- ^ before
124 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
125 -> (c -> m d1) -- ^ on normal stop
126 -> (c -> m d2) -- ^ on GC without normal stop or exception
127 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
128 -> (c -> Stream m b) -- ^ stream generator
129 -> Stream m b
130gbracket bef exc aft gc fexc fnormal =
131 Stream step GBracketIOInit
133 where
135 -- If the stream is never evaluated the "aft" action will never be
136 -- called. For that to occur we will need the user of this API to pass a
137 -- weak pointer to us.
138 {-# INLINE_LATE step #-}
139 step _ GBracketIOInit = do
140 -- We mask asynchronous exceptions to make the execution
141 -- of 'bef' and the registration of 'aft' atomic.
142 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
143 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
144 (r, ref) <- liftBaseOp_ mask_ $ do
145 r <- bef
146 ref <- newIOFinalizer (gc r)
147 return (r, ref)
148 return $ Skip $ GBracketIONormal (fnormal r) r ref
150 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
151 res <- exc $ step1 gst st
152 case res of
153 Right r -> case r of
154 Yield x s ->
155 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
156 Skip s ->
157 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
158 Stop ->
159 clearingIOFinalizer ref (aft v) >> return Stop
160 -- XXX Do not handle async exceptions, just rethrow them.
161 Left e -> do
162 -- Clearing of finalizer and running of exception handler must
163 -- be atomic wrt async exceptions. Otherwise if we have cleared
164 -- the finalizer and have not run the exception handler then we
165 -- may leak the resource.
166 stream <- clearingIOFinalizer ref (fexc 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 :: (MonadIO m, MonadBaseControl IO 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 str =
240 gbracket_ (return ()) (inline MC.try) return
241 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
242 (\_ -> str)
244{-# INLINE_NORMAL _onException #-}
245_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
246_onException action (Stream step state) = Stream step' state
248 where
250 {-# INLINE_LATE step' #-}
251 step' gst st = do
252 res <- step gst st `MC.onException` action
253 case res of
254 Yield x s -> return $ Yield x s
255 Skip s -> return $ Skip s
256 Stop -> return Stop
258-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket_'.
259--
260{-# INLINE_NORMAL bracket_ #-}
261bracket_ :: MonadCatch m
262 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
263bracket_ bef aft bet =
264 gbracket_ bef (inline MC.try) aft
265 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e)) bet
267-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket'.
268--
269{-# INLINE_NORMAL bracket' #-}
270bracket' :: (MonadAsync m, MonadCatch m) =>
271 m b
272 -> (b -> m c)
273 -> (b -> m d)
274 -> (b -> m e)
275 -> (b -> Stream m a)
276 -> Stream m a
277bracket' bef aft exc gc =
278 gbracket bef (inline MC.try) aft gc
279 (\a (e :: SomeException) _ -> exc a >> return (nilM (MC.throwM e)))
281data BracketState s v = BracketInit | BracketRun s v
283-- | Alternate (custom) implementation of 'bracket'.
284--
285{-# INLINE_NORMAL _bracket #-}
286_bracket :: MonadCatch m
287 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
288_bracket bef aft bet = Stream step' BracketInit
290 where
292 {-# INLINE_LATE step' #-}
293 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
295 -- NOTE: It is important to use UnStream instead of the Stream pattern
296 -- here, otherwise we get huge perf degradation, see note in concatMap.
297 step' gst (BracketRun (UnStream step state) v) = do
298 -- res <- step gst state `MC.onException` aft v
299 res <- inline MC.try $ step gst state
300 case res of
301 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
302 Right r -> case r of
303 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
304 Skip s -> return $ Skip (BracketRun (Stream step s) v)
305 Stop -> aft v >> return Stop
307-- | See 'Streamly.Internal.Data.Stream.IsStream.finally_'.
308--
309{-# INLINE finally_ #-}
310finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
311finally_ action xs = bracket_ (return ()) (\_ -> action) (const xs)
313-- | See 'Streamly.Internal.Data.Stream.IsStream.finally'.
314--
315-- finally action xs = after action $ onException action xs
316--
317{-# INLINE finally #-}
318finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
319finally action xs = bracket' (return ()) act act act (const xs)
320 where act _ = action
322-- | See 'Streamly.Internal.Data.Stream.IsStream.ghandle'.
323--
324{-# INLINE_NORMAL ghandle #-}
325ghandle :: (MonadCatch m, Exception e)
326 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
327ghandle f str =
328 gbracket_ (return ()) (inline MC.try) return (\_ -> f) (\_ -> str)
330-- | See 'Streamly.Internal.Data.Stream.IsStream.handle'.
331--
332{-# INLINE_NORMAL handle #-}
333handle :: (MonadCatch m, Exception e)
334 => (e -> Stream m a) -> Stream m a -> Stream m a
335handle f str =
336 gbracket_ (return ()) (inline MC.try) return (\_ e _ -> f e) (\_ -> str)
338-- | Alternate (custom) implementation of 'handle'.
339--
340{-# INLINE_NORMAL _handle #-}
341_handle :: (MonadCatch m, Exception e)
342 => (e -> Stream m a) -> Stream m a -> Stream m a
343_handle f (Stream step state) = Stream step' (Left state)
345 where
347 {-# INLINE_LATE step' #-}
348 step' gst (Left st) = do
349 res <- inline MC.try $ step gst st
350 case res of
351 Left e -> return $ Skip $ Right (f e)
352 Right r -> case r of
353 Yield x s -> return $ Yield x (Left s)
354 Skip s -> return $ Skip (Left s)
355 Stop -> return Stop
357 step' gst (Right (UnStream step1 st)) = do
358 res <- step1 gst st
359 case res of
360 Yield x s -> return $ Yield x (Right (Stream step1 s))
361 Skip s -> return $ Skip (Right (Stream step1 s))
362 Stop -> return Stop
364data RetryState emap s1 s2
365 = RetryWithMap emap s1
366 | RetryDefault s2
368-- | See 'Streamly.Internal.Data.Stream.IsStream.retry'
369--
370{-# INLINE_NORMAL retry #-}
371retry
372 :: forall e m a. (Exception e, Ord e, MonadCatch m)
373 => Map e Int
374 -- ^ map from exception to retry count
375 -> (e -> Stream m a)
376 -- ^ default handler for those exceptions that are not in the map
377 -> Stream m a
378 -> Stream m a
379retry emap0 defaultHandler (Stream step0 state0) = Stream step state
381 where
383 state = RetryWithMap emap0 state0
385 {-# INLINE_LATE step #-}
386 step gst (RetryWithMap emap st) = do
387 eres <- MC.try $ step0 gst st
388 case eres of
389 Left e -> handler e emap st
390 Right res ->
391 return
392 $ case res of
393 Yield x st1 -> Yield x $ RetryWithMap emap st1
394 Skip st1 -> Skip $ RetryWithMap emap st1
395 Stop -> Stop
396 step gst (RetryDefault (UnStream step1 state1)) = do
397 res <- step1 gst state1
398 return
399 $ case res of
400 Yield x st1 -> Yield x $ RetryDefault (Stream step1 st1)
401 Skip st1 -> Skip $ RetryDefault (Stream step1 st1)
402 Stop -> Stop
404 {-# INLINE handler #-}
405 handler e emap st =
406 return
407 $ Skip
408 $ case Map.lookup e emap of
409 Just i
410 | i > 0 ->
411 let emap1 = Map.insert e (i - 1) emap
412 in RetryWithMap emap1 st
413 | otherwise -> RetryDefault $ defaultHandler e
414 Nothing -> RetryDefault $ defaultHandler e