bdce4a0 · streamly · git.iko.soy

git.iko.soy/streamly

Use direct style fromSVar for zip and function apply

Improves concurrent zip perf by 2x. There is modest improvement in zip and apply as well.

commitbdce4a09aa2ee2bbe7d0613ade22805515c7a4c2 authorHarendra Kumar <harendra.kumar@gmail.com> dateFri 22 Nov 2019, 04:32 (7 years ago) parente16a980

6 files changed+399−235

compare with master

unifiedsplit

src/Streamly.hs+4−4 src/Streamly/Internal/Data/SVar.hs+6 src/Streamly/Internal/Prelude.hs+125−15 src/Streamly/Streams/Parallel.hs+15−201 src/Streamly/Streams/SVar.hs−9 src/Streamly/Streams/StreamD.hs+249−6

src/Streamly.hs+4−4

@@ -128,10 +128,10 @@ module Streamly

128128

129129 -- * Parallel Function Application

130130 -- $application

131− , (|$)

132− , (|&)

133− , (|$.)

134− , (|&.)

131+ , (IP.|$)

132+ , (IP.|&)

133+ , (IP.|$.)

134+ , (IP.|&.)

135135 , mkAsync

136136

137137 -- * Merging Streams

src/Streamly/Internal/Data/SVar.hs+6

@@ -107,6 +107,7 @@ module Streamly.Internal.Data.SVar

107107 , toStreamVar

108108 , SVarStats (..)

109109 , dumpSVar

110+ , printSVar

110111 )

111112where

112113

@@ -859,6 +860,11 @@ dumpSVar sv = do

859860 <> "---------STATS-----------\n"

860861 <> stats

861862

863+printSVar :: SVar t m a -> String -> IO ()

864+printSVar sv how = do

865+ svInfo <- dumpSVar sv

866+ hPutStrLn stderr $ "\n" <> how <> "\n" <> svInfo

867+

862868-- MVar diagnostics has some overhead - around 5% on asyncly null benchmark, we

863869-- can keep it on in production to debug problems quickly if and when they

864870-- happen, but it may result in unexpected output when threads are left hanging

src/Streamly/Internal/Prelude.hs+125−15

@@ -131,6 +131,12 @@ module Streamly.Internal.Prelude

131131 , toStream -- XXX rename to write?

132132 , toStreamRev -- XXX rename to writeRev?

133133

134+ -- * Function application

135+ , (|$)

136+ , (|&)

137+ , (|$.)

138+ , (|&.)

139+

134140 -- * Transformation

135141 , transform

136142

@@ -387,6 +393,9 @@ module Streamly.Internal.Prelude

387393 , usingStateT

388394 , runStateT

389395

396+ -- * Concurrency

397+ , mkParallel

398+

390399 -- * Diagnostics

391400 , inspectMode

392401

@@ -435,7 +444,7 @@ import Streamly.Internal.Data.Fold.Types (Fold (..), Fold2 (..))

435444import Streamly.Internal.Data.Unfold.Types (Unfold)

436445import Streamly.Internal.Memory.Array.Types (Array, writeNUnsafe)

437446-- import Streamly.Memory.Ring (Ring)

438−import Streamly.Internal.Data.SVar (MonadAsync, defState, adaptState)

447+import Streamly.Internal.Data.SVar (MonadAsync, State, defState, adaptState)

