68d2bf8 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add parallel stream stop styles

68d2bf8Harendra Kumar7 years ago

src/Streamly/Fold.hs+3−6

compare with master

unifiedsplit

@@ -1123,12 +1123,9 @@ lcatMaybes = lfilter isJust . lmap fromJust

1123sessionsOf

1124 :: (IsStream t, MonadAsync m)

1125 => Double -> Fold m a b -> t m a -> t m b

1126+sessionsOf n f xs =

1126−sessionsOf n f m = splitSuffixBy' isNothing (lcatMaybes f) s

1127+ splitSuffixBy' isNothing (lcatMaybes f)

1127− where

1128− s = S.map Just m `parallel` S.repeatM timeout

1129− timeout = do

1130− liftIO $ threadDelay (round $ n * 1000000)

1131− return Nothing

1128+ (S.intersperseByTime n (return Nothing) (S.map Just xs))

1132

1129

1133------------------------------------------------------------------------------

1130------------------------------------------------------------------------------

1134-- Element Aware APIs

1131-- Element Aware APIs

src/Streamly/Prelude.hs+59−4

@@ -367,6 +367,8 @@ module Streamly.Prelude

367

368 , insertBy

369 , intersperseM

370+ -- , intersperseBySpan

371+ , intersperseByTime

370

372

371 -- ** Reordering

373 -- ** Reordering

372 , reverse

374 , reverse

@@ -470,6 +472,7 @@ module Streamly.Prelude

470 )

472 )

471where

473where

472

474

475+import Control.Concurrent (threadDelay)

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

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

474import Control.Monad.Trans (MonadTrans(..))

477import Control.Monad.Trans (MonadTrans(..))

475import Data.Maybe (isJust, fromJust)

478import Data.Maybe (isJust, fromJust)

@@ -507,6 +510,7 @@ import qualified Streamly.Streams.StreamD as S

507#endif

510#endif

508

511

509import qualified Streamly.Streams.Serial as Serial

512import qualified Streamly.Streams.Serial as Serial

513+import qualified Streamly.Streams.Parallel as Par

510

514

511------------------------------------------------------------------------------

515------------------------------------------------------------------------------

512-- Deconstruction

516-- Deconstruction

@@ -1858,8 +1862,8 @@ mapMaybeMSerial f m = fromStreamD $ D.mapMaybeM f $ toStreamD m

1858-- Transformation by Reordering

1862-- Transformation by Reordering

1859------------------------------------------------------------------------------

1863------------------------------------------------------------------------------

1860

1864

1861−-- XXX to scale this we need to use a slab allocated array backed

1865+-- XXX Use a compact region list to temporarily store the list, in both reverse

1862−-- representation for temporary storage.

1866+-- as well as in reverse'.

1863--

1867--

1864-- |

1868-- |

1865-- > reverse = S.foldlT (flip S.cons) S.nil

1869-- > reverse = S.foldlT (flip S.cons) S.nil

@@ -1886,14 +1890,16 @@ reverse' s = fromStreamD $ D.reverse' $ toStreamD s

1886-- Transformation by Inserting

1890-- Transformation by Inserting

1887------------------------------------------------------------------------------

1891------------------------------------------------------------------------------

1888

1892

1893+-- intersperseM = intersperseBySpan 1

1894+

1889-- | Generate a stream by performing a monadic action between consecutive

1895-- | Generate a stream by performing a monadic action between consecutive

1890-- elements of the given stream.

1896-- elements of the given stream.

1891--

1897--

1892-- /Concurrent (do not use with 'parallely' on infinite streams)/

1898-- /Concurrent (do not use with 'parallely' on infinite streams)/

1893--

1899--

1894-- @

1900-- @

1895−-- > S.toList $ S.intersperseM (putChar \'a' >> return ',') $ S.fromList "hello"

1901+-- > S.toList $ S.intersperseM (return ',') $ S.fromList "hello"

1896−-- aaaa"h,e,l,l,o"

1902+-- "h,e,l,l,o"

1897-- @

1903-- @

1898--

1904--

1899-- @since 0.5.0

1905-- @since 0.5.0

@@ -1901,6 +1907,37 @@ reverse' s = fromStreamD $ D.reverse' $ toStreamD s

1901intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a

1907intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a

1902intersperseM = K.intersperseM

1908intersperseM = K.intersperseM

1903

1909

