on 084cdea help sign in
425 lines 15 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 -- * Finalizers
12 newFinalizedIORef
13 , runIORefFinalizer
14 , withIORefFinalizer
16 -- * Combinators
17 , gbracket_
18 , gbracket
19 , before
20 , after_
21 , after
22 , bracket_
23 , bracket
24 , onException
25 , finally_
26 , finally
27 , ghandle
28 , handle
29 )
30where
32#include "inline.hs"
34import Control.Exception (Exception, SomeException, mask_)
35import Control.Monad (void)
36import Control.Monad.Catch (MonadCatch)
37import Control.Monad.IO.Class (MonadIO(..))
38import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)
39import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
41import qualified Control.Monad.Catch as MC
43import Streamly.Internal.Data.Stream.StreamD.Type
44import Streamly.Internal.Data.SVar
46------------------------------------------------------------------------------
47-- Finalizers
48------------------------------------------------------------------------------
50-- | Make a finalizer from a monadic action @m a@ that can run in IO monad.
51mkIOFinalizer :: MonadBaseControl IO m => m b -> m (IO ())
52mkIOFinalizer f = do
53 mrun <- captureMonadState
54 return $
55 void $ do
56 _ <- runInIO mrun f
57 return ()
59-- | Run an IO action stored in a finalized IORef.
60runIORefFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
61runIORefFinalizerGC ref = do
62 res <- readIORef ref
63 case res of
64 Nothing -> return ()
65 Just f -> f
67-- | Create an IORef holding a finalizer that is called automatically when
68-- the IORef is garbage collected. The IORef can be written to with a 'Nothing'
69-- value to deactivate the finalizer.
70--
71-- Note: The finalizer is always run with the state of the monad captured at
72-- the time of calling newFinalizedIORef. To run it on garbage collection we
73-- have no option but to take a snapshot of the monadic state at some point of
74-- time. For normal case we could run it with the current state of the monad
75-- but we want to keep both the cases consistent.
76--
77newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
78 => m a -> m (IORef (Maybe (IO ())))
79newFinalizedIORef finalizer = do
80 f <- mkIOFinalizer finalizer
81 ref <- liftIO $ newIORef $ Just f
82 _ <- liftIO $ mkWeakIORef ref (runIORefFinalizerGC ref)
83 return ref
85-- | Run the finalizer stored in an IORef and deactivate it so that it is run
86-- only once. Note, the action runs with async exceptions masked.
87--
88runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
89runIORefFinalizer ref = liftIO $ do
90 res <- readIORef ref
91 case res of
92 Nothing -> return ()
93 Just action -> do
94 -- if an async exception comes after writing 'Nothing' then the
95 -- finalizing action will never be run. We need to do this
96 -- atomically wrt async exceptions.
97 mask_ $ do
98 writeIORef ref Nothing
99 action
101-- | Run an action clearing the finalizer IORef atomically wrt async
102-- exceptions. The action is run with async exceptions masked.
103withIORefFinalizer :: MonadBaseControl IO m
104 => IORef (Maybe (IO ())) -> m a -> m a
105withIORefFinalizer ref action = do
106 control $ \runinio ->
107 mask_ $ do
108 writeIORef ref Nothing
109 runinio action
111------------------------------------------------------------------------------
112-- Exception handling combinators.
113------------------------------------------------------------------------------
115data GbracketState s1 s2 v
116 = GBracketInit
117 | GBracketNormal s1 v
118 | GBracketException s2
120-- | Like 'gbracket' but with following differences:
121--
122-- * alloc action @m c@ runs with async exceptions enabled
123-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
124-- after partial evaluation.
125-- * does not require a 'MonadAsync' constraint.
126--
127-- /Inhibits stream fusion/
128--
129-- /Internal/
130--
131{-# INLINE_NORMAL gbracket_ #-}
132gbracket_
133 :: Monad m
134 => m c -- ^ before
135 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
136 -> (c -> m d) -- ^ after, on normal stop
137 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
138 -> (c -> Stream m b) -- ^ stream generator
139 -> Stream m b
140gbracket_ bef exc aft fexc fnormal =
141 Stream step GBracketInit
143 where
145 {-# INLINE_LATE step #-}
146 step _ GBracketInit = do
147 r <- bef
148 return $ Skip $ GBracketNormal (fnormal r) r
150 step gst (GBracketNormal (UnStream step1 st) v) = 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 (GBracketNormal (Stream step1 s) v)
156 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
157 Stop -> aft v >> return Stop
158 -- XXX Do not handle async exceptions, just rethrow them.
159 Left e ->
160 return $ Skip (GBracketException (fexc v e (UnStream step1 st)))
161 step gst (GBracketException (UnStream step1 st)) = do
162 res <- step1 gst st
163 case res of
164 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
165 Skip s -> return $ Skip (GBracketException (Stream step1 s))
166 Stop -> return Stop
168data GbracketIOState s1 s2 v wref
169 = GBracketIOInit
170 | GBracketIONormal s1 v wref
171 | GBracketIOException s2
173-- | Run the alloc action @m c@ with async exceptions disabled but keeping
174-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
175-- output @c@ as input to @c -> Stream m b@ to generate an output stream. When
176-- generating the stream use the supplied @try@ operation @forall s. m s -> m
177-- (Either e s)@ to catch synchronous exceptions. If an exception occurs run
178-- the exception handler @c -> e -> Stream m b -> m (Stream m b)@.
179--
180-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
181-- to a sync or async exception or if it gets garbage collected after a partial
182-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
183--
184-- 'gbracket' can express all other exception handling combinators.
185--
186-- /Inhibits stream fusion/
187--
188-- /Internal/
189{-# INLINE_NORMAL gbracket #-}
190gbracket
191 :: (MonadIO m, MonadBaseControl IO m)
192 => m c -- ^ before
193 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
194 -> (c -> m d) -- ^ after, on normal stop or GC
195 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
196 -> (c -> Stream m b) -- ^ stream generator
197 -> Stream m b
198gbracket bef exc aft fexc fnormal =
199 Stream step GBracketIOInit
201 where
203 -- If the stream is never evaluated the "aft" action will never be
204 -- called. For that to occur we will need the user of this API to pass a
205 -- weak pointer to us.
206 {-# INLINE_LATE step #-}
207 step _ GBracketIOInit = do
208 -- We mask asynchronous exceptions to make the execution
209 -- of 'bef' and the registration of 'aft' atomic.
210 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
211 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
212 (r, ref) <- liftBaseOp_ mask_ $ do
213 r <- bef
214 ref <- newFinalizedIORef (aft r)
215 return (r, ref)
216 return $ Skip $ GBracketIONormal (fnormal r) r ref
218 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
219 res <- exc $ step1 gst st
220 case res of
221 Right r -> case r of
222 Yield x s ->
223 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
224 Skip s ->
225 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
226 Stop -> do
227 runIORefFinalizer ref
228 return Stop
229 -- XXX Do not handle async exceptions, just rethrow them.
230 Left e -> do
231 -- Clearing of finalizer and running of exception handler must
232 -- be atomic wrt async exceptions. Otherwise if we have cleared
233 -- the finalizer and have not run the exception handler then we
234 -- may leak the resource.
235 stream <- withIORefFinalizer ref (fexc v e (UnStream step1 st))
236 return $ Skip (GBracketIOException stream)
237 step gst (GBracketIOException (UnStream step1 st)) = do
238 res <- step1 gst st
239 case res of
240 Yield x s ->
241 return $ Yield x (GBracketIOException (Stream step1 s))
242 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
243 Stop -> return Stop
245-- | See 'Streamly.Internal.Data.Stream.IsStream.before'.
246--
247{-# INLINE_NORMAL before #-}
248before :: Monad m => m b -> Stream m a -> Stream m a
249before action (Stream step state) = Stream step' Nothing
251 where
253 {-# INLINE_LATE step' #-}
254 step' _ Nothing = action >> return (Skip (Just state))
256 step' gst (Just st) = do
257 res <- step gst st
258 case res of
259 Yield x s -> return $ Yield x (Just s)
260 Skip s -> return $ Skip (Just s)
261 Stop -> return Stop
263-- | See 'Streamly.Internal.Data.Stream.IsStream.after_'.
264--
265{-# INLINE_NORMAL after_ #-}
266after_ :: Monad m => m b -> Stream m a -> Stream m a
267after_ action (Stream step state) = Stream step' state
269 where
271 {-# INLINE_LATE step' #-}
272 step' gst st = do
273 res <- step gst st
274 case res of
275 Yield x s -> return $ Yield x s
276 Skip s -> return $ Skip s
277 Stop -> action >> return Stop
279-- | See 'Streamly.Internal.Data.Stream.IsStream.after'.
280--
281{-# INLINE_NORMAL after #-}
282after :: (MonadIO m, MonadBaseControl IO m)
283 => m b -> Stream m a -> Stream m a
284after action (Stream step state) = Stream step' Nothing
286 where
288 {-# INLINE_LATE step' #-}
289 step' _ Nothing = do
290 ref <- newFinalizedIORef action
291 return $ Skip $ Just (state, ref)
292 step' gst (Just (st, ref)) = do
293 res <- step gst st
294 case res of
295 Yield x s -> return $ Yield x (Just (s, ref))
296 Skip s -> return $ Skip (Just (s, ref))
297 Stop -> do
298 runIORefFinalizer ref
299 return Stop
301-- XXX For high performance error checks in busy streams we may need another
302-- Error constructor in step.
303--
304-- | See 'Streamly.Internal.Data.Stream.IsStream.onException'.
305--
306{-# INLINE_NORMAL onException #-}
307onException :: MonadCatch m => m b -> Stream m a -> Stream m a
308onException action str =
309 gbracket_ (return ()) MC.try return
310 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
311 (\_ -> str)
313{-# INLINE_NORMAL _onException #-}
314_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
315_onException action (Stream step state) = Stream step' state
317 where
319 {-# INLINE_LATE step' #-}
320 step' gst st = do
321 res <- step gst st `MC.onException` action
322 case res of
323 Yield x s -> return $ Yield x s
324 Skip s -> return $ Skip s
325 Stop -> return Stop
327-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket_'.
328--
329{-# INLINE_NORMAL bracket_ #-}
330bracket_ :: MonadCatch m
331 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
332bracket_ bef aft bet =
333 gbracket_ bef MC.try aft
334 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e)) bet
336-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket'.
337--
338{-# INLINE_NORMAL bracket #-}
339bracket :: (MonadAsync m, MonadCatch m)
340 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
341bracket bef aft bet =
342 gbracket bef MC.try aft
343 (\a (e :: SomeException) _ -> aft a >> return (nilM (MC.throwM e))) bet
345data BracketState s v = BracketInit | BracketRun s v
347-- | Alternate (custom) implementation of 'bracket'.
348--
349{-# INLINE_NORMAL _bracket #-}
350_bracket :: MonadCatch m
351 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
352_bracket bef aft bet = Stream step' BracketInit
354 where
356 {-# INLINE_LATE step' #-}
357 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
359 -- NOTE: It is important to use UnStream instead of the Stream pattern
360 -- here, otherwise we get huge perf degradation, see note in concatMap.
361 step' gst (BracketRun (UnStream step state) v) = do
362 -- res <- step gst state `MC.onException` aft v
363 res <- MC.try $ step gst state
364 case res of
365 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
366 Right r -> case r of
367 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
368 Skip s -> return $ Skip (BracketRun (Stream step s) v)
369 Stop -> aft v >> return Stop
371-- | See 'Streamly.Internal.Data.Stream.IsStream.finally_'.
372--
373{-# INLINE finally_ #-}
374finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
375finally_ action xs = bracket_ (return ()) (\_ -> action) (const xs)
377-- | See 'Streamly.Internal.Data.Stream.IsStream.finally'.
378--
379-- finally action xs = after action $ onException action xs
380--
381{-# INLINE finally #-}
382finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
383finally action xs = bracket (return ()) (\_ -> action) (const xs)
385-- | See 'Streamly.Internal.Data.Stream.IsStream.ghandle'.
386--
387{-# INLINE_NORMAL ghandle #-}
388ghandle :: (MonadCatch m, Exception e)
389 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
390ghandle f str =
391 gbracket_ (return ()) MC.try return (\_ -> f) (\_ -> str)
393-- | See 'Streamly.Internal.Data.Stream.IsStream.handle'.
394--
395{-# INLINE_NORMAL handle #-}
396handle :: (MonadCatch m, Exception e)
397 => (e -> Stream m a) -> Stream m a -> Stream m a
398handle f str =
399 gbracket_ (return ()) MC.try return (\_ e _ -> f e) (\_ -> str)
401-- | Alternate (custom) implementation of 'handle'.
402--
403{-# INLINE_NORMAL _handle #-}
404_handle :: (MonadCatch m, Exception e)
405 => (e -> Stream m a) -> Stream m a -> Stream m a
406_handle f (Stream step state) = Stream step' (Left state)
408 where
410 {-# INLINE_LATE step' #-}
411 step' gst (Left st) = do
412 res <- MC.try $ step gst st
413 case res of
414 Left e -> return $ Skip $ Right (f e)
415 Right r -> case r of
416 Yield x s -> return $ Yield x (Left s)
417 Skip s -> return $ Skip (Left s)
418 Stop -> return Stop
420 step' gst (Right (UnStream step1 st)) = do
421 res <- step1 gst st
422 case res of
423 Yield x s -> return $ Yield x (Right (Stream step1 s))
424 Skip s -> return $ Skip (Right (Stream step1 s))
425 Stop -> return Stop