503e289 · streamly · git.iko.soy

git.iko.soy/streamly

Capture inner monad state and associate it with the work for async

commit503e2894827101b54218ace38f0f766360e70716 authoradithyaov <adi.obilisetty@gmail.com> dateWed 8 Jul 2020, 17:45 (6 years ago) parente9828b3

4 files changed+57−47

compare with master

unifiedsplit

src/Streamly/Internal/Data/SVar.hs+14−7 src/Streamly/Internal/Data/Stream/Ahead.hs+4−2 src/Streamly/Internal/Data/Stream/Async.hs+38−37 src/Streamly/Internal/Data/Stream/SVar.hs+1−1

src/Streamly/Internal/Data/SVar.hs+14−7

@@ -442,7 +442,7 @@ data SVar t m a = SVar

442442 , yieldRateInfo :: Maybe YieldRateInfo

443443

444444 -- Used only by bounded SVar types

445− , enqueue :: t m a -> IO ()

445+ , enqueue :: (RunInIO m, t m a) -> IO ()

446446 , isWorkDone :: IO Bool

447447 , isQueueDone :: IO Bool

448448 , needDoorBell :: IORef Bool

@@ -1258,7 +1258,8 @@ ringDoorBell sv = do

12581258-- expressions. Large left associated compositions can grow this to a

12591259-- large size

