8a0fdb3 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Refactor/fix returning results from child threads

8a0fdb3Harendra Kumar9 years ago

src/Strands/AsyncT.hs+87−41

compare with master

unifiedsplit

@@ -12,6 +12,10 @@

12module Strands.AsyncT

13 ( AsyncT (..)

14 , waitAsync

15+ , processOneEvent

16+ , drainChildren

17+ , waitForChildren

18+ , getCtxResultDest

15 , (<**)

19 , (<**)

16 , onNothing

20 , onNothing

17 , dbg

21 , dbg

@@ -21,7 +25,7 @@ where

21import Control.Applicative (Alternative (..))

25import Control.Applicative (Alternative (..))

22import Control.Concurrent (ThreadId, killThread)

26import Control.Concurrent (ThreadId, killThread)

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

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

24− readTChan)

28+ readTChan, writeTChan)

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

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

26import Control.Monad.Catch (MonadCatch, MonadThrow, try,

30import Control.Monad.Catch (MonadCatch, MonadThrow, try,

27 throwM, SomeException)

31 throwM, SomeException)

@@ -34,14 +38,15 @@ import Control.Monad.Trans.Control (ComposeSt, MonadBaseControl (..),

34 defaultLiftBaseWith,

38 defaultLiftBaseWith,

35 defaultRestoreM)

39 defaultRestoreM)

36import Data.Dynamic (Typeable)

40import Data.Dynamic (Typeable)

37−import Data.IORef (IORef, newIORef, readIORef,

41+import Data.IORef (IORef, modifyIORef, newIORef, readIORef,

38 writeIORef)

42 writeIORef)

39import Data.List (delete)

43import Data.List (delete)

40−import Data.Maybe (isJust, isNothing)

44+import Data.Maybe (fromJust, isJust, isNothing)

45+import Unsafe.Coerce (unsafeCoerce)

42import Strands.Context

47import Strands.Context

44−-- import Debug.Trace (traceM)

49+--import Debug.Trace (traceM)

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

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

@@ -156,43 +161,82 @@ instance MonadThrow m => MonadThrow (AsyncT m) where

156-- Thread management

161-- Thread management

157------------------------------------------------------------------------------

162------------------------------------------------------------------------------

158

163

159−drainChildren :: (MonadIO m, MonadThrow m)

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

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

166+-- on consumption rate.

167+processOneEvent :: MonadIO m

168+ => ChildEvent a

160− => TChan (ChildEvent a) -> [ThreadId] -> [a] -> m [a]

169+ -> Either (TChan (ChildEvent a)) (IORef [a])

161−drainChildren chan pending res =

170+ -> [ThreadId]

162− if pending == []

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

163− then return res

172+processOneEvent ev dest pending = do

164− else do

173+ -- Collect results unless we have already encountered an exception.

165− ev <- liftIO $ atomically $ readTChan chan

166− case ev of

174+ case ev of

167− ChildDone tid er -> do

175+ ChildDone _ (Just e) -> handleException e

168− dbg $ "drainChildrenTop ChildDone, tid: " ++ show tid

176+ ChildResult (Left e) -> handleException e

169− case er of

177+ ChildDone tid Nothing -> return (delete tid pending, Nothing)

170− Left e -> throwM e

178+ ChildResult (Right r) -> do

171− Right r ->

179+ case dest of

172− drainChildren chan (delete tid pending) (r ++ res)

180+ Left chan -> liftIO $ atomically $ writeTChan chan ev

173− PassOnResult er -> do

181+ Right ref ->

174− dbg $ "drainChildrenTop PassOnResult"

182+ liftIO $ modifyIORef ref $ \rs -> unsafeCoerce r : rs

175− case er of

183+ return (pending, Nothing)

176− Left e -> throwM e

184+ where

177− Right r -> drainChildren chan pending (r ++ res)

185+

178−

186+ handleException e = do

179−waitForChildren :: (MonadIO m, MonadThrow m)

187+ liftIO $ mapM_ killThread pending

188+ return (pending, Just e)

189+

190+drainChildren :: MonadIO m

191+ => Either (TChan (ChildEvent a)) (IORef [a])

180− => TChan (ChildEvent a) -> IORef [ThreadId] -> [a] -> m [a]

192+ -> TChan (ChildEvent a)

193+ -> [ThreadId]

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

195+drainChildren dest cchan pending =

196+ case pending of

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

198+ _ -> do

199+ ev <- liftIO $ atomically $ readTChan cchan

