@@ -302,17 +302,12 @@ module Streamly.Internal.Data.Stream.StreamD
302302where
303303
304304import Control.Concurrent (myThreadId)
305−import Control.Exception
306− (Exception, SomeException, mask_)
307305import Control.Monad (void)
308−import Control.Monad.Catch (MonadCatch)
309306import Control.Monad.IO.Class (MonadIO(..))
310307import Control.Monad.Trans.Reader (ReaderT)
311308import Control.Monad.Trans.State.Strict (StateT)
312309import Control.Monad.Trans.Class (MonadTrans(lift))
313−import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)
314310import Data.Functor.Identity (Identity(..))
315−import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
316311import Data.Maybe (fromJust, isJust)
317312import Streamly.Internal.Data.Time.Units
318313 (TimeUnit64, toRelTime64, diffAbsTime64)
@@ -322,7 +317,6 @@ import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
322317import qualified Streamly.Internal.Data.Pipe.Types as Pipe
323318import qualified Streamly.Internal.Data.Fold as FL
324319import qualified Streamly.Internal.Data.Stream.StreamK as K
325−import qualified Control.Monad.Catch as MC
326320import qualified Control.Monad.Trans.Reader as Reader
327321import qualified Control.Monad.Trans.State.Strict as State
328322
@@ -336,6 +330,7 @@ import Streamly.Internal.Data.Stream.StreamD.Type
336330import Streamly.Internal.Data.Stream.StreamD.Common (takeWhileM, takeWhile)
337331import Streamly.Internal.Data.Stream.StreamD.Generate
338332import Streamly.Internal.Data.Stream.StreamD.Eliminate
333+import Streamly.Internal.Data.Stream.StreamD.Exceptions
339334import Streamly.Internal.Data.Stream.StreamD.Flatten
340335import Streamly.Internal.Data.Stream.StreamD.SplitGroup
341336import Streamly.Internal.Data.SVar
@@ -522,389 +517,6 @@ stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
522517mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
523518mapM_ m = drain . mapM m
524519
525−------------------------------------------------------------------------------
526−-- Exceptions
527−------------------------------------------------------------------------------
528−
529−data GbracketState s1 s2 v
530− = GBracketInit
531− | GBracketNormal s1 v
532− | GBracketException s2
533−
534−-- | Like 'gbracket' but with following differences:
535−--
536−-- * alloc action @m c@ runs with async exceptions enabled
537−-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
538−-- after partial evaluation.
539−-- * does not require a 'MonadAsync' constraint.
540−--
541−-- /Inhibits stream fusion/
542−--
543−-- /Internal/
544−--
545−{-# INLINE_NORMAL gbracket_ #-}
546−gbracket_
547− :: Monad m
548− => m c -- ^ before
549− -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
550− -> (c -> m d) -- ^ after, on normal stop
551− -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
552− -> (c -> Stream m b) -- ^ stream generator
553− -> Stream m b
554−gbracket_ bef exc aft fexc fnormal =
555− Stream step GBracketInit
556−
557− where
558−
559− {-# INLINE_LATE step #-}
560− step _ GBracketInit = do
561− r <- bef
562− return $ Skip $ GBracketNormal (fnormal r) r
563−
564− step gst (GBracketNormal (UnStream step1 st) v) = do
565− res <- exc $ step1 gst st
566− case res of
567− Right r -> case r of
568− Yield x s ->
569− return $ Yield x (GBracketNormal (Stream step1 s) v)
570− Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
571− Stop -> aft v >> return Stop
572− -- XXX Do not handle async exceptions, just rethrow them.
573− Left e ->
574− return $ Skip (GBracketException (fexc v e (UnStream step1 st)))
575− step gst (GBracketException (UnStream step1 st)) = do
576− res <- step1 gst st
577− case res of
578− Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
579− Skip s -> return $ Skip (GBracketException (Stream step1 s))
580− Stop -> return Stop
581−
582−------------------------------------------------------------------------------
583−-- Finalizers
584−------------------------------------------------------------------------------
585−
586−-- | Make a finalizer from a monadic action @m a@ that can run in IO monad.
587−mkIOFinalizer :: MonadBaseControl IO m => m b -> m (IO ())
588−mkIOFinalizer f = do
589− mrun <- captureMonadState
590− return $
591− void $ do
592− _ <- runInIO mrun f
593− return ()
594−
595−-- | Run an IO action stored in a finalized IORef.
596−runIORefFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
597−runIORefFinalizerGC ref = do
598− res <- readIORef ref
599− case res of
600− Nothing -> return ()
601− Just f -> f
602−
603−-- | Create an IORef holding a finalizer that is called automatically when
604−-- the IORef is garbage collected. The IORef can be written to with a 'Nothing'
605−-- value to deactivate the finalizer.
606−--
607−-- Note: The finalizer is always run with the state of the monad captured at
608−-- the time of calling newFinalizedIORef. To run it on garbage collection we
609−-- have no option but to take a snapshot of the monadic state at some point of
610−-- time. For normal case we could run it with the current state of the monad
611−-- but we want to keep both the cases consistent.
612−--
613−newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
614− => m a -> m (IORef (Maybe (IO ())))
615−newFinalizedIORef finalizer = do
616− f <- mkIOFinalizer finalizer
617− ref <- liftIO $ newIORef $ Just f
618− _ <- liftIO $ mkWeakIORef ref (runIORefFinalizerGC ref)
619− return ref
620−
621−-- | Run the finalizer stored in an IORef and deactivate it so that it is run
622−-- only once. Note, the action runs with async exceptions masked.
623−--
624−runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
625−runIORefFinalizer ref = liftIO $ do
626− res <- readIORef ref
627− case res of
628− Nothing -> return ()
629− Just action -> do
630− -- if an async exception comes after writing 'Nothing' then the
631− -- finalizing action will never be run. We need to do this
632− -- atomically wrt async exceptions.
633− mask_ $ do
634− writeIORef ref Nothing
635− action
636−
637−-- | Run an action clearing the finalizer IORef atomically wrt async
638−-- exceptions. The action is run with async exceptions masked.
639−withIORefFinalizer :: MonadBaseControl IO m
640− => IORef (Maybe (IO ())) -> m a -> m a
641−withIORefFinalizer ref action = do
642− control $ \runinio ->
643− mask_ $ do
644− writeIORef ref Nothing
645− runinio action
646−
647−------------------------------------------------------------------------------
648−
649−data GbracketIOState s1 s2 v wref
650− = GBracketIOInit
651− | GBracketIONormal s1 v wref
652− | GBracketIOException s2
653−
654−-- | Run the alloc action @m c@ with async exceptions disabled but keeping
655−-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
656−-- output @c@ as input to @c -> Stream m b@ to generate an output stream. When
657−-- generating the stream use the supplied @try@ operation @forall s. m s -> m
658−-- (Either e s)@ to catch synchronous exceptions. If an exception occurs run
659−-- the exception handler @c -> e -> Stream m b -> m (Stream m b)@.
660−--
661−-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
662−-- to a sync or async exception or if it gets garbage collected after a partial
663−-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
664−--
665−-- 'gbracket' can express all other exception handling combinators.
666−--
667−-- /Inhibits stream fusion/
668−--
669−-- /Internal/
670−{-# INLINE_NORMAL gbracket #-}
671−gbracket
672− :: (MonadIO m, MonadBaseControl IO m)
673− => m c -- ^ before
674− -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
675− -> (c -> m d) -- ^ after, on normal stop or GC
676− -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
677− -> (c -> Stream m b) -- ^ stream generator
678− -> Stream m b
679−gbracket bef exc aft fexc fnormal =
680− Stream step GBracketIOInit
681−
682− where
683−
684− -- If the stream is never evaluated the "aft" action will never be
685− -- called. For that to occur we will need the user of this API to pass a
686− -- weak pointer to us.
687− {-# INLINE_LATE step #-}
688− step _ GBracketIOInit = do
689− -- We mask asynchronous exceptions to make the execution
690− -- of 'bef' and the registration of 'aft' atomic.
691− -- A similar thing is done in the resourcet package: https://git.io/JvKV3
692− -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
693− (r, ref) <- liftBaseOp_ mask_ $ do
694− r <- bef
695− ref <- newFinalizedIORef (aft r)
696− return (r, ref)
697− return $ Skip $ GBracketIONormal (fnormal r) r ref
698−
699− step gst (GBracketIONormal (UnStream step1 st) v ref) = do
700− res <- exc $ step1 gst st
701− case res of
702− Right r -> case r of
703− Yield x s ->
704− return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
705− Skip s ->
706− return $ Skip (GBracketIONormal (Stream step1 s) v ref)
707− Stop -> do
708− runIORefFinalizer ref
709− return Stop
710− -- XXX Do not handle async exceptions, just rethrow them.
711− Left e -> do
712− -- Clearing of finalizer and running of exception handler must
713− -- be atomic wrt async exceptions. Otherwise if we have cleared
714− -- the finalizer and have not run the exception handler then we
715− -- may leak the resource.
716− stream <- withIORefFinalizer ref (fexc v e (UnStream step1 st))
717− return $ Skip (GBracketIOException stream)
718− step gst (GBracketIOException (UnStream step1 st)) = do
719− res <- step1 gst st
720− case res of
721− Yield x s ->
722− return $ Yield x (GBracketIOException (Stream step1 s))
723− Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
724− Stop -> return Stop
725−
726−-- | See 'Streamly.Internal.Data.Stream.IsStream.after_'.
727−--
728−{-# INLINE_NORMAL before #-}
729−before :: Monad m => m b -> Stream m a -> Stream m a
730−before action (Stream step state) = Stream step' Nothing
731−
732− where
733−
734− {-# INLINE_LATE step' #-}
735− step' _ Nothing = action >> return (Skip (Just state))
736−
737− step' gst (Just st) = do
738− res <- step gst st
739− case res of
740− Yield x s -> return $ Yield x (Just s)
741− Skip s -> return $ Skip (Just s)
742− Stop -> return Stop
743−
744−-- | See 'Streamly.Internal.Data.Stream.IsStream.after_'.
745−--
746−{-# INLINE_NORMAL after_ #-}
747−after_ :: Monad m => m b -> Stream m a -> Stream m a
748−after_ action (Stream step state) = Stream step' state
749−
750− where
751−
752− {-# INLINE_LATE step' #-}
753− step' gst st = do
754− res <- step gst st
755− case res of
756− Yield x s -> return $ Yield x s
757− Skip s -> return $ Skip s
758− Stop -> action >> return Stop
759−
760−-- | See 'Streamly.Internal.Data.Stream.IsStream.after'.
761−--
762−{-# INLINE_NORMAL after #-}
763−after :: (MonadIO m, MonadBaseControl IO m)
764− => m b -> Stream m a -> Stream m a
765−after action (Stream step state) = Stream step' Nothing
766−
767− where
768−
769− {-# INLINE_LATE step' #-}
770− step' _ Nothing = do
771− ref <- newFinalizedIORef action
772− return $ Skip $ Just (state, ref)
773− step' gst (Just (st, ref)) = do
774− res <- step gst st
775− case res of
776− Yield x s -> return $ Yield x (Just (s, ref))
777− Skip s -> return $ Skip (Just (s, ref))
778− Stop -> do
779− runIORefFinalizer ref
780− return Stop
781−
782−-- XXX For high performance error checks in busy streams we may need another
783−-- Error constructor in step.
784−--
785−-- | See 'Streamly.Internal.Data.Stream.IsStream.onException'.
786−--
787−{-# INLINE_NORMAL onException #-}
788−onException :: MonadCatch m => m b -> Stream m a -> Stream m a
789−onException action str =
790− gbracket_ (return ()) MC.try return
791− (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
792− (\_ -> str)
793−
794−{-# INLINE_NORMAL _onException #-}
795−_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
796−_onException action (Stream step state) = Stream step' state
797−
798− where
799−
800− {-# INLINE_LATE step' #-}
801− step' gst st = do
802− res <- step gst st `MC.onException` action
803− case res of
804− Yield x s -> return $ Yield x s
805− Skip s -> return $ Skip s
806− Stop -> return Stop
807−
808−-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket_'.
809−--
810−{-# INLINE_NORMAL bracket_ #-}
811−bracket_ :: MonadCatch m
812− => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
813−bracket_ bef aft bet =
814− gbracket_ bef MC.try aft
815− (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e)) bet
816−
817−-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket'.
818−--
819−{-# INLINE_NORMAL bracket #-}
820−bracket :: (MonadAsync m, MonadCatch m)
821− => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
822−bracket bef aft bet =
823− gbracket bef MC.try aft
824− (\a (e :: SomeException) _ -> aft a >> return (nilM (MC.throwM e))) bet
825−
826−data BracketState s v = BracketInit | BracketRun s v
827−
828−-- | Alternate (custom) implementation of 'bracket'.
829−--
830−{-# INLINE_NORMAL _bracket #-}
831−_bracket :: MonadCatch m
832− => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
833−_bracket bef aft bet = Stream step' BracketInit
834−
835− where
836−
837− {-# INLINE_LATE step' #-}
838− step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
839−
840− -- NOTE: It is important to use UnStream instead of the Stream pattern
841− -- here, otherwise we get huge perf degradation, see note in concatMap.
842− step' gst (BracketRun (UnStream step state) v) = do
843− -- res <- step gst state `MC.onException` aft v
844− res <- MC.try $ step gst state
845− case res of
846− Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
847− Right r -> case r of
848− Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
849− Skip s -> return $ Skip (BracketRun (Stream step s) v)
850− Stop -> aft v >> return Stop
851−
852−-- | See 'Streamly.Internal.Data.Stream.IsStream.finally_'.
853−--
854−{-# INLINE finally_ #-}
855−finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
856−finally_ action xs = bracket_ (return ()) (\_ -> action) (const xs)
857−
858−-- | See 'Streamly.Internal.Data.Stream.IsStream.finally'.
859−--
860−-- finally action xs = after action $ onException action xs
861−--
862−{-# INLINE finally #-}
863−finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
864−finally action xs = bracket (return ()) (\_ -> action) (const xs)
865−
866−-- | See 'Streamly.Internal.Data.Stream.IsStream.ghandle'.
867−--
868−{-# INLINE_NORMAL ghandle #-}
869−ghandle :: (MonadCatch m, Exception e)
870− => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
871−ghandle f str =
872− gbracket_ (return ()) MC.try return (\_ -> f) (\_ -> str)
873−
874−-- | See 'Streamly.Internal.Data.Stream.IsStream.handle'.
875−--
876−{-# INLINE_NORMAL handle #-}
877−handle :: (MonadCatch m, Exception e)
878− => (e -> Stream m a) -> Stream m a -> Stream m a
879−handle f str =
880− gbracket_ (return ()) MC.try return (\_ e _ -> f e) (\_ -> str)
881−
882−-- | Alternate (custom) implementation of 'handle'.
883−--
884−{-# INLINE_NORMAL _handle #-}
885−_handle :: (MonadCatch m, Exception e)
886− => (e -> Stream m a) -> Stream m a -> Stream m a
887−_handle f (Stream step state) = Stream step' (Left state)
888−
889− where
890−
891− {-# INLINE_LATE step' #-}
892− step' gst (Left st) = do
893− res <- MC.try $ step gst st
894− case res of
895− Left e -> return $ Skip $ Right (f e)
896− Right r -> case r of
897− Yield x s -> return $ Yield x (Left s)
898− Skip s -> return $ Skip (Left s)
899− Stop -> return Stop
900−
901− step' gst (Right (UnStream step1 st)) = do
902− res <- step1 gst st
903− case res of
904− Yield x s -> return $ Yield x (Right (Stream step1 s))
905− Skip s -> return $ Skip (Right (Stream step1 s))
906− Stop -> return Stop
907−
908520-------------------------------------------------------------------------------
909521-- General transformation
910522-------------------------------------------------------------------------------