2e4fc99 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Use IORefs with CAS and MVar instead of TBQueue

2e4fc99Harendra Kumar9 years ago

asyncly.cabal+1

compare with master

unifiedsplit

@@ -31,6 +31,7 @@ library

31 -fwarn-incomplete-uni-patterns -fwarn-tabs

33 build-depends: base >= 4.8 && < 5

34+ , atomic-primops >= 0.8 && < 0.9

34 , containers >= 0.5 && < 0.6

35 , containers >= 0.5 && < 0.6

35 , exceptions >= 0.8 && < 0.9

36 , exceptions >= 0.8 && < 0.9

36 , lifted-base >= 0.2 && < 0.3

37 , lifted-base >= 0.2 && < 0.3

src/Asyncly/AsyncT.hs+91−76

@@ -37,10 +37,8 @@ where

37import Control.Applicative (Alternative (..), liftA2)

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

39 myThreadId, threadDelay)

40−import Control.Concurrent.STM (TBQueue, atomically, newTBQueue,

40+import Control.Concurrent.MVar (MVar, newEmptyMVar, tryTakeMVar,

41− tryReadTBQueue, writeTBQueue,

41+ tryPutMVar, takeMVar)

42− isEmptyTBQueue, isFullTBQueue,

43− peekTBQueue)

44import Control.Exception (SomeException (..))

42import Control.Exception (SomeException (..))

45import qualified Control.Exception.Lifted as EL

43import qualified Control.Exception.Lifted as EL

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

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

@@ -56,7 +54,9 @@ import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)

56import Data.Functor (void)

54import Data.Functor (void)

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

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

58 writeIORef, readIORef)

56 writeIORef, readIORef)

59−import Data.Maybe (fromJust)

57+import Data.Atomics (atomicModifyIORefCAS,

58+ atomicModifyIORefCAS_)

59+import Data.Maybe (isNothing)

60import Data.Monoid ((<>))

61import Data.Set (Set)

62import qualified Data.Set as S

@@ -93,9 +93,12 @@ data ChildEvent a =

93-- State threaded around the monad for thread management

94------------------------------------------------------------------------------

96+ -- XXX keep just a list of [a] and handle other events out of band as

97+ -- exceptions or using atomicModify

96data Context m a =

98data Context m a =