200+ (p, e) <- processOneEvent ev dest pending

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

202+

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

181−waitForChildren chan pendingRef results = do

204+waitForChildren ctx = do

205+ let pendingRef = pendingThreads ctx

182 pending <- liftIO $ readIORef pendingRef

206 pending <- liftIO $ readIORef pendingRef

183− r <- drainChildren chan pending results

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

184− liftIO $ writeIORef pendingRef []

208+ liftIO $ writeIORef pendingRef p

185− return r

209+ return e

210+

211+getCtxResultDest :: Context -> Either (TChan (ChildEvent a)) (IORef [a])

212+getCtxResultDest ctx =

213+ maybe (Right $ unsafeCoerce $ fromJust $ accumResults ctx)

214+ (Left . unsafeCoerce) (parentChannel ctx)

186

215

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

216------------------------------------------------------------------------------

188-- Running the monad

217-- Running the monad

189------------------------------------------------------------------------------

218------------------------------------------------------------------------------

190

219

220+collectResult :: MonadIO m => a -> StateT Context m ()

221+collectResult r = do

222+ ctx <- get

223+ case parentChannel ctx of

224+ Nothing -> do

225+ let ref = fromJust $ accumResults ctx

226+ liftIO $ modifyIORef ref $ \rs -> unsafeCoerce r : rs

227+ Just chan -> do

228+ -- XXX can we pass the result directly to the root thread

229+ -- instead of passing through all the parents? We can let the

230+ -- parent go away and handle the ChildDone events as well in

231+ -- the root thread.

232+ liftIO $ atomically $ writeTChan chan

233+ (ChildResult (Right (unsafeCoerce r)))

234+

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

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

192-- the computation when the computation is done

236-- the computation when the computation is done

193−finishComputation :: Monad m => a -> AsyncT m b

237+finishComputation :: MonadIO m => a -> AsyncT m b

194finishComputation x = AsyncT $ do

238finishComputation x = AsyncT $ do

195− contextSaveResult x

239+ collectResult x

196 return Nothing

240 return Nothing

197

241

198-- XXX pass a collector function and return a Traversable.

242-- XXX pass a collector function and return a Traversable.

@@ -204,21 +248,23 @@ waitAsync :: forall m a. (MonadIO m, MonadCatch m)

204waitAsync m = do

248waitAsync m = do

205 childChan <- liftIO $ atomically newTChan

249 childChan <- liftIO $ atomically newTChan

206 pendingRef <- liftIO $ newIORef []

250 pendingRef <- liftIO $ newIORef []

251+ resultsRef <- liftIO $ newIORef []

207 credit <- liftIO $ newIORef maxBound

252 credit <- liftIO $ newIORef maxBound

208

253

209− -- XXX this should be moved to Context.hs and then we can make m

210− -- existential and remove the unsafeCoerces

211− r <- try $ runStateT (runAsyncT m) $ initContext

212− (empty :: AsyncT m a) childChan pendingRef credit

254+ let ctx = initContext (empty :: AsyncT m a) childChan pendingRef credit

255+ finishComputation resultsRef

256+

213− finishComputation

257+ r <- try $ runStateT (runAsyncT $ m >>= finishComputation) ctx

214

258

215− xs <- case r of

259+ case r of

216 Left (exc :: SomeException) -> do

260 Left (exc :: SomeException) -> do

217 liftIO $ readIORef pendingRef >>= mapM_ killThread

261 liftIO $ readIORef pendingRef >>= mapM_ killThread

218 throwM exc

262 throwM exc

219− Right (Nothing, ctx) -> return $ contextGetResult ctx

263+ Right _ -> do

220− Right ((Just x), ctx) -> return $ x : contextGetResult ctx

264+ e <- waitForChildren ctx

221− waitForChildren childChan pendingRef xs

265+ case e of

266+ Just (exc :: SomeException) -> throwM exc

267+ Nothing -> liftIO $ readIORef resultsRef

222

268

223------------------------------------------------------------------------------

269------------------------------------------------------------------------------

224-- * Extensible State: Session Data Management

270-- * Extensible State: Session Data Management

src/Strands/Context.hs+10−16

@@ -9,8 +9,6 @@ module Strands.Context

9 , saveContext

10 , restoreContext

11 , resumeContext

12− , contextSaveResult

13− , contextGetResult

14 , setContextMailBox

12 , setContextMailBox

15 , takeContextMailBox

13 , takeContextMailBox

16 , Location(..)

14 , Location(..)

@@ -40,8 +38,8 @@ import GHC.Prim (Any)

