5073377 · streamly · git.iko.soy

git.iko.soy/streamly

Add modular stream stats counting

* pollCounts to poll the element count in another thread * delayPost to introduce a delay in polling * rollingMap to compute diff of successive elements These combinators can be used to compute and report the element processing rate in a stream.

commit5073377ba90db28017aa27e3180ea95f2b071308 authorHarendra Kumar <harendra.kumar@gmail.com> dateSat 28 Dec 2019, 23:01 (7 years ago) parent135063f

5 files changed+168−12

compare with master

unifiedsplit

benchmark/Linear.hs+1 benchmark/Streamly/Benchmark/Prelude.hs+7−1 src/Streamly/Internal/Data/SVar.hs+15−2 src/Streamly/Internal/Data/Stream/StreamD.hs+69−3 src/Streamly/Internal/Prelude.hs+76−6

benchmark/Linear.hs+1

@@ -323,6 +323,7 @@ main =

323323 , benchIOSink "foldrTMap" (Ops.foldrTMap 1)

324324 , benchIOSink "tap" (Ops.tap 1)

325325 , benchIOSink "tapRate 1 second" (Ops.tapRate 1)

326+ , benchIOSink "pollCounts 1 second" (Ops.pollCounts 1)

326327 , benchIOSink "tapAsync" (Ops.tapAsync 1)

327328 , benchIOSink "tapAsyncS" (Ops.tapAsyncS 1)

328329 ]

benchmark/Streamly/Benchmark/Prelude.hs+7−1

@@ -25,7 +25,7 @@ module Streamly.Benchmark.Prelude where

2525

2626import Control.DeepSeq (NFData)

2727import Control.Monad (when)

28−import Control.Monad.IO.Class (MonadIO)

28+import Control.Monad.IO.Class (MonadIO(..))

2929import Control.Monad.State.Strict (StateT, get, put)

3030import Data.Functor.Identity (Identity, runIdentity)

