2c0a6fb · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

cancel parallel tasks when take completes

2c0a6fbHarendra Kumar8 years ago

src/Streamly/SVar.hs+27−6

compare with master

unifiedsplit

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

52 , atomicModifyIORefCAS

53 , WorkerInfo (..)

54 , YieldRateInfo (..)

55+ , ThreadAbort (..)

55 , ChildEvent (..)

56 , ChildEvent (..)

56 , AheadHeapEntry (..)

57 , AheadHeapEntry (..)

57 , send

58 , send

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

74 , workerRateControl

75 , workerRateControl

75 , updateYieldCount

76 , updateYieldCount

76 , decrementYieldLimit

77 , decrementYieldLimit

78+ , decrementYieldLimitPost

77 , incrementYieldLimit

79 , incrementYieldLimit

78 , postProcessBounded

80 , postProcessBounded

79 , postProcessPaced

81 , postProcessPaced

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

96 (ThreadId, myThreadId, threadDelay, getNumCapabilities)

98 (ThreadId, myThreadId, threadDelay, getNumCapabilities)

97import Control.Concurrent.MVar

99import Control.Concurrent.MVar

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

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

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

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

100import Control.Monad (when)

102import Control.Monad (when)

101import Control.Monad.Catch (MonadThrow)

103import Control.Monad.Catch (MonadThrow)

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

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

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

163-- Parent child thread communication type

165-- Parent child thread communication type

164------------------------------------------------------------------------------

166------------------------------------------------------------------------------

165

167

168+data ThreadAbort = ThreadAbort deriving Show

169+

170+instance Exception ThreadAbort

171+

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

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

167data ChildEvent a =

173data ChildEvent a =

168 ChildYield a

174 ChildYield a

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

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

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

656 return $ r >= 1

662 return $ r >= 1

657

663

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+

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

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

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

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

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

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

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

1677 (xs, _) <- readIORef q

1695 (xs, _) <- readIORef q

1678 return $ null xs

1696 return $ null xs

1679

1697

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

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

1681−getParallelSVar = do

1699+getParallelSVar st = do

1682 outQ <- newIORef ([], 0)

1700 outQ <- newIORef ([], 0)

1683 outQMv <- newEmptyMVar

1701 outQMv <- newEmptyMVar

1684 active <- newIORef 0

1702 active <- newIORef 0

1685 running <- newIORef S.empty

1703 running <- newIORef S.empty

1704+ yl <- case getYieldLimit st of

1705+ Nothing -> return Nothing

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

1686

1707

1687 disp <- newIORef 0

1708 disp <- newIORef 0

1688 maxWrk <- newIORef 0

1709 maxWrk <- newIORef 0

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

1696

1717

1697 let sv =

1718 let sv =

1698 SVar { outputQueue = outQ

1719 SVar { outputQueue = outQ

1699− , remainingYields = Nothing

1720+ , remainingYields = yl

1700 , maxBufferLimit = Unlimited

1721 , maxBufferLimit = Unlimited

1701 , maxWorkerLimit = undefined

1722 , maxWorkerLimit = undefined

1702 , yieldRateInfo = Nothing

1723 , yieldRateInfo = Nothing

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

1764 sendFirstWorker sv m

1785 sendFirstWorker sv m

1765

1786

1766{-# INLINABLE newParallelVar #-}

1787{-# INLINABLE newParallelVar #-}

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

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

1768−newParallelVar = liftIO $ getParallelSVar

1789+newParallelVar st = liftIO $ getParallelSVar st

1769

1790

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

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

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

1792-- | 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

34 )

35where

37+import Control.Concurrent (throwTo, myThreadId)

37import Control.Monad (ap)

38import Control.Monad (ap)

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

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

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

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

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

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

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

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

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

45import Data.Functor (void)

46import Data.Functor (void)

46−import Data.IORef (newIORef)

47+import Data.IORef (newIORef, readIORef)

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

47import Data.Maybe (fromJust)

49import Data.Maybe (fromJust)

48import Data.Semigroup (Semigroup(..))

50import Data.Semigroup (Semigroup(..))

49import Prelude hiding (map)

51import Prelude hiding (map)

50import System.Clock

52import System.Clock

54+import qualified Data.Set as S

55+

52import Streamly.Streams.SVar (fromSVar)

56import Streamly.Streams.SVar (fromSVar)

53import Streamly.Streams.Serial (map)

57import Streamly.Streams.Serial (map)

54import Streamly.SVar

58import Streamly.SVar

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

71 yc <- newIORef 0

75 yc <- newIORef 0

72 wls <- newIORef (0, fromNanoSecs 0)

76 wls <- newIORef (0, fromNanoSecs 0)

73 return $ WorkerInfo 0 yc wls

77 return $ WorkerInfo 0 yc wls

78+

74 sv = fromJust $ streamVar st

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

75 stop = liftIO $ winfo >>= sendStop sv

91 stop = liftIO $ winfo >>= sendStop sv

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

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

77− single a = sendit a >> stop

93+ single a = sendit a >> withLimitCheck stop

94+

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

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

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

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

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

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

81 -- output in the buffer.

98 -- 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)

100

84{-# NOINLINE forkSVarPar #-}

101{-# NOINLINE forkSVarPar #-}

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

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

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

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

87− sv <- newParallelVar

104+ sv <- newParallelVar st

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

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

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

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

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

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

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

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

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

98 pushWorkerPar sv (runOne st m1)

115 pushWorkerPar sv (runOne st m1)

99 unStream m2 st stp sng yld

116 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

101

118

102{-# INLINE parallelStream #-}

119{-# INLINE parallelStream #-}

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

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

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

123

140

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

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

125mkParallel m = do

142mkParallel m = do

126− sv <- newParallelVar

143+ sv <- newParallelVar defState

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

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

128 return $ fromSVar sv

145 return $ fromSVar sv

129

146

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

134{-# INLINE applyWith #-}

151{-# INLINE applyWith #-}

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

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

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

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

137− sv <- newParallelVar

154+ sv <- newParallelVar (rstState st)

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

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

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

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

140

157

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

145{-# INLINE runWith #-}

162{-# INLINE runWith #-}

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

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

147runWith f m = do

164runWith f m = do

148− sv <- newParallelVar

165+ sv <- newParallelVar defState

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

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

150 f $ fromSVar sv

167 f $ fromSVar sv

151

168

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

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

36 )

37where

39+import Control.Exception (fromException)

39import Control.Monad.Catch (throwM)

40import Control.Monad.Catch (throwM)

40import Data.Int (Int64)

41import Data.Int (Int64)

41#ifdef DIAGNOSTICS

42#ifdef DIAGNOSTICS

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

100 accountThread sv tid

101 accountThread sv tid

101 case e of

102 case e of

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

103 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

104

109

105{-# INLINE fromSVar #-}

110{-# INLINE fromSVar #-}

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

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

test/Main.hs+49

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

39main = hspec $ do

40 parallelTests

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

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

42 describe "Nested parallel and serial compositions" $ do

44 describe "Nested parallel and serial compositions" $ do

43 let t = timed

45 let t = timed

44 p = wAsyncly

46 p = wAsyncly

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

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

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

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

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

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+

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

103 ---------------------------------------------------------------------------

84 -- Semigroup/Monoidal Composition strict ordering checks

104 -- Semigroup/Monoidal Composition strict ordering checks

85 ---------------------------------------------------------------------------

105 ---------------------------------------------------------------------------

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

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

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

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

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

106

126

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

107

156

108parallelTests :: SpecWith ()

157parallelTests :: SpecWith ()

109parallelTests = H.parallel $ do

158parallelTests = H.parallel $ do