git.iko.soy/streamly
on e01057f help sign in
streamly/src/Streamly/Internal/Data/Stream/Concurrent/Channel/Interleave.hs
250 lines 9.3 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.Concurrent.Channel.Interleave
3-- Copyright : (c) 2017 Composewell Technologies
4-- License : BSD-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
9module Streamly.Internal.Data.Stream.Concurrent.Channel.Interleave
10 (
11 newChannel
12 )
13where
15#include "inline.hs"
17import Control.Concurrent (myThreadId)
18import Control.Concurrent.MVar (newEmptyMVar)
19import Control.Monad.IO.Class (MonadIO(liftIO))
20import Data.Concurrent.Queue.MichaelScott (LinkedQueue, newQ, nullQ, tryPopR, pushL)
21import Data.IORef (newIORef, readIORef)
22import Streamly.Internal.Control.Concurrent
23 (MonadRunInIO, MonadAsync, RunInIO(..), askRunInIO, restoreM)
24import Streamly.Internal.Data.Stream.Channel.Dispatcher (delThread)
26import qualified Data.Set as Set
27import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
29import Streamly.Internal.Data.Stream.Concurrent.Channel.Consumer
30import Streamly.Internal.Data.Stream.Concurrent.Channel.Type
31import Streamly.Internal.Data.Stream.Channel.Types
33------------------------------------------------------------------------------
34-- Creating a channel
35------------------------------------------------------------------------------
37data WorkerStatus = Continue | Suspend
39-- XXX This is not strictly round-robin as the streams that are faster may
40-- yield more elements than the ones that are slower. Also, when streams
41-- suspend due to buffer getting full they get added to the queue in a random
42-- order.
44{-# INLINE enqueueFIFO #-}
45enqueueFIFO ::
46 Channel m a
47 -> LinkedQueue (RunInIO m, K.StreamK m a)
48 -> (RunInIO m, K.StreamK m a)
49 -> IO ()
50enqueueFIFO sv q m = do
51 pushL q m
52 ringDoorBell (doorBellOnWorkQ sv) (outputDoorBell sv)
54{-# INLINE workLoopFIFO #-}
55workLoopFIFO
56 :: MonadRunInIO m
57 => LinkedQueue (RunInIO m, K.StreamK m a)
58 -> Channel m a
59 -> Maybe WorkerInfo
60 -> m ()
61workLoopFIFO q sv winfo = run
63 where
65 run = do
66 work <- liftIO $ tryPopR q
67 case work of
68 Nothing -> liftIO $ stop sv winfo
69 Just (RunInIO runin, m) -> do
70 r <- liftIO
71 $ runin
72 $ K.foldStreamShared
73 undefined yieldk single (return Continue) m
74 res <- restoreM r
75 case res of
76 Continue -> run
77 Suspend -> liftIO $ stop sv winfo
79 single a = do
80 res <- liftIO $ yield sv winfo a
81 return $ if res then Continue else Suspend
83 -- XXX in general we would like to yield "n" elements from a single stream
84 -- before moving on to the next. Single element granularity could be too
85 -- expensive in certain cases. Similarly, we can use time limit for
86 -- yielding.
87 yieldk a r = do
88 res <- liftIO $ yield sv winfo a
89 runInIO <- askRunInIO
90 liftIO $ enqueueFIFO sv q (runInIO, r)
91 return $ if res then Continue else Suspend
93{-# INLINE workLoopFIFOLimited #-}
94workLoopFIFOLimited
95 :: forall m a. MonadRunInIO m
96 => LinkedQueue (RunInIO m, K.StreamK m a)
97 -> Channel m a
98 -> Maybe WorkerInfo
99 -> m ()
100workLoopFIFOLimited q sv winfo = run
102 where
104 incrContinue =
105 liftIO (incrementYieldLimit (remainingWork sv)) >> return Continue
107 run = do
108 work <- liftIO $ tryPopR q
109 case work of
110 Nothing -> liftIO $ stop sv winfo
111 Just (RunInIO runin, m) -> do
112 yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
113 if yieldLimitOk
114 then do
115 r <- liftIO
116 $ runin
117 $ K.foldStreamShared
118 undefined yieldk single incrContinue m
119 res <- restoreM r
120 case res of
121 Continue -> run
122 Suspend -> liftIO $ stop sv winfo
123 else liftIO $ do
124 enqueueFIFO sv q (RunInIO runin, m)
125 incrementYieldLimit (remainingWork sv)
126 stop sv winfo
128 single a = do
129 res <- liftIO $ yield sv winfo a
130 return $ if res then Continue else Suspend
132 yieldk a r = do
133 res <- liftIO $ yield sv winfo a
134 runInIO <- askRunInIO
135 liftIO $ enqueueFIFO sv q (runInIO, r)
136 yieldLimitOk <- liftIO $ decrementYieldLimit (remainingWork sv)
137 if res && yieldLimitOk
138 then return Continue
139 else liftIO $ do
140 incrementYieldLimit (remainingWork sv)
141 return Suspend
143-------------------------------------------------------------------------------
144-- SVar creation
145-------------------------------------------------------------------------------
147-- XXX we have this function in this file because passing runStreamLIFO as a
148-- function argument to this function results in a perf degradation of more
149-- than 10%. Need to investigate what the root cause is.
150-- Interestingly, the same thing does not make any difference for Ahead.
151getFifoSVar :: forall m a. MonadRunInIO m =>
152 RunInIO m -> Config -> IO (Channel m a)
153getFifoSVar mrun cfg = do
154 outQ <- newIORef ([], 0)
155 outQMv <- newEmptyMVar
156 active <- newIORef 0
157 wfw <- newIORef False
158 running <- newIORef Set.empty
159 q <- newQ
160 yl <- case getYieldLimit cfg of
161 Nothing -> return Nothing
162 Just x -> Just <$> newIORef x
163 rateInfo <- newRateInfo cfg
165 stats <- newSVarStats
166 tid <- myThreadId
168 let isWorkFinished _ = nullQ q
169 let isWorkFinishedLimited sv = do
170 yieldsDone <-
171 case remainingWork sv of
172 Just ref -> do
173 n <- readIORef ref
174 return (n <= 0)
175 Nothing -> return False
176 qEmpty <- nullQ q
177 return $ qEmpty || yieldsDone
179 let getSVar :: Channel m a
180 -> (Channel m a -> m [ChildEvent a])
181 -> (Channel m a -> m Bool)
182 -> (Channel m a -> IO Bool)
183 -> (LinkedQueue (RunInIO m, K.StreamK m a)
184 -> Channel m a
185 -> Maybe WorkerInfo
186 -> m())
187 -> Channel m a
188 getSVar sv readOutput postProc workDone wloop = Channel
189 { outputQueue = outQ
190 , remainingWork = yl
191 , maxBufferLimit = getMaxBuffer cfg
192 , maxWorkerLimit = min (getMaxThreads cfg) (getMaxBuffer cfg)
193 , yieldRateInfo = rateInfo
194 , outputDoorBell = outQMv
195 , readOutputQ = readOutput sv
196 , postProcess = postProc sv
197 , workerThreads = running
198 , workLoop = wloop q sv
199 , enqueue = \_ -> enqueueFIFO sv q
200 , eagerDispatch = return ()
201 , isWorkDone = workDone sv
202 , isQueueDone = workDone sv
203 , doorBellOnWorkQ = wfw
204 , svarMrun = mrun
205 , workerCount = active
206 , accountThread = delThread running
207 , workerStopMVar = undefined
208 , svarRef = Nothing
209 , svarInspectMode = getInspectMode cfg
210 , svarCreator = tid
211 , svarStats = stats
212 }
214 let sv =
215 case getStreamRate cfg of
216 Nothing ->
217 case getYieldLimit cfg of
218 Nothing -> getSVar sv (readOutputQBounded False)
219 postProcessBounded
220 isWorkFinished
221 workLoopFIFO
222 Just _ -> getSVar sv (readOutputQBounded False)
223 postProcessBounded
224 isWorkFinishedLimited
225 workLoopFIFOLimited
226 Just _ ->
227 case getYieldLimit cfg of
228 Nothing -> getSVar sv readOutputQPaced
229 postProcessPaced
230 isWorkFinished
231 workLoopFIFO
232 Just _ -> getSVar sv readOutputQPaced
233 postProcessPaced
234 isWorkFinishedLimited
235 workLoopFIFOLimited
236 in return sv
238-- XXX GHC: If instead of MonadAsync we use (MonadIO m, MonadBaseControl IO m)
239-- constraint we get a 2x perf regression. Need to look into that.
240--
241-- | Create a new async style concurrent stream evaluation channel. The monad
242-- state used to run the stream actions is taken from the call site of
243-- newChannel.
244{-# INLINABLE newChannel #-}
245{-# SPECIALIZE newChannel :: (Config -> Config) -> IO (Channel IO a) #-}
246newChannel :: MonadAsync m =>
247 (Config -> Config) -> m (Channel m a)
248newChannel modifier = do
249 mrun <- askRunInIO
250 liftIO $ getFifoSVar mrun (modifier defaultConfig)