bdaefb3 · streamly · git.iko.soy

26

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

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

29− myThreadId, newQSem, QSem,

29+ myThreadId)

30− signalQSem, waitQSem)

31import Control.Concurrent.STM (TBQueue, atomically, newTBQueue,

30import Control.Concurrent.STM (TBQueue, atomically, newTBQueue,

32− tryReadTBQueue, writeTBQueue)

31+ readTBQueue, tryReadTBQueue,

32+ writeTBQueue, isFullTBQueue)

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

34import qualified Control.Exception.Lifted as EL

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

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

36+ when)

36--import Control.Monad.Base (MonadBase (..), liftBaseDefault)

37--import Control.Monad.Base (MonadBase (..), liftBaseDefault)

37import Control.Monad.Catch (MonadThrow, throwM)

38import Control.Monad.Catch (MonadThrow, throwM)

38import Control.Monad.IO.Class (MonadIO(..))

39import Control.Monad.IO.Class (MonadIO(..))

84-- State threaded around the monad for thread management

85-- State threaded around the monad for thread management

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

86------------------------------------------------------------------------------

86

87

87−data Context a =

88+data Context m a =

88 Context { childChannel :: TBQueue (ChildEvent a)

89 Context { childChannel :: TBQueue (ChildEvent a)

89− , dispatchReq :: QSem

90+ , pendingWork :: TBQueue (AsyncT m a)

90 , runningThreads :: IORef (Set ThreadId)

91 , runningThreads :: IORef (Set ThreadId)

91 , doneThreads :: IORef (Set ThreadId)

92 , doneThreads :: IORef (Set ThreadId)

92 }

93 }

102newtype AsyncT m a =

103newtype AsyncT m a =

103 AsyncT {

104 AsyncT {

104 runAsyncT :: forall r.

105 runAsyncT :: forall r.

105− Maybe (Context a) -- state

106+ Maybe (Context m a) -- state

106 -> m r -- stop

107 -> m r -- stop

107− -> (a -> Maybe (Context a) -> Maybe (AsyncT m a) -> m r) -- yield

108+ -> (a -> Maybe (Context m a) -> Maybe (AsyncT m a) -> m r) -- yield

108 -> m r

109 -> m r

109 }

110 }

110

111

181 return ()

182 return ()

182

183

