2c0a6fb · streamly · git.iko.soy

git.iko.soy/streamly

cancel parallel tasks when take completes

commit2c0a6fbd315dd7d824ea86a867e399f6edf4deff authorHarendra Kumar <harendra.kumar@gmail.com> dateTue 14 Aug 2018, 23:19 (8 years ago) parent9617a71

4 files changed+107−15

compare with master

unifiedsplit

src/Streamly/SVar.hs+27−6 src/Streamly/Streams/Parallel.hs+25−8 src/Streamly/Streams/SVar.hs+6−1 test/Main.hs+49

src/Streamly/SVar.hs+27−6

@@ -52,6 +52,7 @@ module Streamly.SVar

5252 , atomicModifyIORefCAS

5353 , WorkerInfo (..)

5454 , YieldRateInfo (..)

55+ , ThreadAbort (..)

5556 , ChildEvent (..)

5657 , AheadHeapEntry (..)

5758 , send

@@ -74,6 +75,7 @@ module Streamly.SVar

7475 , workerRateControl

7576 , updateYieldCount

7677 , decrementYieldLimit

78+ , decrementYieldLimitPost

7779 , incrementYieldLimit

7880 , postProcessBounded

7981 , postProcessPaced

@@ -96,7 +98,7 @@ import Control.Concurrent

9698 (ThreadId, myThreadId, threadDelay, getNumCapabilities)

9799import Control.Concurrent.MVar

98100 (MVar, newEmptyMVar, tryPutMVar, takeMVar, newMVar)

99−import Control.Exception (SomeException(..), catch, mask, assert)

101+import Control.Exception (SomeException(..), catch, mask, assert, Exception)

100102import Control.Monad (when)

101103import Control.Monad.Catch (MonadThrow)

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

@@ -163,6 +165,10 @@ newtype Count = Count Word64

163165-- Parent child thread communication type

164166------------------------------------------------------------------------------

165167

168+data ThreadAbort = ThreadAbort deriving Show

169+

170+instance Exception ThreadAbort

171+

166172-- | Events that a child thread may send to a parent thread.

167173data ChildEvent a =

168174 ChildYield a

@@ -655,6 +661,18 @@ decrementYieldLimit sv =

655661 (if x >= 1 then x - 1 else 0, x)

656662 return $ r >= 1

657663

664+-- decrementYieldLimit returns False when the old limit is 0. This one returns

665+-- False when the old limit is 1.

