cf87d11 · streamly · git.iko.soy

git.iko.soy/streamly

implement maxYields to put a limit on total yields of a stream

maxYields is used to limit the concurrent executions of a stream when it it is immediately followed by a "take" limiting the size of the stream. Also fix the maxBuffer implementation of aheadly.

commitcf87d1132e5c118f90d0462775ee91e4423ebcc2 authorHarendra Kumar <harendra.kumar@gmail.com> dateMon 2 Jul 2018, 08:46 (8 years ago) parent516cab5

7 files changed+109−24

compare with master

unifiedsplit

src/Streamly/Prelude.hs+2−1 src/Streamly/SVar.hs+37−10 src/Streamly/Streams/Ahead.hs+13−6 src/Streamly/Streams/Async.hs+14−6 src/Streamly/Streams/Parallel.hs+1−1 src/Streamly/Streams/SVar.hs+31 test/Main.hs+11

src/Streamly/Prelude.hs+2−1

@@ -165,6 +165,7 @@ import qualified Prelude

165165import qualified System.IO as IO

166166

167167import Streamly.SVar (MonadAsync, defState, rstState)

168+import Streamly.Streams.SVar (maxYields)

168169import Streamly.Streams.StreamK (IsStream(..))

169170import Streamly.Streams.Serial (SerialT)

170171

@@ -688,7 +689,7 @@ filterM p m = fromStreamD $ D.filterM p $ toStreamD m

688689-- @since 0.1.0

689690{-# INLINE take #-}

690691take :: (IsStream t, Monad m) => Int -> t m a -> t m a

691−take n m = fromStreamS $ S.take n $ toStreamS m

692+take n m = fromStreamS $ S.take n $ toStreamS (maxYields (Just n) m)

692693

693694-- | End the stream as soon as the predicate fails on an element.

694695--

src/Streamly/SVar.hs+37−10

@@ -37,7 +37,7 @@ module Streamly.SVar

3737 , atomicModifyIORefCAS

3838 , ChildEvent (..)

3939 , AheadHeapEntry (..)

40− , send

40+ , sendYield

4141 , sendStop

4242 , enqueueLIFO

4343 , workLoopLIFO

@@ -46,7 +46,6 @@ module Streamly.SVar

4646 , enqueueAhead

4747 , pushWorkerPar

4848

49− , maxHeap

5049 , queueEmptyAhead

5150 , dequeueAhead

5251 , dequeueFromHeap

@@ -158,6 +157,7 @@ data SVar t m a =

158157

159158 -- Shared output queue (events, length)

160159 , outputQueue :: IORef ([ChildEvent a], Int)

160+ , maxYieldLimit :: Maybe (IORef Int)

161161 , outputDoorBell :: MVar () -- signal the consumer about output

162162 , readOutputQ :: m [ChildEvent a]

163163 , postProcess :: m Bool

@@ -187,9 +187,10 @@ data SVar t m a =

187187 }

188188

189189data State t m a = State

