a3c32bc · streamly · git.iko.soy

git.iko.soy/streamly

Fix foldStreamSVar space leak

Use tail recursion in the worker loops. Run the work item under the saved monadic state and restore the state after the work is done and before the next work item is picked. Benchmarks: * asyncly/unfoldrM maxThreads 1 * wAsyncly/unfoldrM maxThreads 1 * aheadly/unfoldrM maxThreads 1

commita3c32bc716bd407de9468f30916455b0839ee45b authorHarendra Kumar <harendra.kumar@gmail.com> dateThu 16 Jan 2020, 13:45 (7 years ago) parent37452d9

5 files changed+130−85

compare with master

unifiedsplit

Changelog.md+5−5 src/Streamly/Internal/Data/Stream/Ahead.hs+55−29 src/Streamly/Internal/Data/Stream/Async.hs+68−30 src/Streamly/Internal/Data/Stream/StreamK.hs+1−2 src/Streamly/Internal/Data/Stream/StreamK/Type.hs+1−19

Changelog.md+5−5

@@ -6,11 +6,11 @@

66 cases.

77* Fix a bug in `tap`, `chunksOf` that caused memory consumption to

88 increase in some cases.

9−* Fix a space leak in `async` streams that caused memory consumption

10− to increase with the number of elements in the stream, especially when

11− built with `-threaded` and used with `-N` RTS option. The issue occurs

12− only in cases when a worker thread happens to be used continuously for a

13− long time.

9+* Fix a space leak in concurrent streams (`async`, `wAsync`, and `ahead`) that

10+ caused memory consumption to increase with the number of elements in the

11+ stream, especially when built with `-threaded` and used with `-N` RTS option.

12+ The issue occurs only in cases when a worker thread happens to be used

13+ continuously for a long time.

1414

1515### Behavior change

1616

src/Streamly/Internal/Data/Stream/Ahead.hs+55−29

@@ -36,6 +36,7 @@ import Control.Monad.IO.Class (MonadIO(..))

3636import Control.Monad.Reader.Class (MonadReader(..))

3737import Control.Monad.State.Class (MonadState(..))

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

39+import Control.Monad.Trans.Control (MonadBaseControl (..))

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

4041import Data.IORef (IORef, readIORef, atomicModifyIORef, writeIORef)

4142import Data.Maybe (fromJust)

@@ -49,8 +50,7 @@ import qualified Data.Heap as H

4950import Streamly.Internal.Data.Stream.SVar (fromSVar)

5051import Streamly.Internal.Data.SVar

5152import Streamly.Internal.Data.Stream.StreamK

52− (IsStream(..), Stream, mkStream, foldStream, foldStreamShared,

53− foldStreamSVar)

53+ (IsStream(..), Stream, mkStream, foldStream, foldStreamShared)

5454import qualified Streamly.Internal.Data.Stream.StreamK as K

5555import qualified Streamly.Internal.Data.Stream.StreamD as D

5656

@@ -212,7 +212,8 @@ abortExecution q sv winfo m = do

212212--

213213-- In both cases we give the drainer a chance to run more often.

214214--

215−processHeap :: MonadIO m

215+processHeap

216+ :: (MonadIO m, MonadBaseControl IO m)

216217 => IORef ([Stream m a], Int)

217218 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)), Maybe Int)

218219 -> State Stream m a

@@ -316,7 +317,8 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry

316317 runStreamWithYieldLimit continue seqNo r

317318

318319{-# NOINLINE drainHeap #-}

319−drainHeap :: MonadIO m

320+drainHeap

321+ :: (MonadIO m, MonadBaseControl IO m)

320322 => IORef ([Stream m a], Int)

321323 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)), Maybe Int)

322324 -> State Stream m a

@@ -331,8 +333,10 @@ drainHeap q heap st sv winfo = do

331333 _ -> liftIO $ sendStop sv winfo

332334

333335data HeapStatus = HContinue | HStop

336+data WorkerStatus = Continue | Suspend

334337

335−processWithoutToken :: MonadIO m

338+processWithoutToken

339+ :: (MonadIO m, MonadBaseControl IO m)

336340 => IORef ([Stream m a], Int)

337341 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)), Maybe Int)

338342 -> State Stream m a

@@ -351,12 +355,18 @@ processWithoutToken q heap st sv winfo m seqNo = do

351355 -- put it on heap. So we have to put a null entry on heap even when

352356 -- we stop.

353357 toHeap AheadEntryNull

358+ mrun = runInIO $ svarMrun sv