97− Context { outputQueue :: TBQueue (ChildEvent a)

99+ Context { outputQueue :: IORef [ChildEvent a]

100+ , synchOutQ :: MVar Bool -- wakeup mechanism for outQ

98− , workQueue :: TBQueue (AsyncT m a)

101+ , workQueue :: IORef [AsyncT m a]

99 , runningThreads :: IORef (Set ThreadId)

102 , runningThreads :: IORef (Set ThreadId)

100 , doneThreads :: IORef (Set ThreadId)

103 , doneThreads :: IORef (Set ThreadId)

101 , fairSched :: Bool

104 , fairSched :: Bool

@@ -198,46 +201,40 @@ instance Monad m => Applicative (AsyncT m) where

198{-# INLINE doFork #-}

201{-# INLINE doFork #-}

199doFork :: (MonadIO m, MonadBaseControl IO m)

202doFork :: (MonadIO m, MonadBaseControl IO m)

200 => m ()

203 => m ()

201− -> (SomeException -> IO ())

204+ -> (SomeException -> m ())

202 -> m ThreadId

205 -> m ThreadId

203doFork action exHandler =

206doFork action exHandler =

204 EL.mask $ \restore ->

207 EL.mask $ \restore ->

205 liftBaseWith $ \runInIO -> forkIO $ do

208 liftBaseWith $ \runInIO -> forkIO $ do

206 -- XXX test the exception handling

209 -- XXX test the exception handling

207− _ <- runInIO $ EL.catch (restore action) (liftIO . exHandler)

210+ _ <- runInIO $ EL.catch (restore action) exHandler

208 -- XXX restore state here?

211 -- XXX restore state here?

209 return ()

212 return ()

210

213

211−{-# NOINLINE forkWorker #-}

214+-- XXX exception safety of all atomic/MVar operations

212−forkWorker :: MonadAsync m => Maybe (Context m a) -> m ()

213−forkWorker ctx = do

214− let c = fromJust ctx

215− let q = outputQueue c

216− tid <- doFork (push c) (handleChildException q)

217− liftIO $ atomically $ writeTBQueue q (ChildCreate tid)

218

215

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

216+{-# INLINE send #-}

217+send :: MonadIO m => Context m a -> ChildEvent a -> m ()

218+send ctx msg = liftIO $ do

219+ atomicModifyIORefCAS_ (outputQueue ctx) $ \es -> msg : es

220+ -- XXX need a memory barrier? The wake up must happen only after the

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

222+ void $ tryPutMVar (synchOutQ ctx) True

223+

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

220{-# INLINE queueWork #-}

225{-# INLINE queueWork #-}

221−queueWork :: MonadAsync m

226+queueWork :: MonadAsync m => IORef [AsyncT m a] -> AsyncT m a -> m ()

222− => Maybe (Context m a) -> TBQueue (AsyncT m a) -> AsyncT m a -> m ()

227+queueWork q m = liftIO $ atomicModifyIORefCAS_ q $ \ ms -> m : ms

223−queueWork ctx q m = do

224− -- To guarantee deadlock avoidance we need to dispatch a worker when the

225− -- workQueue goes full. Otherwise all the worker threads might be waiting

226− -- on the queue and never wakeup.

227− --

228− -- TBD If we run out of threads we can also evaluate the action completely

229− -- right here, disallowing any further child workers and turning the

230− -- parallel composition into interleaved serial composition.

231− workQFull <- liftIO $ atomically $ isFullTBQueue q

232− when (workQFull) $ forkWorker ctx

233− liftIO $ atomically $ writeTBQueue q m

234

228

235-- Note: This is performance sensitive code.

229-- Note: This is performance sensitive code.

236{-# INLINE dequeueLoop #-}

230{-# INLINE dequeueLoop #-}

237dequeueLoop :: MonadAsync m

231dequeueLoop :: MonadAsync m

238− => Maybe (Context m a) -> TBQueue (AsyncT m a) -> Bool -> AsyncT m a

232+ => Maybe (Context m a) -> IORef [AsyncT m a] -> Bool -> AsyncT m a

239dequeueLoop ctx q fair = AsyncT $ \_ stp yld -> do

233dequeueLoop ctx q fair = AsyncT $ \_ stp yld -> do

240− work <- liftIO $ atomically $ tryReadTBQueue q

234+ work <- liftIO $ atomicModifyIORefCAS q $ \ ms ->

235+ case ms of

236+ [] -> ([], Nothing)

237+ x : xs -> (xs, Just x)

241 case work of

238 case work of

242 Nothing -> stp

239 Nothing -> stp

243 Just m -> do

240 Just m -> do

@@ -245,7 +242,7 @@ dequeueLoop ctx q fair = AsyncT $ \_ stp yld -> do

245 continue a c = yld a c (Just (dequeueLoop ctx q fair))

242 continue a c = yld a c (Just (dequeueLoop ctx q fair))

246 yield a c Nothing = continue a c

243 yield a c Nothing = continue a c

247 yield a c r | not fair = yld a c r

244 yield a c r | not fair = yld a c r

248− yield a c (Just r) = queueWork ctx q r >> continue a c

245+ yield a c (Just r) = queueWork q r >> continue a c

249 (runAsyncT m) ctx stop yield

246 (runAsyncT m) ctx stop yield

250

247

251{-# NOINLINE push #-}

248{-# NOINLINE push #-}

@@ -255,13 +252,17 @@ push ctx = run (Just ctx) (dequeueLoop (Just ctx) (workQueue ctx)

255

252

256 where

253 where

257

254

258− send msg = atomically $ writeTBQueue (outputQueue ctx) msg

259 stop = do

255 stop = do

260− workQEmpty <- liftIO $ atomically $ isEmptyTBQueue (workQueue ctx)

256+ work <- liftIO $ readIORef (workQueue ctx)

261− if (not workQEmpty) then push ctx

257+ case work of

258+ [] -> do

259+ tid <- liftIO $ myThreadId

262− else liftIO $ myThreadId >>= \tid -> send (ChildStop tid Nothing)

260+ send ctx (ChildStop tid Nothing)

261+ _ -> push ctx

262+ yield a _ Nothing = do

263+ tid <- liftIO myThreadId

263− yield a _ Nothing = liftIO $ myThreadId >>= \tid -> send (ChildDone tid a)

264+ send ctx (ChildDone tid a)

264− yield a c (Just r) = liftIO (send (ChildYield a)) >> run c r

265+ yield a c (Just r) = send ctx (ChildYield a) >> run c r

265 run c m = (runAsyncT m) c stop yield

266 run c m = (runAsyncT m) c stop yield

266

267

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

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

@@ -301,6 +302,13 @@ addThread ctx tid = accountThread tid (doneThreads ctx) (runningThreads ctx)

301delThread :: MonadIO m => Context m a -> ThreadId -> m Bool

302delThread :: MonadIO m => Context m a -> ThreadId -> m Bool

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

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

303

304

305+{-# INLINE allThreadsDone #-}

306+allThreadsDone :: MonadIO m => Context m a -> m Bool

307+allThreadsDone ctx = liftIO $ do

308+ s1 <- readIORef (runningThreads ctx)

309+ s2 <- readIORef (doneThreads ctx)

310+ return $ S.null s1 && S.null s2

311+

304handleException :: (MonadIO m, MonadThrow m)

312handleException :: (MonadIO m, MonadThrow m)

305 => SomeException -> Context m a -> ThreadId -> m r

313 => SomeException -> Context m a -> ThreadId -> m r

306handleException e ctx tid = do

314handleException e ctx tid = do

@@ -308,59 +316,64 @@ handleException e ctx tid = do

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

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

309 throwM e

317 throwM e

310

318

311−{-# NOINLINE sendWorkerWait #-}

319+{-# INLINE sendWorkerWait #-}

312sendWorkerWait :: MonadAsync m => Context m a -> m ()

320sendWorkerWait :: MonadAsync m => Context m a -> m ()

313sendWorkerWait ctx = do

321sendWorkerWait ctx = do

314− liftIO $ threadDelay 4

322+ liftIO $ threadDelay 200

315− let workQ = workQueue ctx

323+ output <- liftIO $ readIORef (outputQueue ctx)

316− outQ = outputQueue ctx

324+ case output of

317− workQEmpty <- liftIO $ atomically $ isEmptyTBQueue workQ

318− outQEmpty <- liftIO $ atomically $ isEmptyTBQueue outQ

319− when (not workQEmpty && outQEmpty) $ pushWorker ctx

325+ [] -> pushWorker ctx

320− void $ liftIO $ atomically $ peekTBQueue (outputQueue ctx)

326+ _ -> return ()

327+ void $ liftIO $ takeMVar (synchOutQ ctx)

321

328

322-- Note: This is performance sensitive code.

329-- Note: This is performance sensitive code.

323{-# NOINLINE pullWorker #-}

330{-# NOINLINE pullWorker #-}

324pullWorker :: MonadAsync m => Context m a -> AsyncT m a

331pullWorker :: MonadAsync m => Context m a -> AsyncT m a

325pullWorker ctx = AsyncT $ \pctx stp yld -> do

332pullWorker ctx = AsyncT $ \pctx stp yld -> do

326− let continue = (runAsyncT (pullWorker ctx)) pctx stp yld

333+ res <- liftIO $ tryTakeMVar (synchOutQ ctx)

327− yield a = yld a pctx (Just (pullWorker ctx))

334+ when (isNothing res) $ sendWorkerWait ctx

328− threadOp tid f finish cont = do

335+ list <- liftIO $ atomicModifyIORefCAS (outputQueue ctx) $ \x -> ([], x)

336+ (runAsyncT $ processEvents list) pctx stp yld

337+ where

338+

339+ {-# INLINE processEvents #-}

340+ processEvents [] = AsyncT $ \pctx stp yld -> do

329− done <- f ctx tid

341+ done <- allThreadsDone ctx

330− if done then finish else cont

342+ if not done

331−

343+ then (runAsyncT (pullWorker ctx)) pctx stp yld

332− res <- liftIO $ atomically $ tryReadTBQueue (outputQueue ctx)

344+ else stp

333− case res of

345+

334− Nothing -> sendWorkerWait ctx >> continue

346+ processEvents (ev : es) = AsyncT $ \pctx stp yld -> do

335− Just ev ->

347+ let continue = (runAsyncT (processEvents es)) pctx stp yld

348+ yield a = yld a pctx (Just (processEvents es))

349+

336− case ev of

350+ case ev of

337− ChildYield a -> yield a

351+ ChildYield a -> yield a

338− ChildDone tid a ->

352+ ChildDone tid a -> delThread ctx tid >> yield a

339− threadOp tid delThread (yld a pctx Nothing) (yield a)

340− ChildStop tid e ->

353+ ChildStop tid e ->

341− case e of

354+ case e of

342− Nothing -> threadOp tid delThread stp continue

355+ Nothing -> delThread ctx tid >> continue

343− Just ex -> handleException ex ctx tid

356+ Just ex -> handleException ex ctx tid

344− ChildCreate tid -> threadOp tid addThread stp continue

357+ ChildCreate tid -> addThread ctx tid >> continue

345

358

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

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

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

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

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

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

349{-# NOINLINE handleChildException #-}

362{-# NOINLINE handleChildException #-}

350−handleChildException :: TBQueue (ChildEvent a) -> SomeException -> IO ()

363+handleChildException :: MonadIO m => Context m a -> SomeException -> m ()

351−handleChildException pchan e = do

364+handleChildException ctx e = do

352− tid <- myThreadId

365+ tid <- liftIO myThreadId

353− atomically $ writeTBQueue pchan (ChildStop tid (Just e))

366+ send ctx (ChildStop tid (Just e))

354

367

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

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

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

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

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

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

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

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

359-- pullWorker will falsely detect that all threads are over.

372-- pullWorker will falsely detect that all threads are over.

360−{-# INLINE pushWorker #-}

373+{-# NOINLINE pushWorker #-}

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

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

362pushWorker ctx = do

375pushWorker ctx = do

363− tid <- doFork (push ctx) (handleChildException (outputQueue ctx))

376+ tid <- doFork (push ctx) (handleChildException ctx)

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

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

365

378

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

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

@@ -370,17 +383,19 @@ pullFork :: MonadAsync m => AsyncT m a -> AsyncT m a -> Bool -> AsyncT m a

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

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

371 ctx <- liftIO $ newContext

384 ctx <- liftIO $ newContext

372 let q = workQueue ctx

385 let q = workQueue ctx

373− queueWork (Just ctx) q m1 >> queueWork (Just ctx) q m2 >> pushWorker ctx

386+ queueWork q m1 >> queueWork q m2 >> pushWorker ctx

374 (runAsyncT (pullWorker ctx)) Nothing stp yld

387 (runAsyncT (pullWorker ctx)) Nothing stp yld

375

388

376 where

389 where

377

390

378 newContext = do

391 newContext = do

379− outQ <- atomically $ newTBQueue 32

392+ outQ <- newIORef []

380− workQ <- atomically $ newTBQueue 32

393+ outQMv <- newEmptyMVar

394+ workQ <- newIORef []

381 running <- newIORef S.empty

395 running <- newIORef S.empty

382 done <- newIORef S.empty

396 done <- newIORef S.empty

383 return $ Context { outputQueue = outQ

397 return $ Context { outputQueue = outQ

398+ , synchOutQ = outQMv

384 , runningThreads = running

399 , runningThreads = running

385 , doneThreads = done

400 , doneThreads = done

386 , workQueue = workQ

401 , workQueue = workQ

@@ -430,7 +445,7 @@ parallel m1 m2 fair = AsyncT $ \ctx stp yld -> do

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

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

431 -- may or may not run the right one.

446 -- may or may not run the right one.

432 let q = workQueue c

447 let q = workQueue c

433− queueWork ctx q m1 >> queueWork ctx q m2

448+ queueWork q m1 >> queueWork q m2

434 (runAsyncT (dequeueLoop ctx q fair)) Nothing stp yld

449 (runAsyncT (dequeueLoop ctx q fair)) Nothing stp yld

435

450

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

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