0d9243f · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Separate the threadstate management layer

0d9243fHarendra Kumar9 years ago

asyncly.cabal+1

compare with master

unifiedsplit

@@ -24,6 +24,7 @@ cabal-version: >= 1.10

24library

25 hs-source-dirs: src

26 exposed-modules: Asyncly

27+ , Asyncly.Threads

27 , Asyncly.AsyncT

28 , Asyncly.AsyncT

28 , Asyncly.RunAsync

29 , Asyncly.RunAsync

src/Asyncly.hs+1

@@ -149,5 +149,6 @@ module Asyncly

149 )

150where

151

152+import Asyncly.Threads

152import Asyncly.AsyncT

153import Asyncly.AsyncT

153import Asyncly.RunAsync

154import Asyncly.RunAsync

src/Asyncly/AsyncT.hs+8−404

@@ -22,18 +22,8 @@

22--

23module Asyncly.AsyncT

24 ( AsyncT (..)

25− , MonadAsync

26− , ChildEvent(..)

27− , Context(..)

28− , initContext

29− , Location(..)

30− , getLocation

31− , setLocation

33 , threads

26 , threads

34− , waitForChildren

35− , handleResult

36−

37 , async

27 , async

38 , makeAsync

28 , makeAsync

39 , each

29 , each

@@ -51,138 +41,26 @@ module Asyncly.AsyncT

51where

41where

53import Control.Applicative (Alternative (..))

43import Control.Applicative (Alternative (..))

54−import Control.Concurrent (ThreadId, forkIO, killThread,

55− myThreadId)

56−import Control.Concurrent.STM (TChan, atomically, newTChan,

57− readTChan, tryReadTChan,

58− writeTChan)

59−import Control.Exception (fromException,

60− SomeException (..))

61−import qualified Control.Exception.Lifted as EL

62import Control.Monad (ap, mzero, when)

44import Control.Monad (ap, mzero, when)

63import Control.Monad.Base (MonadBase (..), liftBaseDefault)

45import Control.Monad.Base (MonadBase (..), liftBaseDefault)

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

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

65import Control.Monad.State (MonadIO (..), MonadPlus (..),

47import Control.Monad.State (MonadIO (..), MonadPlus (..),

66 StateT (..), liftM, runStateT,

48 StateT (..), liftM, runStateT,

67− get, gets, modify)

49+ gets, modify)

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

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

69import Control.Monad.Trans.Control (ComposeSt, MonadBaseControl (..),

51import Control.Monad.Trans.Control (ComposeSt, MonadBaseControl (..),

70 MonadTransControl (..),

52 MonadTransControl (..),

71 defaultLiftBaseWith,

53 defaultLiftBaseWith,

72 defaultRestoreM, liftBaseWith)

54 defaultRestoreM, liftBaseWith)

73−import Data.Atomics (atomicModifyIORefCAS)

74−import Data.IORef (IORef, modifyIORef, newIORef,

75− readIORef, writeIORef)

76−import Data.List (delete)

77−import Data.Maybe (isJust)

78−import GHC.Prim (Any)

79import Unsafe.Coerce (unsafeCoerce)

55import Unsafe.Coerce (unsafeCoerce)

81−import Control.Monad.Trans.Recorder (MonadRecorder(..), Paused(..),

57+import Control.Monad.Trans.Recorder (MonadRecorder(..))

82− Recording)

58+import Data.IORef (newIORef, readIORef, writeIORef)

83--import Debug.Trace (traceM)

59--import Debug.Trace (traceM)

85−newtype AsyncT m a = AsyncT { runAsyncT :: StateT Context m (Maybe a) }

61+import Asyncly.Threads

86−

87−------------------------------------------------------------------------------

88−-- Parent child thread communication types

89−------------------------------------------------------------------------------

90−

91−data ChildEvent = ChildDone ThreadId (Maybe SomeException)

92−

93−------------------------------------------------------------------------------

94−-- State threaded around the monad

95−------------------------------------------------------------------------------

96−

97−data Context = Context

98− {

99− ---------------------------------------------------------------------------

100− -- Execution state

101− ---------------------------------------------------------------------------

102−

103− -- a -> AsyncT m b

104− continuations :: [Any]

105−

106− -- When we suspend we save the logs in this IORef and exit.

107− , logsRef :: Maybe (IORef [Recording])

108− -- XXX this functionality can be in a separate layer?

109− , location :: Location

110−

111− ---------------------------------------------------------------------------

112− -- Thread creation, communication and cleanup

113− ---------------------------------------------------------------------------

114− -- When a new thread is created the parent records it in the

115− -- 'pendingThreads' field and 'threadCredit' is decremented. When a child

116− -- thread is done it sends a done event to the parent on an unbounded

117− -- channel and goes away. Before starting a new thread the parent always

118− -- processes the unbounded channel to clear up the pending zombies. This

119− -- strategy ensures that the zombie queue will never grow more than the

120− -- number of running threads. Before exiting, the parent thread waits on

121− -- the channel until all its children are cleaned up.

122− ---------------------------------------------------------------------------

123−

124− -- XXX setup a cleanup computation to run rather than passing all these

125− -- params.

126−

127− , childChannel :: TChan ChildEvent

128− -- ^ A channel for the immediate children to communicate to us when they

129− -- die. Each thread has its own dedicated channel for its children

130−

131− -- We always track the child threads, otherwise the programmer needs to

132− -- worry about if some threads may remain hanging because they are stuck in

133− -- an infinite loop or a forever blocked IO. Also, it is not clear whether

134− -- there is a significant gain by disabling tracking. Instead of using

135− -- threadIds we can also keep a count instead, that will allow us to wait

136− -- for all children to drain but we cannot kill them.

137− --

138− -- We need these as IORefs since we need any changes to these to be

139− -- reflected back in the calling code that waits for the children.

140−

141− , pendingThreads :: IORef [ThreadId]

142− -- ^ Active immediate child threads spawned by this thread, modified only

143− -- by this thread, therefore atomic modification is not needed.

144−

145− -- By default this limit is shared by the current tree of threads. But at

146− -- any point the 'threads' primitive can be used to start a new limit for

147− -- the computation under that primitive.

148− --

149− -- Shared across many threads, therefore atomic access is required.

150− -- This is incremented by the child thread itelf before it exits.

151− --

152− , threadCredit :: IORef Int

153− -- ^ How many more threads are allowed to be created?

154− }

155−

156−initContext

157− :: TChan ChildEvent

158− -> IORef [ThreadId]

159− -> IORef Int

160− -> (a -> AsyncT m a)

161− -> Maybe (IORef [Recording])

162− -> Context

163−initContext childChan pending credit finalizer lref =

164− Context { continuations = [unsafeCoerce finalizer]

165− , logsRef = lref

166− , location = Worker

167− , childChannel = childChan

168− , pendingThreads = pending

169− , threadCredit = credit }

170

171−------------------------------------------------------------------------------

63+newtype AsyncT m a = AsyncT { runAsyncT :: StateT Context m (Maybe a) }

172−-- Where is the computation running?

173−------------------------------------------------------------------------------

174−

175−-- RemoteNode is used in the Alternative instance. We stop the computation when

176−-- the location is RemoteNode, we do not execute any alternatives.

177−--

178−data Location = Worker | WaitingParent | RemoteNode

179− deriving (Eq, Show)

180−

181−getLocation :: Monad m => StateT Context m Location

182−getLocation = gets location

183−

184−setLocation :: Monad m => Location -> StateT Context m ()

185−setLocation loc = modify $ \Context { .. } -> Context { location = loc, .. }

186

187------------------------------------------------------------------------------

65------------------------------------------------------------------------------

188-- Utilities

66-- Utilities

@@ -194,8 +72,6 @@ setLocation loc = modify $ \Context { .. } -> Context { location = loc, .. }

194dbg :: Monad m => a -> m ()

72dbg :: Monad m => a -> m ()

195dbg _ = return ()

73dbg _ = return ()

196

197−type MonadAsync m = (MonadIO m, MonadBaseControl IO m, MonadThrow m)

198−

199------------------------------------------------------------------------------

75------------------------------------------------------------------------------

200-- Monad

76-- Monad

201------------------------------------------------------------------------------

77------------------------------------------------------------------------------

@@ -233,259 +109,6 @@ instance MonadAsync m => Applicative (AsyncT m) where

233 pure = return

109 pure = return

234 (<*>) = ap

110 (<*>) = ap

235

111

236−------------------------------------------------------------------------------

237−-- Pick up from where we left in the previous thread

238−------------------------------------------------------------------------------

239−

240−{-# SPECIALIZE continue :: [Any] -> a -> StateT Context IO (Maybe a) #-}

241−continue :: Monad m => [Any] -> a -> StateT Context m (Maybe a)

242−continue [] x = do

243− return (Just x)

244−

245−continue [f] x = do

246− modify $ \ctx -> ctx {continuations = []}

247− runAsyncT $ (unsafeCoerce f) x

248−

249−continue (f:fs) x = do

250− modify $ \ctx -> ctx {continuations = fs}

251− y <- runAsyncT $ (unsafeCoerce f) x

252− case y of

253− Nothing -> return y

254− Just z -> continue fs z

255−

256−{-# SPECIALIZE runContext :: Context -> AsyncT IO a -> StateT Context IO () #-}

257−runContext :: MonadAsync m => Context -> AsyncT m a -> StateT Context m ()

258−runContext ctx action = do

259− _ <- lift $ runStateT m ctx

260− return ()

261−

262− where

263−

264− m = do

265− let conts = continuations ctx

266− x <- runAsyncT action

267− case x of

268− Just y -> continue conts y >> return ()

269− Nothing -> return ()

270−

271−------------------------------------------------------------------------------

272−-- Thread Management (creation, reaping and killing)

273−------------------------------------------------------------------------------

274−

275−-- XXX We are using unbounded channels so this will not block on writing to

276−-- pchan. We can use bounded channels to throttle the creation of threads based

277−-- on consumption rate.

278−processOneEvent :: MonadIO m

279− => ChildEvent

280− -> [ThreadId]

281− -> m ([ThreadId], Maybe SomeException)

282−processOneEvent (ChildDone tid e) pending = do

283− when (isJust e) $ liftIO $ mapM_ killThread pending

284− return (delete tid pending, Nothing)

285−

286−drainChildren :: MonadIO m

287− => TChan ChildEvent

288− -> [ThreadId]

289− -> m ([ThreadId], Maybe SomeException)

290−drainChildren cchan pending =

291− case pending of

292− [] -> return (pending, Nothing)

293− _ -> do

294− ev <- liftIO $ atomically $ readTChan cchan

295− (p, e) <- processOneEvent ev pending

296− maybe (drainChildren cchan p) (const $ return (p, e)) e

297−

298−waitForChildren :: MonadIO m => Context -> m (Maybe SomeException)

299−waitForChildren ctx = do

300− let pendingRef = pendingThreads ctx

301− pending <- liftIO $ readIORef pendingRef

302− (p, e) <- drainChildren (childChannel ctx) pending

303− liftIO $ writeIORef pendingRef p

304− return e

305−

306−tryReclaimZombies :: (MonadIO m, MonadThrow m) => Context -> m ()

307−tryReclaimZombies ctx = do

308− let cchan = childChannel ctx

309− pendingRef = pendingThreads ctx

310−

311− pending <- liftIO $ readIORef pendingRef

312− case pending of

313− [] -> return ()

314− _ -> do

315− mev <- liftIO $ atomically $ tryReadTChan cchan

316− case mev of

317− Nothing -> return ()

318− Just ev -> do

319− (p, e) <- processOneEvent ev pending

320− liftIO $ writeIORef pendingRef p

321− maybe (return ()) throwM e

322− tryReclaimZombies ctx

323−

324−{-# NOINLINE waitForOneEvent #-}

325−waitForOneEvent :: (MonadIO m, MonadThrow m) => Context -> m ()

326−waitForOneEvent ctx = do

327− -- XXX assert pending must have at least one element

328− -- assert that the tid is found in our list

329− let cchan = childChannel ctx

330− pendingRef = pendingThreads ctx

331−

332− ev <- liftIO $ atomically $ readTChan cchan

333− pending <- liftIO $ readIORef pendingRef

334− (p, e) <- processOneEvent ev pending

335− liftIO $ writeIORef pendingRef p

336− maybe (return ()) throwM e

337−

338−-- XXX this is not a real semaphore as it does not really block on wait,

339−-- instead it returns whether the value is zero or non-zero.

340−--

341−{-# INLINE waitQSemB #-}

342−waitQSemB :: IORef Int -> IO Bool

343−waitQSemB sem = atomicModifyIORefCAS sem $ \n ->

344− if n > 0

345− then (n - 1, True)

346− else (n, False)

347−

348−signalQSemB :: IORef Int -> IO ()

349−signalQSemB sem = atomicModifyIORefCAS sem $ \n -> (n + 1, ())

350−

351−-- Allocation of threads

352−--

353−-- global thread limit

354−-- thread fan-out i.e. per thread children limit

355−-- min per thread allocation to avoid starvation

356−--

357−-- dynamic adjustment based on the cost, speed of consumption, cpu utilization

358−-- etc. We need to adjust the limits based on cost, throughput and latencies.

359−--

360−-- The event producer thread must put the work on a work-queue and the child

361−-- threads can pick it up from there. But if there is just one consumer then it

362−-- may not make sense to have a separate producer unless the producing cost is

363−-- high.

364−--

365−

366−{-# INLINE handleResult #-}

367−handleResult :: MonadIO m

368− => Context -> Either SomeException a -> m (Maybe SomeException)

369−handleResult ctx res =

370− case res of

371− Left e -> do

372− case fromException e of

373− Just (Paused recording) ->

374− maybe (passException e)

375− (handleRecording recording) (logsRef ctx)

376− Nothing -> passException e

377− Right _ -> waitForChildren ctx

378−

379− where

380−

381− handleRecording recording ref = do

382− liftIO $ atomicModifyIORefCAS ref $ \logs -> (recording : logs, ())

383− waitForChildren ctx

384−

385− passException e = do

386− liftIO $ readIORef (pendingThreads ctx) >>= mapM_ killThread

387− return (Just e)

388−

389−forkFinally1 :: MonadAsync m

390− => Context

391− -> AsyncT m a

392− -> (Either SomeException () -> IO ())

393− -> StateT Context m ThreadId

394−forkFinally1 ctx action preExit =

395− EL.mask $ \restore ->

396− liftBaseWith $ \runInIO -> forkIO $ do

397− _ <- runInIO $ EL.try (restore (runContext ctx action))

398− >>= liftIO . preExit

399− -- XXX restore state here

400− return ()

401−

402−-- | Run a given context in a new thread.

403−--

404−forkContext :: (MonadBaseControl IO m, MonadIO m, MonadThrow m)

405− => AsyncT m a -> StateT Context m ()

406−forkContext action = do

407− parentCtx <- get

408− childCtx <- childContext parentCtx

409− tid <- forkFinally1 childCtx action

410− (beforeExit childCtx (childChannel parentCtx))

411− updatePendingThreads parentCtx tid

412−

413− where

414−

415− updatePendingThreads :: (MonadIO m, MonadThrow m)

416− => Context -> ThreadId -> m ()

417− updatePendingThreads ctx tid = do

418− -- update the new thread before reclaiming zombies so that if it exited

419− -- already reclaim finds it in the list and does not panic.

420− liftIO $ modifyIORef (pendingThreads ctx) $ (\ts -> tid:ts)

421− tryReclaimZombies ctx

422−

423− childContext ctx = do

424− pendingRef <- liftIO $ newIORef []

425− chan <- liftIO $ atomically newTChan

426− -- shares the threadCredit of the parent by default

427− return $ ctx

428− { pendingThreads = pendingRef

429− , childChannel = chan

430− }

431−

432− beforeExit ctx pchan res = do

433− r <- handleResult ctx res

434− signalQSemB (threadCredit ctx)

435− tid <- myThreadId

436− liftIO $ atomically $ writeTChan pchan (ChildDone tid r)

437−

438−-- | Decide whether to resume the context in the same thread or a new thread

439−--

440−canFork :: Context -> IO Bool

441−canFork context = do

442− gotCredit <- waitQSemB (threadCredit context)

443− case gotCredit of

444− False -> do

445− pending <- readIORef $ pendingThreads context

446− case pending of

447− [] -> return False

448− _ -> do

449− -- XXX If we have unreclaimable child threads e.g.

450− -- infinite loop, this is going to deadlock us. We need

451− -- special handling for those cases. Add those to

452− -- unreclaimable list? And always execute them in an

453− -- async thread, cannot use sync for those.

454− --

455− waitForOneEvent context

456− canFork context

457− True -> return True

458−

459−-- Housekeeping, invoked after spawning of all child tasks is done and the

460−-- parent task needs to terminate. Either the task is fully done or we handed

461−-- it over to another thread, in any case the current thread is done.

462−

463−spawningParentDone :: MonadIO m => StateT Context m (Maybe a)

464−spawningParentDone = do

465− loc <- getLocation

466− when (loc /= RemoteNode) $ setLocation WaitingParent

467− return Nothing

468−

469−-- | The task may be run in the same thread or in a new thread depending on the

470−-- forceAsync parameter and the current thread quota.

471−--

472−{-# SPECIALIZE runAsyncTask :: Bool -> AsyncT IO a -> StateT Context IO (Maybe a) #-}

473−{-# INLINE runAsyncTask #-}

474−runAsyncTask :: MonadAsync m

475− => Bool

476− -> AsyncT m a

477− -> StateT Context m (Maybe a)

478−runAsyncTask forceAsync action = do

479− if forceAsync

480− then forkContext action

481− else do

482− context <- get

483− can <- liftIO $ canFork context

484− case can of

485− False -> runContext context action

486− True -> forkContext action

487− return Nothing

488−

489-- The current model is to start a new thread for every task. The input is

112-- The current model is to start a new thread for every task. The input is

490-- provided at the time of the creation and therefore no synchronization is

113-- provided at the time of the creation and therefore no synchronization is

491-- needed compared to a pool of threads contending to get the input from a

114-- needed compared to a pool of threads contending to get the input from a

@@ -521,29 +144,10 @@ runAsyncTask forceAsync action = do

521-- 'async' is not useful as the last action as the last action always runs in

144-- 'async' is not useful as the last action as the last action always runs in

522-- the current thread.

145-- the current thread.

523async :: MonadAsync m => AsyncT m a -> AsyncT m a

146async :: MonadAsync m => AsyncT m a -> AsyncT m a

524−async action = AsyncT $ runAsyncTask True action

147+async action = AsyncT $ runAsyncTask True (runAsyncT action)

525

148

526−-- XXX After setting the callback we should wait for the callback otherwise we

527−-- will just go away. We should use the same finalization mechanism that we use

528−-- in fork.

529−--

530−-- | Makes an asyncly action from a callback setter function; can be used to

531−-- convert existing callback style code into asyncly style code. The first

532−-- parameter is a callback setter function. When we are called back,

533−-- 'makeAsync' behaves as if it was an async computation that just returned a

534−-- value of type 'a'. After the computation is done the result of the action

535−-- in second parameter is returned to the callback.

536−--

537makeAsync :: MonadAsync m => ((a -> m ()) -> m ()) -> AsyncT m a

149makeAsync :: MonadAsync m => ((a -> m ()) -> m ()) -> AsyncT m a

538−makeAsync cbsetter = AsyncT $ do

150+makeAsync = AsyncT . makeCont

539− -- XXX should we force fork a thread or keep running in the context of the

540− -- callback?

541− ctx <- get

542− lift $ cbsetter $ \a -> do

543− let s = runAsyncTask False (return a)

544− _ <- runStateT s ctx

545− return ()

546− spawningParentDone

547

151

548------------------------------------------------------------------------------

152------------------------------------------------------------------------------

549-- Alternative

153-- Alternative

@@ -553,7 +157,7 @@ instance MonadAsync m => Alternative (AsyncT m) where

553 {-# SPECIALIZE instance Alternative (AsyncT IO) #-}

157 {-# SPECIALIZE instance Alternative (AsyncT IO) #-}

554 empty = AsyncT $ return Nothing

158 empty = AsyncT $ return Nothing

555 (<|>) m1 m2 = AsyncT $ do

159 (<|>) m1 m2 = AsyncT $ do

556− r <- runAsyncTask False m1

160+ r <- runAsyncTask False (runAsyncT m1)

557 maybe (runAsyncT m2) (return . Just) r

161 maybe (runAsyncT m2) (return . Just) r

558

162

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

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

src/Asyncly/RunAsync.hs+23−13

@@ -24,13 +24,14 @@ where

24import Control.Concurrent.STM (atomically, newTChan)

25import Control.Monad.Catch (MonadCatch, throwM, try)

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

27−import Control.Monad.State (mzero, runStateT)

27+import Control.Monad.State (StateT, runStateT)

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

29import Data.Atomics (atomicModifyIORefCAS)

30import Data.IORef (IORef, newIORef, readIORef)

32import Control.Monad.Trans.Recorder (MonadRecorder(..), RecorderT,

33 Recording, blank, runRecorderT)

34+import Asyncly.Threads

34import Asyncly.AsyncT

35import Asyncly.AsyncT

36------------------------------------------------------------------------------

37------------------------------------------------------------------------------

@@ -39,10 +40,12 @@ import Asyncly.AsyncT

40-- | Run an 'AsyncT m' computation and collect the results generated by each

41-- | Run an 'AsyncT m' computation and collect the results generated by each

41-- thread of the computation in a list.

42-- thread of the computation in a list.

42−{-# SPECIALIZE waitAsync :: (a -> AsyncT IO a) -> Maybe (IORef [Recording])

43+{-# SPECIALIZE waitAsync :: (a -> StateT Context IO (Maybe a))

43− -> AsyncT IO a -> IO () #-}

44+ -> Maybe (IORef [Recording]) -> AsyncT IO a -> IO () #-}

44waitAsync :: (MonadAsync m, MonadCatch m)

45waitAsync :: (MonadAsync m, MonadCatch m)

46+ => (a -> StateT Context m (Maybe a))

45− => (a -> AsyncT m a) -> Maybe (IORef [Recording]) -> AsyncT m a -> m ()

47+ -> Maybe (IORef [Recording])

48+ -> AsyncT m a -> m ()

46waitAsync finalizer lref m = do

49waitAsync finalizer lref m = do

47 childChan <- liftIO $ atomically newTChan

50 childChan <- liftIO $ atomically newTChan

48 pendingRef <- liftIO $ newIORef []

51 pendingRef <- liftIO $ newIORef []

@@ -50,22 +53,29 @@ waitAsync finalizer lref m = do

51 let ctx = initContext childChan pendingRef credit finalizer lref

54 let ctx = initContext childChan pendingRef credit finalizer lref

53− r <- try $ runStateT (runAsyncT $ m >>= finalizer) ctx

56+ r <- try $ runStateT comp ctx

54 e <- handleResult ctx r

57 e <- handleResult ctx r

55 maybe (return ()) throwM e

58 maybe (return ()) throwM e

60+ where

61+ comp = do

62+ x <- runAsyncT m

63+ maybe (return Nothing) finalizer x

64+

57−-- TBD throttling of producer based on conumption rate.

65+-- TBD throttling of producer based on consumption rate.

59-- | Invoked to store the result of the computation in the context and finish

67-- | Invoked to store the result of the computation in the context and finish

60-- the computation when the computation is done

68-- the computation when the computation is done

61−{-# SPECIALIZE gatherResult :: IORef [a] -> a -> AsyncT IO a #-}

69+{-# SPECIALIZE gatherResult :: IORef [a] -> a -> StateT Context IO (Maybe a) #-}

62−gatherResult :: MonadAsync m => IORef [a] -> a -> AsyncT m a

70+gatherResult :: MonadAsync m => IORef [a] -> a -> StateT Context m (Maybe a)

63gatherResult ref r = do

71gatherResult ref r = do

64 liftIO $ atomicModifyIORefCAS ref $ \rs -> (r : rs, ())

72 liftIO $ atomicModifyIORefCAS ref $ \rs -> (r : rs, ())

65− mzero

73+ return Nothing

67−gather :: (MonadAsync m, MonadCatch m) => AsyncT m a -> AsyncT m [a]

75+gather :: (MonadAsync m, MonadCatch m)

76+ => AsyncT m a

77+ -> StateT Context m (Maybe [a])

68−gather m = AsyncT $ do

78+gather m = do

69 resultsRef <- liftIO $ newIORef []

79 resultsRef <- liftIO $ newIORef []

70 lift $ waitAsync (gatherResult resultsRef) Nothing m

80 lift $ waitAsync (gatherResult resultsRef) Nothing m

71 r <- liftIO $ readIORef resultsRef

81 r <- liftIO $ readIORef resultsRef

@@ -83,7 +93,7 @@ wait m = do

83-- | Run an 'AsyncT m' computation, wait for it to finish and discard the

93-- | Run an 'AsyncT m' computation, wait for it to finish and discard the

84-- results.

94-- results.

85wait_ :: (MonadAsync m, MonadCatch m) => AsyncT m a -> m ()

95wait_ :: (MonadAsync m, MonadCatch m) => AsyncT m a -> m ()

86−wait_ m = waitAsync (const mzero) Nothing m

96+wait_ m = waitAsync (const (return Nothing)) Nothing m

88------------------------------------------------------------------------------

98------------------------------------------------------------------------------

89-- Logging

99-- Logging

@@ -120,6 +130,6 @@ waitRecord_ :: (MonadAsync m, MonadCatch m)

120 => AsyncT (RecorderT m) a -> m [Recording]

130 => AsyncT (RecorderT m) a -> m [Recording]

121waitRecord_ m = do

131waitRecord_ m = do

122 lref <- liftIO $ newIORef []

132 lref <- liftIO $ newIORef []

123− runRecorderT blank (waitAsync (const mzero) (Just lref) m)

133+ runRecorderT blank (waitAsync (const (return Nothing)) (Just lref) m)

124 logs <- liftIO $ readIORef lref

134 logs <- liftIO $ readIORef lref

125 return logs

135 return logs

src/Asyncly/Threads.hsnew file+431

@@ -0,0 +1,431 @@

1+{-# LANGUAGE ConstraintKinds #-}

2+{-# LANGUAGE FlexibleContexts #-}

3+{-# LANGUAGE RecordWildCards #-}

5+-- |

6+-- Module : Asyncly.Threads

8+--

9+-- License : MIT-style

10+-- Maintainer : harendra.kumar@gmail.com

11+-- Stability : experimental

12+-- Portability : GHC

13+--

14+module Asyncly.Threads

15+ ( -- ChildEvent(..)

16+ MonadAsync

17+ , Context(..)

18+ , initContext

19+ , runAsyncTask

20+ , makeCont

21+ -- , Location(..)

22+ -- , getLocation

23+ -- , setLocation

24+

25+ -- , waitForChildren

26+ , handleResult

27+ )

28+where

29+

30+import Control.Concurrent (ThreadId, forkIO, killThread,

31+ myThreadId)

32+import Control.Concurrent.STM (TChan, atomically, newTChan,

33+ readTChan, tryReadTChan,

34+ writeTChan)

35+import Control.Exception (fromException,

36+ SomeException (..))

37+import qualified Control.Exception.Lifted as EL

38+import Control.Monad (when)

39+import Control.Monad.Catch (MonadThrow, throwM)

40+import Control.Monad.State (MonadIO (..),

41+ StateT (..), runStateT,

42+ get, gets, modify)

43+import Control.Monad.Trans.Class (MonadTrans (lift))

44+import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)

45+import Data.Atomics (atomicModifyIORefCAS)

46+import Data.IORef (IORef, modifyIORef, newIORef,

47+ readIORef, writeIORef)

48+import Data.List (delete)

49+import Data.Maybe (isJust)

50+import GHC.Prim (Any)

51+import Unsafe.Coerce (unsafeCoerce)

52+

53+import Control.Monad.Trans.Recorder (Paused (..), Recording)

54+

55+type MonadAsync m = (MonadIO m, MonadBaseControl IO m, MonadThrow m)

56+

57+------------------------------------------------------------------------------

58+-- Parent child thread communication types

59+------------------------------------------------------------------------------

60+

61+data ChildEvent = ChildDone ThreadId (Maybe SomeException)

62+

63+------------------------------------------------------------------------------

64+-- State threaded around the monad

65+------------------------------------------------------------------------------

66+

67+data Context = Context

68+ {

69+ ---------------------------------------------------------------------------

70+ -- Execution state

71+ ---------------------------------------------------------------------------

72+

73+ -- a -> AsyncT m b

74+ continuations :: [Any]

75+

76+ -- When we suspend we save the logs in this IORef and exit.

77+ , logsRef :: Maybe (IORef [Recording])

78+ -- XXX this functionality can be in a separate layer?

79+ , location :: Location

80+

81+ ---------------------------------------------------------------------------

82+ -- Thread creation, communication and cleanup

83+ ---------------------------------------------------------------------------

84+ -- When a new thread is created the parent records it in the

85+ -- 'pendingThreads' field and 'threadCredit' is decremented. When a child

86+ -- thread is done it sends a done event to the parent on an unbounded

87+ -- channel and goes away. Before starting a new thread the parent always

88+ -- processes the unbounded channel to clear up the pending zombies. This

89+ -- strategy ensures that the zombie queue will never grow more than the

90+ -- number of running threads. Before exiting, the parent thread waits on

91+ -- the channel until all its children are cleaned up.

92+ ---------------------------------------------------------------------------

93+

94+ -- XXX setup a cleanup computation to run rather than passing all these

95+ -- params.

96+

97+ , childChannel :: TChan ChildEvent

98+ -- ^ A channel for the immediate children to communicate to us when they

99+ -- die. Each thread has its own dedicated channel for its children

100+

101+ -- We always track the child threads, otherwise the programmer needs to

102+ -- worry about if some threads may remain hanging because they are stuck in

103+ -- an infinite loop or a forever blocked IO. Also, it is not clear whether

104+ -- there is a significant gain by disabling tracking. Instead of using

105+ -- threadIds we can also keep a count instead, that will allow us to wait

106+ -- for all children to drain but we cannot kill them.

107+ --

108+ -- We need these as IORefs since we need any changes to these to be

109+ -- reflected back in the calling code that waits for the children.

110+

111+ , pendingThreads :: IORef [ThreadId]

112+ -- ^ Active immediate child threads spawned by this thread, modified only

113+ -- by this thread, therefore atomic modification is not needed.

114+

115+ -- By default this limit is shared by the current tree of threads. But at

116+ -- any point the 'threads' primitive can be used to start a new limit for

117+ -- the computation under that primitive.

118+ --

119+ -- Shared across many threads, therefore atomic access is required.

120+ -- This is incremented by the child thread itelf before it exits.

121+ --

122+ , threadCredit :: IORef Int

123+ -- ^ How many more threads are allowed to be created?

124+ }

125+

126+initContext

127+ :: TChan ChildEvent

128+ -> IORef [ThreadId]

129+ -> IORef Int

130+ -> (a -> StateT Context m (Maybe a))

131+ -> Maybe (IORef [Recording])

132+ -> Context

133+initContext childChan pending credit finalizer lref =

134+ Context { continuations = [unsafeCoerce finalizer]

135+ , logsRef = lref

136+ , location = Worker

137+ , childChannel = childChan

138+ , pendingThreads = pending

139+ , threadCredit = credit }

140+

141+------------------------------------------------------------------------------

142+-- Where is the computation running?

143+------------------------------------------------------------------------------

144+

145+-- RemoteNode is used in the Alternative instance. We stop the computation when

146+-- the location is RemoteNode, we do not execute any alternatives.

147+--

148+data Location = Worker | WaitingParent | RemoteNode

149+ deriving (Eq, Show)

150+

151+getLocation :: Monad m => StateT Context m Location

152+getLocation = gets location

153+

154+setLocation :: Monad m => Location -> StateT Context m ()

155+setLocation loc = modify $ \Context { .. } -> Context { location = loc, .. }

156+

157+------------------------------------------------------------------------------

158+-- Pick up from where we left in the previous thread

159+------------------------------------------------------------------------------

160+

161+{-# SPECIALIZE continue :: [Any] -> a -> StateT Context IO (Maybe a) #-}

162+continue :: Monad m => [Any] -> a -> StateT Context m (Maybe a)

163+continue [] x = do

164+ return (Just x)

165+

166+continue [f] x = do

167+ modify $ \ctx -> ctx {continuations = []}

168+ (unsafeCoerce f) x

169+

170+continue (f:fs) x = do

171+ modify $ \ctx -> ctx {continuations = fs}

172+ y <- (unsafeCoerce f) x

173+ case y of

174+ Nothing -> return y

175+ Just z -> continue fs z

176+

177+{-# SPECIALIZE runContext :: Context -> StateT Context IO (Maybe a) -> StateT Context IO () #-}

178+runContext :: MonadAsync m => Context -> StateT Context m (Maybe a) -> StateT Context m ()

179+runContext ctx action = do

180+ _ <- lift $ runStateT m ctx

181+ return ()

182+

183+ where

184+

185+ m = do

186+ let conts = continuations ctx

187+ x <- action

188+ case x of

189+ Just y -> continue conts y >> return ()

190+ Nothing -> return ()

191+

192+------------------------------------------------------------------------------

193+-- Thread Management (creation, reaping and killing)

194+------------------------------------------------------------------------------

195+

196+-- XXX We are using unbounded channels so this will not block on writing to

197+-- pchan. We can use bounded channels to throttle the creation of threads based

198+-- on consumption rate.

199+processOneEvent :: MonadIO m

200+ => ChildEvent

201+ -> [ThreadId]

202+ -> m ([ThreadId], Maybe SomeException)

203+processOneEvent (ChildDone tid e) pending = do

204+ when (isJust e) $ liftIO $ mapM_ killThread pending

205+ return (delete tid pending, Nothing)

206+

207+drainChildren :: MonadIO m

208+ => TChan ChildEvent

209+ -> [ThreadId]

210+ -> m ([ThreadId], Maybe SomeException)

211+drainChildren cchan pending =

212+ case pending of

213+ [] -> return (pending, Nothing)

214+ _ -> do

215+ ev <- liftIO $ atomically $ readTChan cchan

216+ (p, e) <- processOneEvent ev pending

217+ maybe (drainChildren cchan p) (const $ return (p, e)) e

218+

219+waitForChildren :: MonadIO m => Context -> m (Maybe SomeException)

220+waitForChildren ctx = do

221+ let pendingRef = pendingThreads ctx

222+ pending <- liftIO $ readIORef pendingRef

223+ (p, e) <- drainChildren (childChannel ctx) pending

224+ liftIO $ writeIORef pendingRef p

225+ return e

226+

227+tryReclaimZombies :: (MonadIO m, MonadThrow m) => Context -> m ()

228+tryReclaimZombies ctx = do

229+ let cchan = childChannel ctx

230+ pendingRef = pendingThreads ctx

231+

232+ pending <- liftIO $ readIORef pendingRef

233+ case pending of

234+ [] -> return ()

235+ _ -> do

236+ mev <- liftIO $ atomically $ tryReadTChan cchan

237+ case mev of

238+ Nothing -> return ()

239+ Just ev -> do

240+ (p, e) <- processOneEvent ev pending

241+ liftIO $ writeIORef pendingRef p

242+ maybe (return ()) throwM e

243+ tryReclaimZombies ctx

244+

245+{-# NOINLINE waitForOneEvent #-}

246+waitForOneEvent :: (MonadIO m, MonadThrow m) => Context -> m ()

247+waitForOneEvent ctx = do

248+ -- XXX assert pending must have at least one element

249+ -- assert that the tid is found in our list

250+ let cchan = childChannel ctx

251+ pendingRef = pendingThreads ctx

252+

253+ ev <- liftIO $ atomically $ readTChan cchan

254+ pending <- liftIO $ readIORef pendingRef

255+ (p, e) <- processOneEvent ev pending

256+ liftIO $ writeIORef pendingRef p

257+ maybe (return ()) throwM e

258+

259+-- XXX this is not a real semaphore as it does not really block on wait,

260+-- instead it returns whether the value is zero or non-zero.

261+--

262+{-# INLINE waitQSemB #-}

263+waitQSemB :: IORef Int -> IO Bool

264+waitQSemB sem = atomicModifyIORefCAS sem $ \n ->

265+ if n > 0

266+ then (n - 1, True)

267+ else (n, False)

268+

269+signalQSemB :: IORef Int -> IO ()

270+signalQSemB sem = atomicModifyIORefCAS sem $ \n -> (n + 1, ())

271+

272+-- Allocation of threads

273+--

274+-- global thread limit

275+-- thread fan-out i.e. per thread children limit

276+-- min per thread allocation to avoid starvation

277+--

278+-- dynamic adjustment based on the cost, speed of consumption, cpu utilization

279+-- etc. We need to adjust the limits based on cost, throughput and latencies.

280+--

281+-- The event producer thread must put the work on a work-queue and the child

282+-- threads can pick it up from there. But if there is just one consumer then it

283+-- may not make sense to have a separate producer unless the producing cost is

284+-- high.

285+--

286+

287+{-# INLINE handleResult #-}

288+handleResult :: MonadIO m

289+ => Context -> Either SomeException a -> m (Maybe SomeException)

290+handleResult ctx res =

291+ case res of

292+ Left e -> do

293+ case fromException e of

294+ Just (Paused recording) ->

295+ maybe (passException e)

296+ (handleRecording recording) (logsRef ctx)

297+ Nothing -> passException e

298+ Right _ -> waitForChildren ctx

299+

300+ where

301+

302+ handleRecording recording ref = do

303+ liftIO $ atomicModifyIORefCAS ref $ \logs -> (recording : logs, ())

304+ waitForChildren ctx

305+

306+ passException e = do

307+ liftIO $ readIORef (pendingThreads ctx) >>= mapM_ killThread

308+ return (Just e)

309+

310+forkFinally1 :: MonadAsync m

311+ => Context

312+ -> StateT Context m (Maybe a)

313+ -> (Either SomeException () -> IO ())

314+ -> StateT Context m ThreadId

315+forkFinally1 ctx action preExit =

316+ EL.mask $ \restore ->

317+ liftBaseWith $ \runInIO -> forkIO $ do

318+ _ <- runInIO $ EL.try (restore (runContext ctx action))

319+ >>= liftIO . preExit

320+ -- XXX restore state here

321+ return ()

322+

323+-- | Run a given context in a new thread.

324+--

325+forkContext :: (MonadBaseControl IO m, MonadIO m, MonadThrow m)

326+ => StateT Context m (Maybe a) -> StateT Context m ()

327+forkContext action = do

328+ parentCtx <- get

329+ childCtx <- childContext parentCtx

330+ tid <- forkFinally1 childCtx action

331+ (beforeExit childCtx (childChannel parentCtx))

332+ updatePendingThreads parentCtx tid

333+

334+ where

335+

336+ updatePendingThreads :: (MonadIO m, MonadThrow m)

337+ => Context -> ThreadId -> m ()

338+ updatePendingThreads ctx tid = do

339+ -- update the new thread before reclaiming zombies so that if it exited

340+ -- already reclaim finds it in the list and does not panic.

341+ liftIO $ modifyIORef (pendingThreads ctx) $ (\ts -> tid:ts)

342+ tryReclaimZombies ctx

343+

344+ childContext ctx = do

345+ pendingRef <- liftIO $ newIORef []

346+ chan <- liftIO $ atomically newTChan

347+ -- shares the threadCredit of the parent by default

348+ return $ ctx

349+ { pendingThreads = pendingRef

350+ , childChannel = chan

351+ }

352+

353+ beforeExit ctx pchan res = do

354+ r <- handleResult ctx res

355+ signalQSemB (threadCredit ctx)

356+ tid <- myThreadId

357+ liftIO $ atomically $ writeTChan pchan (ChildDone tid r)

358+

359+-- | Decide whether to resume the context in the same thread or a new thread

360+--

361+canFork :: Context -> IO Bool

362+canFork context = do

363+ gotCredit <- waitQSemB (threadCredit context)

364+ case gotCredit of

365+ False -> do

366+ pending <- readIORef $ pendingThreads context

367+ case pending of

368+ [] -> return False

369+ _ -> do

370+ -- XXX If we have unreclaimable child threads e.g.

371+ -- infinite loop, this is going to deadlock us. We need

372+ -- special handling for those cases. Add those to

373+ -- unreclaimable list? And always execute them in an

374+ -- async thread, cannot use sync for those.

375+ --

376+ waitForOneEvent context

377+ canFork context

378+ True -> return True

379+

380+-- Housekeeping, invoked after spawning of all child tasks is done and the

381+-- parent task needs to terminate. Either the task is fully done or we handed

382+-- it over to another thread, in any case the current thread is done.

383+

384+spawningParentDone :: MonadIO m => StateT Context m (Maybe a)

385+spawningParentDone = do

386+ loc <- getLocation

387+ when (loc /= RemoteNode) $ setLocation WaitingParent

388+ return Nothing

389+

390+-- | The task may be run in the same thread or in a new thread depending on the

391+-- forceAsync parameter and the current thread quota.

392+--

393+{-# SPECIALIZE runAsyncTask :: Bool -> StateT Context IO (Maybe a)

394+ -> StateT Context IO (Maybe a) #-}

395+{-# INLINE runAsyncTask #-}

396+runAsyncTask :: MonadAsync m

397+ => Bool

398+ -> StateT Context m (Maybe a)

399+ -> StateT Context m (Maybe a)

400+runAsyncTask forceAsync action = do

401+ if forceAsync

402+ then forkContext action

403+ else do

404+ context <- get

405+ can <- liftIO $ canFork context

406+ case can of

407+ False -> runContext context action

408+ True -> forkContext action

409+ return Nothing

410+

411+-- XXX After setting the callback we should wait for the callback otherwise we

412+-- will just go away. We should use the same finalization mechanism that we use

413+-- in fork.

414+--

415+-- | Makes an asyncly action from a callback setter function; can be used to

416+-- convert existing callback style code into asyncly style code. The first

417+-- parameter is a callback setter function. When we are called back,

418+-- 'makeAsync' behaves as if it was an async computation that just returned a

419+-- value of type 'a'. After the computation is done the result of the action

420+-- in second parameter is returned to the callback.

421+--

422+makeCont :: MonadAsync m => ((a -> m ()) -> m ()) -> StateT Context m (Maybe a)

423+makeCont cbsetter = do

424+ -- XXX should we force fork a thread or keep running in the context of the

425+ -- callback?

426+ ctx <- get

427+ lift $ cbsetter $ \a -> do

428+ let s = runAsyncTask False (return (Just a))

429+ _ <- runStateT s ctx

430+ return ()

431+ spawningParentDone