d582037 · streamly · git.iko.soy

git.iko.soy/streamly

Add tapRate combinator.

- Also includes benchmarks for it.

commitd5820372506430d5cb8a853e8d8f71ac070fafbf authorPranay Sashank <pranaysashank@gmail.com> dateFri 13 Dec 2019, 14:44 (7 years ago) parent0d7dad8

5 files changed+83−3

compare with master

unifiedsplit

benchmark/Linear.hs+1 benchmark/Streamly/Benchmark/Prelude.hs+7 src/Streamly/Internal/Data/Stream/StreamD.hs+50−3 src/Streamly/Internal/Prelude.hs+24 src/Streamly/Prelude.hs+1

benchmark/Linear.hs+1

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

319319 , benchIOSink "foldrT" (Ops.foldrT 1)

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

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

322+ , benchIOSink "tapRate 1 second" (Ops.tapRate 1)

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

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

324325 ]

benchmark/Streamly/Benchmark/Prelude.hs+7

@@ -28,6 +28,7 @@ import Control.Monad (when)

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

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

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

31+import Data.IORef (newIORef, modifyIORef')

3132import Data.Maybe (fromJust)

3233import GHC.Generics (Generic)

3334import Prelude

@@ -404,6 +405,12 @@ mapM t n = composeN' n $ t . S.mapM return

404405tap :: MonadIO m => Int -> Stream m Int -> m ()

405406tap n = composeN n $ S.tap FL.sum

406407

408+{-# INLINE tapRate #-}

409+tapRate :: Int -> Stream IO Int -> IO ()

410+tapRate n str = do

411+ cref <- newIORef 0

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

413+

407414{-# INLINE tapAsyncS #-}

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

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

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

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

126126 -- ** Specialized Folds

127127 , tap

128128 , tapAsync

129+ , tapRate

129130 , drain

130131 , null

131132 , head

@@ -300,8 +301,9 @@ module Streamly.Internal.Data.Stream.StreamD

300301 )

301302where

302303

303−import Control.Concurrent (myThreadId, takeMVar)

304−import Control.Exception (Exception, SomeException, fromException)

304+import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)

305+import Control.Exception

306+ (Exception, SomeException, AsyncException, fromException)

305307import Control.Monad (void, when)

306308import Control.Monad.Catch (MonadCatch, throwM)

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

@@ -310,7 +312,7 @@ import Control.Monad.State.Strict (StateT)

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

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

312314import Data.Functor.Identity (Identity(..))

313−import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)

315+import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, modifyIORef')

314316import Data.Maybe (fromJust, isJust, isNothing)

315317import Data.Word (Word32)

316318import Foreign.Ptr (Ptr)

@@ -3218,6 +3220,51 @@ tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing

32183220 void $ extract acc

32193221 return $ Stop

32203222

3223+{-# INLINE_NORMAL tapRate #-}

3224+tapRate ::

3225+ (MonadAsync m, MonadCatch m)

3226+ => Double

3227+ -> (Int -> m b)

3228+ -> Stream m a

3229+ -> Stream m a

3230+tapRate samplingRate action (Stream step state) = Stream step' Nothing

3231+ where

3232+ {-# NOINLINE loop #-}

3233+ loop countRef prev = do

3234+ i <-

3235+ MC.catch

3236+ (do liftIO $ threadDelay (round $ samplingRate * 1000000)

3237+ i <- liftIO $ readIORef countRef

3238+ let !diff = i - prev

3239+ void $ action diff

3240+ return i)

3241+ (\(e :: AsyncException) -> do

3242+ i <- liftIO $ readIORef countRef

3243+ let !diff = i - prev

3244+ void $ action diff

3245+ throwM (MC.toException e))

3246+ loop countRef i

3247+

3248+ {-# INLINE_LATE step' #-}

3249+ step' _ Nothing = do

3250+ countRef <- liftIO $ newIORef 0

3251+ tid <- fork $ loop countRef 0

3252+ ref <- liftIO $ newIORef ()

3253+ _ <- liftIO $ mkWeakIORef ref (killThread tid)

3254+ return $ Skip (Just (countRef, tid, state, ref))

3255+

3256+ step' gst (Just (countRef, tid, st, ref)) = do

3257+ r <- step gst st

3258+ case r of

3259+ Yield x s -> do

3260+ liftIO $ modifyIORef' countRef (+ 1)

3261+ return $ Yield x (Just (countRef, tid, s, ref))

3262+ Skip s -> return $ Skip (Just (countRef, tid, s, ref))

3263+ Stop -> do

3264+ liftIO $ killThread tid

3265+ return Stop

3266+

3267+

32213268-------------------------------------------------------------------------------

32223269-- Filtering

32233270-------------------------------------------------------------------------------

src/Streamly/Internal/Prelude.hs+24

@@ -342,6 +342,7 @@ module Streamly.Internal.Prelude

342342 , trace

343343 , tap

344344 , tapAsync

345+ , tapRate

345346

346347 -- * Windowed Classification

347348

@@ -3602,6 +3603,29 @@ tap f xs = D.fromStreamD $ D.tap f (D.toStreamD xs)

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

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

36043605

3606+-- | Calls the supplied function with the number of elements consumed

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

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

3609+-- stream the thread is killed during the next major GC.

3610+--

3611+-- Note: The action is not guaranteed to run if the main thread exits.

3612+--

3613+-- @

3614+-- > delay n = threadDelay (round $ n * 1000000) >> return n

3615+-- > S.drain $ S.tapRate 2 (\\n -> print $ show n ++ " elements processed") (delay 1 S.|: delay 0.5 S.|: delay 0.5 S.|: S.nil)

3616+-- 2 elements processed

3617+-- 1 elements processed

3618+-- @

3619+--

3620+{-# INLINE tapRate #-}

3621+tapRate ::

3622+ (IsStream t, MonadAsync m, MonadCatch m)

3623+ => Double

3624+ -> (Int -> m b)

3625+ -> t m a

3626+ -> t m a

3627+tapRate n f xs = D.fromStreamD $ D.tapRate n f $ (D.toStreamD xs)

3628+

36053629-- | Apply a monadic function to each element flowing through the stream and

36063630-- discard the results.

36073631--

src/Streamly/Prelude.hs+1

@@ -285,6 +285,7 @@ module Streamly.Prelude

285285 , mapM_

286286 , trace

287287 , tap

288+ , tapRate

288289

289290 -- ** Scanning

290291 --