Skip to content

Commit 3b17ff2

Browse files
adithyaovadithyaov
committed
Add resource management & exception handlers for folds
1 parent 8500b64 commit 3b17ff2

File tree

3 files changed

+202
-0
lines changed

3 files changed

+202
-0
lines changed

‎core/src/Streamly/Internal/Data/Fold.hs

+2
Original file line numberDiff line numberDiff line change
@@ -22,11 +22,13 @@ module Streamly.Internal.Data.Fold
2222
, module Streamly.Internal.Data.Fold.Combinators
2323
, module Streamly.Internal.Data.Fold.Container
2424
, module Streamly.Internal.Data.Fold.Window
25+
, module Streamly.Internal.Data.Fold.Exception
2526
)
2627
where
2728

2829
import Streamly.Internal.Data.Fold.Combinators
2930
import Streamly.Internal.Data.Fold.Container
31+
import Streamly.Internal.Data.Fold.Exception
3032
import Streamly.Internal.Data.Fold.Tee
3133
import Streamly.Internal.Data.Fold.Type
3234
import Streamly.Internal.Data.Fold.Window

‎core/src/Streamly/Internal/Data/Fold/Exception.hs

+199
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
-- |
2+
-- Module : Streamly.Internal.Data.Fold.Exception
3+
-- Copyright : (c) 2025 Composewell Technologies
4+
-- License : BSD-3-Clause
5+
-- Maintainer : streamly@composewell.com
6+
-- Stability : experimental
7+
-- Portability : GHC
8+
--
9+
module Streamly.Internal.Data.Fold.Exception
10+
(
11+
-- * Resources
12+
afterIO
13+
, before
14+
, bracketIO
15+
16+
-- * Exceptions
17+
, onException
18+
)
19+
where
20+
21+
------------------------------------------------------------------------------
22+
-- Imports
23+
------------------------------------------------------------------------------
24+
25+
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
26+
import Control.Monad.IO.Class (MonadIO(..))
27+
import Control.Monad.Catch (MonadCatch)
28+
import Streamly.Internal.Data.IOFinalizer (newIOFinalizer, runIOFinalizer)
29+
30+
import qualified Control.Monad.Catch as MC
31+
32+
import Streamly.Internal.Data.Fold.Step
33+
import Streamly.Internal.Data.Fold.Type
34+
35+
------------------------------------------------------------------------------
36+
-- Exceptions
37+
------------------------------------------------------------------------------
38+
39+
{-
40+
41+
-- | Exception handling states of a fold
42+
data HandleExc s f1 f2 = InitDone !s | InitFailed !f1 | StepFailed !f2
43+
44+
-- | @handle initHandler stepHandler fold@ produces a new fold from a given
45+
-- fold. The new fold executes the original @fold@, if an exception occurs
46+
-- when initializing the fold then @initHandler@ is executed and fold resulting
47+
-- from that starts execution. If an exception occurs while executing the
48+
-- @step@ function of a fold then the @stephandler@ is executed and we start
49+
-- executing the fold resulting from that.
50+
--
51+
-- The exception is caught and handled, not rethrown. If the exception handler
52+
-- itself throws an exception that exception is thrown.
53+
--
54+
-- /Internal/
55+
--
56+
{-# INLINE handle #-}
57+
handle :: (MonadCatch m, Exception e)
58+
=> (e -> m (Fold m a b))
59+
-> (e -> Fold m a b -> m (Fold m a b))
60+
-> Fold m a b
61+
-> Fold m a b
62+
handle initH stepH (Fold step1 initial1 extract1) = Fold step initial extract
63+
64+
where
65+
66+
initial = fmap InitDone initial1 `MC.catch` (fmap InitFailed . initH)
67+
68+
step (InitDone s) a =
69+
let f = Fold step1 (return s) extract1
70+
in fmap InitDone (step1 s a)
71+
`MC.catch` (\e -> fmap StepFailed (stepH e f))
72+
step (InitFailed (Fold step2 initial2 extract2)) a = do
73+
s <- initial2
74+
s1 <- step2 s a
75+
return $ InitFailed $ Fold step2 (return s1) extract2
76+
step (StepFailed (Fold step2 initial2 extract2)) a = do
77+
s <- initial2
78+
s1 <- step2 s a
79+
return $ StepFailed $ Fold step2 (return s1) extract2
80+
81+
extract (InitDone s) = extract1 s
82+
extract (InitFailed (Fold _ initial2 extract2)) = initial2 >>= extract2
83+
extract (StepFailed (Fold _ initial2 extract2)) = initial2 >>= extract2
84+
85+
-}
86+
87+
-- | @onException action fold@ runs @action@ whenever the fold throws an
88+
-- exception. The action is executed on any exception whether it is in
89+
-- initial, step or extract action of the fold.
90+
--
91+
-- The exception is not caught, simply rethrown. If the @action@ itself
92+
-- throws an exception that exception is thrown instead of the original
93+
-- exception.
94+
--
95+
-- /Internal/
96+
--
97+
{-# INLINE onException #-}
98+
onException :: MonadCatch m => m x -> Fold m a b -> Fold m a b
99+
onException action (Fold step1 initial1 extract1 final1) =
100+
Fold step initial extract final
101+
102+
where
103+
104+
initial = initial1 `MC.onException` action
105+
step s a = step1 s a `MC.onException` action
106+
extract s = extract1 s `MC.onException` action
107+
final s = final1 s `MC.onException` action
108+
109+
-- | @bracketIO before after between@ runs @before@ and invokes @between@ using
110+
-- its output, then runs the fold generated by @between@. If the fold ends
111+
-- normally, due to an exception or if it is garbage collected prematurely then
112+
-- @after@ is run with the output of @before@ as argument.
113+
--
114+
-- If @before@ or @after@ throw an exception that exception is thrown.
115+
--
116+
{-# INLINE bracketIO #-}
117+
bracketIO :: (MonadIO m, MonadCatch m)
118+
=> IO x -> (x -> IO c) -> (x -> Fold m a b) -> Fold m a b
119+
bracketIO bef aft bet = Fold step initial extract final
120+
121+
where
122+
123+
initial = do
124+
r <- liftIO $ bef
125+
ref <- liftIO $ newIOFinalizer (aft r)
126+
case bet r of
127+
Fold step1 initial1 extract1 final1 -> do
128+
res <- initial1 `MC.onException` liftIO (runIOFinalizer ref)
129+
case res of
130+
Partial s -> do
131+
let fld1 = Fold step1 (pure (Partial s)) extract1 final1
132+
pure $ Partial $ Tuple' ref fld1
133+
Done b -> do
134+
liftIO $ runIOFinalizer ref
135+
pure $ Done b
136+
137+
step (Tuple' ref (Fold step1 initial1 extract1 final1)) a = do
138+
res <- initial1
139+
case res of
140+
Partial s -> do
141+
s1 <- step1 s a `MC.onException` liftIO (runIOFinalizer ref)
142+
let fld1 = Fold step1 (pure s1) extract1 final1
143+
pure $ Partial $ Tuple' ref fld1
144+
Done b -> do
145+
liftIO $ runIOFinalizer ref
146+
pure $ Done b
147+
148+
extract (Tuple' ref (Fold _ initial1 extract1 _)) = do
149+
res <- initial1
150+
case res of
151+
Partial s -> extract1 s `MC.onException` liftIO (runIOFinalizer ref)
152+
Done b -> pure b
153+
154+
final (Tuple' ref (Fold _ initial1 _ final1)) = do
155+
res <- initial1
156+
case res of
157+
Partial s -> do
158+
val <- final1 s `MC.onException` liftIO (runIOFinalizer ref)
159+
runIOFinalizer ref
160+
pure val
161+
Done b -> pure b
162+
163+
-- | Run a side effect whenever the fold stops normally, aborts due to an
164+
-- exception or is garbage collected.
165+
--
166+
{-# INLINE afterIO #-}
167+
afterIO :: (MonadIO m, MonadCatch m) => IO b -> Fold m a b -> Fold m a b
168+
afterIO aft (Fold step1 initial1 extract1 final1) =
169+
Fold step initial extract final
170+
171+
where
172+
173+
initial = do
174+
ref <- liftIO $ newIOFinalizer aft
175+
res <- initial1 `MC.onException` liftIO (runIOFinalizer ref)
176+
pure $ case res of
177+
Done b -> Done b
178+
Partial s -> Partial $ Tuple' ref s
179+
180+
step (Tuple' ref s) a = do
181+
res <- step1 s a `MC.onException` liftIO (runIOFinalizer ref)
182+
pure $ case res of
183+
Done b -> Done b
184+
Partial s1 -> Partial $ Tuple' ref s1
185+
186+
extract (Tuple' ref s) =
187+
extract1 s `MC.onException` liftIO (runIOFinalizer ref)
188+
189+
final (Tuple' ref s) = do
190+
res <- final1 s `MC.onException` liftIO (runIOFinalizer ref)
191+
liftIO $ runIOFinalizer ref
192+
pure res
193+
194+
195+
-- | Run a side effect before the fold consumes its first element.
196+
--
197+
{-# INLINE before #-}
198+
before :: Monad m => m x -> Fold m a b -> Fold m a b
199+
before effect (Fold s i e f) = Fold s (effect *> i) e f

‎core/streamly-core.cabal

+1
Original file line numberDiff line numberDiff line change
@@ -503,6 +503,7 @@ library
503503
, Streamly.Internal.Data.Fold.Type
504504
, Streamly.Internal.Data.Fold.Combinators
505505
, Streamly.Internal.Data.Fold.Container
506+
, Streamly.Internal.Data.Fold.Exception
506507
, Streamly.Internal.Data.Fold.Tee
507508
, Streamly.Internal.Data.Fold.Window
508509

0 commit comments

Comments
 (0)