183{-# NOINLINE push #-}

184{-# NOINLINE push #-}

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

185+push :: MonadAsync m => Context m a -> m ()

185−push context action = do

186− liftIO $ waitQSem (dispatchReq context)

187− run (Just context) action

186+push context = run (Just context) (dequeueLoop context)

188

187

189 where

188 where

190

189

191 run ctx m = (runAsyncT m) ctx channelStop yielder

190 run ctx m = (runAsyncT m) ctx channelStop yielder

192

191

193− -- XXX make this a bounded channel so that we block if the previous value

194− -- is not consumed yet.

195 chan = childChannel context

192 chan = childChannel context

196 channelYield a = liftIO $ atomically $ writeTBQueue chan (ChildYield a)

193 channelYield a = liftIO $ atomically $ writeTBQueue chan (ChildYield a)

197 channelDone a = do

194 channelDone a = do

235 return False

232 return False

236

233

237{-# NOINLINE addThread #-}

234{-# NOINLINE addThread #-}

238−addThread :: MonadIO m => Context a -> ThreadId -> m Bool

235+addThread :: MonadIO m => Context m a -> ThreadId -> m Bool

239addThread ctx tid = accountThread tid (doneThreads ctx) (runningThreads ctx)

236addThread ctx tid = accountThread tid (doneThreads ctx) (runningThreads ctx)

240

237

241{-# INLINE delThread #-}

238{-# INLINE delThread #-}

242−delThread :: MonadIO m => Context a -> ThreadId -> m Bool

239+delThread :: MonadIO m => Context m a -> ThreadId -> m Bool

243delThread ctx tid = accountThread tid (runningThreads ctx) (doneThreads ctx)

240delThread ctx tid = accountThread tid (runningThreads ctx) (doneThreads ctx)

244

241

245handleException :: (MonadIO m, MonadThrow m)

242handleException :: (MonadIO m, MonadThrow m)

246− => SomeException -> Context a -> ThreadId -> m r

243+ => SomeException -> Context m a -> ThreadId -> m r

247handleException e ctx tid = do

244handleException e ctx tid = do

248 void (delThread ctx tid)

245 void (delThread ctx tid)

249 liftIO $ readIORef (runningThreads ctx) >>= mapM_ killThread

246 liftIO $ readIORef (runningThreads ctx) >>= mapM_ killThread

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

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

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

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

254-- there.

251-- there.

255−{-# NOINLINE pullDrain #-}

252+{-# NOINLINE pullWorker #-}

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

253+pullWorker :: (MonadIO m, MonadThrow m) => Context m a -> AsyncT m a

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

254+pullWorker ctx = AsyncT $ \_ stp yld -> do

258− res <- liftIO $ atomically $ tryReadTBQueue (childChannel ctx)

255+ ev <- liftIO $ atomically $ readTBQueue (childChannel ctx)

259− maybe (dispatch >> continue stp yld)

260− (\ev -> handleEvent ev stp yld) res

256+ handleEvent ev stp yld

261

257

262 where

258 where

263

259

264− {-# INLINE continue #-}

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

260+ continue stp yld = (runAsyncT (pullWorker ctx)) Nothing stp yld

266−

267− {-# INLINE dispatch #-}

268− dispatch = liftIO $ signalQSem (dispatchReq ctx)

269

261

270 {-# INLINE handleEvent #-}

262 {-# INLINE handleEvent #-}

271 handleEvent ev stp yld = do

263 handleEvent ev stp yld = do

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

264+ let yielder a = yld a Nothing (Just (pullWorker ctx))

273 case ev of

265 case ev of

274 ChildYield a -> yielder a

266 ChildYield a -> yielder a

275 ChildDone tid a -> do

267 ChildDone tid a -> do

276− dispatch

277 done <- delThread ctx tid

268 done <- delThread ctx tid

278 if done then (yld a Nothing Nothing) else (yielder a)

269 if done then (yld a Nothing Nothing) else (yielder a)

279 ChildStop tid e -> do

270 ChildStop tid e -> do

280 case e of

271 case e of

281 Nothing -> do

272 Nothing -> do

282− dispatch

283 done <- delThread ctx tid

273 done <- delThread ctx tid

284 if done then stp else continue stp yld

274 if done then stp else continue stp yld

285 Just x -> handleException x ctx tid

275 Just x -> handleException x ctx tid

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

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

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

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

304 ctx <- liftIO $ newContext

294 ctx <- liftIO $ newContext

305− initialFork ctx m1 >> initialFork ctx m2

295+ queueWork ctx m1 >> queueWork ctx m2 >> pushWorker ctx

306− (runAsyncT (pullDrain ctx)) Nothing stp yld

296+ (runAsyncT (pullWorker ctx)) Nothing stp yld

307

297

308 where

298 where

309

299

310− -- This function is different than "fork" because we have to directly

300+ -- This function is different than "forkWorker" because we have to directly

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

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

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

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

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

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

314− -- pullDrain will falsely detect that all threads are over.

304+ -- pullWorker will falsely detect that all threads are over.

315− initialFork ctx m = do

305+ pushWorker ctx = do

316 let chan = childChannel ctx

306 let chan = childChannel ctx

317− tid <- doFork (push ctx m) (handleChildException chan)

307+ tid <- doFork (push ctx) (handleChildException chan)

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

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

319

309

320 newContext = do

310 newContext = do

321 channel <- atomically $ newTBQueue 16

311 channel <- atomically $ newTBQueue 16

312+ work <- atomically $ newTBQueue 16

322 running <- newIORef S.empty

313 running <- newIORef S.empty

323 done <- newIORef S.empty

314 done <- newIORef S.empty

324− count <- newQSem 1

325 return $ Context { childChannel = channel

315 return $ Context { childChannel = channel

326− , dispatchReq = count

327 , runningThreads = running

316 , runningThreads = running

328 , doneThreads = done

317 , doneThreads = done

318+ , pendingWork = work

329 }

319 }

330

320

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

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

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

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

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

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

361

351

362−{-# INLINE fork #-}

352+{-# INLINE forkWorker #-}

363−fork :: MonadAsync m => Context a -> AsyncT m a -> m ()

353+forkWorker :: MonadAsync m => Context m a -> m ()

364−fork ctx m = do

354+forkWorker ctx = do

365 let chan = childChannel ctx

355 let chan = childChannel ctx

366− tid <- doFork (push ctx m) (handleChildException chan)

356+ tid <- doFork (push ctx) (handleChildException chan)

367 liftIO $ atomically $ writeTBQueue chan (ChildCreate tid)

357 liftIO $ atomically $ writeTBQueue chan (ChildCreate tid)

368

358

359+{-# INLINE queueWork #-}

360+queueWork :: MonadAsync m => Context m a -> AsyncT m a -> m ()

361+queueWork ctx m = do

362+ -- To guarantee deadlock avoidance we need to dispatch a worker when the

363+ -- workQueue goes full. Otherwise all the worker threads might be waiting

364+ -- on the queue and never wakeup.

365+ --

366+ -- TBD If we run out of threads we can also evaluate the action completely

367+ -- right here, disallowing any further child workers and turning the

368+ -- parallel composition into interleaved serial composition.

369+ workQFull <- liftIO $ atomically $ isFullTBQueue (pendingWork ctx)

370+ when (workQFull) $ forkWorker ctx

371+ liftIO $ atomically $ writeTBQueue (pendingWork ctx) m

372+

373+{-# INLINE dequeueLoop #-}

374+dequeueLoop :: MonadAsync m => Context m a -> AsyncT m a

375+dequeueLoop ctx = AsyncT $ \_ stp yld -> do

376+ work <- liftIO $ atomically $ tryReadTBQueue (pendingWork ctx)

377+ case work of

378+ Nothing -> stp

379+ Just m -> do

380+ let loop = (runAsyncT (dequeueLoop ctx)) Nothing stp yld

381+ yielder a c Nothing = yld a c (Just (dequeueLoop ctx))

382+ yielder a c r = yld a c r

383+ (runAsyncT m) (Just ctx) loop yielder

384+

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

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

370 empty = mempty

386 empty = mempty

371

387

374 {-# INLINE (<|>) #-}

390 {-# INLINE (<|>) #-}

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

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

376 case ctx of

392 case ctx of

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

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

394+ Just c -> do

395+ -- we open up both the branches fairly but on a given level we

396+ -- go left to right. If the left one keeps producing results we

397+ -- may or may not run the right one.

398+ queueWork c m1 >> queueWork c m2

378− Just c -> fork c m2 >> (runAsyncT m1) ctx stp yld

399+ (runAsyncT (dequeueLoop c)) Nothing stp yld

379

400

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

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

381 mzero = empty

402 mzero = empty