git.iko.soy/streamly
on master help sign in
Capture inner monad state and associate it with the work for ahead
b02925aadithyaov
src/Streamly/Internal/Data/SVar.hs+2−2
compare with master
unifiedsplit
@@ -177,7 +177,7 @@ data ChildEvent a =
177data AheadHeapEntry (t :: (Type -> Type) -> Type -> Type) m a =
177data AheadHeapEntry (t :: (Type -> Type) -> Type -> Type) m a =
178 AheadEntryNull
178 AheadEntryNull
179 | AheadEntryPure a
179 | AheadEntryPure a
180− | AheadEntryStream (t m a)
180+ | AheadEntryStream (RunInIO m, t m a)
181#undef Type
181#undef Type
182
182
183------------------------------------------------------------------------------
183------------------------------------------------------------------------------
@@ -1347,7 +1347,7 @@ enqueueFIFO sv q m = do
1347-- we can even run the already queued items but they will have to be sorted in
1347-- we can even run the already queued items but they will have to be sorted in
1348-- layers in the heap. We can use a list of heaps for that.
1348-- layers in the heap. We can use a list of heaps for that.
1349{-# INLINE enqueueAhead #-}
1349{-# 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 ()
1351enqueueAhead sv q m = do
1351enqueueAhead sv q m = do
1352 atomicModifyIORefCAS_ q $ \ case
1352 atomicModifyIORefCAS_ q $ \ case
1353 ([], n) -> ([snd m], n + 1) -- increment sequence
1353 ([], 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
239 -- transferring available results from heap to outputQueue.
239 -- transferring available results from heap to outputQueue.
240 void $ liftIO $ send sv (ChildYield a)
240 void $ liftIO $ send sv (ChildYield a)
241 nextHeap seqNo
241 nextHeap seqNo
242− AheadEntryStream r ->
242+ AheadEntryStream (RunInIO runin, r) ->
243 if stopping
243 if stopping
244 then stopIfNeeded ent seqNo r
244 then stopIfNeeded ent seqNo r
245+ else do
245− else runStreamWithYieldLimit True seqNo r
246+ res <- liftIO $ runin (runStreamWithYieldLimit True seqNo r)
247+ restoreM res
248+
246
249
247 nextHeap prevSeqNo = do
250 nextHeap prevSeqNo = do
248 res <- liftIO $ dequeueFromHeapSeq heap (prevSeqNo + 1)
251 res <- liftIO $ dequeueFromHeapSeq heap (prevSeqNo + 1)
@@ -299,11 +302,13 @@ processHeap q heap st sv winfo entry sno stopping = loopHeap sno entry
299 (singleStreamFromHeap seqNo)
302 (singleStreamFromHeap seqNo)
300 stop
303 stop
301 r
304 r
302− else liftIO $ do
305+ else do
306+ runIn <- captureMonadState
303− let ent = Entry seqNo (AheadEntryStream r)
307+ let ent = Entry seqNo (AheadEntryStream (runIn, r))
308+ liftIO $ do
304− liftIO $ requeueOnHeapTop heap ent seqNo
309+ requeueOnHeapTop heap ent seqNo
305− incrementYieldLimit sv
310+ incrementYieldLimit sv
306− sendStop sv winfo
311+ sendStop sv winfo
307
312
308 yieldStreamFromHeap seqNo a r = do
313 yieldStreamFromHeap seqNo a r = do
309 continue <- liftIO $ sendYield sv winfo (ChildYield a)
314 continue <- liftIO $ sendYield sv winfo (ChildYield a)
@@ -352,7 +357,9 @@ processWithoutToken q heap st sv winfo m seqNo = do
352
357
353 r <- liftIO $ mrun $
358 r <- liftIO $ mrun $
354 foldStreamShared st
359 foldStreamShared st
360+ (\a r -> do
361+ runIn <- captureMonadState
355− (\a r -> toHeap $ AheadEntryStream $ K.cons a r)
362+ toHeap $ AheadEntryStream (runIn, K.cons a r))
356 (toHeap . AheadEntryPure)
363 (toHeap . AheadEntryPure)
357 stop
364 stop
358 m
365 m
@@ -454,7 +461,8 @@ processWithToken q heap st sv winfo action sno = do
454 stop
461 stop
455 r
462 r
456 else do
463 else do
464+ runIn <- captureMonadState
457− let ent = Entry seqNo (AheadEntryStream r)
465+ let ent = Entry seqNo (AheadEntryStream (runIn, r))
458 liftIO $ requeueOnHeapTop heap ent seqNo
466 liftIO $ requeueOnHeapTop heap ent seqNo
459 liftIO $ incrementYieldLimit sv
467 liftIO $ incrementYieldLimit sv
460 return TokenSuspend
468 return TokenSuspend