@@ -88,8 +88,9 @@ module Streamly.SVar
8888 , workerRateControl
8989 , updateYieldCount
9090 , decrementYieldLimit
91− , decrementYieldLimitPost
9291 , incrementYieldLimit
92+ , decrementBufferLimit
93+ , incrementBufferLimit
9394 , postProcessBounded
9495 , postProcessPaced
9596 , readOutputQBounded
@@ -107,7 +108,7 @@ where
107108import Control.Concurrent
108109 (ThreadId, myThreadId, threadDelay, throwTo)
109110import Control.Concurrent.MVar
110− (MVar, newEmptyMVar, tryPutMVar, takeMVar, newMVar, tryReadMVar)
111+ (MVar, newEmptyMVar, tryPutMVar, takeMVar, tryTakeMVar, newMVar, tryReadMVar)
111112import Control.Exception
112113 (SomeException(..), catch, mask, assert, Exception, catches,
113114 throwIO, Handler(..), BlockedIndefinitelyOnMVar(..),
@@ -350,6 +351,26 @@ data SVarStopStyle =
350351 | StopBy -- stop when a specific stream finishes
351352 deriving (Eq, Show)
352353
354+-- | Buffering policy for persistent push workers (in ParallelT). In a pull
355+-- style SVar (in AsyncT, AheadT etc.), the consumer side dispatches workers on
356+-- demand, workers terminate if the buffer is full or if the consumer is not
357+-- cosuming fast enough. In a push style SVar, a worker is dispatched only
358+-- once, workers are persistent and keep pushing work to the consumer via a
359+-- bounded buffer. If the buffer becomes full the worker either blocks, or it
360+-- can drop an item from the buffer to make space.
361+--
362+-- Pull style SVars are useful in lazy stream evaluation whereas push style
363+-- SVars are useful in strict left Folds.
364+--
365+-- XXX Maybe we can separate the implementation in two different types instead
366+-- of using a common SVar type.
367+--
368+data PushBufferPolicy =
369+ PushBufferDropNew -- drop the latest element and continue
370+ | PushBufferDropOld -- drop the oldest element and continue
371+ | PushBufferBlock -- block the thread until space
372+ -- becomes available
373+
353374-- IMPORTANT NOTE: we cannot update the SVar after generating it as we have
354375-- references to the original SVar stored in several functions which will keep
355376-- pointing to the original data and the new updates won't reflect there.
@@ -392,6 +413,13 @@ data SVar t m a = SVar
392413 -- case exceeding the requested buffer size.
393414 , maxWorkerLimit :: Limit
394415 , maxBufferLimit :: Limit
416+ -- These two are valid and used only when maxBufferLimit is Limited.
417+ , pushBufferSpace :: IORef Count
418+ , pushBufferPolicy :: PushBufferPolicy
419+ -- [LOCKING] The consumer puts this MVar after emptying the buffer, workers
420+ -- block on it when the buffer becomes full. No overhead unless the buffer
421+ -- becomes full.
422+ , pushBufferMVar :: MVar ()
395423
396424 -- [LOCKING] Read only access by consumer when dispatching a worker.
397425 -- Decremented by workers when picking work and undo decrement if the
@@ -920,17 +948,6 @@ decrementYieldLimit sv =
920948 r <- atomicModifyIORefCAS ref $ \x -> (x - 1, x)
921949 return $ r >= 1
922950
923−-- decrementYieldLimit returns False when the old limit is 0. This one returns
924−-- False when the old limit is 1.
925−{-# INLINE decrementYieldLimitPost #-}
926−decrementYieldLimitPost :: SVar t m a -> IO Bool
927−decrementYieldLimitPost sv =
928− case remainingWork sv of
929− Nothing -> return True
930− Just ref -> do
931− r <- atomicModifyIORefCAS ref $ \x -> (x - 1, x)
932− return $ r > 1
933−
934951{-# INLINE incrementYieldLimit #-}
935952incrementYieldLimit :: SVar t m a -> IO ()
936953incrementYieldLimit sv =
@@ -946,15 +963,81 @@ incrementYieldLimit sv =
946963-- XXX Only yields should be counted in the buffer limit and not the Stop
947964-- events.
948965
966+{-# INLINE decrementBufferLimit #-}
967+decrementBufferLimit :: SVar t m a -> IO ()
968+decrementBufferLimit sv =
969+ case maxBufferLimit sv of
970+ Unlimited -> return ()
971+ Limited _ -> do
972+ let ref = pushBufferSpace sv
973+ old <- atomicModifyIORefCAS ref $ \x ->
974+ (if x >= 1 then x - 1 else x, x)
975+ when (old <= 0) $
976+ case pushBufferPolicy sv of
977+ PushBufferBlock -> blockAndRetry
978+ PushBufferDropNew -> do
979+ -- We just drop one item and proceed. It is possible
980+ -- that by the time we drop the item the consumer
981+ -- thread might have run and created space in the
982+ -- buffer, but we do not care about that condition.
983+ -- This is not pedantically correct but it should be
984+ -- fine in practice.
985+ -- XXX we may want to drop only if n == maxBuf
986+ -- otherwise we must have space in the buffer and a
987+ -- decrement should be possible.
988+ block <- atomicModifyIORefCAS (outputQueue sv) $
989+ \(es, n) ->
990+ case es of
991+ [] -> (([],n), True)
992+ _ : xs -> ((xs, n - 1), False)
993+ when block blockAndRetry
994+ -- XXX need a dequeue or ring buffer for this
995+ PushBufferDropOld -> undefined
996+
997+ where
998+
999+ blockAndRetry = do
1000+ let ref = pushBufferSpace sv
1001+ liftIO $ takeMVar (pushBufferMVar sv)
1002+ old <- atomicModifyIORefCAS ref $ \x ->
1003+ (if x >= 1 then x - 1 else x, x)
1004+ -- When multiple threads sleep on takeMVar, the first thread would
1005+ -- wakeup due to a putMVar by the consumer, but the rest of the threads
1006+ -- would have to put back the MVar after taking it and decrementing the
1007+ -- buffer count, otherwise all other threads will remain asleep.
1008+ if old >= 1
1009+ then void $ liftIO $ tryPutMVar (pushBufferMVar sv) ()
1010+ -- We do not put the MVar back in this case, instead we
1011+ -- wait for the consumer to put it.
1012+ else blockAndRetry
1013+
1014+{-# INLINE incrementBufferLimit #-}
1015+incrementBufferLimit :: SVar t m a -> IO ()
1016+incrementBufferLimit sv =
1017+ case maxBufferLimit sv of
1018+ Unlimited -> return ()
1019+ Limited _ -> do
1020+ atomicModifyIORefCAS_ (pushBufferSpace sv) (+ 1)
1021+ writeBarrier
1022+ void $ liftIO $ tryPutMVar (pushBufferMVar sv) ()
1023+
1024+{-# INLINE resetBufferLimit #-}
1025+resetBufferLimit :: SVar t m a -> IO ()
1026+resetBufferLimit sv =
1027+ case maxBufferLimit sv of
1028+ Unlimited -> return ()
1029+ Limited n -> atomicModifyIORefCAS_ (pushBufferSpace sv)
1030+ (const (fromIntegral n))
1031+
9491032-- | This function is used by the producer threads to queue output for the
9501033-- consumer thread to consume. Returns whether the queue has more space.
951−send :: SVar t m a -> ChildEvent a -> IO Bool
1034+send :: SVar t m a -> ChildEvent a -> IO Int
9521035send sv msg = do
9531036 -- XXX can the access to outputQueue and maxBufferLimit be made faster
9541037 -- somehow?
955− len <- atomicModifyIORefCAS (outputQueue sv) $ \(es, n) ->
1038+ oldlen <- atomicModifyIORefCAS (outputQueue sv) $ \(es, n) ->
9561039 ((msg : es, n + 1), n)
957− when (len <= 0) $ do
1040+ when (oldlen <= 0) $ do
9581041 -- The wake up must happen only after the store has finished otherwise
9591042 -- we can have lost wakeup problems.
9601043 writeBarrier
@@ -965,13 +1048,7 @@ send sv msg = do
9651048 -- The important point is that the consumer is guaranteed to receive a
9661049 -- doorbell if something was added to the queue after it empties it.
9671050 void $ tryPutMVar (outputDoorBell sv) ()
968−
969− let limit = maxBufferLimit sv
970− case limit of
971− Unlimited -> return True
972− Limited lim -> do
973− active <- readIORef (workerCount sv)
974− return $ len < (fromIntegral lim - active)
1051+ return oldlen
9751052
9761053workerCollectLatency :: WorkerInfo -> IO (Maybe (Count, NanoSecond64))
9771054workerCollectLatency winfo = do
@@ -1062,10 +1139,18 @@ workerRateControl sv yinfo winfo = do
10621139-- XXX we should do rate control here but not latency update in case of ahead
10631140-- streams. latency update must be done when we yield directly to outputQueue
10641141-- or when we yield to heap.
1142+--
1143+-- returns whether the worker should continue (True) or stop (False).
10651144{-# INLINE sendYield #-}
10661145sendYield :: SVar t m a -> Maybe WorkerInfo -> ChildEvent a -> IO Bool
10671146sendYield sv mwinfo msg = do
1068− r <- send sv msg
1147+ oldlen <- send sv msg
1148+ let limit = maxBufferLimit sv
1149+ bufferSpaceOk <- case limit of
1150+ Unlimited -> return True
1151+ Limited lim -> do
1152+ active <- readIORef (workerCount sv)
1153+ return $ (oldlen + 1) < (fromIntegral lim - active)
10691154 rateLimitOk <-
10701155 case mwinfo of
10711156 Just winfo ->
@@ -1073,7 +1158,7 @@ sendYield sv mwinfo msg = do
10731158 Nothing -> return True
10741159 Just yinfo -> workerRateControl sv yinfo winfo
10751160 Nothing -> return True
1076− return $ r && rateLimitOk
1161+ return $ bufferSpaceOk && rateLimitOk
10771162
10781163{-# INLINE workerStopUpdate #-}
10791164workerStopUpdate :: WorkerInfo -> YieldRateInfo -> IO ()
@@ -2102,6 +2187,9 @@ getAheadSVar st f mrun = do
21022187 { outputQueue = outQ
21032188 , remainingWork = yl
21042189 , maxBufferLimit = getMaxBuffer st
2190+ , pushBufferSpace = undefined
2191+ , pushBufferPolicy = undefined
2192+ , pushBufferMVar = undefined
21052193 , maxWorkerLimit = min (getMaxThreads st) (getMaxBuffer st)
21062194 , yieldRateInfo = rateInfo
21072195 , outputDoorBell = outQMv
@@ -2173,6 +2261,12 @@ getParallelSVar ss st mrun = do
21732261 Nothing -> return Nothing
21742262 Just x -> Just <$> newIORef x
21752263 rateInfo <- getYieldRateInfo st
2264+ let bufLim =
2265+ case getMaxBuffer st of
2266+ Unlimited -> undefined
2267+ Limited x -> (fromIntegral x)
2268+ remBuf <- newIORef bufLim
2269+ pbMVar <- newMVar ()
21762270
21772271 stats <- newSVarStats
21782272 tid <- myThreadId
@@ -2184,8 +2278,11 @@ getParallelSVar ss st mrun = do
21842278
21852279 let sv =
21862280 SVar { outputQueue = outQ
2187− , remainingWork = yl
2188− , maxBufferLimit = Unlimited
2281+ , remainingWork = yl
2282+ , maxBufferLimit = getMaxBuffer st
2283+ , pushBufferSpace = remBuf
2284+ , pushBufferPolicy = PushBufferBlock
2285+ , pushBufferMVar = pbMVar
21892286 , maxWorkerLimit = Unlimited
21902287 -- Used only for diagnostics
21912288 , yieldRateInfo = rateInfo
@@ -2222,7 +2319,13 @@ getParallelSVar ss st mrun = do
22222319 case yieldRateInfo sv of
22232320 Nothing -> return ()
22242321 Just yinfo -> void $ collectLatency sv yinfo False
2225− fst `fmap` readOutputQRaw sv
2322+ r <- fst `fmap` readOutputQRaw sv
2323+ liftIO $ do
2324+ void $ tryTakeMVar (pushBufferMVar sv)
2325+ resetBufferLimit sv
2326+ writeBarrier
2327+ void $ tryPutMVar (pushBufferMVar sv) ()
2328+ return r
22262329
22272330sendFirstWorker :: MonadAsync m => SVar t m a -> t m a -> m (SVar t m a)
22282331sendFirstWorker sv m = do