439448import Streamly.Streams.Async (mkAsync')

440449import Streamly.Streams.Combinators (inspectMode, maxYields)

441450import Streamly.Streams.Prelude

@@ -1473,6 +1482,101 @@ toPure = foldr K.cons K.nil

14731482toPureRev :: Monad m => SerialT m a -> m (SerialT Identity a)

14741483toPureRev = foldl' (flip K.cons) K.nil

14751484

1485+------------------------------------------------------------------------------

1486+-- Concurrent Application

1487+------------------------------------------------------------------------------

1488+

1489+-- XXX mkParallel is usually used to make serial streams concurrent by

1490+-- evaluating the stream in another thread. Ideally, we should not be creating

1491+-- an SVar if the stream itself is concurrent. For example, if we are zipping

1492+-- two streams that are concurrently composed, we should be able to pull from

1493+-- the SVar of that concurrent stream instead of creating another redundant

1494+-- SVar on top of it.

1495+--

1496+{-# INLINE_NORMAL mkParallel #-}

1497+mkParallel :: (IsStream t, MonadAsync m)

1498+ => State K.Stream m a -> t m a -> m (t m a)

1499+mkParallel st m = fmap fromStreamD $ D.mkParallel st (toStreamD m)

1500+-- mkParallel = Par.mkParallel

1501+

1502+infixr 0 |$

1503+infixr 0 |$.

1504+

1505+infixl 1 |&

1506+infixl 1 |&.

1507+

1508+-- | Parallel function application operator for streams; just like the regular

1509+-- function application operator '$' except that it is concurrent. The

1510+-- following code prints a value every second even though each stage adds a 1

1511+-- second delay.

1512+--

1513+--

1514+-- @

1515+-- drain $

1516+-- S.mapM (\\x -> threadDelay 1000000 >> print x)

1517+-- |$ S.repeatM (threadDelay 1000000 >> return 1)

1518+-- @

1519+--

1520+-- /Concurrent/

1521+--

1522+-- @since 0.3.0

1523+{-# INLINE (|$) #-}

1524+(|$) :: (IsStream t, MonadAsync m) => (t m a -> t m b) -> t m a -> t m b

1525+f |$ x = D.fromStreamD $

1526+ D.applyParallel (D.toStreamD . f . D.fromStreamD) (D.toStreamD x)

1527+-- (|$) = Par.applyParallel

1528+

1529+-- | Parallel reverse function application operator for streams; just like the

1530+-- regular reverse function application operator '&' except that it is

1531+-- concurrent.

1532+--

1533+-- @

1534+-- drain $

1535+-- S.repeatM (threadDelay 1000000 >> return 1)

1536+-- |& S.mapM (\\x -> threadDelay 1000000 >> print x)

1537+-- @

1538+--

1539+-- /Concurrent/

1540+--

1541+-- @since 0.3.0

1542+{-# INLINE (|&) #-}

1543+(|&) :: (IsStream t, MonadAsync m) => t m a -> (t m a -> t m b) -> t m b

1544+x |& f = f |$ x

1545+

1546+-- | Parallel function application operator; applies a @run@ or @fold@ function

1547+-- to a stream such that the fold consumer and the stream producer run in

1548+-- parallel. A @run@ or @fold@ function reduces the stream to a value in the

1549+-- underlying monad. The @.@ at the end of the operator is a mnemonic for

1550+-- termination of the stream.

1551+--

1552+-- @

1553+-- S.foldlM' (\\_ a -> threadDelay 1000000 >> print a) ()

1554+-- |$. S.repeatM (threadDelay 1000000 >> return 1)

1555+-- @

1556+--

1557+-- /Concurrent/

1558+--

1559+-- @since 0.3.0

1560+{-# INLINE (|$.) #-}

1561+(|$.) :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

1562+(|$.) = D.foldParallel

1563+-- (|$.) = Par.foldParallel

1564+

1565+-- | Parallel reverse function application operator for applying a run or fold

1566+-- functions to a stream. Just like '|$.' except that the operands are reversed.

1567+--

1568+-- @

1569+-- S.repeatM (threadDelay 1000000 >> return 1)

1570+-- |&. S.foldlM' (\\_ a -> threadDelay 1000000 >> print a) ()

1571+-- @

1572+--

1573+-- /Concurrent/

1574+--

1575+-- @since 0.3.0

1576+{-# INLINE (|&.) #-}

1577+(|&.) :: (IsStream t, MonadAsync m) => t m a -> (t m a -> m b) -> m b

1578+x |&. f = f |$. x

1579+

14761580------------------------------------------------------------------------------

14771581-- General Transformation

14781582------------------------------------------------------------------------------

@@ -2036,29 +2140,35 @@ zipWith f m1 m2 = fromStreamS $ S.zipWith f (toStreamS m1) (toStreamS m2)

20362140-- Parallel Zipping

20372141------------------------------------------------------------------------------

20382142

2039−-- | Like 'zipWith' but zips concurrently i.e. both the streams being zipped

2040−-- are generated concurrently.

2041−--

2042−-- @since 0.1.0

2043−{-# INLINE zipAsyncWith #-}

2044−zipAsyncWith :: (IsStream t, MonadAsync m)

2045− => (a -> b -> c) -> t m a -> t m b -> t m c

2046−zipAsyncWith f m1 m2 = K.mkStream $ \st stp sng yld -> do

2047− ma <- Par.mkParallel (adaptState st) m1

2048− mb <- Par.mkParallel (adaptState st) m2

2049− K.foldStream st stp sng yld $ zipWith f ma mb

2143+-- The CPS version and the direct version of zipAsyncWithM below seem to have

2144+-- identical performane. However, we need to use the StreamD version of

2145+-- mkParallel which uses a direct implementation of fromSVar. In comparison to

2146+-- CPS version of fromSVar the direct version gives a 2x improvement.

20502147

20512148-- | Like 'zipWithM' but zips concurrently i.e. both the streams being zipped

20522149-- are generated concurrently.

20532150--

20542151-- @since 0.4.0

2055−{-# INLINABLE zipAsyncWithM #-}

2152+{-# INLINE zipAsyncWithM #-}

20562153zipAsyncWithM :: (IsStream t, MonadAsync m)

20572154 => (a -> b -> m c) -> t m a -> t m b -> t m c

20582155zipAsyncWithM f m1 m2 = K.mkStream $ \st stp sng yld -> do

2059− ma <- Par.mkParallel (adaptState st) m1

2060− mb <- Par.mkParallel (adaptState st) m2

2156+ ma <- mkParallel (adaptState st) m1

2157+ mb <- mkParallel (adaptState st) m2

20612158 K.foldStream st stp sng yld $ zipWithM f ma mb

2159+{-

2160+zipAsyncWithM f m1 m2 =

2161+ fromStreamD $ D.zipAsyncWithM f (toStreamD m1) (toStreamD m2)

2162+-}

2163+

2164+-- | Like 'zipWith' but zips concurrently i.e. both the streams being zipped

2165+-- are generated concurrently.

2166+--

2167+-- @since 0.1.0

2168+{-# INLINE zipAsyncWith #-}

2169+zipAsyncWith :: (IsStream t, MonadAsync m)

2170+ => (a -> b -> c) -> t m a -> t m b -> t m c

2171+zipAsyncWith f = zipAsyncWithM (\a b -> return (f a b))

20622172

20632173------------------------------------------------------------------------------

20642174-- Comparison

src/Streamly/Streams/Parallel.hs+15−201

@@ -28,13 +28,9 @@ module Streamly.Streams.Parallel

2828 , parallelFst

2929 , parallelMin

3030 , tapAsync

31−

32− -- * Function application

3331 , mkParallel

34− , (|$)

35− , (|&)

36− , (|$.)

37− , (|&.)

32+ , applyParallel

33+ , foldParallel

3834 )

3935where

4036

@@ -49,7 +45,7 @@ import Control.Monad.Reader.Class (MonadReader(..))

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

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

5147import Data.Functor (void)

52−import Data.IORef (readIORef, writeIORef, newIORef)

48+import Data.IORef (readIORef, writeIORef)

5349import Data.Maybe (fromJust)

5450#if __GLASGOW_HASKELL__ < 808

5551import Data.Semigroup (Semigroup(..))

@@ -63,14 +59,10 @@ import Streamly.Streams.Serial (map)

6359import Streamly.Streams.StreamK (IsStream(..), Stream, mkStream, foldStream,

6460 foldStreamShared, adapt)

6561

66−import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)

67−import Streamly.Internal.Data.Time.Clock (Clock(..), getTime)

6862import Streamly.Internal.Data.SVar hiding (handleChildException)

6963

7064import qualified Streamly.Streams.StreamK as K

7165import qualified Streamly.Streams.StreamD as D

72−import qualified Streamly.Internal.Data.SVar as SVar

73−import qualified Streamly.Internal.Data.Fold as FL

7466

7567#include "Instances.hs"

7668

@@ -133,71 +125,6 @@ runOneLimited st m0 winfo = go m0

133125 single a = sendit a >> (liftIO $ sendStop sv winfo)

134126 yieldk a r = sendit a >> go r

135127

136−-------------------------------------------------------------------------------

137−-- StreamD based worker routines

138−-------------------------------------------------------------------------------

139−

140−-- Using StreamD the worker stream producing code can fuse with the code to

141−-- queue output to the SVar giving some perf boost.

142−--

143−-- Note that StreamD can only be used in limited situations, specifically, we

144−-- cannot implement joinStreamVarPar using this.

145−

146−{-# INLINE_NORMAL pushWorkerParD #-}

147−pushWorkerParD :: MonadAsync m

148− => State t m a -> SVar t m a -> D.Stream m a -> m ()

149−pushWorkerParD st sv xs =

150− if svarInspectMode sv

151− then forkWithDiag

152− else do

153− tid <-

154− case getYieldLimit st of

155− Nothing -> doFork (work Nothing)

156− (svarMrun sv)

157− (SVar.handleChildException sv)

158− Just _ -> doFork (workLim Nothing)

159− (svarMrun sv)

160− (SVar.handleChildException sv)

161− modifyThread sv tid

162−

163− where

164−

165− {-# NOINLINE work #-}

166− work info = (D.runFold (FL.toParallelSVar sv info) xs)

167−

168− {-# NOINLINE workLim #-}

169− workLim info = D.runFold (FL.toParallelSVarLimited sv info) xs

170−

171− {-# NOINLINE forkWithDiag #-}

172− forkWithDiag = do

173− -- We do not use workerCount in case of ParallelVar but still there is

174− -- no harm in maintaining it correctly.

175− liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1

176− recordMaxWorkers sv

177− -- This allocation matters when significant number of workers are being

178− -- sent. We allocate it only when needed. The overhead increases by 4x.

179− winfo <-

180− case yieldRateInfo sv of

181− Nothing -> return Nothing

182− Just _ -> liftIO $ do

183− cntRef <- newIORef 0

184− t <- getTime Monotonic

185− lat <- newIORef (0, t)

186− return $ Just WorkerInfo

187− { workerYieldMax = 0

188− , workerYieldCount = cntRef

189− , workerLatencyStart = lat

190− }

191− tid <-

192− case getYieldLimit st of

193− Nothing -> doFork (work winfo)

194− (svarMrun sv)

195− (SVar.handleChildException sv)

196− Just _ -> doFork (workLim winfo)

197− (svarMrun sv)

198− (SVar.handleChildException sv)

199− modifyThread sv tid

200−

201128-------------------------------------------------------------------------------

202129-- Consing and appending a stream in parallel style

203130-------------------------------------------------------------------------------

@@ -276,72 +203,38 @@ parallelMin = joinStreamVarPar ParallelVar StopAny

276203-- Convert a stream to parallel

277204------------------------------------------------------------------------------

278205

279−_mkParallelK :: (IsStream t, MonadAsync m) => t m a -> m (t m a)

280−_mkParallelK m = do

281− sv <- newParallelVar StopNone defState

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

283− return $ fromSVar sv

284−

285−-- We may have to use it in higher order functions like concatMap, so use the

286−-- Normal phase for inlining.

287−{-# INLINE_NORMAL mkParallel #-}

288206mkParallel :: (IsStream t, MonadAsync m)

289207 => State Stream m a -> t m a -> m (t m a)

290208mkParallel st m = do

291209 sv <- newParallelVar StopNone defState

292− pushWorkerParD st sv (D.toStreamD m)

210+ -- pushWorkerPar sv (runOne st{streamVar = Just sv} $ toStream m)

211+ D.toSVarParallel st sv $ D.toStreamD m

293212 return $ fromSVar sv

294213

295214------------------------------------------------------------------------------

296215-- Stream to stream concurrent function application

297216------------------------------------------------------------------------------

298217

299−{-# INLINE _applyParallelK #-}

300−_applyParallelK :: (IsStream t, MonadAsync m)

301− => (t m a -> t m b) -> t m a -> t m b

302−_applyParallelK f m = mkStream $ \st yld sng stp -> do

303− sv <- newParallelVar StopNone (adaptState st)

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

305− foldStream st yld sng stp $ f $ fromSVar sv

306−

307−-- XXX need to implement fromSVar in direct style for fused generation from

308−-- SVar.

309−{-# INLINE applyParallelD #-}

310−applyParallelD :: (IsStream t, MonadAsync m)

311− => (t m a -> t m b) -> t m a -> t m b

312−applyParallelD f m = mkStream $ \st yld sng stp -> do

313− sv <- newParallelVar StopNone (adaptState st)

314− pushWorkerParD (adaptState st) sv (D.toStreamD m)

315− foldStream st yld sng stp $ f $ fromSVar sv

316−

317218{-# INLINE applyParallel #-}

318219applyParallel :: (IsStream t, MonadAsync m)

319220 => (t m a -> t m b) -> t m a -> t m b

320−applyParallel = applyParallelD

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

222+ sv <- newParallelVar StopNone (adaptState st)

223+ -- pushWorkerPar sv (runOne st{streamVar = Just sv} (toStream m))

224+ D.toSVarParallel (adaptState st) sv $ D.toStreamD m

225+ foldStream st yld sng stp $ f $ fromSVar sv

321226

322227------------------------------------------------------------------------------

323228-- Stream runner concurrent function application

324229------------------------------------------------------------------------------

325230

326−{-# INLINE _foldParallelK #-}

327−_foldParallelK :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

328−_foldParallelK f m = do

329− sv <- newParallelVar StopNone defState

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

331− f $ fromSVar sv

332−

333−-- XXX need to implement fromSVar in direct style for fused generation from

334−-- SVar.

335−{-# INLINE foldParallelD #-}

336−foldParallelD :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

337−foldParallelD f m = do

338− sv <- newParallelVar StopNone defState

339− pushWorkerParD defState sv (D.toStreamD m)

340− f $ fromSVar sv

341−

342231{-# INLINE foldParallel #-}

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

344−foldParallel = foldParallelD

233+foldParallel f m = do

234+ sv <- newParallelVar StopNone defState

235+ -- pushWorkerPar sv (runOne defState{streamVar = Just sv} $ toStream m)

236+ D.toSVarParallel defState sv $ D.toStreamD m

237+ f $ fromSVar sv

345238

346239------------------------------------------------------------------------------

347240-- Parallel tee

@@ -427,85 +320,6 @@ tapAsync f m = mkStream $ \st yld sng stp -> do

427320 (svarMrun sv) (handleChildException sv)

428321 foldStreamShared st yld sng stp (teeToSVar sv m)

429322

430−------------------------------------------------------------------------------

431−-- Concurrent Application

432−------------------------------------------------------------------------------

433−

434−infixr 0 |$

435−infixr 0 |$.

436−

437−infixl 1 |&

438−infixl 1 |&.

439−

440−-- | Parallel function application operator for streams; just like the regular

441−-- function application operator '$' except that it is concurrent. The

442−-- following code prints a value every second even though each stage adds a 1

443−-- second delay.

444−--

445−--

446−-- @

447−-- drain $

448−-- S.mapM (\\x -> threadDelay 1000000 >> print x)

449−-- |$ S.repeatM (threadDelay 1000000 >> return 1)

450−-- @

451−--

452−-- /Concurrent/

453−--

454−-- @since 0.3.0

455−{-# INLINE (|$) #-}

456−(|$) :: (IsStream t, MonadAsync m) => (t m a -> t m b) -> t m a -> t m b

457−f |$ x = applyParallel f x

458−

459−-- | Parallel reverse function application operator for streams; just like the

460−-- regular reverse function application operator '&' except that it is

461−-- concurrent.

462−--

463−-- @

464−-- drain $

465−-- S.repeatM (threadDelay 1000000 >> return 1)

466−-- |& S.mapM (\\x -> threadDelay 1000000 >> print x)

467−-- @

468−--

469−-- /Concurrent/

470−--

471−-- @since 0.3.0

472−{-# INLINE (|&) #-}

473−(|&) :: (IsStream t, MonadAsync m) => t m a -> (t m a -> t m b) -> t m b

474−x |& f = f |$ x

475−

476−-- | Parallel function application operator; applies a @run@ or @fold@ function

477−-- to a stream such that the fold consumer and the stream producer run in

478−-- parallel. A @run@ or @fold@ function reduces the stream to a value in the

479−-- underlying monad. The @.@ at the end of the operator is a mnemonic for

480−-- termination of the stream.

481−--

482−-- @

483−-- S.foldlM' (\\_ a -> threadDelay 1000000 >> print a) ()

484−-- |$. S.repeatM (threadDelay 1000000 >> return 1)

485−-- @

486−--

487−-- /Concurrent/

488−--

489−-- @since 0.3.0

490−{-# INLINE (|$.) #-}

491−(|$.) :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

492−f |$. x = foldParallel f x

493−

494−-- | Parallel reverse function application operator for applying a run or fold

495−-- functions to a stream. Just like '|$.' except that the operands are reversed.

496−--

497−-- @

498−-- S.repeatM (threadDelay 1000000 >> return 1)

499−-- |&. S.foldlM' (\\_ a -> threadDelay 1000000 >> print a) ()

500−-- @

501−--

502−-- /Concurrent/

503−--

504−-- @since 0.3.0

505−{-# INLINE (|&.) #-}

506−(|&.) :: (IsStream t, MonadAsync m) => t m a -> (t m a -> m b) -> m b

507−x |&. f = f |$. x

508−

509323------------------------------------------------------------------------------

510324-- ParallelT

511325------------------------------------------------------------------------------

src/Streamly/Streams/SVar.hs−9

@@ -24,21 +24,12 @@ import Control.Monad.Catch (throwM)

2424import Control.Monad.IO.Class (MonadIO(liftIO))

2525import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

2626import Data.Maybe (isNothing)

27−#if __GLASGOW_HASKELL__ < 808

28−import Data.Semigroup ((<>))

29−#endif

30−import System.IO (hPutStrLn, stderr)

3127import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

3228import System.Mem (performMajorGC)

3329

3430import Streamly.Internal.Data.SVar

3531import Streamly.Streams.StreamK hiding (reverse)

3632

37−printSVar :: SVar t m a -> String -> IO ()

38−printSVar sv how = do

39− svInfo <- dumpSVar sv

40− hPutStrLn stderr $ "\n" <> how <> "\n" <> svInfo

41−

4233-- | Pull a stream from an SVar.

4334{-# NOINLINE fromStreamVar #-}

4435fromStreamVar :: MonadAsync m => SVar Stream m a -> Stream m a

src/Streamly/Streams/StreamD.hs+249−6

@@ -95,6 +95,7 @@ module Streamly.Streams.StreamD

9595 , fromListM

9696 , fromStreamK

9797 , fromStreamD

98+ , fromSVar

9899

99100 -- * Elimination

100101 -- ** General Folds

@@ -136,6 +137,7 @@ module Streamly.Streams.StreamD

136137 , findM

137138 , find

138139 , (!!)

140+ , toSVarParallel

139141

140142 -- ** Flattening nested streams

141143 , concatMapM

@@ -260,6 +262,8 @@ module Streamly.Streams.StreamD

260262 , indexedR

261263 , zipWith

262264 , zipWithM

265+ , zipAsyncWith

266+ , zipAsyncWithM

263267

264268 -- * Comparisons

265269 , eqBy

@@ -292,19 +296,25 @@ module Streamly.Streams.StreamD

292296 , resumeDecodeUtf8Either

293297 , decodeUtf8Arrays

294298 , decodeUtf8ArraysLenient

299+

300+ -- * Concurrent Application

301+ , mkParallel

302+ , applyParallel

303+ , foldParallel

295304 )

296305where

297306

298−import Control.Exception (Exception, SomeException)

299−import Control.Monad (void)

300−import Control.Monad.Catch (MonadCatch)

307+import Control.Exception (Exception, SomeException, fromException)

308+import Control.Monad (void, when)

309+import Control.Monad.Catch (MonadCatch, throwM)

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

302311import Control.Monad.Reader (ReaderT)

303312import Control.Monad.State.Strict (StateT)

304313import Control.Monad.Trans (MonadTrans(lift))

305314import Data.Bits (shiftR, shiftL, (.|.), (.&.))

306315import Data.Functor.Identity (Identity(..))

307−import Data.Maybe (fromJust, isJust)

316+import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

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

308318import Data.Word (Word32)

309319import Foreign.Ptr (Ptr)

310320import Foreign.Storable (Storable(..))

@@ -315,6 +325,7 @@ import GHC.Ptr (Ptr (..))

315325import GHC.Types (SPEC(..))

316326import GHC.Word (Word8(..))

317327import System.IO.Unsafe (unsafePerformIO)

328+import System.Mem (performMajorGC)

318329import Prelude

319330 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,

320331 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,

@@ -325,18 +336,22 @@ import Prelude

325336import qualified Control.Monad.Catch as MC

326337import qualified Control.Monad.Reader as Reader

327338import qualified Control.Monad.State.Strict as State

339+import qualified Prelude

328340

341+import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)

329342import Streamly.Internal.Memory.Array.Types (Array(..))

330343import Streamly.Internal.Data.Fold.Types (Fold(..))

331344import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))

332−import Streamly.Internal.Data.SVar (MonadAsync, defState, adaptState)

345+import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)

333346import Streamly.Internal.Data.Unfold.Types (Unfold(..))

334347import Streamly.Internal.Data.Strict (Tuple'(..))

335348

336349import Streamly.Internal.Data.Stream.StreamD.Type

350+import Streamly.Internal.Data.SVar

337351

338352import qualified Streamly.Internal.Data.Pipe.Types as Pipe

339353import qualified Streamly.Internal.Memory.Array.Types as A

354+import qualified Streamly.Internal.Data.Fold as FL

340355import qualified Streamly.Memory.Ring as RB

341356import qualified Streamly.Streams.StreamK as K

342357

@@ -648,10 +663,100 @@ fromListM = Stream step

648663toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a

649664toStreamD = fromStreamK . K.toStream

650665

666+-------------------------------------------------------------------------------

667+-- Generation from SVar

668+-------------------------------------------------------------------------------

669+

670+data FromSVarState t m a =

671+ FromSVarInit

672+ | FromSVarRead (SVar t m a)

673+ | FromSVarLoop (SVar t m a) [ChildEvent a]

674+ | FromSVarDone (SVar t m a)

675+

676+{-# INLINE_NORMAL fromSVar #-}

677+fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a

678+fromSVar svar = Stream step FromSVarInit

679+ where

680+

681+ {-# INLINE_LATE step #-}

682+ step _ FromSVarInit = do

683+ ref <- liftIO $ newIORef ()

684+ _ <- liftIO $ mkWeakIORef ref hook

685+ -- when this copy of svar gets garbage collected "ref" will get

686+ -- garbage collected and our GC hook will be called.

687+ let sv = svar{svarRef = Just ref}

688+ return $ Skip (FromSVarRead sv)

689+

690+ where

691+

692+ {-# NOINLINE hook #-}

693+ hook = do

694+ when (svarInspectMode svar) $ do

695+ r <- liftIO $ readIORef (svarStopTime (svarStats svar))

696+ when (isNothing r) $

697+ printSVar svar "SVar Garbage Collected"

698+ cleanupSVar svar

699+ -- If there are any SVars referenced by this SVar a GC will prompt

700+ -- them to be cleaned up quickly.

701+ when (svarInspectMode svar) performMajorGC

702+

703+ step _ (FromSVarRead sv) = do

704+ list <- readOutputQ sv

705+ -- Reversing the output is important to guarantee that we process the

706+ -- outputs in the same order as they were generated by the constituent

707+ -- streams.

708+ return $ Skip $ FromSVarLoop sv (Prelude.reverse list)

709+

710+ step _ (FromSVarLoop sv []) = do

711+ done <- postProcess sv

712+ return $ Skip $ if done

713+ then (FromSVarDone sv)

714+ else (FromSVarRead sv)

715+

716+ step _ (FromSVarLoop sv (ev : es)) = do

717+ case ev of

718+ ChildYield a -> return $ Yield a (FromSVarLoop sv es)

719+ ChildStop tid e -> do

720+ accountThread sv tid

721+ case e of

722+ Nothing -> do

723+ stop <- shouldStop tid

724+ if stop

725+ then do

726+ liftIO (cleanupSVar sv)

727+ return $ Skip (FromSVarDone sv)

728+ else return $ Skip (FromSVarLoop sv es)

729+ Just ex ->

730+ case fromException ex of

731+ Just ThreadAbort ->

732+ return $ Skip (FromSVarLoop sv es)

733+ Nothing -> liftIO (cleanupSVar sv) >> throwM ex

734+ where

735+

736+ shouldStop tid =

737+ case svarStopStyle sv of

738+ StopNone -> return False

739+ StopAny -> return True

740+ StopBy -> do

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

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

743+

744+ step _ (FromSVarDone sv) = do

745+ when (svarInspectMode sv) $ do

746+ t <- liftIO $ getTime Monotonic

747+ liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)

748+ liftIO $ printSVar sv "SVar Done"

749+ return Stop

750+

751+-------------------------------------------------------------------------------

752+-- Hoisting the inner monad

753+-------------------------------------------------------------------------------

754+

651755{-# INLINE_NORMAL hoist #-}

652756hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a

653757hoist f (Stream step state) = (Stream step' state)

654758 where

759+ {-# INLINE_LATE step' #-}

655760 step' gst st = do

656761 r <- f $ step (adaptState gst) st

657762 return $ case r of

@@ -659,7 +764,7 @@ hoist f (Stream step state) = (Stream step' state)

659764 Skip s -> Skip s

660765 Stop -> Stop

661766

662−{-# INLINE_NORMAL generally #-}

767+{-# INLINE generally #-}

663768generally :: Monad m => Stream Identity a -> Stream m a

664769generally = hoist (return . runIdentity)

665770

@@ -668,6 +773,7 @@ liftInner :: (Monad m, MonadTrans t, Monad (t m))

668773 => Stream m a -> Stream (t m) a

669774liftInner (Stream step state) = Stream step' state

670775 where

776+ {-# INLINE_LATE step' #-}

671777 step' gst st = do

672778 r <- lift $ step (adaptState gst) st

673779 return $ case r of

@@ -679,6 +785,7 @@ liftInner (Stream step state) = Stream step' state

679785runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a

680786runReaderT sval (Stream step state) = Stream step' state

681787 where

788+ {-# INLINE_LATE step' #-}

682789 step' gst st = do

683790 r <- Reader.runReaderT (step (adaptState gst) st) sval

684791 return $ case r of

@@ -690,6 +797,7 @@ runReaderT sval (Stream step state) = Stream step' state

690797evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a

691798evalStateT sval (Stream step state) = Stream step' (state, sval)

692799 where

800+ {-# INLINE_LATE step' #-}

693801 step' gst (st, sv) = do

694802 (r, sv') <- State.runStateT (step (adaptState gst) st) sv

695803 return $ case r of

@@ -701,6 +809,7 @@ evalStateT sval (Stream step state) = Stream step' (state, sval)

701809runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)

702810runStateT sval (Stream step state) = Stream step' (state, sval)

703811 where

812+ {-# INLINE_LATE step' #-}

704813 step' gst (st, sv) = do

705814 (r, sv') <- State.runStateT (step (adaptState gst) st) sv

706815 return $ case r of

@@ -3426,6 +3535,37 @@ zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)

34263535zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c

34273536zipWith f = zipWithM (\a b -> return (f a b))

34283537

3538+-- | Like 'zipWithM' but zips concurrently i.e. both the streams being zipped

3539+-- are generated concurrently.

3540+--

3541+{-# INLINE_NORMAL zipAsyncWithM #-}

3542+zipAsyncWithM :: MonadAsync m

3543+ => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c

3544+zipAsyncWithM f m1 m2 = Stream step Nothing

3545+

3546+ where

3547+

3548+ {-# INLINE_LATE step #-}

3549+ step gst Nothing = do

3550+ ma <- mkParallel (adaptState gst) m1

3551+ mb <- mkParallel (adaptState gst) m2

3552+ return $ Skip (Just $ zipWithM f ma mb)

3553+

3554+ step gst (Just (UnStream step1 st)) = do

3555+ r <- step1 gst st

3556+ return $ case r of

3557+ Yield a s -> Yield a (Just $ Stream step1 s)

3558+ Skip s -> Skip (Just $ Stream step1 s)

3559+ Stop -> Stop

3560+

3561+-- | Like 'zipWith' but zips concurrently i.e. both the streams being zipped

3562+-- are generated concurrently.

3563+--

3564+{-# INLINE zipAsyncWith #-}

3565+zipAsyncWith :: MonadAsync m

3566+ => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c

3567+zipAsyncWith f = zipAsyncWithM (\a b -> return (f a b))

3568+

34293569------------------------------------------------------------------------------

34303570-- Merging

34313571------------------------------------------------------------------------------

@@ -3983,3 +4123,106 @@ encodeUtf8 (Stream step state) = Stream step' (EncodeState state WNil)

39834123 Skip s -> Skip (EncodeState s WNil)

39844124 Stop -> Stop

39854125 step' _ (EncodeState s (WCons x xs)) = return $ Yield x (EncodeState s xs)

4126+

4127+-------------------------------------------------------------------------------

4128+-- Concurrent application and fold

4129+-------------------------------------------------------------------------------

4130+

4131+-- Using StreamD the worker stream producing code can fuse with the code to

4132+-- queue output to the SVar giving some perf boost.

4133+--

4134+-- Note that StreamD can only be used in limited situations, specifically, we

4135+-- cannot implement joinStreamVarPar using this.

4136+--

4137+-- XXX make sure that the SVar passed is a Parallel style SVar.

4138+

4139+-- | Fold the supplied stream to the SVar asynchronously using Parallel

4140+-- concurrency style.

4141+-- {-# INLINE_NORMAL toSVarParallel #-}

4142+{-# INLINE toSVarParallel #-}

4143+toSVarParallel :: MonadAsync m

4144+ => State t m a -> SVar t m a -> Stream m a -> m ()

4145+toSVarParallel st sv xs =

4146+ if svarInspectMode sv

4147+ then forkWithDiag

4148+ else do

4149+ tid <-

4150+ case getYieldLimit st of

4151+ Nothing -> doFork (work Nothing)

4152+ (svarMrun sv)

4153+ (handleChildException sv)

4154+ Just _ -> doFork (workLim Nothing)

4155+ (svarMrun sv)

4156+ (handleChildException sv)

4157+ modifyThread sv tid

4158+

4159+ where

4160+

4161+ {-# NOINLINE work #-}

4162+ work info = (runFold (FL.toParallelSVar sv info) xs)

4163+

4164+ {-# NOINLINE workLim #-}

4165+ workLim info = runFold (FL.toParallelSVarLimited sv info) xs

4166+

4167+ {-# NOINLINE forkWithDiag #-}

4168+ forkWithDiag = do

4169+ -- We do not use workerCount in case of ParallelVar but still there is

4170+ -- no harm in maintaining it correctly.

4171+ liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1

4172+ recordMaxWorkers sv

4173+ -- This allocation matters when significant number of workers are being

4174+ -- sent. We allocate it only when needed. The overhead increases by 4x.

4175+ winfo <-

4176+ case yieldRateInfo sv of

4177+ Nothing -> return Nothing

4178+ Just _ -> liftIO $ do

4179+ cntRef <- newIORef 0

4180+ t <- getTime Monotonic

4181+ lat <- newIORef (0, t)

4182+ return $ Just WorkerInfo

4183+ { workerYieldMax = 0

4184+ , workerYieldCount = cntRef

4185+ , workerLatencyStart = lat

4186+ }

4187+ tid <-

4188+ case getYieldLimit st of

4189+ Nothing -> doFork (work winfo)

4190+ (svarMrun sv)

4191+ (handleChildException sv)

4192+ Just _ -> doFork (workLim winfo)

4193+ (svarMrun sv)

4194+ (handleChildException sv)

4195+ modifyThread sv tid

4196+

4197+{-# INLINE_NORMAL mkParallel #-}

4198+mkParallel :: MonadAsync m => State t m a -> Stream m a -> m (Stream m a)

4199+mkParallel st m = do

4200+ sv <- newParallelVar StopNone defState

4201+ toSVarParallel st sv m

4202+ return $ fromSVar sv

4203+

4204+{-# INLINE_NORMAL applyParallel #-}

4205+applyParallel :: MonadAsync m

4206+ => (Stream m a -> Stream m b) -> Stream m a -> Stream m b

4207+applyParallel f m = Stream step Nothing

4208+ where

4209+

4210+ {-# INLINE_LATE step #-}

4211+ step gst Nothing = do

4212+ sv <- newParallelVar StopNone (adaptState gst)

4213+ toSVarParallel (adaptState gst) sv m

4214+ return $ Skip $ Just $ f (fromSVar sv)

4215+

4216+ step gst (Just (UnStream step1 st)) = do

4217+ r <- step1 gst st

4218+ return $ case r of

4219+ Yield a s -> Yield a (Just $ Stream step1 s)

4220+ Skip s -> Skip (Just $ Stream step1 s)

4221+ Stop -> Stop

4222+

4223+{-# INLINE foldParallel #-}

4224+foldParallel :: (K.IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

4225+foldParallel f m = do

4226+ sv <- newParallelVar StopNone defState

4227+ toSVarParallel defState sv (toStreamD m)

4228+ f $ fromStreamD $ fromSVar sv