e41e315 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Simplify result passing, direct to the IORef

e41e315Harendra Kumar9 years ago

src/Strands/Context.hs+4−20

compare with master

unifiedsplit

@@ -11,7 +11,6 @@ module Strands.Context

11 , composeContext

12 , setContextMailBox

13 , takeContextMailBox

14− , getCtxResultDest

15 , Location(..)

14 , Location(..)

16 , getLocation

15 , getLocation

17 , setLocation

16 , setLocation

@@ -25,7 +24,6 @@ import Control.Monad.State (StateT, MonadPlus(..), MonadIO(..))

25import qualified Control.Monad.Trans.State.Lazy as Lazy (get, gets, modify, put)

24import qualified Control.Monad.Trans.State.Lazy as Lazy (get, gets, modify, put)

26import Data.Dynamic (Typeable)

25import Data.Dynamic (Typeable)

27import Data.IORef (IORef)

26import Data.IORef (IORef)

28−import Data.Maybe (fromJust)

29import Unsafe.Coerce (unsafeCoerce)

27import Unsafe.Coerce (unsafeCoerce)

30import GHC.Prim (Any)

28import GHC.Prim (Any)

@@ -33,9 +31,7 @@ import GHC.Prim (Any)

33-- State of a continuation

31-- State of a continuation

34------------------------------------------------------------------------------

32------------------------------------------------------------------------------

36−data ChildEvent a =

37− ChildDone ThreadId (Maybe SomeException) -- A child is finished

34+data ChildEvent a = ChildDone ThreadId (Maybe SomeException)

38− | ChildResult (Either SomeException a) -- Pass on the result of a child

40-- | Describes the context of a computation.

36-- | Describes the context of a computation.

41data Context = Context

37data Context = Context

@@ -84,14 +80,9 @@ data Context = Context

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

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

86 -- params.

82 -- params.

87− -- XXX use Either parentChannel accumResults

88− , accumResults :: Maybe (IORef [Any])

89− -- ^ Accumulated results, only the top level thread context accumulates.

90− -- Child threads just pass on results via the parent channels.

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

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

93− -- ^ Parent thread's channel to communicate to when a child thread dies or

85+ -- ^ Parent thread's channel to communicate to when a child thread dies

94− -- yields a result value

96 , childChannel :: TChan (ChildEvent Any)

87 , childChannel :: TChan (ChildEvent Any)

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

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

@@ -128,9 +119,8 @@ initContext

128 -> IORef [ThreadId]

119 -> IORef [ThreadId]

129 -> IORef Int

120 -> IORef Int

130 -> (b -> m a)

121 -> (b -> m a)

131− -> IORef [a]

132 -> Context

122 -> Context

133−initContext x childChan pending credit finalizer results =

123+initContext x childChan pending credit finalizer =

134 Context { mailBox = Nothing

124 Context { mailBox = Nothing

135 , currentm = unsafeCoerce x

125 , currentm = unsafeCoerce x

136 , fstack = [unsafeCoerce finalizer]

126 , fstack = [unsafeCoerce finalizer]

@@ -139,8 +129,7 @@ initContext x childChan pending credit finalizer results =

139 , parentChannel = Nothing

129 , parentChannel = Nothing

140 , childChannel = unsafeCoerce childChan

130 , childChannel = unsafeCoerce childChan

141 , pendingThreads = pending

131 , pendingThreads = pending

142− , threadCredit = credit

132+ , threadCredit = credit }

143− , accumResults = Just (unsafeCoerce results) }

144

133

145------------------------------------------------------------------------------

134------------------------------------------------------------------------------

146-- Where is the computation running?

135-- Where is the computation running?

@@ -229,8 +218,3 @@ takeContextMailBox = do

229 Just x -> do

218 Just x -> do

230 Lazy.put ctx { mailBox = Nothing }

219 Lazy.put ctx { mailBox = Nothing }

231 return $ Right (unsafeCoerce x)

220 return $ Right (unsafeCoerce x)

232−

233−getCtxResultDest :: Context -> Either (TChan (ChildEvent a)) (IORef [a])

234−getCtxResultDest ctx =

235− maybe (Right $ unsafeCoerce $ fromJust $ accumResults ctx)

236− (Left . unsafeCoerce) (parentChannel ctx)

src/Strands/Threads.hs+27−63

@@ -36,7 +36,7 @@ import Data.IORef (IORef, atomicModifyIORef,

36 modifyIORef, newIORef, readIORef,

37 writeIORef)

38import Data.List (delete)

39−import Data.Maybe (fromJust)

39+import Data.Maybe (fromJust, isJust)

