packages feed

streamly-0.11.1: src/Streamly/Internal/Data/Stream/Channel/Interleave.hs

-- |
-- Module      : Streamly.Internal.Data.Stream.Channel.Interleave
-- Copyright   : (c) 2017 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC

module Streamly.Internal.Data.Stream.Channel.Interleave
    (
      newInterleaveChannel
    )
where

#include "inline.hs"

import Control.Concurrent (myThreadId)
import Control.Concurrent.MVar (newEmptyMVar, newMVar)
import Control.Monad.IO.Class (MonadIO(liftIO))
import Data.Concurrent.Queue.MichaelScott (LinkedQueue, newQ, nullQ, tryPopR, pushL)
import Data.IORef (newIORef, readIORef)
import Streamly.Internal.Control.Concurrent
    (MonadRunInIO, MonadAsync, RunInIO(..), askRunInIO, restoreM)
import Streamly.Internal.Data.Channel.Dispatcher (delThread)

import qualified Data.Set as Set
import qualified Streamly.Internal.Data.StreamK as K

import Streamly.Internal.Data.Channel.Types
import Streamly.Internal.Data.Stream.Channel.Consumer
import Streamly.Internal.Data.Stream.Channel.Type

------------------------------------------------------------------------------
-- Creating a channel
------------------------------------------------------------------------------

data WorkerStatus = Continue | Suspend

-- XXX This is not strictly round-robin as the streams that are faster may
-- yield more elements than the ones that are slower. Also, when streams
-- suspend due to buffer getting full they get added to the queue in a random
-- order. Document this under interleaved config option or fix it.