40------------------------------------------------------------------------------

38------------------------------------------------------------------------------

42data ChildEvent a =

40data ChildEvent a =

43− ChildDone ThreadId (Either SomeException [a]) -- A child is finished

41+ ChildDone ThreadId (Maybe SomeException) -- A child is finished

44− | PassOnResult (Either SomeException [a]) -- Pass on the result of a child

42+ | ChildResult (Either SomeException a) -- Pass on the result of a child

46-- | Describes the context of a computation.

44-- | Describes the context of a computation.

47data Context = Context

45data Context = Context

@@ -63,7 +61,12 @@ data Context = Context

63 , mfData :: M.Map TypeRep Any -- untyped, type coerced

61 , mfData :: M.Map TypeRep Any -- untyped, type coerced

64 -- ^ State data accessed with get or put operations

62 -- ^ State data accessed with get or put operations

64+ -- XXX use Either parentChannel accumResults

65+ , accumResults :: Maybe (IORef [Any])

66+ -- ^ Accumulated results, only the top level thread context accumulates.

67+ -- Child threads just pass on results via the parent channels.

66 -- XXX we can pass this at the time of fork rather than keeping it here.

68 -- XXX we can pass this at the time of fork rather than keeping it here.

69+

67 , parentChannel :: Maybe (TChan (ChildEvent Any))

70 , parentChannel :: Maybe (TChan (ChildEvent Any))

68 -- ^ Our parent thread's channel to communicate to when we die

71 -- ^ Our parent thread's channel to communicate to when we die

@@ -104,8 +107,6 @@ data Context = Context

104 --

107 --

105 , threadCredit :: IORef Int

108 , threadCredit :: IORef Int

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

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

107− , accumResults :: [Any]

108− -- ^ Accumulated results when running synchronously

109 } deriving Typeable

110 } deriving Typeable

110

111

111initContext

112initContext

@@ -114,8 +115,9 @@ initContext

114 -> IORef [ThreadId]

115 -> IORef [ThreadId]

115 -> IORef Int

116 -> IORef Int

116 -> (b -> m a)

117 -> (b -> m a)

118+ -> IORef [a]

117 -> Context

119 -> Context

118−initContext x childChan pending credit finalizer =

120+initContext x childChan pending credit finalizer results =

119 Context { mailBox = Nothing

121 Context { mailBox = Nothing

120 , currentm = unsafeCoerce x

122 , currentm = unsafeCoerce x

121 , fstack = [unsafeCoerce finalizer]

123 , fstack = [unsafeCoerce finalizer]

@@ -125,7 +127,7 @@ initContext x childChan pending credit finalizer =

125 , childChannel = unsafeCoerce childChan

127 , childChannel = unsafeCoerce childChan

126 , pendingThreads = pending

128 , pendingThreads = pending

127 , threadCredit = credit

129 , threadCredit = credit

128− , accumResults = [] }

130+ , accumResults = Just (unsafeCoerce results) }

129

131

130------------------------------------------------------------------------------

132------------------------------------------------------------------------------

131-- Where is the computation running?

133-- Where is the computation running?

