git.iko.soy/streamly
on master help sign in
Add bracketIO for StreamK
fc45e9bHarendra Kumar
core/src/Streamly/Internal/Data/IOFinalizer.hs+12−2
compare with master
unifiedsplit
@@ -66,6 +66,9 @@ newIOFinalizer finalizer = liftIO $ do
66-- never runs again. Note, the finalizing action runs with async exceptions
66-- never runs again. Note, the finalizing action runs with async exceptions
67-- masked.
67-- masked.
68--
68--
69+-- If this function is called multiple times, the action is guaranteed to run
70+-- once and only once.
71+--
69-- /Pre-release/
72-- /Pre-release/
70runIOFinalizer :: MonadIO m => IOFinalizer -> m ()
73runIOFinalizer :: MonadIO m => IOFinalizer -> m ()
71runIOFinalizer (IOFinalizer ref) = liftIO $ do
74runIOFinalizer (IOFinalizer ref) = liftIO $ do
@@ -83,9 +86,16 @@ runIOFinalizer (IOFinalizer ref) = liftIO $ do
83-- | Run an action clearing the finalizer atomically wrt async exceptions. The
86-- | Run an action clearing the finalizer atomically wrt async exceptions. The
84-- action is run with async exceptions masked.
87-- action is run with async exceptions masked.
85--
88--
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+--
86-- /Pre-release/
92-- /Pre-release/
87clearingIOFinalizer :: MonadIO m => IOFinalizer -> IO a -> m a
93clearingIOFinalizer :: MonadIO m => IOFinalizer -> IO a -> m a
88clearingIOFinalizer (IOFinalizer ref) action = do
94clearingIOFinalizer (IOFinalizer ref) action = do
89 liftIO $ mask_ $ do
95 liftIO $ mask_ $ do
96+ res <- readIORef ref
97+ case res of
98+ Just _ -> do
90− writeIORef ref Nothing
99+ writeIORef ref Nothing
91− action
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
183 , append
183 , append
184 , interleave
184 , interleave
185
185
186+ -- * Resource Management
187+ , bracketIO
188+
186 -- * Utilities
189 -- * Utilities
187 , consM
190 , consM
188 , mfix
191 , mfix
@@ -193,7 +196,10 @@ where
193#include "inline.hs"
196#include "inline.hs"
194#include "assert.hs"
197#include "assert.hs"
195
198
199+import Control.Exception (mask_)
196import Control.Monad (void, join)
200import Control.Monad (void, join)
201+import Control.Monad.IO.Class (MonadIO(..))
202+import Control.Monad.Catch (MonadCatch)
197import Data.Proxy (Proxy(..))
203import Data.Proxy (Proxy(..))
198import GHC.Types (SPEC(..))
204import GHC.Types (SPEC(..))
199import Streamly.Internal.Data.Array.Type (Array(..))
205import Streamly.Internal.Data.Array.Type (Array(..))
@@ -202,7 +208,10 @@ import Streamly.Internal.Data.Producer.Type (Producer(..))
202import Streamly.Internal.Data.SVar.Type (adaptState, defState)
208import Streamly.Internal.Data.SVar.Type (adaptState, defState)
203import Streamly.Internal.Data.Unbox (sizeOf, Unbox)
209import Streamly.Internal.Data.Unbox (sizeOf, Unbox)
204import Streamly.Internal.Data.Parser.ParserK.Type (ParserK)
210import Streamly.Internal.Data.Parser.ParserK.Type (ParserK)
211+import Streamly.Internal.Data.IOFinalizer
212+ (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
205
213
214+import qualified Control.Monad.Catch as MC
206import qualified Streamly.Internal.Data.Array.Type as Array
215import qualified Streamly.Internal.Data.Array.Type as Array
207import qualified Streamly.Internal.Data.Fold.Type as FL
216import qualified Streamly.Internal.Data.Fold.Type as FL
208import qualified Streamly.Internal.Data.Parser as Parser
217import qualified Streamly.Internal.Data.Parser as Parser
@@ -917,6 +926,50 @@ mapMaybe f = go
917 Nothing -> foldStream (adaptState st) yieldk single stp r
926 Nothing -> foldStream (adaptState st) yieldk single stp r
918 in foldStream (adaptState st) yieldk single stp m1
927 in foldStream (adaptState st) yieldk single stp m1
919
928
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+
920------------------------------------------------------------------------------
973------------------------------------------------------------------------------
921-- Serial Zipping
974-- Serial Zipping
922------------------------------------------------------------------------------
975------------------------------------------------------------------------------