{-# INLINE enqueueFIFO #-}
enqueueFIFO ::
       Channel m a
    -> LinkedQueue (RunInIO m, K.StreamK m a)
    -> (RunInIO m, K.StreamK m a)
    -> IO ()
enqueueFIFO sv q m = do
    pushL q m
    ringDoorBell (doorBellOnWorkQ sv) (outputDoorBell sv)

{-# INLINE workLoopFIFO #-}
workLoopFIFO
    :: MonadRunInIO m
    => LinkedQueue (RunInIO m, K.StreamK m a)
    -> Channel m a
    -> Maybe WorkerInfo
    -> m ()
workLoopFIFO q sv winfo = run

    where

    run = do
        work <- liftIO $ tryPopR q
        case work of
            Nothing -> return ()
            Just (RunInIO runin, m) -> do
                r <- liftIO
                        $ runin
                        $ K.foldStreamShared
                            undefined yieldk single (return Continue) m
                res <- restoreM r
                case res of
                    Continue -> run
                    Suspend -> return ()

    single a = do
        res <- liftIO $ yieldWith winfo sv a
        return $ if res then Continue else Suspend

    -- XXX in general we would like to yield "n" elements from a single stream
    -- before moving on to the next. Single element granularity could be too
    -- expensive in certain cases. Similarly, we can use time limit for
    -- yielding.
    yieldk a r = do
        res <- liftIO $ yieldWith winfo sv a
        runInIO <- askRunInIO
        -- XXX If the queue is empty we do not need to enqueue. We can just
        -- continue evaluating the stream.
        liftIO $ enqueueFIFO sv q (runInIO, r)
        return $ if res then Continue else Suspend

{-# INLINE workLoopFIFOLimited #-}
workLoopFIFOLimited
    :: forall m a. MonadRunInIO m
    => LinkedQueue (RunInIO m, K.StreamK m a)
    -> Channel m a
    -> Maybe WorkerInfo
    -> m ()
workLoopFIFOLimited q sv winfo = run

    where

    incrContinue =
        liftIO (incrementYieldLimit (remainingWork sv)) >> return Continue

    run = do
        work <- liftIO $ tryPopR q
        case work of
            Nothing -> return ()
            Just (RunInIO runin, m) -> do
                yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
                if yieldLimitOk
                then do
                    r <- liftIO
                            $ runin
                            $ K.foldStreamShared
                                undefined yieldk single incrContinue m
                    res <- restoreM r
                    case res of
                        Continue -> run
                        Suspend -> return ()
                else liftIO $ do
                    enqueueFIFO sv q (RunInIO runin, m)
                    incrementYieldLimit (remainingWork sv)

    single a = do
        res <- liftIO $ yieldWith winfo sv a
        return $ if res then Continue else Suspend

    yieldk a r = do
        res <- liftIO $ yieldWith winfo sv a
        runInIO <- askRunInIO
        liftIO $ enqueueFIFO sv q (runInIO, r)
        yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
        if res && yieldLimitOk
        then return Continue
        else liftIO $ do
            incrementYieldLimit (remainingWork sv)
            return Suspend

-------------------------------------------------------------------------------
-- SVar creation
-------------------------------------------------------------------------------

-- XXX we have this function in this file because passing runStreamLIFO as a
-- function argument to this function results in a perf degradation of more
-- than 10%.  Need to investigate what the root cause is.
-- Interestingly, the same thing does not make any difference for Ahead.
getFifoSVar :: forall m a. MonadRunInIO m =>
    RunInIO m -> Config -> IO (Channel m a)
getFifoSVar mrun cfg = do
    outQ    <- newIORef ([], 0)
    outQMv  <- newEmptyMVar
    active  <- newIORef 0
    wfw     <- newIORef False
    running <- newIORef Set.empty
    q       <- newQ
    yl      <- case getYieldLimit cfg of
                Nothing -> return Nothing
                Just x -> Just <$> newIORef x
    stoppingRef <- newIORef False
    stoppedMVar <- newMVar False
    rateInfo <- newRateInfo cfg

    stats <- newSVarStats
    tid <- myThreadId

    let isWorkFinished _ = nullQ q
    let isWorkFinishedLimited sv = do
            yieldsDone <-
                    case remainingWork sv of
                        Just ref -> do
                            n <- readIORef ref
                            return (n <= 0)
                        Nothing -> return False
            qEmpty <- nullQ q
            return $ qEmpty || yieldsDone

    let getSVar :: Channel m a
            -> (Channel m a -> m [ChildEvent a])
            -> (Channel m a -> m Bool)
            -> (Channel m a -> IO Bool)
            -> (LinkedQueue (RunInIO m, K.StreamK m a)
                -> Channel m a
                -> Maybe WorkerInfo
                -> m())
            -> Channel m a
        getSVar sv readOutput postProc workDone wloop = Channel
            { outputQueue      = outQ
            , remainingWork    = yl
            , maxBufferLimit   = getMaxBuffer cfg
            , maxWorkerLimit   = min (getMaxThreads cfg) (getMaxBuffer cfg)
            , yieldRateInfo    = rateInfo
            , outputDoorBell   = outQMv
            , readOutputQ      = readOutput sv
            , postProcess      = postProc sv
            , workerThreads    = running
            , workLoop         = wloop q sv
            , channelStopping   = stoppingRef
            , channelStopped   = stoppedMVar
            , enqueue          = enqueueFIFO sv q
            , eagerDispatch    = return ()
            , isWorkDone       = workDone sv
            , isQueueDone      = workDone sv
            , doorBellOnWorkQ  = wfw
            , svarMrun         = mrun
            , workerCount      = active
            , accountThread    = delThread running
            , workerStopMVar   = undefined
            , svarRef          = Nothing
            , svarInspectMode  = getInspectMode cfg
            , svarCreator      = tid
            , svarStats        = stats
            }

    let sv =
            case getStreamRate cfg of
                Nothing ->
                    case getYieldLimit cfg of
                        Nothing -> getSVar sv (readOutputQBounded False)
                                              postProcessBounded
                                              isWorkFinished
                                              workLoopFIFO
                        Just _  -> getSVar sv (readOutputQBounded False)
                                              postProcessBounded
                                              isWorkFinishedLimited
                                              workLoopFIFOLimited
                Just _  ->
                    case getYieldLimit cfg of
                        Nothing -> getSVar sv readOutputQPaced
                                              postProcessPaced
                                              isWorkFinished
                                              workLoopFIFO
                        Just _  -> getSVar sv readOutputQPaced
                                              postProcessPaced
                                              isWorkFinishedLimited
                                              workLoopFIFOLimited
     in return sv

-- XXX GHC: If instead of MonadAsync we use (MonadIO m, MonadBaseControl IO m)
-- constraint we get a 2x perf regression. Need to look into that.

-- | Create a new 'interleaved' style concurrent stream evaluation channel. The
-- monad state used to run the stream actions is taken from the call site of
-- newInterleaveChannel.
--
-- This is a low level API, use newChannel instead.
{-# INLINABLE newInterleaveChannel #-}
{-# SPECIALIZE newInterleaveChannel :: (Config -> Config) -> IO (Channel IO a) #-}
newInterleaveChannel :: MonadAsync m =>
    (Config -> Config) -> m (Channel m a)
newInterleaveChannel modifier = do
    mrun <- askRunInIO
    liftIO $ getFifoSVar mrun (modifier defaultConfig)