666+{-# INLINE decrementYieldLimitPost #-}

667+decrementYieldLimitPost :: SVar t m a -> IO Bool

668+decrementYieldLimitPost sv =

669+ case remainingYields sv of

670+ Nothing -> return True

671+ Just ref -> do

672+ r <- atomicModifyIORefCAS ref $ \x ->

673+ (if x >= 1 then x - 1 else 0, x)

674+ return $ r > 1

675+

658676-- XXX increment is not safe when the yield limit was already 0

659677-- and therefore was not decremented. We need yield limit to be

660678-- int and negative if we want this to work.

@@ -1677,12 +1695,15 @@ getAheadSVar st f = do

16771695 (xs, _) <- readIORef q

16781696 return $ null xs

16791697

1680−getParallelSVar :: MonadIO m => IO (SVar t m a)

1681−getParallelSVar = do

1698+getParallelSVar :: MonadIO m => State t m a -> IO (SVar t m a)

1699+getParallelSVar st = do

16821700 outQ <- newIORef ([], 0)

16831701 outQMv <- newEmptyMVar

16841702 active <- newIORef 0

16851703 running <- newIORef S.empty

1704+ yl <- case getYieldLimit st of

1705+ Nothing -> return Nothing

1706+ Just x -> Just <$> newIORef x

16861707

16871708 disp <- newIORef 0

16881709 maxWrk <- newIORef 0

@@ -1696,7 +1717,7 @@ getParallelSVar = do

16961717

16971718 let sv =

16981719 SVar { outputQueue = outQ

1699− , remainingYields = Nothing

1720+ , remainingYields = yl

17001721 , maxBufferLimit = Unlimited

17011722 , maxWorkerLimit = undefined

17021723 , yieldRateInfo = Nothing

@@ -1764,8 +1785,8 @@ newAheadVar st m wloop = do

17641785 sendFirstWorker sv m

17651786

17661787{-# INLINABLE newParallelVar #-}

1767−newParallelVar :: MonadAsync m => m (SVar t m a)

1768−newParallelVar = liftIO $ getParallelSVar

1788+newParallelVar :: MonadAsync m => State t m a -> m (SVar t m a)

1789+newParallelVar st = liftIO $ getParallelSVar st

17691790

17701791-- XXX this errors out for Parallel/Ahead SVars

17711792-- | Write a stream to an 'SVar' in a non-blocking manner. The stream can then

src/Streamly/Streams/Parallel.hs+25−8

@@ -34,6 +34,7 @@ module Streamly.Streams.Parallel

3434 )

3535where

3636

37+import Control.Concurrent (throwTo, myThreadId)

3738import Control.Monad (ap)

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

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

@@ -43,12 +44,15 @@ import Control.Monad.Reader.Class (MonadReader(..))

4344import Control.Monad.State.Class (MonadState(..))

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

4546import Data.Functor (void)

46−import Data.IORef (newIORef)

47+import Data.IORef (newIORef, readIORef)

48+import Data.List ((\\))

4749import Data.Maybe (fromJust)

4850import Data.Semigroup (Semigroup(..))

4951import Prelude hiding (map)

5052import System.Clock

5153

54+import qualified Data.Set as S

55+

5256import Streamly.Streams.SVar (fromSVar)

5357import Streamly.Streams.Serial (map)

5458import Streamly.SVar

@@ -71,20 +75,33 @@ runOne st m = unStream m st stop single yieldk

7175 yc <- newIORef 0

7276 wls <- newIORef (0, fromNanoSecs 0)

7377 return $ WorkerInfo 0 yc wls

78+

7479 sv = fromJust $ streamVar st

80+

81+ withLimitCheck action = do

82+ yieldLimitOk <- liftIO $ decrementYieldLimitPost sv

83+ if yieldLimitOk

84+ then action

85+ else liftIO $ do

86+ workers <- readIORef (workerThreads sv)

87+ self <- myThreadId

88+ mapM_ (\tid -> throwTo tid ThreadAbort)

89+ (S.toList workers \\ [self])

90+

7591 stop = liftIO $ winfo >>= sendStop sv

7692 sendit a = liftIO $ winfo >>= \x -> sendYield sv x (ChildYield a)

77− single a = sendit a >> stop

93+ single a = sendit a >> withLimitCheck stop

94+

7895 -- XXX there is no flow control in parallel case. We should perhaps use a

7996 -- queue and queue it back on that and exit the thread when the outputQueue

8097 -- overflows. Parallel is dangerous because it can accumulate unbounded

8198 -- output in the buffer.

82− yieldk a r = void (sendit a) >> runOne st r

99+ yieldk a r = void (sendit a) >> withLimitCheck (runOne st r)

83100

84101{-# NOINLINE forkSVarPar #-}

85102forkSVarPar :: MonadAsync m => Stream m a -> Stream m a -> Stream m a

86103forkSVarPar m r = Stream $ \st stp sng yld -> do

87− sv <- newParallelVar

104+ sv <- newParallelVar st

88105 pushWorkerPar sv (runOne st{streamVar = Just sv} m)

89106 pushWorkerPar sv (runOne st{streamVar = Just sv} r)

90107 (unStream (fromSVar sv)) (rstState st) stp sng yld

@@ -97,7 +114,7 @@ joinStreamVarPar style m1 m2 = Stream $ \st stp sng yld ->

97114 Just sv | svarStyle sv == style -> do

98115 pushWorkerPar sv (runOne st m1)

99116 unStream m2 st stp sng yld

100− _ -> unStream (forkSVarPar m1 m2) (rstState st) stp sng yld

117+ _ -> unStream (forkSVarPar m1 m2) st stp sng yld

101118

102119{-# INLINE parallelStream #-}

103120parallelStream :: MonadAsync m => Stream m a -> Stream m a -> Stream m a

@@ -123,7 +140,7 @@ parallel m1 m2 = fromStream $ parallelStream (toStream m1) (toStream m2)

123140

124141mkParallel :: (IsStream t, MonadAsync m) => t m a -> m (t m a)

125142mkParallel m = do

126− sv <- newParallelVar

143+ sv <- newParallelVar defState

127144 pushWorkerPar sv (runOne defState{streamVar = Just sv} $ toStream m)

128145 return $ fromSVar sv

129146

@@ -134,7 +151,7 @@ mkParallel m = do

134151{-# INLINE applyWith #-}

135152applyWith :: (IsStream t, MonadAsync m) => (t m a -> t m b) -> t m a -> t m b

136153applyWith f m = fromStream $ Stream $ \st stp sng yld -> do

137− sv <- newParallelVar

154+ sv <- newParallelVar (rstState st)

138155 pushWorkerPar sv (runOne st{streamVar = Just sv} (toStream m))

139156 unStream (toStream $ f $ fromSVar sv) (rstState st) stp sng yld

140157

@@ -145,7 +162,7 @@ applyWith f m = fromStream $ Stream $ \st stp sng yld -> do

145162{-# INLINE runWith #-}

146163runWith :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

147164runWith f m = do

148− sv <- newParallelVar

165+ sv <- newParallelVar defState

149166 pushWorkerPar sv (runOne defState{streamVar = Just sv} $ toStream m)

150167 f $ fromSVar sv

151168

src/Streamly/Streams/SVar.hs+6−1

@@ -36,6 +36,7 @@ module Streamly.Streams.SVar

3636 )

3737where

3838

39+import Control.Exception (fromException)

3940import Control.Monad.Catch (throwM)

4041import Data.Int (Int64)

4142#ifdef DIAGNOSTICS

@@ -100,7 +101,11 @@ fromStreamVar sv = Stream $ \st stp sng yld -> do

100101 accountThread sv tid

101102 case e of

102103 Nothing -> unStream rest (rstState st) stp sng yld

103− Just ex -> throwM ex

104+ Just ex ->

105+ case fromException ex of

106+ Just ThreadAbort ->

107+ unStream rest (rstState st) stp sng yld

108+ Nothing -> throwM ex

104109

105110{-# INLINE fromSVar #-}

106111fromSVar :: (MonadAsync m, IsStream t) => SVar Stream m a -> t m a

test/Main.hs+49

@@ -39,6 +39,8 @@ main :: IO ()

3939main = hspec $ do

4040 parallelTests

4141

42+ -- These are not run parallely because the timing gets affected

43+ -- unpredictably when other tests are running on the same machine.

4244 describe "Nested parallel and serial compositions" $ do

4345 let t = timed

4446 p = wAsyncly

@@ -80,6 +82,24 @@ main = hspec $ do

8082 <> ((t 4 <> t 8) <> (t 0 <> t 2)))

8183 `shouldReturn` ([0,0,2,2,4,4,8,8])

8284

85+ describe "restricts concurrency and cleans up extra tasks" $ do

86+ it "take 1 asyncly" $ checkCleanup asyncly (S.take 1)

87+ it "take 1 wAsyncly" $ checkCleanup wAsyncly (S.take 1)

88+ it "take 1 aheadly" $ checkCleanup aheadly (S.take 1)

89+ it "take 1 parallely" $ checkCleanup parallely (S.take 1)

90+

91+{-

92+ it "takeWhile (< 0) asyncly" $ checkCleanup asyncly (S.takeWhile (< 0))

93+ it "takeWhile (< 0) wAsyncly" $ checkCleanup wAsyncly (S.takeWhile (< 0))

94+ it "takeWhile (< 0) aheadly" $ checkCleanup aheadly (S.takeWhile (< 0))

95+ it "takeWhile (< 0) parallely" $ checkCleanup parallely (S.takeWhile (< 0))

96+

97+ it "head asyncly" $ checkCleanupFold asyncly (S.head)

98+ it "head wAsyncly" $ checkCleanupFold wAsyncly (S.head)

99+ it "head aheadly" $ checkCleanupFold aheadly (S.head)

100+ it "head parallely" $ checkCleanupFold parallely (S.head)

101+-}

102+

83103 ---------------------------------------------------------------------------

84104 -- Semigroup/Monoidal Composition strict ordering checks

85105 ---------------------------------------------------------------------------

@@ -104,6 +124,35 @@ main = hspec $ do

104124 describe "Parallel (<>) time order check" $ parallelCheck parallely (<>)

105125 describe "Parallel mappend time order check" $ parallelCheck parallely mappend

106126

127+checkCleanup :: IsStream t

128+ => (t IO Int -> SerialT IO Int)

129+ -> (t IO Int -> t IO Int)

130+ -> IO ()

131+checkCleanup t op = do

132+ r <- newIORef (-1 :: Int)

133+ runStream . serially $ do

134+ _ <- t $ op $ delay r 0 S.|: delay r 1 S.|: delay r 2 S.|: S.nil

135+ return ()

136+ threadDelay 500000

137+ res <- readIORef r

138+ res `shouldBe` 0

139+ where

140+ delay ref i = threadDelay (i*100000) >> writeIORef ref i >> return i

141+

142+{-

143+checkCleanupFold :: IsStream t

144+ => (t IO Int -> SerialT IO Int)

145+ -> (SerialT IO Int -> IO (Maybe Int))

146+ -> IO ()

147+checkCleanupFold t op = do

148+ r <- newIORef (-1 :: Int)

149+ _ <- op $ t $ delay r 0 S.|: delay r 1 S.|: delay r 2 S.|: S.nil

150+ threadDelay 500000

151+ res <- readIORef r

152+ res `shouldBe` 0

153+ where

154+ delay ref i = threadDelay (i*100000) >> writeIORef ref i >> return i

155+-}

107156

108157parallelTests :: SpecWith ()

109158parallelTests = H.parallel $ do