3131import Data.IORef (newIORef, modifyIORef')

@@ -411,6 +411,12 @@ tapRate n str = do

411411 cref <- newIORef 0

412412 composeN n (Internal.tapRate 1 (\c -> modifyIORef' cref (c +))) str

413413

414+{-# INLINE pollCounts #-}

415+pollCounts :: Int -> Stream IO Int -> IO ()

416+pollCounts n str = do

417+ composeN n (Internal.pollCounts f FL.drain) str

418+ where f = Internal.rollingMap (P.-) . Internal.delayPost 1

419+

414420{-# INLINE tapAsyncS #-}

415421tapAsyncS :: S.MonadAsync m => Int -> Stream m Int -> m ()

416422tapAsyncS n = composeN n $ Par.tapAsync S.sum

src/Streamly/Internal/Data/SVar.hs+15−2

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

108108 , modifyThread

109109 , doFork

110110 , fork

111+ , forkManaged

111112

112113 , toStreamVar

113114 , SVarStats (..)

@@ -118,7 +119,7 @@ module Streamly.Internal.Data.SVar

118119where

119120

120121import Control.Concurrent

121− (ThreadId, myThreadId, threadDelay, throwTo, forkIO)

122+ (ThreadId, myThreadId, threadDelay, throwTo, forkIO, killThread)

122123import Control.Concurrent.MVar

123124 (MVar, newEmptyMVar, tryPutMVar, takeMVar, tryTakeMVar, newMVar,

124125 tryReadMVar)

@@ -151,11 +152,13 @@ import Data.Set (Set)

151152import GHC.Conc (ThreadId(..))

152153import GHC.Exts

153154import GHC.IO (IO(..))

155+import System.IO (hPutStrLn, stderr)

156+import System.Mem.Weak (addFinalizer)

157+

154158import Streamly.Internal.Data.Time.Clock (Clock(..), getTime)

155159import Streamly.Internal.Data.Time.Units

156160 (AbsTime, NanoSecond64(..), MicroSecond64(..), diffAbsTime64,

157161 fromRelTime64, toRelTime64, showNanoSecond64, showRelTime64)

158−import System.IO (hPutStrLn, stderr)

159162

160163import qualified Data.Heap as H

161164import qualified Data.Set as S

@@ -958,6 +961,16 @@ doFork action (RunInIO mrun) exHandler =

958961fork :: MonadBaseControl IO m => m () -> m ThreadId

959962fork = liftBaseDiscard forkIO

960963

964+-- | Fork a thread that is automatically killed as soon as the reference to the

965+-- returned threadId is garbage collected.

966+--

967+{-# INLINABLE forkManaged #-}

968+forkManaged :: (MonadIO m, MonadBaseControl IO m) => m () -> m ThreadId

969+forkManaged action = do

970+ tid <- fork action

971+ liftIO $ addFinalizer tid (killThread tid)

972+ return tid

973+

961974------------------------------------------------------------------------------

962975-- Collecting results from child workers in a streamed fashion

963976------------------------------------------------------------------------------

src/Streamly/Internal/Data/Stream/StreamD.hs+69−3

@@ -130,6 +130,7 @@ module Streamly.Internal.Data.Stream.StreamD

130130 , tap

131131 , tapAsync

132132 , tapRate

133+ , pollCounts

133134 , drain

134135 , null

135136 , head

@@ -254,6 +255,8 @@ module Streamly.Internal.Data.Stream.StreamD

254255 , map

255256 , mapM

256257 , sequence

258+ , rollingMap

259+ , rollingMapM

257260

258261 -- * Inserting

259262 , intersperseM

@@ -335,7 +338,7 @@ import qualified Control.Monad.State.Strict as State

335338import qualified Prelude

336339

337340import Streamly.Internal.Mutable.Prim.Var

338− (MonadMut, Prim, Var, readVar, newVar, modifyVar')

341+ (Prim, Var, readVar, newVar, modifyVar')

339342import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)

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

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

@@ -668,11 +671,11 @@ toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a

668671toStreamD = fromStreamK . K.toStream

669672

670673{-# INLINE_NORMAL fromPrimVar #-}

671−fromPrimVar :: (MonadMut m, Prim a) => Var m a -> Stream m a

674+fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a

672675fromPrimVar var = Stream step ()

673676 where

674677 {-# INLINE_LATE step #-}

675− step _ () = readVar var >>= \x -> return $ Yield x ()

678+ step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()

676679

677680-------------------------------------------------------------------------------

678681-- Generation from SVar

@@ -3221,6 +3224,40 @@ scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)

32213224scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a

32223225scanl1' f = scanl1M' (\x y -> return (f x y))

32233226

3227+------------------------------------------------------------------------------

3228+-- Stateful map/scan

3229+------------------------------------------------------------------------------

3230+

3231+data RollingMapState s a = RollingMapInit s | RollingMapGo s a

3232+

3233+{-# INLINE rollingMapM #-}

3234+rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b

3235+rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)

3236+ where

3237+ step gst (RollingMapInit st) = do

3238+ r <- step1 (adaptState gst) st

3239+ return $ case r of

3240+ Yield x s -> Skip $ RollingMapGo s x

3241+ Skip s -> Skip $ RollingMapInit s

3242+ Stop -> Stop

3243+

3244+ step gst (RollingMapGo s1 x1) = do

3245+ r <- step1 (adaptState gst) s1

3246+ case r of

3247+ Yield x s -> do

3248+ !res <- f x x1

3249+ return $ Yield res $ RollingMapGo s x

3250+ Skip s -> return $ Skip $ RollingMapGo s x1

3251+ Stop -> return $ Stop

3252+

3253+{-# INLINE rollingMap #-}

3254+rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b

3255+rollingMap f = rollingMapM (\x y -> return $ f x y)

3256+

3257+------------------------------------------------------------------------------

3258+-- Tapping/Distributing

3259+------------------------------------------------------------------------------

3260+

32243261{-# INLINE tap #-}

32253262tap :: Monad m => Fold m a b -> Stream m a -> Stream m a

32263263tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing

@@ -3242,6 +3279,35 @@ tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing

32423279 void $ extract acc

32433280 return $ Stop

32443281

3282+{-# INLINE_NORMAL pollCounts #-}

3283+pollCounts

3284+ :: MonadAsync m

3285+ => (Stream m Int -> Stream m Int)

3286+ -> Fold m Int b

3287+ -> Stream m a

3288+ -> Stream m a

3289+pollCounts transf fld (Stream step state) = Stream step' Nothing

3290+ where

3291+

3292+ {-# INLINE_LATE step' #-}

3293+ step' _ Nothing = do

3294+ countVar <- liftIO $ newVar (0 :: Int)

3295+ tid <- forkManaged

3296+ $ void $ runFold fld

3297+ $ transf $ fromPrimVar countVar

3298+ return $ Skip (Just (countVar, tid, state))

3299+

3300+ step' gst (Just (countVar, tid, st)) = do

3301+ r <- step gst st

3302+ case r of

3303+ Yield x s -> do

3304+ liftIO $ modifyVar' countVar (+ 1)

3305+ return $ Yield x (Just (countVar, tid, s))

3306+ Skip s -> return $ Skip (Just (countVar, tid, s))

3307+ Stop -> do

3308+ liftIO $ killThread tid

3309+ return Stop

3310+

32453311{-# INLINE_NORMAL tapRate #-}

32463312tapRate ::

32473313 (MonadAsync m, MonadCatch m)

src/Streamly/Internal/Prelude.hs+76−6

@@ -205,6 +205,8 @@ module Streamly.Internal.Prelude

205205 -- ** Mapping Filters

206206 , mapMaybe

207207 , mapMaybeM

208+ , rollingMapM

209+ , rollingMap

208210

209211 -- ** Scanning Filters

210212 , findIndices

@@ -219,6 +221,7 @@ module Streamly.Internal.Prelude

219221 , intersperseSuffixBySpan

220222 -- , intersperseBySpan

221223 , interjectSuffix

224+ , delayPost

222225

223226 -- ** Reordering

224227 , reverse

@@ -344,6 +347,7 @@ module Streamly.Internal.Prelude

344347 , tap

345348 , tapAsync

346349 , tapRate

350+ , pollCounts

347351

348352 -- * Windowed Classification

349353

@@ -467,7 +471,7 @@ import Streamly.Internal.Data.Pipe.Types (Pipe (..))

467471import Streamly.Internal.Data.Time.Units

468472 (AbsTime, MilliSecond64(..), addToAbsTime, diffAbsTime, toRelTime,

469473 toAbsTime)

470−import Streamly.Internal.Mutable.Prim.Var (MonadMut, Prim, Var)

474+import Streamly.Internal.Mutable.Prim.Var (Prim, Var)

471475

472476import Streamly.Internal.Data.Strict

473477

@@ -879,7 +883,7 @@ fromHandle h = go

879883-- /Internal/

880884--

881885{-# INLINE fromPrimVar #-}

882−fromPrimVar :: (IsStream t, MonadMut m, Prim a) => Var m a -> t m a

886+fromPrimVar :: (IsStream t, MonadIO m, Prim a) => Var IO a -> t m a

883887fromPrimVar = fromStreamD . D.fromPrimVar

884888

885889------------------------------------------------------------------------------

@@ -1786,6 +1790,27 @@ scan (Fold step begin done) = P.scanlMx' step begin done

17861790postscan :: (IsStream t, Monad m) => Fold m a b -> t m a -> t m b

17871791postscan (Fold step begin done) = P.postscanlMx' step begin done

17881792

1793+------------------------------------------------------------------------------

1794+-- Stateful Transformations

1795+------------------------------------------------------------------------------

1796+

1797+-- | Apply a function on every two successive elements of a stream. If the

1798+-- stream consists of a single element the output is an empty stream.

1799+--

1800+-- /Internal/

1801+--

1802+{-# INLINE rollingMap #-}

1803+rollingMap :: (IsStream t, Monad m) => (a -> a -> b) -> t m a -> t m b

1804+rollingMap f m = fromStreamD $ D.rollingMap f $ toStreamD m

1805+

1806+-- | Like 'rollingMap' but with an effectful map function.

1807+--

1808+-- /Internal/

1809+--

1810+{-# INLINE rollingMapM #-}

1811+rollingMapM :: (IsStream t, Monad m) => (a -> a -> m b) -> t m a -> t m b

1812+rollingMapM f m = fromStreamD $ D.rollingMapM f $ toStreamD m

1813+

17891814------------------------------------------------------------------------------

17901815-- Transformation by Filtering

17911816------------------------------------------------------------------------------

@@ -2004,10 +2029,9 @@ reverse' s = fromStreamD $ D.reverse' $ toStreamD s

20042029

20052030-- intersperseM = intersperseBySpan 1

20062031

2007−-- | Generate a stream by performing a monadic action between consecutive

2008−-- elements of the given stream.

2009−--

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

2032+-- | Generate a stream by inserting the result of a monadic action between

2033+-- consecutive elements of the given stream. Note that the monadic action is

2034+-- performed after the stream action before which its result is inserted.

20112035--

20122036-- @

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

@@ -2039,6 +2063,28 @@ intersperse a = fromStreamS . S.intersperse a . toStreamS

20392063intersperseSuffix :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a

20402064intersperseSuffix m = fromStreamD . D.intersperseSuffix m . toStreamD

20412065

2066+-- | Perform a side effect after each element of a stream. The output of the

2067+-- effectful action is discarded, therefore, the input stream remains

2068+-- unchanged.

2069+--

2070+-- @

2071+-- > S.mapM_ putChar $ S.intersperseSuffix_ (threadDelay 1000000) $ S.fromList "hello"

2072+-- @

2073+--

2074+-- /Internal/

2075+--

2076+{-# INLINE intersperseSuffix_ #-}

2077+intersperseSuffix_ :: (IsStream t, Monad m) => m b -> t m a -> t m a

2078+intersperseSuffix_ m = Serial.mapM (\x -> void m >> return x)

2079+

2080+-- | Introduces a delay of specified seconds after each element of a stream.

2081+--

2082+-- /Internal/

2083+--

2084+{-# INLINE delayPost #-}

2085+delayPost :: (IsStream t, MonadIO m) => Double -> t m a -> t m a

2086+delayPost n = intersperseSuffix_ $ liftIO $ threadDelay $ round $ n * 1000000

2087+

20422088-- | Like 'intersperseSuffix' but intersperses a monadic action into the input

20432089-- stream after every @n@ elements and after the last element.

20442090--

@@ -3614,6 +3660,30 @@ tap f xs = D.fromStreamD $ D.tap f (D.toStreamD xs)

36143660tapAsync :: (IsStream t, MonadAsync m) => FL.Fold m a b -> t m a -> t m a

36153661tapAsync f xs = D.fromStreamD $ D.tapAsync f (D.toStreamD xs)

36163662

3663+-- | Maintain the count of elements flowing in the stream and poll the count

3664+-- asynchronously from another thread. The count stream is transformed using

3665+-- the supplied transform and then folded using the supplied fold. The thread

3666+-- is automatically cleaned up if the stream stops or aborts due to exception.

3667+--

3668+-- For example, to print the count of elements processed every second:

3669+--

3670+-- @

3671+-- > S.drain $ S.pollCounts (rollingMap (-) . delayPost 1) (FL.drainBy print)

3672+-- $ S.enumerateFrom 0

3673+-- @

3674+--

3675+{-# INLINE pollCounts #-}

3676+pollCounts ::

3677+ (IsStream t, MonadAsync m)

3678+ => (t m Int -> t m Int)

3679+ -> Fold m Int b

3680+ -> t m a

3681+ -> t m a

3682+pollCounts transf f xs =

3683+ D.fromStreamD

3684+ $ D.pollCounts (D.toStreamD . transf . D.fromStreamD) f

3685+ $ (D.toStreamD xs)

3686+

36173687-- | Calls the supplied function with the number of elements consumed

36183688-- every @n@ seconds. The given function is run in a separate thread

36193689-- until the end of the stream. In case there is an exception in the