@@ -183,6 +183,9 @@ module Streamly.Internal.Data.StreamK
183183 , append
184184 , interleave
185185
186+ -- * Exceptions
187+ , handle
188+
186189 -- * Resource Management
187190 , bracketIO
188191
@@ -196,20 +199,19 @@ where
196199#include "inline.hs"
197200#include "assert.hs"
198201
199−import Control.Exception (mask_)
202+import Control.Exception (mask_, Exception)
200203import Control.Monad (void, join)
201−import Control.Monad.IO.Class (MonadIO(..))
202204import Control.Monad.Catch (MonadCatch)
205+import Control.Monad.IO.Class (MonadIO(..))
203206import Data.Proxy (Proxy(..))
204207import GHC.Types (SPEC(..))
205208import Streamly.Internal.Data.Array.Type (Array(..))
206209import Streamly.Internal.Data.Fold.Type (Fold(..))
210+import Streamly.Internal.Data.IOFinalizer (newIOFinalizer, runIOFinalizer)
211+import Streamly.Internal.Data.Parser.ParserK.Type (ParserK)
207212import Streamly.Internal.Data.Producer.Type (Producer(..))
208213import Streamly.Internal.Data.SVar.Type (adaptState, defState)
209214import Streamly.Internal.Data.Unbox (sizeOf, Unbox)
210−import Streamly.Internal.Data.Parser.ParserK.Type (ParserK)
211−import Streamly.Internal.Data.IOFinalizer
212− (newIOFinalizer, runIOFinalizer, clearingIOFinalizer)
213215
214216import qualified Control.Monad.Catch as MC
215217import qualified Streamly.Internal.Data.Array.Type as Array
@@ -926,6 +928,38 @@ mapMaybe f = go
926928 Nothing -> foldStream (adaptState st) yieldk single stp r
927929 in foldStream (adaptState st) yieldk single stp m1
928930
931+-------------------------------------------------------------------------------
932+-- Exception Handling
933+-------------------------------------------------------------------------------
934+
935+-- | Like 'Streamly.Data.Stream.handle' but with one significant difference,
936+-- this function observes exceptions from the consumer of the stream as well.
937+--
938+-- You can also convert 'StreamK' to 'Stream' and use exception handling from
939+-- 'Stream' module:
940+--
941+-- >>> handle f s = StreamK.fromStream $ Stream.handle (\e -> StreamK.toStream (f e)) (StreamK.toStream s)
942+--
943+{-# INLINABLE handle #-}
944+handle :: (MonadCatch m, Exception e)
945+ => (e -> m (StreamK m a)) -> StreamK m a -> StreamK m a
946+handle f stream = go stream
947+
948+ where
949+
950+ go m1 = mkStream $ \st yld sng stp ->
951+ let yieldk a r = yld a $ go r
952+ in do
953+ res <- MC.try (foldStream (adaptState st) yieldk sng stp m1)
954+ case res of
955+ Right r -> return r
956+ Left e -> do
957+ r <- f e
958+ foldStream (adaptState st) yld sng stp r
959+
960+-------------------------------------------------------------------------------
961+-- Resource Management
962+-------------------------------------------------------------------------------
929963
930964-- If we are folding the stream and we do not drain the entire stream (e.g. if
931965-- the fold terminates before the stream) then the finalizer will run on GC.
@@ -935,13 +969,16 @@ mapMaybe f = go
935969-- the entire chain can be invoked when the stream ends voluntarily or if
936970-- someone decides to abandon the stream.
937971
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−
972+-- | Like 'Streamly.Data.Stream.bracketIO' but with one significant difference,
973+-- this function observes exceptions from the consumer of the stream as well.
974+-- Therefore, it cleans up the resource promptly when the consumer encounters
975+-- an exception.
976+--
977+-- You can also convert 'StreamK' to 'Stream' and use resource handling from
978+-- 'Stream' module:
979+--
980+-- >>> bracketIO bef aft bet = StreamK.fromStream $ Stream.bracketIO bef aft (StreamK.toStream . bet)
981+--
945982{-# INLINABLE bracketIO #-}
946983bracketIO :: (MonadIO m, MonadCatch m)
947984 => IO b -> (b -> IO c) -> (b -> StreamK m a) -> StreamK m a
@@ -957,6 +994,8 @@ bracketIO bef aft bet =
957994
958995 go ref m1 = mkStream $ \st yld sng stp ->
959996 let
997+ -- We can discard exceptions on continuations to make it equivalent
998+ -- to StreamD, but it seems like a desirable behavior.
960999 stop = liftIO (runIOFinalizer ref) >> stp
9611000 single a = liftIO (runIOFinalizer ref) >> sng a
9621001 yieldk a r = yld a $ go ref r