git.iko.soy/streamly
on master help sign in
Capture inner monad state and associate it with the work for ahead
commitb02925aa17df2e8099ea2786b80e54e1b560a8b2 authoradithyaov <adi.obilisetty@gmail.com> dateTue 14 Jul 2020, 14:26 (6 years ago) parent503e289
2 files changed+19−11
compare with master
unifiedsplit
src/Streamly/Internal/Data/SVar.hs+2−2 src/Streamly/Internal/Data/Stream/Ahead.hs+17−9
src/Streamly/Internal/Data/SVar.hs+2−2
@@ -177,7 +177,7 @@ data ChildEvent a =
177177data AheadHeapEntry (t :: (Type -> Type) -> Type -> Type) m a =
178178 AheadEntryNull
179179 | AheadEntryPure a
180− | AheadEntryStream (t m a)
180+ | AheadEntryStream (RunInIO m, t m a)
181181#undef Type
182182
183183------------------------------------------------------------------------------
@@ -1347,7 +1347,7 @@ enqueueFIFO sv q m = do
13471347-- we can even run the already queued items but they will have to be sorted in
13481348-- layers in the heap. We can use a list of heaps for that.
13491349{-# INLINE enqueueAhead #-}
1350−enqueueAhead :: SVar t m a -> IORef ([t m a], Int) -> StreamWR t m a -> IO ()
1350+enqueueAhead :: SVar t m a -> IORef ([t m a], Int) -> (RunInIO m, t m a) -> IO ()
13511351enqueueAhead sv q m = do
13521352 atomicModifyIORefCAS_ q $ \ case
13531353 ([], n) -> ([snd m], n + 1) -- increment sequence
src/Streamly/Internal/Data/Stream/Ahead.hs+17−9
@@ -239,10 +239,13 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry
239239 -- transferring available results from heap to outputQueue.
240240 void $ liftIO $ send sv (ChildYield a)
241241 nextHeap seqNo
242− AheadEntryStream r ->
242+ AheadEntryStream (RunInIO runin, r) ->
243243 if stopping
244244 then stopIfNeeded ent seqNo r
245− else runStreamWithYieldLimit True seqNo r
245+ else do
246+ res <- liftIO $ runin (runStreamWithYieldLimit True seqNo r)
247+ restoreM res
248+
246249
247250 nextHeap prevSeqNo = do
248251 res <- liftIO $ dequeueFromHeapSeq heap (prevSeqNo + 1)
@@ -299,11 +302,13 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry
299302 (singleStreamFromHeap seqNo)
300303 stop
301304 r
302− else liftIO $ do
303− let ent = Entry seqNo (AheadEntryStream r)
304− liftIO $ requeueOnHeapTop heap ent seqNo
305− incrementYieldLimit sv
306− sendStop sv winfo
305+ else do
306+ runIn <- captureMonadState
307+ let ent = Entry seqNo (AheadEntryStream (runIn, r))
308+ liftIO $ do
309+ requeueOnHeapTop heap ent seqNo
310+ incrementYieldLimit sv
311+ sendStop sv winfo
307312
308313 yieldStreamFromHeap seqNo a r = do
309314 continue <- liftIO $ sendYield sv winfo (ChildYield a)
@@ -352,7 +357,9 @@ processWithoutToken q heap st sv winfo m seqNo = do
352357
353358 r <- liftIO $ mrun $
354359 foldStreamShared st
355− (\a r -> toHeap $ AheadEntryStream $ K.cons a r)
360+ (\a r -> do
361+ runIn <- captureMonadState
362+ toHeap $ AheadEntryStream (runIn, K.cons a r))
356363 (toHeap . AheadEntryPure)
357364 stop
358365 m
@@ -454,7 +461,8 @@ processWithToken q heap st sv winfo action sno = do
454461 stop
455462 r
456463 else do
457− let ent = Entry seqNo (AheadEntryStream r)
464+ runIn <- captureMonadState
465+ let ent = Entry seqNo (AheadEntryStream (runIn, r))
458466 liftIO $ requeueOnHeapTop heap ent seqNo
459467 liftIO $ incrementYieldLimit sv
460468 return TokenSuspend