d582037 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add tapRate combinator.

d582037Pranay Sashank7 years ago

benchmark/Linear.hs+1

compare with master

unifiedsplit

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

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

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

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

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

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

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

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

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

324 ]

325 ]

benchmark/Streamly/Benchmark/Prelude.hs+7

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

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

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

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

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

31import Data.Maybe (fromJust)

32import Data.Maybe (fromJust)

32import GHC.Generics (Generic)

33import GHC.Generics (Generic)

33import Prelude

34import Prelude

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

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

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

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

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

406

407

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+

407{-# INLINE tapAsyncS #-}

414{-# INLINE tapAsyncS #-}

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

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

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

416tapAsyncS 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

126 -- ** Specialized Folds

127 , tap

128 , tapAsync

129+ , tapRate

129 , drain

130 , drain

130 , null

131 , null

131 , head

132 , head

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

300 )

301 )

301where

302where

302

303

303−import Control.Concurrent (myThreadId, takeMVar)

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

305+import Control.Exception

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

306+ (Exception, SomeException, AsyncException, fromException)

305import Control.Monad (void, when)

307import Control.Monad (void, when)

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

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

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

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

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

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

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

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

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

312import Data.Functor.Identity (Identity(..))

314import Data.Functor.Identity (Identity(..))

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

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

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

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

315import Data.Word (Word32)

317import Data.Word (Word32)

316import Foreign.Ptr (Ptr)

318import Foreign.Ptr (Ptr)

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

3218 void $ extract acc

3220 void $ extract acc

3219 return $ Stop

3221 return $ Stop

3220

3222

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+

3221-------------------------------------------------------------------------------

3268-------------------------------------------------------------------------------

3222-- Filtering

3269-- Filtering

3223-------------------------------------------------------------------------------

3270-------------------------------------------------------------------------------

src/Streamly/Internal/Prelude.hs+24

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

342 , trace

343 , tap

344 , tapAsync

345+ , tapRate

345

346

346 -- * Windowed Classification

347 -- * Windowed Classification

347

348

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

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

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

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

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

3604

3605

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+

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

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

3606-- discard the results.

3630-- discard the results.

3607--

3631--

src/Streamly/Prelude.hs+1

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

285 , mapM_

286 , trace

287 , tap

288+ , tapRate

288

289

289 -- ** Scanning

290 -- ** Scanning

290 --

291 --