354359

355− foldStreamSVar sv st

356− (\a r -> toHeap $ AheadEntryStream $ K.cons a r)

357− (toHeap . AheadEntryPure)

358− stop

359− m

360+ r <- liftIO $ mrun $

361+ foldStreamShared st

362+ (\a r -> toHeap $ AheadEntryStream $ K.cons a r)

363+ (toHeap . AheadEntryPure)

364+ stop

365+ m

366+ res <- restoreM r

367+ case res of

368+ Continue -> workLoopAhead q heap st sv winfo

369+ Suspend -> drainHeap q heap st sv winfo

360370

361371 where

362372

@@ -377,8 +387,6 @@ processWithoutToken q heap st sv winfo m seqNo = do

377387 writeIORef (maxHeapSize $ svarStats sv) (H.size newHp)

378388

379389 heapOk <- liftIO $ underMaxHeap sv newHp

380− let drainAndStop = drainHeap q heap st sv winfo

381− mainLoop = workLoopAhead q heap st sv winfo

382390 status <-

383391 case yieldRateInfo sv of

384392 Nothing -> return HContinue

@@ -394,11 +402,14 @@ processWithoutToken q heap st sv winfo m seqNo = do

394402 if heapOk

395403 then

396404 case status of

397− HContinue -> mainLoop

398− HStop -> drainAndStop

399− else drainAndStop

405+ HContinue -> return Continue

406+ HStop -> return Suspend

407+ else return Suspend

408+

409+data TokenWorkerStatus = TokenContinue Int | TokenSuspend

400410

401−processWithToken :: MonadIO m

411+processWithToken

412+ :: (MonadIO m, MonadBaseControl IO m)

402413 => IORef ([Stream m a], Int)

403414 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)), Maybe Int)

404415 -> State Stream m a

@@ -412,19 +423,26 @@ processWithToken q heap st sv winfo action sno = do

412423 -- XXX deduplicate stop in all invocations

413424 let stop = do

414425 liftIO (incrementYieldLimit sv)

415− loopWithToken (sno + 1)

426+ return $ TokenContinue (sno + 1)

427+ mrun = runInIO $ svarMrun sv

416428

417− foldStreamSVar sv st (yieldOutput sno) (singleOutput sno) stop action

429+ r <- liftIO $ mrun $

430+ foldStreamShared st (yieldOutput sno) (singleOutput sno) stop action

431+

432+ res <- restoreM r

433+ case res of

434+ TokenContinue seqNo -> loopWithToken seqNo

435+ TokenSuspend -> drainHeap q heap st sv winfo

418436

419437 where

420438

421439 singleOutput seqNo a = do

422440 continue <- liftIO $ sendYield sv winfo (ChildYield a)

423441 if continue

424− then loopWithToken (seqNo + 1)

442+ then return $ TokenContinue (seqNo + 1)

425443 else do

426444 liftIO $ updateHeapSeq heap (seqNo + 1)

427− drainHeap q heap st sv winfo

445+ return TokenSuspend

428446

429447 -- XXX use a wrapper function around stop so that we never miss

430448 -- incrementing the yield in a stop continuation. Essentiatlly all

@@ -436,7 +454,7 @@ processWithToken q heap st sv winfo action sno = do

436454 then do

437455 let stop = do

438456 liftIO (incrementYieldLimit sv)

439− loopWithToken (seqNo + 1)

457+ return $ TokenContinue (seqNo + 1)

440458 foldStreamShared st

441459 (yieldOutput seqNo)

442460 (singleOutput seqNo)

@@ -446,7 +464,7 @@ processWithToken q heap st sv winfo action sno = do

446464 let ent = Entry seqNo (AheadEntryStream r)

447465 liftIO $ requeueOnHeapTop heap ent seqNo

448466 liftIO $ incrementYieldLimit sv

449− drainHeap q heap st sv winfo

467+ return TokenSuspend

450468

451469 loopWithToken nextSeqNo = do

452470 work <- dequeueAhead q

@@ -467,12 +485,19 @@ processWithToken q heap st sv winfo action sno = do

467485 then do

468486 let stop = do

469487 liftIO (incrementYieldLimit sv)

470− loopWithToken (seqNo + 1)

471− foldStreamShared st

472− (yieldOutput seqNo)

473− (singleOutput seqNo)

474− stop

475− m

488+ return $ TokenContinue (seqNo + 1)

489+ mrun = runInIO $ svarMrun sv

490+ r <- liftIO $ mrun $