12601260{-# INLINE enqueueLIFO #-}

1261−enqueueLIFO :: SVar t m a -> IORef [t m a] -> t m a -> IO ()

1261+enqueueLIFO ::

1262+ SVar t m a -> IORef [(RunInIO m, t m a)] -> (RunInIO m, t m a) -> IO ()

12621263enqueueLIFO sv q m = do

12631264 atomicModifyIORefCAS_ q $ \ms -> m : ms

12641265 ringDoorBell sv

@@ -1272,7 +1273,11 @@ enqueueLIFO sv q m = do

12721273-- first as long as possible.

12731274

12741275{-# INLINE enqueueFIFO #-}

1275−enqueueFIFO :: SVar t m a -> LinkedQueue (t m a) -> t m a -> IO ()

1276+enqueueFIFO ::

1277+ SVar t m a

1278+ -> LinkedQueue (RunInIO m, t m a)

1279+ -> (RunInIO m, t m a)

1280+ -> IO ()

12761281enqueueFIFO sv q m = do

12771282 pushL q m

12781283 ringDoorBell sv

@@ -1342,10 +1347,10 @@ enqueueFIFO sv q m = do

13421347-- we can even run the already queued items but they will have to be sorted in

13431348-- layers in the heap. We can use a list of heaps for that.

13441349{-# INLINE enqueueAhead #-}

1345−enqueueAhead :: SVar t m a -> IORef ([t m a], Int) -> t m a -> IO ()

1350+enqueueAhead :: SVar t m a -> IORef ([t m a], Int) -> StreamWR t m a -> IO ()

13461351enqueueAhead sv q m = do

13471352 atomicModifyIORefCAS_ q $ \ case

1348− ([], n) -> ([m], n + 1) -- increment sequence

1353+ ([], n) -> ([snd m], n + 1) -- increment sequence

13491354 _ -> error "enqueueAhead: queue is not empty"

13501355 ringDoorBell sv

13511356

@@ -2399,7 +2404,8 @@ sendFirstWorker sv m = do

23992404 -- Note: We must have all the work on the queue before sending the

24002405 -- pushworker, otherwise the pushworker may exit before we even get a

24012406 -- chance to push.

2402− liftIO $ enqueue sv m

2407+ runIn <- captureMonadState

2408+ liftIO $ enqueue sv (runIn, m)

24032409 case yieldRateInfo sv of

24042410 Nothing -> pushWorker 0 sv

24052411 Just yinfo ->

@@ -2440,7 +2446,8 @@ newParallelVar ss st = do

24402446-- be read back from the SVar using 'fromSVar'.

24412447toStreamVar :: MonadAsync m => SVar t m a -> t m a -> m ()

24422448toStreamVar sv m = do

2443− liftIO $ enqueue sv m

2449+ runIn <- captureMonadState

2450+ liftIO $ enqueue sv (runIn, m)

24442451 done <- allThreadsDone sv

24452452 -- XXX This is safe only when called from the consumer thread or when no

24462453 -- consumer is present. There may be a race if we are not running in the

src/Streamly/Internal/Data/Stream/Ahead.hs+4−2

@@ -579,7 +579,8 @@ forkSVarAhead m1 m2 = mkStream $ \st yld sng stp -> do

579579 foldStream st yld sng stp (fromSVar sv)

580580 where

581581 concurrently ma mb = mkStream $ \st yld sng stp -> do

582− liftIO $ enqueue (fromJust $ streamVar st) mb

582+ runInIO <- captureMonadState

583+ liftIO $ enqueue (fromJust $ streamVar st) (runInIO, mb)

583584 foldStream st yld sng stp ma

584585

585586infixr 6 `ahead`

@@ -593,7 +594,8 @@ ahead :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a

593594ahead m1 m2 = mkStream $ \st yld sng stp ->

594595 case streamVar st of

595596 Just sv | svarStyle sv == AheadVar -> do

596− liftIO $ enqueue sv (toStream m2)

597+ runInIO <- captureMonadState

598+ liftIO $ enqueue sv (runInIO, toStream m2)

597599 -- Always run the left side on a new SVar to avoid complexity in

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

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

src/Streamly/Internal/Data/Stream/Async.hs+38−37

@@ -68,7 +68,7 @@ data WorkerStatus = Continue | Suspend

6868{-# INLINE workLoopLIFO #-}

6969workLoopLIFO

7070 :: (MonadIO m, MonadBaseControl IO m)

71− => IORef [Stream m a]

71+ => IORef [(RunInIO m, Stream m a)]

7272 -> State Stream m a

7373 -> SVar Stream m a

7474 -> Maybe WorkerInfo

@@ -77,13 +77,12 @@ workLoopLIFO q st sv winfo = run

7777

7878 where

7979

80− mrun = runInIO $ svarMrun sv

80+ stop = liftIO $ sendStop sv winfo

8181 run = do

8282 work <- dequeue

83− let stop = liftIO $ sendStop sv winfo

8483 case work of

8584 Nothing -> stop

86− Just m -> do

85+ Just (RunInIO runin, m) -> do

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

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

8988 -- and return the state in the stop continuation.

@@ -91,7 +90,7 @@ workLoopLIFO q st sv winfo = run

9190 -- Instead of using the run function we can just restore the

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

9392 -- distributed case as well.

94− r <- liftIO $ mrun $

93+ r <- liftIO $ runin $

9594 foldStreamShared st yieldk single (return Continue) m

9695 res <- restoreM r

9796 case res of

@@ -106,9 +105,9 @@ workLoopLIFO q st sv winfo = run

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

107106 if res

108107 then foldStreamShared st yieldk single (return Continue) r

109− else liftIO $ do

110− -- XXX we also need to save the monadic state here

111− enqueueLIFO sv q r

108+ else do

109+ runInIO <- captureMonadState

110+ liftIO $ enqueueLIFO sv q (runInIO, r)

112111 return Suspend

113112

114113 dequeue = liftIO $ atomicModifyIORefCAS q $ \case

@@ -122,8 +121,8 @@ workLoopLIFO q st sv winfo = run

122121-- make a check every time.

123122{-# INLINE workLoopLIFOLimited #-}

124123workLoopLIFOLimited

125− :: (MonadIO m, MonadBaseControl IO m)

126− => IORef [Stream m a]

124+ :: forall m a. (MonadIO m, MonadBaseControl IO m)

125+ => IORef [(RunInIO m, Stream m a)]

127126 -> State Stream m a

128127 -> SVar Stream m a

129128 -> Maybe WorkerInfo

@@ -132,14 +131,13 @@ workLoopLIFOLimited q st sv winfo = run

132131

133132 where

134133

135− mrun = runInIO $ svarMrun sv

136134 incrContinue = liftIO (incrementYieldLimit sv) >> return Continue

135+ stop = liftIO $ sendStop sv winfo

137136 run = do

138137 work <- dequeue

139− let stop = liftIO $ sendStop sv winfo

140138 case work of

141139 Nothing -> stop

142− Just m -> do

140+ Just (RunInIO runin, m) -> do

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

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

145143 -- streams we do not reserve the yield limit in the constituent

@@ -149,7 +147,7 @@ workLoopLIFOLimited q st sv winfo = run

149147 yieldLimitOk <- liftIO $ decrementYieldLimit sv

150148 if yieldLimitOk

151149 then do

152− r <- liftIO $ mrun $

150+ r <- liftIO $ runin $

153151 foldStreamShared st yieldk single incrContinue m

154152 res <- restoreM r

155153 case res of

@@ -158,7 +156,7 @@ workLoopLIFOLimited q st sv winfo = run

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

159157 -- never yielded anything.

160158 else liftIO $ do

161− enqueueLIFO sv q m

159+ enqueueLIFO sv q (RunInIO runin, m)

162160 incrementYieldLimit sv

163161 sendStop sv winfo

164162

@@ -173,9 +171,10 @@ workLoopLIFOLimited q st sv winfo = run

173171 yieldLimitOk <- liftIO $ decrementYieldLimit sv

174172 if res && yieldLimitOk

175173 then foldStreamShared st yieldk single incrContinue r

176− else liftIO $ do

177− incrementYieldLimit sv

178− enqueueLIFO sv q r

174+ else do

175+ runInIO <- captureMonadState

176+ liftIO $ incrementYieldLimit sv

177+ liftIO $ enqueueLIFO sv q (runInIO, r)

179178 return Suspend

180179

181180 dequeue = liftIO $ atomicModifyIORefCAS q $ \case

@@ -191,7 +190,7 @@ workLoopLIFOLimited q st sv winfo = run

191190{-# INLINE workLoopFIFO #-}

192191workLoopFIFO

193192 :: (MonadIO m, MonadBaseControl IO m)

194− => LinkedQueue (Stream m a)

193+ => LinkedQueue (RunInIO m, Stream m a)

195194 -> State Stream m a

196195 -> SVar Stream m a

197196 -> Maybe WorkerInfo

@@ -200,14 +199,13 @@ workLoopFIFO q st sv winfo = run

200199

201200 where

202201

203− mrun = runInIO $ svarMrun sv

202+ stop = liftIO $ sendStop sv winfo

204203 run = do

205204 work <- liftIO $ tryPopR q

206− let stop = liftIO $ sendStop sv winfo

207205 case work of

208206 Nothing -> stop

209− Just m -> do

210− r <- liftIO $ mrun $

207+ Just (RunInIO runin, m) -> do

208+ r <- liftIO $ runin $

211209 foldStreamShared st yieldk single (return Continue) m

212210 res <- restoreM r

213211 case res of

@@ -224,13 +222,14 @@ workLoopFIFO q st sv winfo = run

224222 -- yielding.

225223 yieldk a r = do

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

227− liftIO $ enqueueFIFO sv q r

225+ runInIO <- captureMonadState

226+ liftIO $ enqueueFIFO sv q (runInIO, r)

228227 return $ if res then Continue else Suspend

229228

230229{-# INLINE workLoopFIFOLimited #-}

231230workLoopFIFOLimited

232− :: (MonadIO m, MonadBaseControl IO m)

233− => LinkedQueue (Stream m a)

231+ :: forall m a. (MonadIO m, MonadBaseControl IO m)

232+ => LinkedQueue (RunInIO m, Stream m a)

234233 -> State Stream m a

235234 -> SVar Stream m a

236235 -> Maybe WorkerInfo

@@ -239,25 +238,24 @@ workLoopFIFOLimited q st sv winfo = run

239238

240239 where

241240

242− mrun = runInIO $ svarMrun sv

241+ stop = liftIO $ sendStop sv winfo

243242 incrContinue = liftIO (incrementYieldLimit sv) >> return Continue

244243 run = do

245244 work <- liftIO $ tryPopR q

246− let stop = liftIO $ sendStop sv winfo

247245 case work of

248246 Nothing -> stop

249− Just m -> do

247+ Just (RunInIO runin, m) -> do

250248 yieldLimitOk <- liftIO $ decrementYieldLimit sv

251249 if yieldLimitOk

252250 then do

253− r <- liftIO $ mrun $

251+ r <- liftIO $ runin $

254252 foldStreamShared st yieldk single incrContinue m

255253 res <- restoreM r

256254 case res of

257255 Continue -> run

258256 Suspend -> stop

259257 else liftIO $ do

260− enqueueFIFO sv q m

258+ enqueueFIFO sv q (RunInIO runin, m)

261259 incrementYieldLimit sv

262260 sendStop sv winfo

263261

@@ -267,7 +265,8 @@ workLoopFIFOLimited q st sv winfo = run

267265

268266 yieldk a r = do

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

270− liftIO $ enqueueFIFO sv q r

268+ runInIO <- captureMonadState

269+ liftIO $ enqueueFIFO sv q (runInIO, r)

271270 yieldLimitOk <- liftIO $ decrementYieldLimit sv

272271 if res && yieldLimitOk

273272 then return Continue

@@ -292,7 +291,7 @@ getLifoSVar st mrun = do

292291 active <- newIORef 0

293292 wfw <- newIORef False

294293 running <- newIORef S.empty

295− q <- newIORef []

294+ q <- newIORef ([] :: [(RunInIO m, Stream m a)])

296295 yl <- case getYieldLimit st of

297296 Nothing -> return Nothing

298297 Just x -> Just <$> newIORef x

@@ -317,7 +316,7 @@ getLifoSVar st mrun = do

317316 -> (SVar Stream m a -> m [ChildEvent a])

318317 -> (SVar Stream m a -> m Bool)

319318 -> (SVar Stream m a -> IO Bool)

320− -> (IORef [Stream m a]

319+ -> (IORef [(RunInIO m, Stream m a)]

321320 -> State Stream m a

322321 -> SVar Stream m a

323322 -> Maybe WorkerInfo

@@ -414,7 +413,7 @@ getFifoSVar st mrun = do

414413 -> (SVar Stream m a -> m [ChildEvent a])

415414 -> (SVar Stream m a -> m Bool)

416415 -> (SVar Stream m a -> IO Bool)

417− -> (LinkedQueue (Stream m a)

416+ -> (LinkedQueue (RunInIO m, Stream m a)

418417 -> State Stream m a

419418 -> SVar Stream m a

420419 -> Maybe WorkerInfo

@@ -611,7 +610,8 @@ forkSVarAsync style m1 m2 = mkStream $ \st yld sng stp -> do

611610 foldStream st yld sng stp $ fromSVar sv

612611 where

613612 concurrently ma mb = mkStream $ \st yld sng stp -> do

614− liftIO $ enqueue (fromJust $ streamVar st) mb

613+ runInIO <- captureMonadState

614+ liftIO $ enqueue (fromJust $ streamVar st) $ (runInIO, mb)

615615 foldStreamShared st yld sng stp ma

616616

617617{-# INLINE joinStreamVarAsync #-}

@@ -620,7 +620,8 @@ joinStreamVarAsync :: (IsStream t, MonadAsync m)

620620joinStreamVarAsync style m1 m2 = mkStream $ \st yld sng stp ->

621621 case streamVar st of

622622 Just sv | svarStyle sv == style -> do

623− liftIO $ enqueue sv (toStream m2)

623+ runInIO <- captureMonadState

624+ liftIO $ enqueue sv $ (runInIO, toStream m2)

624625 foldStreamShared st yld sng stp m1

625626 _ -> foldStreamShared st yld sng stp (forkSVarAsync style m1 m2)

626627

src/Streamly/Internal/Data/Stream/SVar.hs+1−1

@@ -146,7 +146,7 @@ fromSVar sv =

146146-- | Write a stream to an 'SVar' in a non-blocking manner. The stream can then

147147-- be read back from the SVar using 'fromSVar'.

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

149−toSVar sv m = toStreamVar sv (toStream m)

149+toSVar sv m = toStreamVar sv $ toStream m

150150

151151-------------------------------------------------------------------------------

152152-- Process events received by a fold consumer from a stream producer