d8c399e · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Implement tracking and killing of child threads

d8c399eHarendra Kumar9 years ago

src/Asyncly/AsyncT.hs+96−58

compare with master

unifiedsplit

@@ -25,14 +25,12 @@ module Asyncly.AsyncT

25where

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

28−import Control.Concurrent (ThreadId, forkIO,

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

29− myThreadId, threadDelay)

29+ myThreadId)

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

31 readTChan, tryReadTChan,

32 writeTChan)

33−import Control.Exception (fromException, try,

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

34− SomeException (..),

35− BlockedIndefinitelyOnSTM(..))

36import qualified Control.Exception.Lifted as EL

34import qualified Control.Exception.Lifted as EL

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

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

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

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

@@ -46,10 +44,13 @@ import Control.Monad.Trans.Control (ComposeSt, MonadBaseControl (..),

46 defaultLiftBaseWith,

44 defaultLiftBaseWith,

47 defaultRestoreM, liftBaseWith)

45 defaultRestoreM, liftBaseWith)

48 -}

46 -}

47+import Data.Functor (void)

49−import Data.IORef (IORef, writeIORef, newIORef,

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

50− readIORef)

49+ writeIORef, readIORef)

51import Data.Maybe (maybe)

50import Data.Maybe (maybe)

52import Data.Monoid ((<>))

51import Data.Monoid ((<>))

52+import Data.Set (Set)

53+import qualified Data.Set as S

54import Control.Monad.Trans.Recorder (MonadRecorder(..))

55import Control.Monad.Trans.Recorder (MonadRecorder(..))

@@ -77,6 +78,7 @@ data ChildEvent a =

77 ChildYield a

78 ChildYield a

78 | ChildDone ThreadId a

79 | ChildDone ThreadId a

79 | ChildStop ThreadId (Maybe SomeException)

80 | ChildStop ThreadId (Maybe SomeException)

81+ | ChildCreate ThreadId

81------------------------------------------------------------------------------

83------------------------------------------------------------------------------

82-- State threaded around the monad for thread management

84-- State threaded around the monad for thread management

@@ -85,6 +87,8 @@ data ChildEvent a =

85data Context a =

87data Context a =

86 Context { childChannel :: TChan (ChildEvent a)

88 Context { childChannel :: TChan (ChildEvent a)

87 , pullSide :: Bool

89 , pullSide :: Bool

90+ , runningThreads :: IORef (Set ThreadId)

91+ , doneThreads :: IORef (Set ThreadId)

88 }

92 }

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

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

@@ -98,7 +102,7 @@ data Context a =

98newtype AsyncT m a =

102newtype AsyncT m a =