491+ foldStreamShared st

492+ (yieldOutput seqNo)

493+ (singleOutput seqNo)

494+ stop

495+ m

496+ res <- restoreM r

497+ case res of

498+ TokenContinue seqNo1 -> loopWithToken seqNo1

499+ TokenSuspend -> drainHeap q heap st sv winfo

500+

476501 else

477502 -- To avoid a race when another thread puts something

478503 -- on the heap and goes away, the consumer will not get

@@ -496,7 +521,8 @@ processWithToken q heap st sv winfo action sno = do

496521

497522-- XXX we can remove the sv parameter as it can be derived from st

498523

499−workLoopAhead :: MonadIO m

524+workLoopAhead

525+ :: (MonadIO m, MonadBaseControl IO m)

500526 => IORef ([Stream m a], Int)

501527 -> IORef (Heap (Entry Int (AheadHeapEntry Stream m a)), Maybe Int)

502528 -> State Stream m a

src/Streamly/Internal/Data/Stream/Async.hs+68−30

@@ -39,6 +39,7 @@ where

3939import Control.Concurrent (myThreadId)

4040import Control.Monad.Base (MonadBase(..), liftBaseDefault)

4141import Control.Monad.Catch (MonadThrow, throwM)

42+import Control.Monad.Trans.Control (MonadBaseControl (..))

4243import Control.Concurrent.MVar (newEmptyMVar)

4344-- import Control.Monad.Error.Class (MonadError(..))

4445import Control.Monad.IO.Class (MonadIO(..))

@@ -59,8 +60,7 @@ import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS)

5960import Streamly.Internal.Data.Stream.SVar (fromSVar)

6061import Streamly.Internal.Data.SVar

6162import Streamly.Internal.Data.Stream.StreamK

62− (IsStream(..), Stream, mkStream, foldStream, adapt, foldStreamShared,

63− foldStreamSVar)

63+ (IsStream(..), Stream, mkStream, foldStream, adapt, foldStreamShared)

6464import qualified Streamly.Internal.Data.Stream.StreamK as K

6565import qualified Streamly.Internal.Data.Stream.StreamD as D

6666

@@ -70,9 +70,11 @@ import qualified Streamly.Internal.Data.Stream.StreamD as D

7070-- Async

7171-------------------------------------------------------------------------------

7272

73+data WorkerStatus = Continue | Suspend

74+

