Flushing source conduit

FlushSource.hs

{-# LANGUAGE LambdaCase #-}

-- | Sources that flush when their monadic action takes longer than n
-- microseconds.

module Data.Conduit.Flush where

import Data.Conduit
import Data.Conduit.Internal as Pipe (ConduitT(..), Pipe(..))
import UnliftIO

-- | Each monadic action in the given conduit will run in an async
-- thread, and the conduit will yield a 'Flush' if the action takes longer
-- than N microseconds.
flushSource ::
     MonadUnliftIO m => Int -> ConduitT () o m r -> ConduitT () (Flush o) m r
flushSource delay (ConduitT c0) =
  ConduitT $ \rest ->
    let go =
          \case
            Leftover p l -> Leftover (go p) l
            Done r -> rest r
            HaveOutput p o -> HaveOutput (go p) (Chunk o)
            NeedInput p c -> NeedInput (go . p) (go . c)
            PipeM m ->
              PipeM
                (do original <- async m
                    result <- timeout delay (wait original)
                    case result of
                      Nothing ->
                        pure
                          (HaveOutput (PipeM (fmap go (wait original))) Flush)
                      Just p -> pure (go p))
     in go (c0 Done)

-- | Join nested flushes into one.
joinFlush :: Flush (Flush a) -> Flush a
joinFlush =
  \case
    Flush -> Flush
    Chunk (Chunk a) -> Chunk a
    Chunk Flush -> Flush