git.iko.soy/streamly
on e01057f help sign in
streamly/src/Streamly/Internal/Data/Stream/Exception/Lifted.hs
331 lines 11 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.Exception.Lifted
3-- Copyright : (c) 2019 Composewell Technologies
4-- License : BSD-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
9module Streamly.Internal.Data.Stream.Exception.Lifted
10 (
11 after
12 , bracket
13 , bracket3
14 , finally
15 , retry
17 -- For IsStream module
18 , afterD
19 , bracket3D
20 , retryD
21 )
22where
24#include "inline.hs"
26import Control.Exception (Exception, SomeException, mask_)
27import Control.Monad.Catch (MonadCatch)
28#ifdef USE_UNLIFTIO
29import Control.Monad.IO.Unlift (MonadUnliftIO)
30#else
31import Control.Monad.IO.Class (MonadIO)
32import Control.Monad.Trans.Control (MonadBaseControl)
33#endif
34import Data.Map.Strict (Map)
35import GHC.Exts (inline)
36import Streamly.Internal.Control.Concurrent
37 (MonadRunInIO, MonadAsync, withRunInIO)
38import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
39import Streamly.Internal.Data.IOFinalizer.Lifted
40 (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
41import Streamly.Internal.Data.Stream.StreamD (Step(..))
43import qualified Control.Monad.Catch as MC
44import qualified Data.Map.Strict as Map
45import qualified Streamly.Internal.Data.Stream.StreamD as D
47-- $setup
48-- >>> :m
49-- >>> import qualified Streamly.Internal.Data.Stream.Exception.Lifted as Stream
51-- XXX Implement in terms of the corresponding IO operation (gbracketIO).
53data GbracketIOState s1 s2 v wref
54 = GBracketIOInit
55 | GBracketIONormal s1 v wref
56 | GBracketIOException s2
58{-# INLINE_NORMAL gbracket #-}
59gbracket
60 :: MonadRunInIO m
61 => m c -- ^ before
62 -> (c -> m d1) -- ^ on normal stop
63 -> (c -> e -> D.Stream m b -> m (D.Stream m b)) -- ^ on exception
64 -> (c -> m d2) -- ^ on GC without normal stop or exception
65 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
66 -> (c -> D.Stream m b) -- ^ stream generator
67 -> D.Stream m b
68gbracket bef aft onExc onGC ftry action =
69 D.Stream step GBracketIOInit
71 where
73 -- If the stream is never evaluated the "aft" action will never be
74 -- called. For that to occur we will need the user of this API to pass a
75 -- weak pointer to us.
76 {-# INLINE_LATE step #-}
77 step _ GBracketIOInit = do
78 -- We mask asynchronous exceptions to make the execution
79 -- of 'bef' and the registration of 'aft' atomic.
80 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
81 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
82 (r, ref) <- withRunInIO $ \run -> mask_ $ run $ do
83 r <- bef
84 ref <- newIOFinalizer (onGC r)
85 return (r, ref)
86 return $ Skip $ GBracketIONormal (action r) r ref
88 step gst (GBracketIONormal (D.UnStream step1 st) v ref) = do
89 res <- ftry $ step1 gst st
90 case res of
91 Right r -> case r of
92 Yield x s ->
93 return $ Yield x (GBracketIONormal (D.Stream step1 s) v ref)
94 Skip s ->
95 return $ Skip (GBracketIONormal (D.Stream step1 s) v ref)
96 Stop ->
97 clearingIOFinalizer ref (aft v) >> return Stop
98 -- XXX Do not handle async exceptions, just rethrow them.
99 Left e -> do
100 -- Clearing of finalizer and running of exception handler must
101 -- be atomic wrt async exceptions. Otherwise if we have cleared
102 -- the finalizer and have not run the exception handler then we
103 -- may leak the resource.
104 stream <-
105 clearingIOFinalizer ref (onExc v e (D.UnStream step1 st))
106 return $ Skip (GBracketIOException stream)
107 step gst (GBracketIOException (D.UnStream step1 st)) = do
108 res <- step1 gst st
109 case res of
110 Yield x s ->
111 return $ Yield x (GBracketIOException (D.Stream step1 s))
112 Skip s -> return $ Skip (GBracketIOException (D.Stream step1 s))
113 Stop -> return Stop
115{-# INLINE_NORMAL bracket3D #-}
116bracket3D :: (MonadAsync m, MonadCatch m) =>
117 m b
118 -> (b -> m c)
119 -> (b -> m d)
120 -> (b -> m e)
121 -> (b -> D.Stream m a)
122 -> D.Stream m a
123bracket3D bef aft onExc onGC =
124 gbracket
125 bef
126 aft
127 (\a (e :: SomeException) _ -> onExc a >> return (D.nilM (MC.throwM e)))
128 onGC
129 (inline MC.try)
131-- For a use case of this see the "streamly-process" package. It needs to kill
132-- the process in case of exception or garbage collection, but waits for the
133-- process to terminate in normal cases.
135-- | Like 'bracket' but can use 3 separate cleanup actions depending on the
136-- mode of termination:
137--
138-- 1. When the stream stops normally
139-- 2. When the stream is garbage collected
140-- 3. When the stream encounters an exception
141--
142-- @bracket3 before onStop onGC onException action@ runs @action@ using the
143-- result of @before@. If the stream stops, @onStop@ action is executed, if the
144-- stream is abandoned @onGC@ is executed, if the stream encounters an
145-- exception @onException@ is executed.
146--
147-- /Pre-release/
148{-# INLINE bracket3 #-}
149bracket3 :: (MonadAsync m, MonadCatch m)
150 => m b
151 -> (b -> m c)
152 -> (b -> m d)
153 -> (b -> m e)
154 -> (b -> Stream m a)
155 -> Stream m a
156bracket3 = bracket3D
158-- | Run the alloc action @m b@ with async exceptions disabled but keeping
159-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
160-- output @b@ as input to @b -> Stream m a@ to generate an output stream.
161--
162-- @b@ is usually a resource under the state of monad @m@, e.g. a file
163-- handle, that requires a cleanup after use. The cleanup action @b -> m c@,
164-- runs whenever the stream ends normally, due to a sync or async exception or
165-- if it gets garbage collected after a partial lazy evaluation.
166--
167-- 'bracket' only guarantees that the cleanup action runs, and it runs with
168-- async exceptions enabled. The action must ensure that it can successfully
169-- cleanup the resource in the face of sync or async exceptions.
170--
171-- When the stream ends normally or on a sync exception, cleanup action runs
172-- immediately in the current thread context, whereas in other cases it runs in
173-- the GC context, therefore, cleanup may be delayed until the GC gets to run.
174--
175-- /See also: 'bracket_'/
176--
177-- /Inhibits stream fusion/
178--
179{-# INLINE bracket #-}
180bracket :: (MonadAsync m, MonadCatch m)
181 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
182bracket bef aft = bracket3 bef aft aft aft
184-- | Run the action @m b@ whenever the stream @Stream m a@ stops normally,
185-- aborts due to an exception or if it is garbage collected after a partial
186-- lazy evaluation.
187--
188-- The semantics of running the action @m b@ are similar to the cleanup action
189-- semantics described in 'bracket'.
190--
191-- >>> finally action xs = Stream.bracket (return ()) (const action) (const xs)
192--
193-- /See also 'finally_'/
194--
195-- /Inhibits stream fusion/
196--
197{-# INLINE finally #-}
198finally :: (MonadAsync m, MonadCatch m) =>
199 m b -> Stream m a -> Stream m a
200finally action xs = bracket (return ()) (const action) (const xs)
202{-# INLINE_NORMAL afterD #-}
203afterD :: MonadRunInIO m
204 => m b -> D.Stream m a -> D.Stream m a
205afterD action (D.Stream step state) = D.Stream step' Nothing
207 where
209 {-# INLINE_LATE step' #-}
210 step' _ Nothing = do
211 ref <- newIOFinalizer action
212 return $ Skip $ Just (state, ref)
213 step' gst (Just (st, ref)) = do
214 res <- step gst st
215 case res of
216 Yield x s -> return $ Yield x (Just (s, ref))
217 Skip s -> return $ Skip (Just (s, ref))
218 Stop -> do
219 runIOFinalizer ref
220 return Stop
222-- | Run the action @m b@ whenever the stream @Stream m a@ stops normally, or
223-- if it is garbage collected after a partial lazy evaluation.
224--
225-- The semantics of the action @m b@ are similar to the semantics of cleanup
226-- action in 'bracket'.
227--
228-- /See also 'after_'/
229--
230{-# INLINE after #-}
231after ::
232#ifdef USE_UNLIFTIO
233 MonadUnliftIO m
234#else
235 (MonadIO m, MonadBaseControl IO m)
236#endif
237 => m b -> Stream m a -> Stream m a
238after = afterD
240data RetryState emap s1 s2
241 = RetryWithMap emap s1
242 | RetryDefault s2
244-- | See 'Streamly.Internal.Data.Stream.retry'
245--
246{-# INLINE_NORMAL retryD #-}
247retryD
248 :: forall e m a. (Exception e, Ord e, MonadCatch m)
249 => Map e Int
250 -- ^ map from exception to retry count
251 -> (e -> D.Stream m a)
252 -- ^ default handler for those exceptions that are not in the map
253 -> D.Stream m a
254 -> D.Stream m a
255retryD emap0 defaultHandler (D.Stream step0 state0) = D.Stream step state
257 where
259 state = RetryWithMap emap0 state0
261 {-# INLINE_LATE step #-}
262 step gst (RetryWithMap emap st) = do
263 eres <- MC.try $ step0 gst st
264 case eres of
265 Left e -> handler e emap st
266 Right res ->
267 return
268 $ case res of
269 Yield x st1 -> Yield x $ RetryWithMap emap st1
270 Skip st1 -> Skip $ RetryWithMap emap st1
271 Stop -> Stop
272 step gst (RetryDefault (D.UnStream step1 state1)) = do
273 res <- step1 gst state1
274 return
275 $ case res of
276 Yield x st1 -> Yield x $ RetryDefault (D.Stream step1 st1)
277 Skip st1 -> Skip $ RetryDefault (D.Stream step1 st1)
278 Stop -> Stop
280 {-# INLINE handler #-}
281 handler e emap st =
282 return
283 $ Skip
284 $ case Map.lookup e emap of
285 Just i
286 | i > 0 ->
287 let emap1 = Map.insert e (i - 1) emap
288 in RetryWithMap emap1 st
289 | otherwise -> RetryDefault $ defaultHandler e
290 Nothing -> RetryDefault $ defaultHandler e
292-- | @retry@ takes 3 arguments
293--
294-- 1. A map @m@ whose keys are exceptions and values are the number of times to
295-- retry the action given that the exception occurs.
296--
297-- 2. A handler @han@ that decides how to handle an exception when the exception
298-- cannot be retried.
299--
300-- 3. The stream itself that we want to run this mechanism on.
301--
302-- When evaluating a stream if an exception occurs,
303--
304-- 1. The stream evaluation aborts
305--
306-- 2. The exception is looked up in @m@
307--
308-- a. If the exception exists and the mapped value is > 0 then,
309--
310-- i. The value is decreased by 1.
311--
312-- ii. The stream is resumed from where the exception was called, retrying
313-- the action.
314--
315-- b. If the exception exists and the mapped value is == 0 then the stream
316-- evaluation stops.
317--
318-- c. If the exception does not exist then we handle the exception using
319-- @han@.
320--
321-- /Internal/
322--
323{-# INLINE retry #-}
324retry :: (MonadCatch m, Exception e, Ord e)
325 => Map e Int
326 -- ^ map from exception to retry count
327 -> (e -> Stream m a)
328 -- ^ default handler for those exceptions that are not in the map
329 -> Stream m a
330 -> Stream m a
331retry = retryD