7375{-# INLINE workLoopLIFO #-}

7476workLoopLIFO

75− :: MonadIO m

77+ :: (MonadIO m, MonadBaseControl IO m)

7678 => IORef [Stream m a]

7779 -> State Stream m a

7880 -> SVar Stream m a

@@ -82,23 +84,39 @@ workLoopLIFO q st sv winfo = run

8284

8385 where

8486

87+ mrun = runInIO $ svarMrun sv

8588 run = do

8689 work <- dequeue

90+ let stop = liftIO $ sendStop sv winfo

8791 case work of

88− Nothing -> liftIO $ sendStop sv winfo

89− Just m -> foldStreamSVar sv st yieldk single run m

92+ Nothing -> stop

93+ Just m -> do

94+ -- XXX when we finish we need to send the monadic state back to

95+ -- the parent so that the state can be merged back. We capture

96+ -- and return the state in the stop continuation.

97+ --

98+ -- Instead of using the run function we can just restore the

99+ -- monad state here. That way it can work easily for

100+ -- distributed case as well.

101+ r <- liftIO $ mrun $

102+ foldStreamShared st yieldk single (return Continue) m

103+ res <- restoreM r

104+ case res of

105+ Continue -> run

106+ Suspend -> stop

90107

91108 single a = do

92109 res <- liftIO $ sendYield sv winfo (ChildYield a)

93− if res then run else liftIO $ sendStop sv winfo

110+ return $ if res then Continue else Suspend

94111

95112 yieldk a r = do

96113 res <- liftIO $ sendYield sv winfo (ChildYield a)

97114 if res

98− then foldStreamShared st yieldk single run r

115+ then foldStreamShared st yieldk single (return Continue) r

99116 else liftIO $ do

117+ -- XXX we also need to save the monadic state here

100118 enqueueLIFO sv q r

101− sendStop sv winfo

119+ return Suspend

102120

103121 dequeue = liftIO $ atomicModifyIORefCAS q $ \case

104122 [] -> ([], Nothing)

@@ -111,7 +129,7 @@ workLoopLIFO q st sv winfo = run

111129-- make a check every time.

112130{-# INLINE workLoopLIFOLimited #-}

113131workLoopLIFOLimited

114− :: MonadIO m

132+ :: (MonadIO m, MonadBaseControl IO m)

115133 => IORef [Stream m a]

116134 -> State Stream m a

117135 -> SVar Stream m a

@@ -121,10 +139,13 @@ workLoopLIFOLimited q st sv winfo = run

121139

122140 where

123141

142+ mrun = runInIO $ svarMrun sv

143+ incrContinue = liftIO (incrementYieldLimit sv) >> return Continue

124144 run = do

125145 work <- dequeue

146+ let stop = liftIO $ sendStop sv winfo

126147 case work of

127− Nothing -> liftIO $ sendStop sv winfo

148+ Nothing -> stop

128149 Just m -> do

129150 -- XXX This is just a best effort minimization of concurrency

130151 -- to the yield limit. If the stream is made of concurrent

@@ -135,8 +156,12 @@ workLoopLIFOLimited q st sv winfo = run

135156 yieldLimitOk <- liftIO $ decrementYieldLimit sv

136157 if yieldLimitOk

137158 then do

138− let stop = liftIO (incrementYieldLimit sv) >> run

139− foldStreamSVar sv st yieldk single stop m

159+ r <- liftIO $ mrun $

160+ foldStreamShared st yieldk single incrContinue m

161+ res <- restoreM r

162+ case res of

163+ Continue -> run

164+ Suspend -> stop

140165 -- Avoid any side effects, undo the yield limit decrement if we

141166 -- never yielded anything.

142167 else liftIO $ do

@@ -146,20 +171,19 @@ workLoopLIFOLimited q st sv winfo = run

146171

147172 single a = do

148173 res <- liftIO $ sendYield sv winfo (ChildYield a)

149− if res then run else liftIO $ sendStop sv winfo

174+ return $ if res then Continue else Suspend

150175

151176 -- XXX can we pass on the yield limit downstream to limit the concurrency

152177 -- of constituent streams.

153178 yieldk a r = do

154179 res <- liftIO $ sendYield sv winfo (ChildYield a)

155180 yieldLimitOk <- liftIO $ decrementYieldLimit sv

156− let stop = liftIO (incrementYieldLimit sv) >> run

157181 if res && yieldLimitOk

158− then foldStreamShared st yieldk single stop r

182+ then foldStreamShared st yieldk single incrContinue r

159183 else liftIO $ do

160184 incrementYieldLimit sv

161185 enqueueLIFO sv q r

162− sendStop sv winfo

186+ return Suspend

163187

164188 dequeue = liftIO $ atomicModifyIORefCAS q $ \case

165189 [] -> ([], Nothing)

@@ -173,7 +197,7 @@ workLoopLIFOLimited q st sv winfo = run

173197

174198{-# INLINE workLoopFIFO #-}

175199workLoopFIFO

176− :: MonadIO m

200+ :: (MonadIO m, MonadBaseControl IO m)

177201 => LinkedQueue (Stream m a)

178202 -> State Stream m a

179203 -> SVar Stream m a

@@ -183,27 +207,35 @@ workLoopFIFO q st sv winfo = run

183207

184208 where

185209

210+ mrun = runInIO $ svarMrun sv

186211 run = do

187212 work <- liftIO $ tryPopR q

213+ let stop = liftIO $ sendStop sv winfo

188214 case work of

189− Nothing -> liftIO $ sendStop sv winfo

190− Just m -> foldStreamSVar sv st yieldk single run m

215+ Nothing -> stop

216+ Just m -> do

217+ r <- liftIO $ mrun $

218+ foldStreamShared st yieldk single (return Continue) m

219+ res <- restoreM r

220+ case res of

221+ Continue -> run

222+ Suspend -> stop

191223

192224 single a = do

193225 res <- liftIO $ sendYield sv winfo (ChildYield a)

194− if res then run else liftIO $ sendStop sv winfo

226+ return $ if res then Continue else Suspend

195227

196228 yieldk a r = do

197229 res <- liftIO $ sendYield sv winfo (ChildYield a)

198230 if res

199− then foldStreamShared st yieldk single run r

231+ then foldStreamShared st yieldk single (return Continue) r

200232 else liftIO $ do

201233 enqueueFIFO sv q r

202− sendStop sv winfo

234+ return Suspend

203235

204236{-# INLINE workLoopFIFOLimited #-}

205237workLoopFIFOLimited

206− :: MonadIO m

238+ :: (MonadIO m, MonadBaseControl IO m)

207239 => LinkedQueue (Stream m a)

208240 -> State Stream m a

209241 -> SVar Stream m a

@@ -213,16 +245,23 @@ workLoopFIFOLimited q st sv winfo = run

213245

214246 where

215247

248+ mrun = runInIO $ svarMrun sv

249+ incrContinue = liftIO (incrementYieldLimit sv) >> return Continue

216250 run = do

217251 work <- liftIO $ tryPopR q

252+ let stop = liftIO $ sendStop sv winfo

218253 case work of

219− Nothing -> liftIO $ sendStop sv winfo

254+ Nothing -> stop

220255 Just m -> do

221256 yieldLimitOk <- liftIO $ decrementYieldLimit sv

222257 if yieldLimitOk

223258 then do

224− let stop = liftIO (incrementYieldLimit sv) >> run

225− foldStreamSVar sv st yieldk single stop m

259+ r <- liftIO $ mrun $

260+ foldStreamShared st yieldk single incrContinue m

261+ res <- restoreM r

262+ case res of

263+ Continue -> run

264+ Suspend -> stop

226265 else liftIO $ do

227266 enqueueFIFO sv q m

228267 incrementYieldLimit sv

@@ -230,18 +269,17 @@ workLoopFIFOLimited q st sv winfo = run

230269

231270 single a = do

232271 res <- liftIO $ sendYield sv winfo (ChildYield a)

233− if res then run else liftIO $ sendStop sv winfo

272+ return $ if res then Continue else Suspend

234273

235274 yieldk a r = do

236275 res <- liftIO $ sendYield sv winfo (ChildYield a)

237276 yieldLimitOk <- liftIO $ decrementYieldLimit sv

238− let stop = liftIO (incrementYieldLimit sv) >> run

239277 if res && yieldLimitOk

240− then foldStreamShared st yieldk single stop r

278+ then foldStreamShared st yieldk single incrContinue r

241279 else liftIO $ do

242280 incrementYieldLimit sv

243281 enqueueFIFO sv q r

244− sendStop sv winfo

282+ return Suspend

245283

246284-------------------------------------------------------------------------------

247285-- SVar creation

src/Streamly/Internal/Data/Stream/StreamK.hs+1−2

@@ -47,7 +47,6 @@ module Streamly.Internal.Data.Stream.StreamK

4747 -- * Elimination Primitives

4848 , foldStream

4949 , foldStreamShared

50− , foldStreamSVar

5150

5251 -- * Transformation Primitives

5352 , unShare

@@ -67,7 +66,7 @@ module Streamly.Internal.Data.Stream.StreamK

6766 , replicateM

6867 , fromIndices

6968 , fromIndicesM

70− , iterate

69+ , iterate

7170 , iterateM

7271

7372 -- ** Conversions

src/Streamly/Internal/Data/Stream/StreamK/Type.hs+1−19

@@ -47,7 +47,6 @@ module Streamly.Internal.Data.Stream.StreamK.Type

4747 -- * Elimination

4848 , foldStream

4949 , foldStreamShared

50− , foldStreamSVar

5150

5251 -- * foldr/build

5352 , foldrM

@@ -85,8 +84,7 @@ module Streamly.Internal.Data.Stream.StreamK.Type

8584 )

8685where

8786

88−import Control.Monad (void, ap, (>=>))

89−import Control.Monad.IO.Class (MonadIO(liftIO))

87+import Control.Monad (ap, (>=>))

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

9189#if __GLASGOW_HASKELL__ >= 800

9290import Data.Kind (Type)

@@ -415,22 +413,6 @@ foldStream st yld sng stp m =

415413 MkStream k = toStream m

416414 in k (adaptState st) yieldk sng stp

417415

418−-- Run the stream using a run function associated with the SVar that runs the

419−-- streams with a captured snapshot of the monadic state.

420−{-# INLINE foldStreamSVar #-}

421−foldStreamSVar

422− :: (IsStream t, MonadIO m)

423− => SVar Stream m a

424− -> State Stream m a -- state

425− -> (a -> t m a -> m r) -- yield

426− -> (a -> m r) -- singleton

427− -> m r -- stop

428− -> t m a

429− -> m ()

430−foldStreamSVar sv st yld sng stp m =

431− let mrun = runInIO $ svarMrun sv

432− in void $ liftIO $ mrun $ foldStreamShared st yld sng stp m

433−

434416-------------------------------------------------------------------------------

435417-- Instances

436418-------------------------------------------------------------------------------