b3faafa · streamly · git.iko.soy

26

27import Control.Applicative (Alternative (..))

28import Control.Concurrent (ThreadId, forkIO, killThread,

29− myThreadId)

29+ myThreadId, newQSem, QSem,

30+ signalQSem, waitQSem)

30import Control.Concurrent.STM (TChan, atomically, newTChan,

31import Control.Concurrent.STM (TChan, atomically, newTChan,

31− readTChan, tryReadTChan,

32− writeTChan)

32+ tryReadTChan, writeTChan)

33import Control.Exception (SomeException (..))

34import qualified Control.Exception.Lifted as EL

35import Control.Monad (ap, liftM, MonadPlus(..), mzero)

85------------------------------------------------------------------------------

86

87data Context a =

88− Context { childChannel :: TChan (ChildEvent a)

88+ Context { childChannel :: TChan (ChildEvent a)

89− , pullSide :: Bool

89+ , dispatchReq :: QSem

90− , runningThreads :: IORef (Set ThreadId)

90+ , runningThreads :: IORef (Set ThreadId)

91− , doneThreads :: IORef (Set ThreadId)

91+ , doneThreads :: IORef (Set ThreadId)

92 }

93

94-- The 'Maybe (AsyncT m a)' is redundant as we can use 'stop' value for the

182

183{-# NOINLINE push #-}

184push :: MonadIO m => Context a -> AsyncT m a -> m ()

185+push context action = do

186+ liftIO $ waitQSem (dispatchReq context)

185−push context action = run (Just context) action

187+ run (Just context) action

186

188

187 where

189 where

188

190

203 continue a ctx m = channelYield a >> run ctx m

205 continue a ctx m = channelYield a >> run ctx m

204 yielder a ctx r = maybe (done a) (\rx -> continue a ctx rx) r

206 yielder a ctx r = maybe (done a) (\rx -> continue a ctx rx) r

205

207

206−-- If an exception occurs we push it to the channel so that it can handled by

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

207−-- the parent. 'Paused' exceptions are to be collected at the top level.

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

208−-- XXX Paused exceptions should only bubble up to the runRecorder computation

210+--

209−{-# NOINLINE handleChildException #-}

211+-- 1) Killing threads on exceptions. Threads may not be allowed to go away by

210−handleChildException :: TChan (ChildEvent a) -> SomeException -> IO ()

212+-- themselves because they may run for significant times before going away or

211−handleChildException pchan e = do

213+-- worse they may be stuck in IO and never go away.

212− tid <- myThreadId

214+--

213− atomically $ writeTChan pchan (ChildStop tid (Just e))

215+-- 2) To know when all threads are done. This can be acheived by detecting a

214−

216+-- BlockedIndefinitelyOnSTM exception too. But we will have to trigger a GC to

215−{-# NOINLINE pushSideDispatch #-}

217+-- make sure that we detect it promptly.

216−pushSideDispatch :: MonadAsync m

218+--

217− => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a

218−pushSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do

219− let chan = childChannel ctx

220− tid <- doFork (push ctx m1) (handleChildException chan)

221− liftIO $ atomically $ writeTChan chan (ChildCreate tid)

222− (runAsyncT m2) (Just ctx) stp yld

223−

224-- This is a bit messy because ChildCreate and ChildDone events can arrive out

219-- This is a bit messy because ChildCreate and ChildDone events can arrive out

225-- of order in case of pushSideDispatch. Returns whether we are done draining

220-- of order in case of pushSideDispatch. Returns whether we are done draining

226-- threads.

221-- threads.

258-- We re-raise any exceptions received from the child threads, that way

253-- We re-raise any exceptions received from the child threads, that way

259-- exceptions get propagated to the top level computation and can be handled

254-- exceptions get propagated to the top level computation and can be handled

260-- there.

255-- there.

261−{-# NOINLINE pullDispatch #-}

256+{-# NOINLINE pullDrain #-}

262−pullDispatch :: (MonadIO m, MonadThrow m)

257+pullDrain :: (MonadIO m, MonadThrow m) => Context a -> AsyncT m a

263− => Context a -> AsyncT m a -> Bool -> AsyncT m a

264−pullDispatch ctx m dispatch = AsyncT $ \_ stp yld -> do

258+pullDrain ctx = AsyncT $ \_ stp yld -> do

265− if dispatch then do

266− (runAsyncT m) (Just ctx) stp yld

267− else do

268− res <- liftIO $ atomically $ tryReadTChan (childChannel ctx)

259+ res <- liftIO $ atomically $ tryReadTChan (childChannel ctx)

260+ maybe (dispatch >> continue stp yld)

269− maybe (continue stp yld) (\ev -> handleEvent ev stp yld) res

261+ (\ev -> handleEvent ev stp yld) res

270

262

271 where

263 where

272

264

273 {-# INLINE continue #-}

265 {-# INLINE continue #-}

274− continue stp yld = (runAsyncT (pullDispatch ctx m True)) (Just ctx) stp yld

266+ continue stp yld = (runAsyncT (pullDrain ctx)) Nothing stp yld

267+

268+ {-# INLINE dispatch #-}

269+ dispatch = liftIO $ signalQSem (dispatchReq ctx)

275

270

276 {-# INLINE handleEvent #-}

271 {-# INLINE handleEvent #-}

277− handleEvent ev stp yld =

272+ handleEvent ev stp yld = do

273+ let yielder a = yld a Nothing (Just (pullDrain ctx))

278− case ev of

274+ case ev of

279− ChildYield a -> yld a Nothing (Just (pullDispatch ctx m False))

275+ ChildYield a -> yielder a

280 ChildDone tid a -> do

276 ChildDone tid a -> do

277+ dispatch

281− void $ delThread ctx tid

278+ done <- delThread ctx tid

282− yld a Nothing (Just (pullDispatch ctx m True))

279+ if done then (yld a Nothing Nothing) else (yielder a)

283 ChildStop tid e -> do

280 ChildStop tid e -> do

284 case e of

281 case e of

282+ Nothing -> do

283+ dispatch

285− Nothing -> void (delThread ctx tid) >> continue stp yld

284+ done <- delThread ctx tid

285+ if done then stp else continue stp yld

286 Just x -> handleException x ctx tid

287+ ChildCreate tid -> do

287− ChildCreate tid -> void (addThread ctx tid) >> continue stp yld

288+ done <- addThread ctx tid

288−

289+ if done then stp else continue stp yld

289−-- | run m1 in a new thread, pushing its results to a pull channel and then run

290−-- m2 in the parent thread. Any exceptions are also pushed to the channel.

291−{-# INLINE pullSideDispatch #-}

292−pullSideDispatch :: MonadAsync m

293− => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a

294−pullSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do

295− let chan = childChannel ctx

296− tid <- doFork (push (ctx {pullSide = False}) m1)

297− (handleChildException chan)

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

299− (runAsyncT (pullDispatch ctx m2 False)) Nothing stp yld

300

290

301−{-# NOINLINE pullDrain #-}

291+-- If an exception occurs we push it to the channel so that it can handled by

302−pullDrain :: (MonadIO m, MonadThrow m) => Context a -> AsyncT m a

292+-- the parent. 'Paused' exceptions are to be collected at the top level.

303−pullDrain ctx = AsyncT $ \_ stp yld -> do

293+-- XXX Paused exceptions should only bubble up to the runRecorder computation

304− let yielder a = yld a Nothing (Just (pullDrain ctx))

294+{-# NOINLINE handleChildException #-}

305− continue = (runAsyncT (pullDrain ctx)) Nothing stp yld

295+handleChildException :: TChan (ChildEvent a) -> SomeException -> IO ()

306−

296+handleChildException pchan e = do

307− res <- liftIO $ atomically $ readTChan (childChannel ctx)

297+ tid <- myThreadId

308− case res of

309− ChildYield a -> yielder a

310− ChildDone tid a -> do

311− done <- delThread ctx tid

312− if done then (yld a Nothing Nothing) else (yielder a)

313− ChildStop tid e -> do

298+ atomically $ writeTChan pchan (ChildStop tid (Just e))

314− case e of

315− Nothing -> do

316− done <- delThread ctx tid

317− if done then stp else continue

318− Just x -> handleException x ctx tid

319− ChildCreate tid -> do

320− done <- addThread ctx tid

321− if done then stp else continue

322−

323−pullDrainStart :: (MonadIO m, MonadThrow m) => Context a -> AsyncT m a

324−pullDrainStart ctx = AsyncT $ \_ stp yld -> do

325− r <- liftIO $ readIORef (runningThreads ctx)

326− d <- liftIO $ readIORef (doneThreads ctx)

327− if (S.null r && S.null d)

328− then stp

329− else (runAsyncT (pullDrain ctx)) Nothing stp yld

330

299

331-- | Split the original computation in a pull-push pair. The original

300-- | Split the original computation in a pull-push pair. The original

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

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

334pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a

303pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> AsyncT m a

335pullFork m1 m2 = AsyncT $ \_ stp yld -> do

304pullFork m1 m2 = AsyncT $ \_ stp yld -> do

336 ctx <- liftIO $ newContext

305 ctx <- liftIO $ newContext

337− let m = pullDispatch ctx ((m1 <|> m2) <> pullDrainStart ctx) True

306+ initialFork ctx m1 >> initialFork ctx m2

338− (runAsyncT m) Nothing stp yld

307+ (runAsyncT (pullDrain ctx)) Nothing stp yld

339

308

340 where

309 where

341

310

311+ -- This function is different than "fork" because we have to directly

312+ -- insert the threadIds here and cannot use the channel to send ChildCreate

313+ -- unlike on the push side. If we do that, the first thread's done message

314+ -- may arrive even before the second thread is forked, in that case

315+ -- pullDrain will falsely detect that all threads are over.

316+ initialFork ctx m = do

317+ let chan = childChannel ctx

318+ tid <- doFork (push ctx m) (handleChildException chan)

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

320+

342 newContext = do

321 newContext = do

343 channel <- atomically newTChan

322 channel <- atomically newTChan

344 running <- liftIO $ newIORef S.empty

323 running <- liftIO $ newIORef S.empty

345− done <- liftIO $ newIORef S.empty

324+ done <- liftIO $ newIORef S.empty

325+ count <- liftIO $ newQSem 1

346− return $ Context { childChannel = channel

326+ return $ Context { childChannel = channel

347− , pullSide = True

327+ , dispatchReq = count

348 , runningThreads = running

328 , runningThreads = running

349− , doneThreads = done

329+ , doneThreads = done

350 }

330 }

351

331

352-- Concurrency rate control. Our objective is to create more threads on

332-- Concurrency rate control. Our objective is to create more threads on

379-- XXX to rate control left folded structrues we will have to return the

359-- XXX to rate control left folded structrues we will have to return the

380-- residual work back to the dispatcher. It will also consume a lot of

360-- residual work back to the dispatcher. It will also consume a lot of

381-- memory due to queueing of all the work before execution starts.

361-- memory due to queueing of all the work before execution starts.

382−--

362+

363+{-# INLINE fork #-}

364+fork :: MonadAsync m => Context a -> AsyncT m a -> m ()

365+fork ctx m = do

366+ let chan = childChannel ctx

367+ tid <- doFork (push ctx m) (handleChildException chan)

368+ liftIO $ atomically $ writeTChan chan (ChildCreate tid)

369+

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

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

384 empty = mempty

371 empty = mempty

385

372

389 m1 <|> m2 = AsyncT $ \ctx stp yld -> do

376 m1 <|> m2 = AsyncT $ \ctx stp yld -> do

390 case ctx of

377 case ctx of

391 Nothing -> (runAsyncT (pullFork m1 m2)) ctx stp yld

378 Nothing -> (runAsyncT (pullFork m1 m2)) ctx stp yld

392− Just c -> do

393− if pullSide c then

394− (runAsyncT (pullSideDispatch c m1 m2)) ctx stp yld

395− else

396− -- for left associated compositions

397− (runAsyncT (pushSideDispatch c m1 m2)) ctx stp yld

379+ Just c -> fork c m2 >> (runAsyncT m1) ctx stp yld

398

380

399instance MonadAsync m => MonadPlus (AsyncT m) where

381instance MonadAsync m => MonadPlus (AsyncT m) where

400 mzero = empty

382 mzero = empty