@@ -199,14 +201,6 @@ resumeContext Context { currentm = m

199 composefStack [] _ = error "Bug: this should never be reached"

201 composefStack [] _ = error "Bug: this should never be reached"

200 composefStack (f:ff) x = f x >>= composefStack ff

202 composefStack (f:ff) x = f x >>= composefStack ff

201

203

202−contextSaveResult :: Monad m => a -> StateT Context m ()

203−contextSaveResult r =

204− modify $ \ Context {accumResults = rs, ..} ->

205− Context {accumResults = unsafeCoerce r : rs, ..}

206−

207−contextGetResult :: Context -> [a]

208−contextGetResult ctx = unsafeCoerce $ accumResults ctx

209−

210setContextMailBox :: Context -> a -> Context

204setContextMailBox :: Context -> a -> Context

211setContextMailBox ctx mbdata = ctx { mailBox = Just $ unsafeCoerce mbdata }

205setContextMailBox ctx mbdata = ctx { mailBox = Just $ unsafeCoerce mbdata }

212

206

src/Strands/Threads.hs+17−99

@@ -69,74 +69,20 @@ import Strands.Context

69-- | Continue execution of the closure that we were executing when we migrated

70-- to a new thread.

72−resume :: MonadIO m => Context -> StateT Context m ()

72+runContext :: MonadIO m => Context -> StateT Context m ()

73−resume ctx = do

73+runContext ctx = do

74 -- XXX rename to buildContext or buildState?

75 let s = runAsyncT (resumeContext ctx)

76− -- The returned value is always 'Nothing', we just discard it

77− (_, c) <- lift $ runStateT s ctx

76+ _ <- lift $ runStateT s ctx

78−

77+ return ()

79− -- XXX can we pass the result directly to the root thread instead of

80− -- passing through all the parents? We can let the parent go away and

81− -- handle the ChildDone events as well in the root thread.

82− case parentChannel c of

83− Nothing -> modify $ \x ->

84− x { accumResults = accumResults c ++ accumResults x }

85− Just chan -> do

86− let r = accumResults c

87− when (length r /= 0) $

88− -- there is only one result in case of a non-root thread

89− -- XXX change the return type to 'a' instead of '[a]'

90− liftIO $ atomically $ writeTChan chan (PassOnResult (Right r))

92------------------------------------------------------------------------------

79------------------------------------------------------------------------------

93-- Thread Management (creation, reaping and killing)

80-- Thread Management (creation, reaping and killing)

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

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

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

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

98−-- on consumption rate.

99−processOneEvent :: MonadIO m

100− => ChildEvent a

101− -> TChan (ChildEvent a)

102− -> [ThreadId]

103− -> Maybe SomeException

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

105−processOneEvent ev pchan pending exc = do

106− e <- case exc of

107− Nothing ->

108− case ev of

109− ChildDone tid res -> do

110− dbg $ "processOneEvent ChildDone: " ++ show tid

111− handlePass res

112− PassOnResult res -> do

113− dbg $ "processOneEvent PassOnResult"

114− handlePass res

115− Just _ -> return exc

116− let p = case ev of

117− ChildDone tid _ -> delete tid pending

118− _ -> pending

119− return (p, e)

120−

121− where

122−

123− handlePass :: MonadIO m

124− => Either SomeException [a] -> m (Maybe SomeException)

125− handlePass res =

126− case res of

127− Left e -> do

128− dbg $ "handlePass: caught exception"

129− liftIO $ mapM_ killThread pending

130− return (Just e)

131− Right [] -> return Nothing

132− Right _ -> do

133− liftIO $ atomically $ writeTChan pchan

134− (PassOnResult (unsafeCoerce res))

135− return Nothing

136−

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

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

138tryReclaimZombies ctx = do

84tryReclaimZombies ctx = do

139− let pchan = fromJust (parentChannel ctx)

85+ let dest = getCtxResultDest ctx

140 cchan = childChannel ctx

86 cchan = childChannel ctx

141 pendingRef = pendingThreads ctx

87 pendingRef = pendingThreads ctx

142

@@ -148,7 +94,7 @@ tryReclaimZombies ctx = do

148 case mev of

94 case mev of

149 Nothing -> return ()

95 Nothing -> return ()

150 Just ev -> do

96 Just ev -> do

151− (p, e) <- processOneEvent ev pchan pending Nothing

97+ (p, e) <- processOneEvent ev dest pending

152 liftIO $ writeIORef pendingRef p

98 liftIO $ writeIORef pendingRef p

153 maybe (return ()) throwM e

99 maybe (return ()) throwM e

154 tryReclaimZombies ctx

100 tryReclaimZombies ctx

@@ -157,40 +103,16 @@ waitForOneEvent :: (MonadIO m, MonadThrow m) => Context -> m ()

157waitForOneEvent ctx = do

103waitForOneEvent ctx = do

158 -- XXX assert pending must have at least one element

104 -- XXX assert pending must have at least one element

159 -- assert that the tid is found in our list

105 -- assert that the tid is found in our list

160− let pchan = fromJust (parentChannel ctx)

106+ let dest = getCtxResultDest ctx

161 cchan = childChannel ctx

107 cchan = childChannel ctx

162 pendingRef = pendingThreads ctx

108 pendingRef = pendingThreads ctx

163

109

164 ev <- liftIO $ atomically $ readTChan cchan

110 ev <- liftIO $ atomically $ readTChan cchan

165 pending <- liftIO $ readIORef pendingRef

111 pending <- liftIO $ readIORef pendingRef

166− (p, e) <- processOneEvent ev pchan pending Nothing

112+ (p, e) <- processOneEvent ev dest pending

167 liftIO $ writeIORef pendingRef p

113 liftIO $ writeIORef pendingRef p

168 maybe (return ()) throwM e

114 maybe (return ()) throwM e

169

115

170−drainChildren :: MonadIO m

171− => TChan (ChildEvent a)

172− -> TChan (ChildEvent a)

173− -> [ThreadId]

174− -> Maybe SomeException

175− -> m (Maybe SomeException)

176−drainChildren pchan cchan pending exc =

177− if pending == []

178− then return exc

179− else do

180− ev <- liftIO $ atomically $ readTChan cchan

181− (p, e) <- processOneEvent ev pchan pending exc

182− drainChildren pchan cchan p e

183−

184−waitForChildren :: MonadIO m

185− => Context -> Maybe SomeException -> m (Maybe SomeException)

186−waitForChildren ctx exc = do

187− let pendingRef = pendingThreads ctx

188− pchan = fromJust (parentChannel ctx)

189− pending <- liftIO $ readIORef pendingRef

190− e <- drainChildren pchan (childChannel ctx) pending exc

191− liftIO $ writeIORef pendingRef []

192− return e

193−

194-- | kill all the child threads associated with the continuation context

116-- | kill all the child threads associated with the continuation context

195killChildren :: Context -> IO ()

117killChildren :: Context -> IO ()

196killChildren ctx = do

118killChildren ctx = do

@@ -233,7 +155,8 @@ forkFinally1 :: (MonadIO m, MonadBaseControl IO m)

233forkFinally1 ctx preExit =

155forkFinally1 ctx preExit =

234 EL.mask $ \restore ->

156 EL.mask $ \restore ->

235 liftBaseWith $ \runInIO -> forkIO $ do

157 liftBaseWith $ \runInIO -> forkIO $ do

236− _ <- runInIO $ EL.try (restore (resume ctx)) >>= liftIO . preExit

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

159+ >>= liftIO . preExit

237 return ()

160 return ()

238

161

239-- | Run a given context in a new thread.

162-- | Run a given context in a new thread.

@@ -263,27 +186,22 @@ forkContext context = do

263 { parentChannel = Just (childChannel ctx)

186 { parentChannel = Just (childChannel ctx)

264 , pendingThreads = pendingRef

187 , pendingThreads = pendingRef

265 , childChannel = chan

188 , childChannel = chan

266− , accumResults = []

189+ , accumResults = Nothing

267 }

190 }

268

191

269 beforeExit ctx res = do

192 beforeExit ctx res = do

270 tid <- myThreadId

193 tid <- myThreadId

271− exc <- case res of

194+ r <- case res of

272− Left e -> do

195+ Left e -> do

273 dbg $ "beforeExit: " ++ show tid ++ " caught exception"

196 dbg $ "beforeExit: " ++ show tid ++ " caught exception"

274 liftIO $ killChildren ctx

197 liftIO $ killChildren ctx

275 return (Just e)

198 return (Just e)

276− Right _ -> return Nothing

277−

278− e <- waitForChildren ctx exc

199+ Right _ -> waitForChildren ctx

279

200

280− -- We are guaranteed to have a parent because we have been explicitly

201+ -- We are guaranteed to have a parent because we are forked.

281− -- forked by some parent.

282 let p = fromJust (parentChannel ctx)

202 let p = fromJust (parentChannel ctx)

283 signalQSemB (threadCredit ctx)

203 signalQSemB (threadCredit ctx)

284− -- XXX change the return value type to Maybe SomeException

285− liftIO $ atomically $ writeTChan p

204+ liftIO $ atomically $ writeTChan p (ChildDone tid r)

286− (ChildDone tid (maybe (Right []) Left e))

287

205

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

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

289--

207--

@@ -318,7 +236,7 @@ resumeContextWith context synch action = do

318 let ctx = setContextMailBox context (action context)

236 let ctx = setContextMailBox context (action context)

319 can <- liftIO $ canFork context

237 can <- liftIO $ canFork context

320 case can && (not synch) of

238 case can && (not synch) of

321− False -> resume ctx -- run synchronously

239+ False -> runContext ctx -- run synchronously

322 True -> forkContext ctx

240 True -> forkContext ctx

323

241

324-- | 'StreamData' represents a task in a task stream being generated.

242-- | 'StreamData' represents a task in a task stream being generated.

test/2-waitAsync.hs+5−2

@@ -1,5 +1,8 @@

1import Strands

2import Control.Monad.IO.Class (liftIO)

4+main = do

4−main = waitAsync $ do

5+ xs <- waitAsync $ do

5− liftIO $ putStrLn "hello"

6+ liftIO $ putStrLn "hello"

7+ return 5

8+ print xs