e06b319 · streamly · git.iko.soy

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

53import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)

54import Data.Concurrent.Queue.MichaelScott (LinkedQueue, newQ, pushL,

55− tryPopR, nullQ)

55+ tryPopR)

56import Data.Functor (void)

57import Data.IORef (IORef, modifyIORef, newIORef,

58 writeIORef, readIORef)

100data Context m a =

101 Context { outputQueue :: IORef [ChildEvent a]

102 , synchOutQ :: MVar Bool -- wakeup mechanism for outQ

103− , dequeueLoop :: AsyncT m a

104− , enqueueWork :: AsyncT m a -> IO ()

103+ , enqueue :: AsyncT m a -> IO ()

105− , hasNoWork :: IO Bool

104+ , runqueue :: m ()

106 , runningThreads :: IORef (Set ThreadId)

105 , runningThreads :: IORef (Set ThreadId)

107 , doneThreads :: IORef (Set ThreadId)

106 , doneThreads :: IORef (Set ThreadId)

108 }

107 }

224 -- store has finished otherwise we can have lost wakeup problems.

223 -- store has finished otherwise we can have lost wakeup problems.

225 void $ tryPutMVar (synchOutQ ctx) True

224 void $ tryPutMVar (synchOutQ ctx) True

226

225

226+{-# INLINE sendStop #-}

227+sendStop :: MonadIO m => Context m a -> m ()

228+sendStop ctx = liftIO myThreadId >>= \tid -> send ctx (ChildStop tid Nothing)

229+

227-- Note: Left associated operations can grow this queue to a large size

230-- Note: Left associated operations can grow this queue to a large size

228{-# INLINE enqueueLIFO #-}

231{-# INLINE enqueueLIFO #-}

229enqueueLIFO :: IORef [AsyncT m a] -> AsyncT m a -> IO ()

232enqueueLIFO :: IORef [AsyncT m a] -> AsyncT m a -> IO ()

230enqueueLIFO q m = atomicModifyIORefCAS_ q $ \ ms -> m : ms

233enqueueLIFO q m = atomicModifyIORefCAS_ q $ \ ms -> m : ms

231

234

232−{-# INLINE dequeueLIFO #-}

235+runqueueLIFO :: MonadIO m => Context m a -> IORef [AsyncT m a] -> m ()

233−dequeueLIFO :: IORef [AsyncT m a] -> IO (Maybe (AsyncT m a))

236+runqueueLIFO ctx q = run

234−dequeueLIFO q =

237+

238+ where

239+

240+ run = do

241+ work <- dequeue

242+ case work of

243+ Nothing -> sendStop ctx

244+ Just m -> (runAsyncT m) (Just ctx) run yield

245+

246+ sendit a = send ctx (ChildYield a)

247+ yield a _ Nothing = sendit a >> run

248+ yield a c (Just r) = sendit a >> (runAsyncT r) c run yield

249+

235− atomicModifyIORefCAS q $ \ ms ->

250+ dequeue = liftIO $ atomicModifyIORefCAS q $ \ ms ->

236 case ms of

251 case ms of

237 [] -> ([], Nothing)

252 [] -> ([], Nothing)

238 x : xs -> (xs, Just x)

253 x : xs -> (xs, Just x)

239

254

240−{-# INLINE isEmptyLIFO #-}

255+{-# INLINE enqueueFIFO #-}

241−isEmptyLIFO :: IORef [AsyncT m a] -> IO Bool

256+enqueueFIFO :: LinkedQueue (AsyncT m a) -> AsyncT m a -> IO ()

242−isEmptyLIFO q = do

257+enqueueFIFO = pushL

243− work <- readIORef q

244− return $ case work of

245− [] -> True

246− _ -> False

247

258

248−-- Note: This is performance sensitive code.

249−{-# INLINE dequeueLoopLIFO #-}

250−dequeueLoopLIFO :: MonadAsync m

251− => Maybe (Context m a)

252− -> IORef [AsyncT m a]

253− -> AsyncT m a

254−dequeueLoopLIFO ctx q = AsyncT $ \_ stp yld -> do

255− work <- liftIO $ dequeueLIFO q

256− case work of

257− Nothing -> stp

258− Just m -> do

259− let continue = dequeueLoopLIFO ctx q

260− stop = (runAsyncT continue) Nothing stp yld

261− yield a c Nothing = yld a c (Just continue)

262− yield a c r = yld a c r

263− (runAsyncT m) ctx stop yield

264−

265−{-# INLINE dequeueLoopFIFO #-}

266−dequeueLoopFIFO :: MonadAsync m

267− => Maybe (Context m a)

268− -> LinkedQueue (AsyncT m a)

259+runqueueFIFO :: MonadIO m => Context m a -> LinkedQueue (AsyncT m a) -> m ()

269− -> AsyncT m a

260+runqueueFIFO ctx q = run

270−dequeueLoopFIFO ctx q = AsyncT $ \_ stp yld -> do

271− work <- liftIO $ tryPopR q

272− case work of

273− Nothing -> stp

274− Just m -> do

275− let continue = dequeueLoopFIFO ctx q

276− stop = (runAsyncT continue) Nothing stp yld

277− yield a c Nothing = yld a c (Just continue)

278− yield a c (Just r) = do

279− liftIO $ pushL q r

280− yld a c (Just continue)

281− (runAsyncT m) ctx stop yield

282−

283−{-# NOINLINE push #-}

284−push :: MonadAsync m => Context m a -> m ()

285−push ctx = run (Just ctx) (dequeueLoop ctx)

286

261

287 where

262 where

288

263

289− stop = do

264+ run = do

290− done <- liftIO $ hasNoWork ctx

265+ work <- dequeue

291− if done then do

266+ case work of

292− tid <- liftIO $ myThreadId

267+ Nothing -> sendStop ctx

293− send ctx (ChildStop tid Nothing)

268+ Just m -> (runAsyncT m) (Just ctx) run yield

294− else push ctx

269+

270+ dequeue = liftIO $ tryPopR q

271+ sendit a = send ctx (ChildYield a)

295− yield a _ Nothing = do

272+ yield a _ Nothing = sendit a >> run

296− tid <- liftIO myThreadId

273+ -- XXX do we need to save the context as well?

297− send ctx (ChildDone tid a)

274+ -- we can enqueue "(runAsyncT m) ctx run yield" instead

298− yield a c (Just r) = send ctx (ChildYield a) >> run c r

275+ yield a c (Just r) = sendit a >> liftIO (enqueueFIFO q r) >> run

299− run c m = (runAsyncT m) c stop yield

300

276

301-- Thread tracking has a significant performance overhead (~20% on empty

277-- Thread tracking has a significant performance overhead (~20% on empty

302-- threads, it will be lower for heavy threads). It is needed for two reasons:

278-- threads, it will be lower for heavy threads). It is needed for two reasons:

354{-# NOINLINE pushWorker #-}

330{-# NOINLINE pushWorker #-}

355pushWorker :: MonadAsync m => Context m a -> m ()

331pushWorker :: MonadAsync m => Context m a -> m ()

356pushWorker ctx = do

332pushWorker ctx = do

357− tid <- doFork (push ctx) (handleChildException ctx)

333+ tid <- doFork (runqueue ctx) (handleChildException ctx)

358 liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)

334 liftIO $ modifyIORef (runningThreads ctx) $ (\s -> S.insert tid s)

359

335

360{-# INLINE sendWorkerWait #-}

336{-# INLINE sendWorkerWait #-}

407-- computation pulls from a Channel while m1 and m2 push to the channel.

383-- computation pulls from a Channel while m1 and m2 push to the channel.

408{-# NOINLINE pullFork #-}

384{-# NOINLINE pullFork #-}

409pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a

385pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a

410−pullFork m1 m2 fair = AsyncT $ \_ stp yld -> do

386+pullFork m1 m2 fifo = AsyncT $ \_ stp yld -> do

411 ctx <- liftIO $ newContext

387 ctx <- liftIO $ newContext

412 pushWorker ctx >> (runAsyncT (pullWorker ctx)) Nothing stp yld

388 pushWorker ctx >> (runAsyncT (pullWorker ctx)) Nothing stp yld

413

389

419 running <- newIORef S.empty

395 running <- newIORef S.empty

420 done <- newIORef S.empty

396 done <- newIORef S.empty

421

397

422− case fair of

398+ case fifo of

423 True -> do

399 True -> do

424 q <- newQ

400 q <- newQ

425 pushL q m1 >> pushL q m2

401 pushL q m1 >> pushL q m2

428 , synchOutQ = outQMv

404 , synchOutQ = outQMv

429 , runningThreads = running

405 , runningThreads = running

430 , doneThreads = done

406 , doneThreads = done

431− , dequeueLoop = dequeueLoopFIFO (Just ctx) q

407+ , runqueue = runqueueFIFO ctx q

432− , enqueueWork = pushL q

408+ , enqueue = pushL q

433− , hasNoWork = nullQ q

434 }

409 }

435 in return ctx

410 in return ctx

436 False -> do

411 False -> do

441 , synchOutQ = outQMv

416 , synchOutQ = outQMv

442 , runningThreads = running

417 , runningThreads = running

443 , doneThreads = done

418 , doneThreads = done

444− , dequeueLoop = dequeueLoopLIFO (Just ctx) q

419+ , runqueue = runqueueLIFO ctx q

445− , enqueueWork = enqueueLIFO q

420+ , enqueue = enqueueLIFO q

446− , hasNoWork = isEmptyLIFO q

447 }

421 }

448 in return ctx

422 in return ctx

449

423

482-- therefore always use a right fold for folding bigger structures.

456-- therefore always use a right fold for folding bigger structures.

483{-# INLINE parallel #-}

457{-# INLINE parallel #-}

484parallel :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a

458parallel :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a

485−parallel m1 m2 fair = AsyncT $ \ctx stp yld -> do

459+parallel m1 m2 fifo = AsyncT $ \ctx stp yld -> do

486 case ctx of

460 case ctx of

487− Nothing -> (runAsyncT (pullFork m1 m2 fair)) Nothing stp yld

461+ Nothing -> (runAsyncT (pullFork m1 m2 fifo)) Nothing stp yld

488 Just c -> do

462 Just c -> do

489− let dqloop = (runAsyncT (dequeueLoop c)) Nothing stp yld

490 -- XXX We can statically decide which function to call here, that

463 -- XXX We can statically decide which function to call here, that

491 -- will inline the call and improve perf a little bit.

464 -- will inline the call and improve perf a little bit.

492− liftIO $ (enqueueWork c) m2

465+ liftIO $ (enqueue c) m2

493− (runAsyncT m1) ctx dqloop yld

466+ (runAsyncT m1) ctx stp yld

494

467

495instance MonadAsync m => Alternative (AsyncT m) where

468instance MonadAsync m => Alternative (AsyncT m) where

496 empty = mempty

469 empty = mempty