5073377 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add modular stream stats counting

5073377Harendra Kumar7 years ago

benchmark/Linear.hs+1

compare with master

unifiedsplit

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

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

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

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

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

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

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

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

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

328 ]

329 ]

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

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

26import Control.DeepSeq (NFData)

27import Control.Monad (when)

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

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

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

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

31import Data.IORef (newIORef, modifyIORef')

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

411 cref <- newIORef 0

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

413

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+

414{-# INLINE tapAsyncS #-}

420{-# INLINE tapAsyncS #-}

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

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

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

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

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

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

108 , modifyThread

109 , doFork

110 , fork

111+ , forkManaged

111

112

112 , toStreamVar

113 , toStreamVar

113 , SVarStats (..)

114 , SVarStats (..)

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

118where

119where

119

120

120import Control.Concurrent

121import Control.Concurrent

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

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

122import Control.Concurrent.MVar

123import Control.Concurrent.MVar

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

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

124 tryReadMVar)

125 tryReadMVar)

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

151import GHC.Conc (ThreadId(..))

152import GHC.Conc (ThreadId(..))

152import GHC.Exts

153import GHC.Exts

153import GHC.IO (IO(..))

154import GHC.IO (IO(..))

155+import System.IO (hPutStrLn, stderr)

156+import System.Mem.Weak (addFinalizer)

157+

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

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

155import Streamly.Internal.Data.Time.Units

159import Streamly.Internal.Data.Time.Units

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

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

157 fromRelTime64, toRelTime64, showNanoSecond64, showRelTime64)

161 fromRelTime64, toRelTime64, showNanoSecond64, showRelTime64)

158−import System.IO (hPutStrLn, stderr)

159

162

160import qualified Data.Heap as H

163import qualified Data.Heap as H

161import qualified Data.Set as S

164import qualified Data.Set as S

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

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

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

959fork = liftBaseDiscard forkIO

962fork = liftBaseDiscard forkIO

960

963

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+

961------------------------------------------------------------------------------

974------------------------------------------------------------------------------

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

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

963------------------------------------------------------------------------------

976------------------------------------------------------------------------------

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

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

130 , tap

131 , tapAsync

132 , tapRate

133+ , pollCounts

133 , drain

134 , drain

134 , null

135 , null

135 , head

136 , head

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

254 , map

255 , map

255 , mapM

256 , mapM

256 , sequence

257 , sequence

258+ , rollingMap

259+ , rollingMapM

257

260

258 -- * Inserting

261 -- * Inserting

259 , intersperseM

262 , intersperseM

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

335import qualified Prelude

338import qualified Prelude

336

339

337import Streamly.Internal.Mutable.Prim.Var

340import Streamly.Internal.Mutable.Prim.Var

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

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

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

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

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

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

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

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

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

668toStreamD = fromStreamK . K.toStream

671toStreamD = fromStreamK . K.toStream

669

672

670{-# INLINE_NORMAL fromPrimVar #-}

673{-# 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

672fromPrimVar var = Stream step ()

675fromPrimVar var = Stream step ()

673 where

676 where

674 {-# INLINE_LATE step #-}

677 {-# INLINE_LATE step #-}

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

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

676

679

677-------------------------------------------------------------------------------

680-------------------------------------------------------------------------------

678-- Generation from SVar

681-- Generation from SVar

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

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

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

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

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

3223

3226

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+

3224{-# INLINE tap #-}

3261{-# INLINE tap #-}

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

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

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

3263tap (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

3242 void $ extract acc

3279 void $ extract acc

3243 return $ Stop

3280 return $ Stop

3244

3281

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+

3245{-# INLINE_NORMAL tapRate #-}

3311{-# INLINE_NORMAL tapRate #-}

3246tapRate ::

3312tapRate ::

3247 (MonadAsync m, MonadCatch m)

3313 (MonadAsync m, MonadCatch m)

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

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

205 -- ** Mapping Filters

206 , mapMaybe

207 , mapMaybeM

208+ , rollingMapM

209+ , rollingMap

208

210

209 -- ** Scanning Filters

211 -- ** Scanning Filters

210 , findIndices

212 , findIndices

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

219 , intersperseSuffixBySpan

221 , intersperseSuffixBySpan

220 -- , intersperseBySpan

222 -- , intersperseBySpan

221 , interjectSuffix

223 , interjectSuffix

224+ , delayPost

222

225

223 -- ** Reordering

226 -- ** Reordering

224 , reverse

227 , reverse

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

344 , tap

347 , tap

345 , tapAsync

348 , tapAsync

346 , tapRate

349 , tapRate

350+ , pollCounts

347

351

348 -- * Windowed Classification

352 -- * Windowed Classification

349

353

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

467import Streamly.Internal.Data.Time.Units

471import Streamly.Internal.Data.Time.Units

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

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

469 toAbsTime)

473 toAbsTime)

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

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

471

475

472import Streamly.Internal.Data.Strict

476import Streamly.Internal.Data.Strict

473

477

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

879-- /Internal/

883-- /Internal/

880--

884--

881{-# INLINE fromPrimVar #-}

885{-# 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

883fromPrimVar = fromStreamD . D.fromPrimVar

887fromPrimVar = fromStreamD . D.fromPrimVar

884

888

885------------------------------------------------------------------------------

889------------------------------------------------------------------------------

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

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

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

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

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

1788

1792

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+

1789------------------------------------------------------------------------------

1814------------------------------------------------------------------------------

1790-- Transformation by Filtering

1815-- Transformation by Filtering

1791------------------------------------------------------------------------------

1816------------------------------------------------------------------------------

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

2004

2029

2005-- intersperseM = intersperseBySpan 1

2030-- intersperseM = intersperseBySpan 1

2006

2031

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

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

2008−-- elements of the given stream.

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

2009−--

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

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

2011--

2035--

2012-- @

2036-- @

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

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

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

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

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

2040intersperseSuffix m = fromStreamD . D.intersperseSuffix m . toStreamD

2064intersperseSuffix m = fromStreamD . D.intersperseSuffix m . toStreamD

2041

2065

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+

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

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

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

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

2044--

2090--

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

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

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

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

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

3616

3662

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+

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

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

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

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

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

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