git.iko.soy/streamly
on master help sign in
Add bracketIO for StreamK
commitfc45e9b96a1046d5079e34696c25b063bcc21472 authorHarendra Kumar <harendra@composewell.com> dateWed 2 Aug 2023, 03:22 (3 years ago) parent87cb22c
2 files changed+65−2
compare with master
unifiedsplit
core/src/Streamly/Internal/Data/IOFinalizer.hs+12−2 core/src/Streamly/Internal/Data/StreamK.hs+53
core/src/Streamly/Internal/Data/IOFinalizer.hs+12−2
@@ -66,6 +66,9 @@ newIOFinalizer finalizer = liftIO $ do
6666-- never runs again. Note, the finalizing action runs with async exceptions
6767-- masked.
6868--
69+-- If this function is called multiple times, the action is guaranteed to run
70+-- once and only once.
71+--
6972-- /Pre-release/
7073runIOFinalizer :: MonadIO m => IOFinalizer -> m ()
7174runIOFinalizer (IOFinalizer ref) = liftIO $ do
@@ -83,9 +86,16 @@ runIOFinalizer (IOFinalizer ref) = liftIO $ do
8386-- | Run an action clearing the finalizer atomically wrt async exceptions. The
8487-- action is run with async exceptions masked.
8588--
89+-- This function can be called at most once after setting the finalizer. If the
90+-- finalizer is not set it is considered a bug.
91+--
8692-- /Pre-release/
8793clearingIOFinalizer :: MonadIO m => IOFinalizer -> IO a -> m a
8894clearingIOFinalizer (IOFinalizer ref) action = do
8995 liftIO $ mask_ $ do
90− writeIORef ref Nothing
91− action
96+ res <- readIORef ref
97+ case res of
98+ Just _ -> do
99+ writeIORef ref Nothing
100+ action
101+ Nothing -> error "clearingIOFinalizer: finalizer not set"
core/src/Streamly/Internal/Data/StreamK.hs+53
@@ -183,6 +183,9 @@ module Streamly.Internal.Data.StreamK
183183 , append
184184 , interleave
185185
186+ -- * Resource Management
187+ , bracketIO
188+
186189 -- * Utilities
187190 , consM
188191 , mfix
@@ -193,7 +196,10 @@ where
193196#include "inline.hs"
194197#include "assert.hs"
195198
199+import Control.Exception (mask_)
196200import Control.Monad (void, join)
201+import Control.Monad.IO.Class (MonadIO(..))
202+import Control.Monad.Catch (MonadCatch)
197203import Data.Proxy (Proxy(..))
198204import GHC.Types (SPEC(..))
199205import Streamly.Internal.Data.Array.Type (Array(..))
@@ -202,7 +208,10 @@ import Streamly.Internal.Data.Producer.Type (Producer(..))
202208import Streamly.Internal.Data.SVar.Type (adaptState, defState)
203209import Streamly.Internal.Data.Unbox (sizeOf, Unbox)
204210import Streamly.Internal.Data.Parser.ParserK.Type (ParserK)
211+import Streamly.Internal.Data.IOFinalizer
212+ (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
205213
214+import qualified Control.Monad.Catch as MC
206215import qualified Streamly.Internal.Data.Array.Type as Array
207216import qualified Streamly.Internal.Data.Fold.Type as FL
208217import qualified Streamly.Internal.Data.Parser as Parser
@@ -917,6 +926,50 @@ mapMaybe f = go
917926 Nothing -> foldStream (adaptState st) yieldk single stp r
918927 in foldStream (adaptState st) yieldk single stp m1
919928
929+
930+-- If we are folding the stream and we do not drain the entire stream (e.g. if
931+-- the fold terminates before the stream) then the finalizer will run on GC.
932+--
933+-- XXX To implement a prompt cleanup, we will have to yield a cleanup function
934+-- via the yield continuation. A chain of cleanup functions can be built and
935+-- the entire chain can be invoked when the stream ends voluntarily or if
936+-- someone decides to abandon the stream.
937+
938+{-
939+bracketIO :: (MonadIO m, MonadCatch m)
940+ => IO b -> (b -> IO c) -> (b -> StreamK m a) -> StreamK m a
941+bracketIO bef aft bet =
942+ fromStream $ Stream.bracketIO bef aft (\x -> toStream $ bet x)
943+-}
944+
945+{-# INLINABLE bracketIO #-}
946+bracketIO :: (MonadIO m, MonadCatch m)
947+ => IO b -> (b -> IO c) -> (b -> StreamK m a) -> StreamK m a
948+bracketIO bef aft bet =
949+ concatEffect $ do
950+ (r, ref) <- liftIO $ mask_ $ do
951+ r <- bef
952+ ref <- newIOFinalizer (aft r)
953+ return (r, ref)
954+ return $ go ref (bet r)
955+
956+ where
957+
958+ go ref m1 = mkStream $ \st yld sng stp ->
959+ let
960+ stop = liftIO (runIOFinalizer ref) >> stp
961+ single a = liftIO (runIOFinalizer ref) >> sng a
962+ yieldk a r = yld a $ go ref r
963+ in do
964+ -- Do not call the finalizer twice if it has already been
965+ -- called via stop continuation and stop continuation itself
966+ -- generated an exception. runIOFinalizer takes care of that.
967+ res <- MC.try (foldStream (adaptState st) yieldk single stop m1)
968+ case res of
969+ Right r -> return r
970+ Left (e :: MC.SomeException) ->
971+ liftIO (runIOFinalizer ref) >> MC.throwM e
972+
920973------------------------------------------------------------------------------
921974-- Serial Zipping
922975------------------------------------------------------------------------------