1910+{-

1911+-- | Intersperse a monadic action into the input stream after every @n@

1912+-- elements.

1913+--

1914+-- @

1915+-- > S.toList $ S.intersperseBySpan 2 (return ',') $ S.fromList "hello"

1916+-- "he,ll,o"

1917+-- @

1918+--

1919+-- @since 0.7.0

1920+{-# INLINE intersperseBySpan #-}

1921+intersperseBySpan :: IsStream t => Int -> m a -> t m a -> t m a

1922+intersperseBySpan _n _f _xs = undefined

1923+-}

1924+

1925+-- | Intersperse a monadic action into the input stream after every @n@

1926+-- seconds.

1927+--

1928+-- @

1929+-- > S.drain $ S.intersperseByTime 1 (putChar ',') $ S.mapM (\\x -> threadDelay 1000000 >> putChar x) $ S.fromList "hello"

1930+-- "h,e,l,l,o"

1931+-- @

1932+--

1933+-- @since 0.7.0

1934+{-# INLINE intersperseByTime #-}

1935+intersperseByTime

1936+ :: (IsStream t, MonadAsync m)

1937+ => Double -> m a -> t m a -> t m a

1938+intersperseByTime n f xs = xs `Par.parallelEndByFirst` repeatM timed

1939+ where timed = liftIO (threadDelay (round $ n * 1000000)) >> f

1940+

1904-- | @insertBy cmp elem stream@ inserts @elem@ before the first element in

1941-- | @insertBy cmp elem stream@ inserts @elem@ before the first element in

1905-- @stream@ that is less than @elem@ when compared using @cmp@.

1942-- @stream@ that is less than @elem@ when compared using @cmp@.

1906--

1943--

@@ -2071,6 +2108,24 @@ mergeByM

2071 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a

2108 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a

2072mergeByM f m1 m2 = fromStreamS $ S.mergeByM f (toStreamS m1) (toStreamS m2)

2109mergeByM f m1 m2 = fromStreamS $ S.mergeByM f (toStreamS m1) (toStreamS m2)

2073

2110

2111+{-

2112+-- | Like 'mergeByM' but stops merging as soon as any of the two streams stops.

2113+{-# INLINABLE mergeUptoShortest #-}

2114+mergeEndByAny

2115+ :: (IsStream t, Monad m)

2116+ => (a -> a -> m Ordering) -> t m a -> t m a -> t m a

2117+mergeEndByAny f m1 m2 = fromStreamD $

2118+ D.mergeEndByAny f (toStreamD m1) (toStreamD m2)

2119+

2120+-- Like 'mergeByM' but stops merging as soon as the first stream stops.

2121+{-# INLINABLE mergeEndByFirst #-}

2122+mergeEndByFirst

2123+ :: (IsStream t, Monad m)

2124+ => (a -> a -> m Ordering) -> t m a -> t m a -> t m a

2125+mergeEndByFirst f m1 m2 = fromStreamS $

2126+ D.mergeEndByFirst f (toStreamD m1) (toStreamD m2)

2127+-}

2128+

2074-- Holding this back for now, we may want to use the name "merge" differently

2129-- Holding this back for now, we may want to use the name "merge" differently

2075{-

2130{-

2076-- | Same as @'mergeBy' 'compare'@.

2131-- | Same as @'mergeBy' 'compare'@.

src/Streamly/SVar.hs+31−5

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

25 (

26 MonadAsync

27 , SVarStyle (..)

28+ , SVarStopStyle (..)

28 , SVar (..)

29 , SVar (..)

30 -- State threaded around the stream

31 -- State threaded around the stream

@@ -314,11 +315,25 @@ data SVarStats = SVarStats {

314

315

315data Limit = Unlimited | Limited Word deriving Show

316data Limit = Unlimited | Limited Word deriving Show

316

317

318+-- When to stop the composed stream.

319+data SVarStopStyle =

320+ StopNone -- stops only when all streams are finished

321+ | StopAny -- stop when any stream finishes

322+ | StopBy -- stop when a specific stream finishes

323+ deriving (Eq, Show)

324+

325+-- IMPORTANT NOTE: we cannot update the SVar after generating it as we have

326+-- references to the original SVar stored in several functions which will keep

327+-- pointing to the original data and the new updates won't reflect there.

328+-- Any updateable parts must be kept in mutable references (IORef).

329+--

317data SVar t m a = SVar

330data SVar t m a = SVar

318 {

331 {

319 -- Read only state

332 -- Read only state

320 svarStyle :: SVarStyle

333 svarStyle :: SVarStyle

321 , svarMrun :: RunInIO m

334 , svarMrun :: RunInIO m

335+ , svarStopStyle :: SVarStopStyle

336+ , svarStopBy :: IORef ThreadId

322

337

323 -- Shared output queue (events, length)

338 -- Shared output queue (events, length)

324 -- XXX For better efficiency we can try a preallocated array type (perhaps

339 -- XXX For better efficiency we can try a preallocated array type (perhaps

@@ -2045,6 +2060,8 @@ getAheadSVar st f mrun = do

2045 , isQueueDone = isQueueDoneAhead sv q

2060 , isQueueDone = isQueueDoneAhead sv q

2046 , needDoorBell = wfw

2061 , needDoorBell = wfw

2047 , svarStyle = AheadVar

2062 , svarStyle = AheadVar

2063+ , svarStopStyle = StopNone

2064+ , svarStopBy = undefined

2048 , svarMrun = mrun

2065 , svarMrun = mrun

2049 , workerCount = active

2066 , workerCount = active

2050 , accountThread = delThread sv

2067 , accountThread = delThread sv

@@ -2091,8 +2108,9 @@ getAheadSVar st f mrun = do

2091 (xs, _) <- readIORef q

2108 (xs, _) <- readIORef q

2092 return $ null xs

2109 return $ null xs

2093

2110

2094−getParallelSVar :: MonadIO m => State t m a -> RunInIO m -> IO (SVar t m a)

2111+getParallelSVar :: MonadIO m

2112+ => SVarStopStyle -> State t m a -> RunInIO m -> IO (SVar t m a)

2095−getParallelSVar st mrun = do

2113+getParallelSVar ss st mrun = do

2096 outQ <- newIORef ([], 0)

2114 outQ <- newIORef ([], 0)

2097 outQMv <- newEmptyMVar

2115 outQMv <- newEmptyMVar

2098 active <- newIORef 0

2116 active <- newIORef 0

@@ -2105,6 +2123,11 @@ getParallelSVar st mrun = do

2105 stats <- newSVarStats

2123 stats <- newSVarStats

2106 tid <- myThreadId

2124 tid <- myThreadId

2107

2125

2126+ stopBy <-

2127+ case ss of

2128+ StopBy -> liftIO $ newIORef undefined

2129+ _ -> return undefined

2130+

2108 let sv =

2131 let sv =

2109 SVar { outputQueue = outQ

2132 SVar { outputQueue = outQ

2110 , remainingWork = yl

2133 , remainingWork = yl

@@ -2122,6 +2145,8 @@ getParallelSVar st mrun = do

2122 , isQueueDone = undefined

2145 , isQueueDone = undefined

2123 , needDoorBell = undefined

2146 , needDoorBell = undefined

2124 , svarStyle = ParallelVar

2147 , svarStyle = ParallelVar

2148+ , svarStopStyle = ss

2149+ , svarStopBy = stopBy

2125 , svarMrun = mrun

2150 , svarMrun = mrun

2126 , workerCount = active

2151 , workerCount = active

2127 , accountThread = modifyThread sv

2152 , accountThread = modifyThread sv

@@ -2176,10 +2201,11 @@ newAheadVar st m wloop = do

2176 sendFirstWorker sv m

2201 sendFirstWorker sv m

2177

2202

2178{-# INLINABLE newParallelVar #-}

2203{-# INLINABLE newParallelVar #-}

2179−newParallelVar :: MonadAsync m => State t m a -> m (SVar t m a)

2204+newParallelVar :: MonadAsync m

2205+ => SVarStopStyle -> State t m a -> m (SVar t m a)

2180−newParallelVar st = do

2206+newParallelVar ss st = do

2181 mrun <- captureMonadState

2207 mrun <- captureMonadState

2182− liftIO $ getParallelSVar st mrun

2208+ liftIO $ getParallelSVar ss st mrun

2183

2209

2184-------------------------------------------------------------------------------

2210-------------------------------------------------------------------------------

2185-- Write a stream to an SVar

2211-- Write a stream to an SVar

src/Streamly/Streams/Async.hs+4

@@ -306,6 +306,8 @@ getLifoSVar st mrun = do

306 , isQueueDone = workDone sv

307 , needDoorBell = wfw

308 , svarStyle = AsyncVar

309+ , svarStopStyle = StopNone

310+ , svarStopBy = undefined

309 , svarMrun = mrun

311 , svarMrun = mrun

310 , workerCount = active

312 , workerCount = active

311 , accountThread = delThread sv

313 , accountThread = delThread sv

@@ -396,6 +398,8 @@ getFifoSVar st mrun = do

396 , isQueueDone = workDone sv

398 , isQueueDone = workDone sv

397 , needDoorBell = wfw

399 , needDoorBell = wfw

398 , svarStyle = WAsyncVar

400 , svarStyle = WAsyncVar

401+ , svarStopStyle = StopNone

402+ , svarStopBy = undefined

399 , svarMrun = mrun

403 , svarMrun = mrun

400 , workerCount = active

404 , workerCount = active

401 , accountThread = delThread sv

405 , accountThread = delThread sv

src/Streamly/Streams/Parallel.hs+43−11

@@ -23,6 +23,8 @@ module Streamly.Streams.Parallel

23 , Parallel

24 , parallely

25 , parallel

26+ , parallelEndByFirst

27+ , parallelEndByAny

27 -- * Function application

29 -- * Function application

28 , mkParallel

30 , mkParallel

@@ -42,10 +44,13 @@ import Control.Monad.Reader.Class (MonadReader(..))

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

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

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

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

44import Data.Functor (void)

46import Data.Functor (void)

47+import Data.IORef (readIORef, writeIORef)

45import Data.Maybe (fromJust)

48import Data.Maybe (fromJust)

46import Data.Semigroup (Semigroup(..))

49import Data.Semigroup (Semigroup(..))

47import Prelude hiding (map)

50import Prelude hiding (map)

52+import qualified Data.Set as Set

53+

49import Streamly.Streams.SVar (fromSVar)

54import Streamly.Streams.SVar (fromSVar)

50import Streamly.Streams.Serial (map)

55import Streamly.Streams.Serial (map)

51import Streamly.SVar

56import Streamly.SVar

@@ -88,22 +93,28 @@ runOne st m winfo = foldStreamShared st yieldk single stop m

88 >> withLimitCheck (void $ liftIO $ mrun $ runOne st r winfo)

93 >> withLimitCheck (void $ liftIO $ mrun $ runOne st r winfo)

90{-# NOINLINE forkSVarPar #-}

95{-# NOINLINE forkSVarPar #-}

91−forkSVarPar :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a

96+forkSVarPar :: (IsStream t, MonadAsync m)

97+ => SVarStopStyle -> t m a -> t m a -> t m a

92−forkSVarPar m r = mkStream $ \st yld sng stp -> do

98+forkSVarPar ss m r = mkStream $ \st yld sng stp -> do

93− sv <- newParallelVar st

99+ sv <- newParallelVar ss st

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

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

101+ case ss of

102+ StopBy -> liftIO $ do

103+ set <- readIORef (workerThreads sv)

104+ writeIORef (svarStopBy sv) $ Set.elemAt 0 set

105+ _ -> return ()

95 pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream r)

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

96 foldStream st yld sng stp (fromSVar sv)

107 foldStream st yld sng stp (fromSVar sv)

108

98{-# INLINE joinStreamVarPar #-}

109{-# INLINE joinStreamVarPar #-}

99joinStreamVarPar :: (IsStream t, MonadAsync m)

110joinStreamVarPar :: (IsStream t, MonadAsync m)

100− => SVarStyle -> t m a -> t m a -> t m a

111+ => SVarStyle -> SVarStopStyle -> t m a -> t m a -> t m a

101−joinStreamVarPar style m1 m2 = mkStream $ \st yld sng stp ->

112+joinStreamVarPar style ss m1 m2 = mkStream $ \st yld sng stp ->

102 case streamVar st of

113 case streamVar st of

103− Just sv | svarStyle sv == style -> do

114+ Just sv | svarStyle sv == style && svarStopStyle sv == ss -> do

104 pushWorkerPar sv (runOne st $ toStream m1)

115 pushWorkerPar sv (runOne st $ toStream m1)

105 foldStreamShared st yld sng stp m2

116 foldStreamShared st yld sng stp m2

106− _ -> foldStreamShared st yld sng stp (forkSVarPar m1 m2)

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

107

118

108-- | XXX we can implement it more efficienty by directly implementing instead

119-- | XXX we can implement it more efficienty by directly implementing instead

109-- of combining streams using parallel.

120-- of combining streams using parallel.

@@ -118,7 +129,28 @@ consMParallel m r = fromStream $ K.yieldM m `parallel` (toStream r)

118-- @since 0.2.0

129-- @since 0.2.0

119{-# INLINE parallel #-}

130{-# INLINE parallel #-}

120parallel :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a

131parallel :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a

121−parallel = joinStreamVarPar ParallelVar

132+parallel = joinStreamVarPar ParallelVar StopNone

133+

134+-- This is a co-parallel like combinator for streams, where first stream is the

135+-- main stream and the rest are just supporting it, when the first ends

136+-- everything ends.

137+--

138+-- | Like `parallel` but stops the output as soon as the first stream stops.

139+--

140+-- @since 0.7.0

141+{-# INLINE parallelEndByFirst #-}

142+parallelEndByFirst :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a

143+parallelEndByFirst = joinStreamVarPar ParallelVar StopBy

144+

145+-- This is a race like combinator for streams.

146+--

147+-- | Like `parallel` but stops the output as soon as any of the two streams

148+-- stops.

149+--

150+-- @since 0.7.0

151+{-# INLINE parallelEndByAny #-}

152+parallelEndByAny :: (IsStream t, MonadAsync m) => t m a -> t m a -> t m a

153+parallelEndByAny = joinStreamVarPar ParallelVar StopAny

122

154

123------------------------------------------------------------------------------

155------------------------------------------------------------------------------

124-- Convert a stream to parallel

156-- Convert a stream to parallel

@@ -126,7 +158,7 @@ parallel = joinStreamVarPar ParallelVar

126

158

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

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

128mkParallel m = do

160mkParallel m = do

129− sv <- newParallelVar defState

161+ sv <- newParallelVar StopNone defState

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

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

131 return $ fromSVar sv

163 return $ fromSVar sv

132

164

@@ -137,7 +169,7 @@ mkParallel m = do

137{-# INLINE applyParallel #-}

169{-# INLINE applyParallel #-}

138applyParallel :: (IsStream t, MonadAsync m) => (t m a -> t m b) -> t m a -> t m b

170applyParallel :: (IsStream t, MonadAsync m) => (t m a -> t m b) -> t m a -> t m b

139applyParallel f m = mkStream $ \st yld sng stp -> do

171applyParallel f m = mkStream $ \st yld sng stp -> do

140− sv <- newParallelVar (adaptState st)

172+ sv <- newParallelVar StopNone (adaptState st)

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

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

142 foldStream st yld sng stp $ f $ fromSVar sv

174 foldStream st yld sng stp $ f $ fromSVar sv

143

175

@@ -148,7 +180,7 @@ applyParallel f m = mkStream $ \st yld sng stp -> do

148{-# INLINE foldParallel #-}

180{-# INLINE foldParallel #-}

149foldParallel :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

181foldParallel :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

150foldParallel f m = do

182foldParallel f m = do

151− sv <- newParallelVar defState

183+ sv <- newParallelVar StopNone defState

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

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

153 f $ fromSVar sv

185 f $ fromSVar sv

154

186

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

@@ -68,12 +68,23 @@ fromStreamVar sv = mkStream $ \st yld sng stp -> do

68 ChildStop tid e -> do

69 accountThread sv tid

70 case e of

71+ Nothing -> do

72+ stop <- shouldStop tid

73+ if stop

74+ then liftIO (cleanupSVar sv) >> allDone stp

71− Nothing -> foldStream st yld sng stp rest

75+ else foldStream st yld sng stp rest

72 Just ex ->

76 Just ex ->

73 case fromException ex of

77 case fromException ex of

74 Just ThreadAbort ->

78 Just ThreadAbort ->

75 foldStream st yld sng stp rest

79 foldStream st yld sng stp rest

76 Nothing -> liftIO (cleanupSVar sv) >> throwM ex

80 Nothing -> liftIO (cleanupSVar sv) >> throwM ex

81+ shouldStop tid =

82+ case svarStopStyle sv of

83+ StopNone -> return False

84+ StopAny -> return True

85+ StopBy -> do

86+ sid <- liftIO $ readIORef (svarStopBy sv)

87+ return $ if tid == sid then True else False

78{-# INLINE fromSVar #-}

89{-# INLINE fromSVar #-}

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

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