40import Unsafe.Coerce (unsafeCoerce)

42import Strands.AsyncT

@@ -88,52 +88,35 @@ runContext ctx = do

88-- on consumption rate.

89processOneEvent :: MonadIO m

90 => ChildEvent a

91− -> Either (TChan (ChildEvent a)) (IORef [a])

92 -> [ThreadId]

91 -> [ThreadId]

93 -> m ([ThreadId], Maybe SomeException)

92 -> m ([ThreadId], Maybe SomeException)

94−processOneEvent ev dest pending = do

93+processOneEvent (ChildDone tid e) pending = do

95− -- Collect results unless we have already encountered an exception.

96− case ev of

97− ChildDone _ (Just e) -> handleException e

98− ChildResult (Left e) -> handleException e

99− ChildDone tid Nothing -> return (delete tid pending, Nothing)

100− ChildResult (Right r) -> do

101− case dest of

102− Left chan -> liftIO $ atomically $ writeTChan chan ev

103− Right ref ->

104− liftIO $ modifyIORef ref $ \rs -> unsafeCoerce r : rs

105− return (pending, Nothing)

106− where

107−

108− handleException e = do

109− liftIO $ mapM_ killThread pending

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

110− return (pending, Just e)

95+ return (delete tid pending, Nothing)

111

112drainChildren :: MonadIO m

97drainChildren :: MonadIO m

113− => Either (TChan (ChildEvent a)) (IORef [a])

114− -> TChan (ChildEvent a)

98+ => TChan (ChildEvent a)

115 -> [ThreadId]

99 -> [ThreadId]

116 -> m ([ThreadId], Maybe SomeException)

100 -> m ([ThreadId], Maybe SomeException)

117−drainChildren dest cchan pending =

101+drainChildren cchan pending =

118 case pending of

102 case pending of

119 [] -> return (pending, Nothing)

103 [] -> return (pending, Nothing)

120 _ -> do

104 _ -> do

121 ev <- liftIO $ atomically $ readTChan cchan

105 ev <- liftIO $ atomically $ readTChan cchan

122− (p, e) <- processOneEvent ev dest pending

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

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

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

124

108

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

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

126waitForChildren ctx = do

110waitForChildren ctx = do

127 let pendingRef = pendingThreads ctx

111 let pendingRef = pendingThreads ctx

128 pending <- liftIO $ readIORef pendingRef

112 pending <- liftIO $ readIORef pendingRef

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

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

130 liftIO $ writeIORef pendingRef p

114 liftIO $ writeIORef pendingRef p

131 return e

115 return e

132

116

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

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

134tryReclaimZombies ctx = do

118tryReclaimZombies ctx = do

135− let dest = getCtxResultDest ctx

136− cchan = childChannel ctx

119+ let cchan = childChannel ctx

137 pendingRef = pendingThreads ctx

120 pendingRef = pendingThreads ctx

138

121

139 pending <- liftIO $ readIORef pendingRef

122 pending <- liftIO $ readIORef pendingRef

@@ -144,7 +127,7 @@ tryReclaimZombies ctx = do

144 case mev of

127 case mev of

145 Nothing -> return ()

128 Nothing -> return ()

146 Just ev -> do

129 Just ev -> do

147− (p, e) <- processOneEvent ev dest pending

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

148 liftIO $ writeIORef pendingRef p

131 liftIO $ writeIORef pendingRef p

149 maybe (return ()) throwM e

132 maybe (return ()) throwM e

150 tryReclaimZombies ctx

133 tryReclaimZombies ctx

@@ -153,13 +136,12 @@ waitForOneEvent :: (MonadIO m, MonadThrow m) => Context -> m ()

153waitForOneEvent ctx = do

136waitForOneEvent ctx = do

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

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

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

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

156− let dest = getCtxResultDest ctx

157− cchan = childChannel ctx

139+ let cchan = childChannel ctx

158 pendingRef = pendingThreads ctx

140 pendingRef = pendingThreads ctx

159

141

160 ev <- liftIO $ atomically $ readTChan cchan

142 ev <- liftIO $ atomically $ readTChan cchan

161 pending <- liftIO $ readIORef pendingRef

143 pending <- liftIO $ readIORef pendingRef

162− (p, e) <- processOneEvent ev dest pending

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

163 liftIO $ writeIORef pendingRef p

145 liftIO $ writeIORef pendingRef p

164 maybe (return ()) throwM e

146 maybe (return ()) throwM e

165

147

@@ -226,7 +208,6 @@ forkContext context = do