99 AsyncT {

103 AsyncT {

100 runAsyncT :: forall r.

104 runAsyncT :: forall r.

101− Maybe (Context a) -- state

105+ Maybe (Context a) -- state

102 -> m r -- stop

106 -> m r -- stop

103 -> (a -> Maybe (Context a) -> Maybe (AsyncT m a) -> m r) -- yield

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

104 -> m r

108 -> m r

@@ -168,11 +172,11 @@ doFork :: (MonadIO m, MonadBaseControl IO m)

168 => m ()

172 => m ()

169 -> (SomeException -> IO ())

173 -> (SomeException -> IO ())

170 -> m ThreadId

174 -> m ThreadId

171−doFork action preExit =

175+doFork action exHandler =

172 EL.mask $ \restore ->

176 EL.mask $ \restore ->

173 liftBaseWith $ \runInIO -> forkIO $ do

177 liftBaseWith $ \runInIO -> forkIO $ do

174 -- XXX test the exception handling

178 -- XXX test the exception handling

175− _ <- runInIO $ EL.catch (restore action) (liftIO . preExit)

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

176 -- XXX restore state here?

180 -- XXX restore state here?

177 return ()

181 return ()

178

182

@@ -202,9 +206,9 @@ push context action = run (Just context) action

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

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

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

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

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

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

205−{-# NOINLINE handlePushException #-}

209+{-# NOINLINE handleChildException #-}

206−handlePushException :: TChan (ChildEvent a) -> SomeException -> IO ()

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

207−handlePushException pchan e = do

211+handleChildException pchan e = do

208 tid <- myThreadId

212 tid <- myThreadId

209 atomically $ writeTChan pchan (ChildStop tid (Just e))

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

210

214

@@ -213,9 +217,43 @@ pushSideDispatch :: MonadAsync m

213 => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a

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

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

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

215 let chan = childChannel ctx

219 let chan = childChannel ctx

216− _ <- doFork (push ctx m1) (handlePushException chan)

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

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

217 (runAsyncT m2) (Just ctx) stp yld

222 (runAsyncT m2) (Just ctx) stp yld

218

223

224+-- 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

226+-- threads.

227+{-# INLINE accountThread #-}

228+accountThread :: MonadIO m

229+ => ThreadId -> IORef (Set ThreadId) -> IORef (Set ThreadId) -> m Bool

230+accountThread tid ref1 ref2 = liftIO $ do

231+ s1 <- readIORef ref1

232+ s2 <- readIORef ref2

233+

234+ if (S.member tid s1) then do

235+ let r = S.delete tid s1

236+ writeIORef ref1 r

237+ return $ S.null r && S.null s2

238+ else do

239+ liftIO $ writeIORef ref2 (S.insert tid s2)

240+ return False

241+

242+{-# NOINLINE addThread #-}

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

244+addThread ctx tid = accountThread tid (doneThreads ctx) (runningThreads ctx)

245+

246+{-# INLINE delThread #-}

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

248+delThread ctx tid = accountThread tid (runningThreads ctx) (doneThreads ctx)

249+

250+handleException :: (MonadIO m, MonadThrow m)

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

252+handleException e ctx tid = do

253+ void (delThread ctx tid)

254+ liftIO $ readIORef (runningThreads ctx) >>= mapM_ killThread

255+ throwM e

256+

219-- XXX the TBQueue size should be proportional to the pending threads.

257-- XXX the TBQueue size should be proportional to the pending threads.

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

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

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

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

@@ -224,24 +262,14 @@ pushSideDispatch ctx m1 m2 = AsyncT $ \_ stp yld -> do

224pullDispatch :: (MonadIO m, MonadThrow m)

262pullDispatch :: (MonadIO m, MonadThrow m)

225 => Context a -> AsyncT m a -> Bool -> AsyncT m a

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

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

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

227− let chan = childChannel ctx

228− eHandler e = handleException e stp

229− evHandler ev = handleEvent ev stp yld

230− maybeEvHandler ev = maybe (continue stp yld) evHandler ev

231−

232 if dispatch then do

265 if dispatch then do

233 (runAsyncT m) (Just ctx) stp yld

266 (runAsyncT m) (Just ctx) stp yld

234 else do

267 else do

235− res <- liftIO $ try $ atomically $ tryReadTChan chan

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

236− either eHandler maybeEvHandler res

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

237

270

238 where

271 where

239

272

240− handleException e stp =

241− case fromException e of

242− Just BlockedIndefinitelyOnSTM -> stp

243− Nothing -> throwM e

244−

245 {-# INLINE continue #-}

273 {-# INLINE continue #-}

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

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

247

275

@@ -249,9 +277,14 @@ pullDispatch ctx m dispatch = AsyncT $ \_ stp yld -> do

249 handleEvent ev stp yld =

277 handleEvent ev stp yld =

250 case ev of

278 case ev of

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

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

280+ ChildDone tid a -> do

281+ void $ delThread ctx tid

252− ChildDone _ a -> yld a Nothing (Just (pullDispatch ctx m True))

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

253− -- XXX kill threads

283+ ChildStop tid e -> do

254− ChildStop _ e -> maybe (continue stp yld) throwM e

284+ case e of

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

286+ Just x -> handleException x ctx tid

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

255

288

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

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

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

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

@@ -260,37 +293,32 @@ pullSideDispatch :: MonadAsync m

260 => Context a -> AsyncT m a -> AsyncT m a -> AsyncT m a

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

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

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

262 let chan = childChannel ctx

295 let chan = childChannel ctx

263− _ <- doFork (push (ctx {pullSide = False}) m1) (handlePushException chan)

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

264− -- liftIO $ threadDelay 0

297+ (handleChildException chan)

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

265 (runAsyncT (pullDispatch ctx m2 False)) Nothing stp yld

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

266

300

267{-# NOINLINE pullDrain #-}

301{-# NOINLINE pullDrain #-}

268−pullDrain :: (MonadIO m, MonadThrow m) => TChan (ChildEvent a) -> AsyncT m a

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

269−pullDrain chan = AsyncT $ \_ stp yld -> do

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

270− let eHandler e = handleException e stp

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

271− evHandler ev = handleEvent ev stp yld

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

272−

306+

273− res <- liftIO $ try (atomically $ readTChan chan)

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

274− either eHandler evHandler res

275−

276− where

277−

278− handleException e stp =

279− case fromException e of

280− Just BlockedIndefinitelyOnSTM -> stp

281− Nothing -> throwM e

282−

283− {-# INLINE continue #-}

284− continue stp yld = (runAsyncT (pullDrain chan)) Nothing stp yld

285−

286− {-# INLINE handleEvent #-}

287− handleEvent ev stp yld =

288− let yielder a = yld a Nothing (Just (pullDrain chan))

289− in case ev of

308+ case res of

290− ChildYield a -> yielder a

309+ ChildYield a -> yielder a

291− ChildDone _ a -> yielder a

310+ ChildDone tid a -> do

292− -- XXX kill threads

311+ done <- delThread ctx tid

293− ChildStop _ e -> maybe (continue stp yld) throwM e

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

313+ ChildStop tid e -> do

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

294

322

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

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

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

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

@@ -298,15 +326,25 @@ pullDrain chan = AsyncT $ \_ stp yld -> do

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

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

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

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

300 ctx <- liftIO $ newContext

328 ctx <- liftIO $ newContext

301− let m = pullDispatch ctx ((m1 <|> m2) <> pullDrain (childChannel ctx)) True

329+ let m = pullDispatch ctx ((m1 <|> m2) <> finalizer ctx) True

302 (runAsyncT m) Nothing stp yld

330 (runAsyncT m) Nothing stp yld

303

331

304 where

332 where

305

333

334+ finalizer ctx = AsyncT $ \_ stp yld -> do

335+ running <- liftIO $ readIORef (runningThreads ctx)

336+ done <- liftIO $ readIORef (runningThreads ctx)

337+ if (S.null running && S.null done) then stp

338+ else (runAsyncT (pullDrain ctx)) Nothing stp yld

339+

306 newContext = do

340 newContext = do

307 channel <- atomically newTChan

341 channel <- atomically newTChan

342+ running <- liftIO $ newIORef S.empty

343+ done <- liftIO $ newIORef S.empty

308 return $ Context { childChannel = channel

344 return $ Context { childChannel = channel

309 , pullSide = True

345 , pullSide = True

346+ , runningThreads = running

347+ , doneThreads = done

310 }

348 }

311

349

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

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

test/Main.hs+4−1

@@ -44,8 +44,11 @@ main = hspec $ do

44 ((toList $ (((return 0 <|> return 1) <|> return 2) <|> return 3)) >>= return . sort) `shouldReturn` ([0, 1, 2, 3] :: [Int])

45 it "Alternative composition - right associated" $

46 ((toList $ (return 0 <|> (return 1 <|> (return 2 <|> return 3)))) >>= return . sort) `shouldReturn` ([0, 1, 2, 3] :: [Int])

47− it "Alternative composition with bind" $

47+ it "Alternative composition (right fold) with bind" $

48 (toList (for [1..10 :: Int] $ \x -> return x >>= return . id) >>= return . sort) `shouldReturn` ([1..10] :: [Int])

49+ it "Alternative composition (left fold) with bind" $

50+ let forL xs f = foldl (<|>) empty $ map f xs

51+ in (toList (forL [1..10 :: Int] $ \x -> return x >>= return . id) >>= return . sort) `shouldReturn` ([1..10] :: [Int])

49 {-

52 {-

50 it "Alternative composition of async and sync tasks" $

53 it "Alternative composition of async and sync tasks" $

51 ((wait (threads 0 ((async (return 0) <|> return 1)))) >>= return . sort)

54 ((wait (threads 0 ((async (return 0) <|> return 1)))) >>= return . sort)