190− { streamVar :: Maybe (SVar t m a)

190+ { streamVar :: Maybe (SVar t m a)

191+ , yieldLimit :: Maybe Int

191192 , threadsHigh :: Int

192− , bufferHigh :: Int

193+ , bufferHigh :: Int

193194 }

194195

195196defaultMaxThreads, defaultMaxBuffer :: Int

@@ -199,15 +200,23 @@ defaultMaxBuffer = 1500

199200defState :: State t m a

200201defState = State

201202 { streamVar = Nothing

203+ , yieldLimit = Nothing

202204 , threadsHigh = defaultMaxThreads

203205 , bufferHigh = defaultMaxBuffer

204206 }

205207

208+-- XXX if perf gets affected we can have all the Nothing params in a single

209+-- structure so that we reset is fast. We can also use rewrite rules such that

210+-- reset occurs only in concurrent streams to reduce the impact on serial

211+-- streams.

206212-- We can optimize this so that we clear it only if it is a Just value, it

207213-- results in slightly better perf for zip/zipM but the performance of scan

208214-- worsens a lot, it does not fuse.

209215rstState :: State t m a -> State t m b

210−rstState st = st {streamVar = Nothing}

216+rstState st = st

217+ { streamVar = Nothing

218+ , yieldLimit = Nothing

219+ }

211220

212221#ifdef DIAGNOSTICS

213222{-# NOINLINE dumpSVar #-}

@@ -337,7 +346,6 @@ doFork action exHandler =

337346

338347-- | This function is used by the producer threads to queue output for the

339348-- consumer thread to consume. Returns whether the queue has more space.

340−{-# NOINLINE send #-}

341349send :: Int -> SVar t m a -> ChildEvent a -> IO Bool

342350send maxOutputQLen sv msg = do

343351 len <- atomicModifyIORefCAS (outputQueue sv) $ \(es, n) ->

@@ -355,6 +363,15 @@ send maxOutputQLen sv msg = do

355363 void $ tryPutMVar (outputDoorBell sv) ()

356364 return (len < maxOutputQLen || maxOutputQLen < 0)

357365

366+{-# NOINLINE sendYield #-}

367+sendYield :: Int -> SVar t m a -> ChildEvent a -> IO Bool

368+sendYield maxOutputQLen sv msg = do

369+ ylimit <- case maxYieldLimit sv of

370+ Nothing -> return True

371+ Just ref -> atomicModifyIORefCAS ref $ \x -> (x - 1, x > 1)

372+ r <- send maxOutputQLen sv msg

373+ return $ r && ylimit

374+

358375{-# NOINLINE sendStop #-}

359376sendStop :: SVar t m a -> IO ()

360377sendStop sv = do

@@ -534,9 +551,6 @@ enqueueAhead sv q m = do

534551--

535552-- XXX review for livelock

536553--

537−maxHeap :: Int

538−maxHeap = 1500

539−

540554{-# INLINE queueEmptyAhead #-}

541555queueEmptyAhead :: MonadIO m => IORef ([t m a], Int) -> m Bool

542556queueEmptyAhead q = liftIO $ do

@@ -675,7 +689,15 @@ dispatchWorker maxWorkerLimit sv = do

675689 -- executing. In that case we should either configure the maxWorker

676690 -- count to higher or use parallel style instead of ahead or async

677691 -- style.

678− when (cnt < maxWorkerLimit || maxWorkerLimit < 0) $ pushWorker sv

692+ limit <- case maxYieldLimit sv of

693+ Nothing -> return maxWorkerLimit

694+ Just x -> do

695+ lim <- liftIO $ readIORef x

696+ return $

697+ if maxWorkerLimit > 0

698+ then min maxWorkerLimit lim

699+ else lim

700+ when (cnt < limit || limit < 0) $ pushWorker sv

679701

680702{-# NOINLINE sendWorkerWait #-}

681703sendWorkerWait :: MonadAsync m => Int -> SVar t m a -> m ()

@@ -816,6 +838,9 @@ getAheadSVar st f = do

816838 wfw <- newIORef False

817839 running <- newIORef S.empty

818840 q <- newIORef ([], -1)

841+ yl <- case yieldLimit st of

842+ Nothing -> return Nothing

843+ Just x -> Just <$> newIORef x

819844

820845#ifdef DIAGNOSTICS

821846 disp <- newIORef 0

@@ -826,6 +851,7 @@ getAheadSVar st f = do

826851#endif

827852 let sv =

828853 SVar { outputQueue = outQ

854+ , maxYieldLimit = yl

829855 , outputDoorBell = outQMv

830856 , readOutputQ = readOutputQBounded (threadsHigh st) sv

831857 , postProcess = postProcessBounded sv

@@ -879,6 +905,7 @@ getParallelSVar = do

879905#endif

880906 let sv =

881907 SVar { outputQueue = outQ

908+ , maxYieldLimit = Nothing

882909 , outputDoorBell = outQMv

883910 , readOutputQ = readOutputQPar sv

884911 , postProcess = allThreadsDone sv

src/Streamly/Streams/Ahead.hs+13−6

@@ -37,7 +37,7 @@ import Control.Monad.State.Class (MonadState(..))

3737import Control.Monad.Trans.Class (MonadTrans(lift))

3838import Data.Atomics (atomicModifyIORefCAS_)

3939import Data.Heap (Heap, Entry(..))

40−import Data.IORef (IORef)

40+import Data.IORef (IORef, readIORef)

4141import Data.Maybe (fromJust)

4242import Data.Semigroup (Semigroup(..))

4343

@@ -51,7 +51,7 @@ import qualified Streamly.Streams.StreamK as K

5151

5252#ifdef DIAGNOSTICS

5353import Control.Monad (when)

54−import Data.IORef (writeIORef, readIORef)

54+import Data.IORef (writeIORef)

5555#endif

5656import Prelude hiding (map)

5757

@@ -128,7 +128,14 @@ workLoopAhead st q heap = runHeap

128128 toHeap seqNo ent = do

129129 hp <- liftIO $ atomicModifyIORefCAS heap $ \(h, snum) ->

130130 ((H.insert (Entry seqNo ent) h, snum), h)

131− if H.size hp <= maxHeap

131+ (_, len) <- liftIO $ readIORef (outputQueue sv)

132+ let maxHeap = maxBuf - len

133+ limit <- case maxYieldLimit sv of

134+ Nothing -> return maxHeap

135+ Just ref -> do

136+ r <- liftIO $ readIORef ref

137+ return $ if r >= 0 then r else maxHeap

138+ if H.size hp <= limit

132139 then runHeap

133140 else liftIO $ sendStop sv

134141

@@ -136,7 +143,7 @@ workLoopAhead st q heap = runHeap

136143 yieldToHeap seqNo a r = toHeap seqNo (AheadEntryStream (a `K.cons` r))

137144

138145 singleOutput seqNo a = do

139− continue <- liftIO $ send maxBuf sv (ChildYield a)

146+ continue <- liftIO $ sendYield maxBuf sv (ChildYield a)

140147 if continue

141148 then runQueueToken seqNo

142149 else liftIO $ do

@@ -144,7 +151,7 @@ workLoopAhead st q heap = runHeap

144151 sendStop sv

145152

146153 yieldOutput seqNo a r = do

147− continue <- liftIO $ send maxBuf sv (ChildYield a)

154+ continue <- liftIO $ sendYield maxBuf sv (ChildYield a)

148155 if continue

149156 then unStream r st (runQueueToken seqNo)

150157 (singleOutput seqNo)

@@ -241,7 +248,7 @@ aheadS m1 m2 = Stream $ \st stp sng yld -> do

241248 -- sequencing results. This means the left side cannot further

242249 -- split into more ahead computations on the same SVar.

243250 unStream m1 (rstState st) stp sng yld

244− _ -> unStream (forkSVarAhead m1 m2) (rstState st) stp sng yld

251+ _ -> unStream (forkSVarAhead m1 m2) st stp sng yld

245252

246253-- | XXX we can implement it more efficienty by directly implementing instead

247254-- of combining streams using ahead.

src/Streamly/Streams/Async.hs+14−6

@@ -72,10 +72,10 @@ runStreamLIFO st q m stop = unStream m st stop single yieldk

7272 sv = fromJust $ streamVar st

7373 maxBuf = bufferHigh st

7474 single a = do

75− res <- liftIO $ send maxBuf sv (ChildYield a)

75+ res <- liftIO $ sendYield maxBuf sv (ChildYield a)

7676 if res then stop else liftIO $ sendStop sv

7777 yieldk a r = do

78− res <- liftIO $ send maxBuf sv (ChildYield a)

78+ res <- liftIO $ sendYield maxBuf sv (ChildYield a)

7979 if res

8080 then (unStream r) st stop single yieldk

8181 else liftIO $ enqueueLIFO sv q r >> sendStop sv

@@ -97,10 +97,10 @@ runStreamFIFO st q m stop = unStream m st stop single yieldk

9797 sv = fromJust $ streamVar st

9898 maxBuf = bufferHigh st

9999 single a = do

100− res <- liftIO $ send maxBuf sv (ChildYield a)

100+ res <- liftIO $ sendYield maxBuf sv (ChildYield a)

101101 if res then stop else liftIO $ sendStop sv

102102 yieldk a r = do

103− res <- liftIO $ send maxBuf sv (ChildYield a)

103+ res <- liftIO $ sendYield maxBuf sv (ChildYield a)

104104 liftIO (enqueueFIFO sv q r)

105105 if res then stop else liftIO $ sendStop sv

106106

@@ -121,6 +121,9 @@ getLifoSVar st = do

121121 wfw <- newIORef False

122122 running <- newIORef S.empty

123123 q <- newIORef []

124+ yl <- case yieldLimit st of

125+ Nothing -> return Nothing

126+ Just x -> Just <$> newIORef x

124127#ifdef DIAGNOSTICS

125128 disp <- newIORef 0

126129 maxWrk <- newIORef 0

@@ -131,6 +134,7 @@ getLifoSVar st = do

131134 let checkEmpty = null <$> readIORef q

132135 let sv =

133136 SVar { outputQueue = outQ

137+ , maxYieldLimit = yl

134138 , outputDoorBell = outQMv

135139 , readOutputQ = readOutputQBounded (threadsHigh st) sv

136140 , postProcess = postProcessBounded sv

@@ -163,6 +167,9 @@ getFifoSVar st = do

163167 wfw <- newIORef False

164168 running <- newIORef S.empty

165169 q <- newQ

170+ yl <- case yieldLimit st of

171+ Nothing -> return Nothing

172+ Just x -> Just <$> newIORef x

166173#ifdef DIAGNOSTICS

167174 disp <- newIORef 0

168175 maxWrk <- newIORef 0

@@ -172,6 +179,7 @@ getFifoSVar st = do

172179#endif

173180 let sv =

174181 SVar { outputQueue = outQ

182+ , maxYieldLimit = yl

175183 , outputDoorBell = outQMv

176184 , readOutputQ = readOutputQBounded (threadsHigh st) sv

177185 , postProcess = postProcessBounded sv

@@ -307,11 +315,11 @@ forkSVarAsync style m1 m2 = Stream $ \st stp sng yld -> do

307315{-# INLINE joinStreamVarAsync #-}

308316joinStreamVarAsync :: MonadAsync m

309317 => SVarStyle -> Stream m a -> Stream m a -> Stream m a

310−joinStreamVarAsync style m1 m2 = Stream $ \st stp sng yld ->

318+joinStreamVarAsync style m1 m2 = Stream $ \st stp sng yld -> do

311319 case streamVar st of

312320 Just sv | svarStyle sv == style ->

313321 liftIO (enqueue sv m2) >> unStream m1 st stp sng yld

314− _ -> unStream (forkSVarAsync style m1 m2) (rstState st) stp sng yld

322+ _ -> unStream (forkSVarAsync style m1 m2) st stp sng yld

315323

316324------------------------------------------------------------------------------

317325-- Semigroup and Monoid style compositions for parallel actions

src/Streamly/Streams/Parallel.hs+1−1

@@ -67,7 +67,7 @@ runOne st m = unStream m st stop single yieldk

6767

6868 sv = fromJust $ streamVar st

6969 stop = liftIO $ sendStop sv

70− sendit a = liftIO $ send (-1) sv (ChildYield a)

70+ sendit a = liftIO $ sendYield (-1) sv (ChildYield a)

7171 single a = sendit a >> stop

7272 -- XXX there is no flow control in parallel case. We should perhaps use a

7373 -- queue and queue it back on that and exit the thread when the outputQueue

src/Streamly/Streams/SVar.hs+31

@@ -9,6 +9,8 @@

99{-# LANGUAGE UnboxedTuples #-}

1010{-# LANGUAGE UndecidableInstances #-} -- XXX

1111

12+#include "inline.h"

13+

1214-- |

1315-- Module : Streamly.Streams.SVar

@@ -25,6 +27,7 @@ module Streamly.Streams.SVar

2527 , toSVar

2628 , maxThreads

2729 , maxBuffer

30+ , maxYields

2831 )

2932where

3033

@@ -32,6 +35,7 @@ import Control.Monad.Catch (throwM)

3235

3336import Streamly.SVar

3437import Streamly.Streams.StreamK

38+import Streamly.Streams.Serial (SerialT)

3539

3640-- MVar diagnostics has some overhead - around 5% on asyncly null benchmark, we

3741-- can keep it on in production to debug problems quickly if and when they

@@ -85,6 +89,10 @@ fromSVar sv = fromStream $ fromStreamVar sv

8589toSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> m ()

8690toSVar sv m = toStreamVar sv (toStream m)

8791

92+-------------------------------------------------------------------------------

93+-- Concurrency control

94+-------------------------------------------------------------------------------

95+--

8896-- | Specify the maximum number of threads that can be spawned concurrently

8997-- when using concurrent streams. This is not the grand total number of threads

9098-- but maximum threads at each point of concurrency.

@@ -92,11 +100,16 @@ toSVar sv m = toStreamVar sv (toStream m)

92100-- there is no limit. The default value is 1500.

93101--

94102-- @since 0.4.0

103+{-# INLINE_NORMAL maxThreads #-}

95104maxThreads :: IsStream t => Int -> t m a -> t m a

96105maxThreads n m = fromStream $ Stream $ \st stp sng yld -> do

97106 let n' = if n == 0 then defaultMaxThreads else n

98107 unStream (toStream m) (st {threadsHigh = n'}) stp sng yld

99108

109+{-# RULES "maxThreadsSerial serial" maxThreads = maxThreadsSerial #-}

110+maxThreadsSerial :: Int -> SerialT m a -> SerialT m a

111+maxThreadsSerial _ = id

112+

100113-- | Specify the maximum size of the buffer for storing the results from

101114-- concurrent computations. If the buffer becomes full we stop spawning more

102115-- concurrent tasks until there is space in the buffer.

@@ -104,7 +117,25 @@ maxThreads n m = fromStream $ Stream $ \st stp sng yld -> do

104117-- there is no limit. The default value is 1500.

105118--

106119-- @since 0.4.0

120+{-# INLINE_NORMAL maxBuffer #-}

107121maxBuffer :: IsStream t => Int -> t m a -> t m a

108122maxBuffer n m = fromStream $ Stream $ \st stp sng yld -> do

109123 let n' = if n == 0 then defaultMaxBuffer else n

110124 unStream (toStream m) (st {bufferHigh = n'}) stp sng yld

125+

126+{-# RULES "maxBuffer serial" maxBuffer = maxBufferSerial #-}

127+maxBufferSerial :: Int -> SerialT m a -> SerialT m a

128+maxBufferSerial _ = id

129+

130+-- Stop concurrent dispatches after this limit. This is useful in API's like

131+-- "take" where we want to dispatch only upto the number of elements "take"

132+-- needs. This value applies only to the immediate next level and is not

133+-- inherited by everything in enclosed scope.

134+{-# INLINE_NORMAL maxYields #-}

135+maxYields :: IsStream t => Maybe Int -> t m a -> t m a

136+maxYields n m = fromStream $ Stream $ \st stp sng yld -> do

137+ unStream (toStream m) (st {yieldLimit = n}) stp sng yld

138+

139+{-# RULES "maxYields serial" maxYields = maxYieldsSerial #-}

140+maxYieldsSerial :: Maybe Int -> SerialT m a -> SerialT m a

141+maxYieldsSerial _ = id

test/Main.hs+11

@@ -380,6 +380,10 @@ main = hspec $ do

380380 describe "Composed MonadThrow parallely" $ composeWithMonadThrow parallely

381381 describe "Composed MonadThrow aheadly" $ composeWithMonadThrow aheadly

382382

383+ describe "take on infinite concurrent stream" $ takeInfinite asyncly

384+ describe "take on infinite concurrent stream" $ takeInfinite wAsyncly

385+ describe "take on infinite concurrent stream" $ takeInfinite aheadly

386+

383387 it "asyncly crosses thread limit (2000 threads)" $

384388 runStream (asyncly $ fold $

385389 replicate 2000 $ S.yieldM $ threadDelay 1000000)

@@ -390,6 +394,13 @@ main = hspec $ do

390394 replicate 4000 $ S.yieldM $ threadDelay 1000000)

391395 `shouldReturn` ()

392396

397+takeInfinite :: IsStream t => (t IO Int -> SerialT IO Int) -> Spec

398+takeInfinite t = do

399+ it "take 1" $

400+ (runStream $ t $

401+ S.take 1 $ S.repeatM (print "hello" >> return (1::Int)))

402+ `shouldReturn` ()

403+

393404-- XXX need to test that we have promptly cleaned up everything after the error

394405-- XXX We can also check the output that we are expected to get before the

395406-- error occurs.