226 { parentChannel = Just (childChannel ctx)

208 { parentChannel = Just (childChannel ctx)

227 , pendingThreads = pendingRef

209 , pendingThreads = pendingRef

228 , childChannel = chan

210 , childChannel = chan

229− , accumResults = Nothing

230 }

211 }

231

212

232 beforeExit ctx res = do

213 beforeExit ctx res = do

@@ -361,7 +342,6 @@ loopContextWith ioaction context = AsyncT $ do

361 { parentChannel = parentChannel curCtx

342 { parentChannel = parentChannel curCtx

362 , pendingThreads = pendingThreads curCtx

343 , pendingThreads = pendingThreads curCtx

363 , childChannel = childChannel curCtx

344 , childChannel = childChannel curCtx

364− , accumResults = accumResults curCtx

365 }

345 }

366

346

367 where

347 where

@@ -518,36 +498,17 @@ noTrans x = AsyncT $ x >>= return . Just

518-- Running the monad

498-- Running the monad

519------------------------------------------------------------------------------

499------------------------------------------------------------------------------

520

500

521−collectResult :: MonadIO m => a -> StateT Context m ()

522−collectResult r = do

523− ctx <- get

524− case parentChannel ctx of

525− Nothing -> do

526− let ref = fromJust $ accumResults ctx

527− liftIO $ modifyIORef ref $ \rs -> unsafeCoerce r : rs

528− Just chan -> do

529− -- XXX can we pass the result directly to the root thread

530− -- instead of passing through all the parents? We can let the

531− -- parent go away and handle the ChildDone events as well in

532− -- the root thread.

533− liftIO $ atomically $ writeTChan chan

534− (ChildResult (Right (unsafeCoerce r)))

535−

536−-- XXX pass a collector function and return a Traversable?

537−-- XXX Ideally it should be a non-empty list.

538−

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

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

540-- thread of the computation in a list.

502-- thread of the computation in a list.

541waitAsync :: forall m a b. (MonadIO m, MonadCatch m)

503waitAsync :: forall m a b. (MonadIO m, MonadCatch m)

542− => (a -> AsyncT m a) -> AsyncT m a -> m [a]

504+ => (a -> AsyncT m a) -> AsyncT m a -> m ()

543waitAsync finalizer m = do

505waitAsync finalizer m = do

544 childChan <- liftIO $ atomically newTChan

506 childChan <- liftIO $ atomically newTChan

545 pendingRef <- liftIO $ newIORef []

507 pendingRef <- liftIO $ newIORef []

546− resultsRef <- liftIO $ newIORef []

547 credit <- liftIO $ newIORef maxBound

508 credit <- liftIO $ newIORef maxBound

548

509

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

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

550− finalizer resultsRef

511+ finalizer

551

512

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

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

553

514

@@ -559,25 +520,28 @@ waitAsync finalizer m = do

559 e <- waitForChildren ctx

520 e <- waitForChildren ctx

560 case e of

521 case e of

561 Just (exc :: SomeException) -> throwM exc

522 Just (exc :: SomeException) -> throwM exc

562− Nothing -> liftIO $ readIORef resultsRef

523+ Nothing -> return ()

524+

525+-- TBD throttling of producer based on conumption rate.

563

526

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

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

565-- the computation when the computation is done

528-- the computation when the computation is done

566−gatherResult :: MonadIO m => a -> AsyncT m a

529+gatherResult :: MonadIO m => IORef [a] -> a -> AsyncT m a

567−gatherResult x = AsyncT $ do

530+gatherResult ref r = do

568− collectResult x

531+ liftIO $ atomicModifyIORef ref $ \rs -> (r : rs, rs)

569− return Nothing

532+ mzero

570

533

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

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

572-- thread of the computation in a list.

535-- thread of the computation in a list.

573gather :: forall m a. (MonadIO m, MonadCatch m)

536gather :: forall m a. (MonadIO m, MonadCatch m)

574 => AsyncT m a -> m [a]

537 => AsyncT m a -> m [a]

538+gather m = do

539+ resultsRef <- liftIO $ newIORef []

575−gather m = waitAsync gatherResult m

540+ waitAsync (gatherResult resultsRef) m

541+ liftIO $ readIORef resultsRef

576

542

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

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

578-- results.

544-- results.

579wait :: forall m a. (MonadIO m, MonadCatch m)

545wait :: forall m a. (MonadIO m, MonadCatch m)

580 => AsyncT m a -> m ()

546 => AsyncT m a -> m ()

581−wait m = do

582− _ <- waitAsync (const mzero) m

547+wait m = waitAsync (const mzero) m

583− return ()