@@ -347,6 +347,7 @@ import qualified Prelude
347347import Streamly.Internal.Mutable.Prim.Var
348348 (Prim, Var, readVar, newVar, modifyVar')
349349import Streamly.Internal.Data.Time.Units
350+ (TimeUnit64, toRelTime64, diffAbsTime64)
350351
351352import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
352353import Streamly.Internal.Memory.Array.Types (Array(..))
@@ -4045,51 +4046,80 @@ lastN n = Fold step initial done
40454046-- Time related
40464047------------------------------------------------------------------------------
40474048
4048−-- XXX This does not fuse. Removing the if function results in fusion though.
4049−-- XXX This function will fuse with the use of fusion-plugin.
4050−-- XXX This function also fuses if compiled with -funfolding-use-threshold=500
4049+-- XXX using getTime in the loop can be pretty expensive especially for
4050+-- computations where iterations are lightweight. We have the following
4051+-- options:
4052+--
4053+-- 1) Run a timeout thread updating a flag asynchronously and check that
4054+-- flag here, that way we can have a cheap termination check.
4055+--
4056+-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4057+--
4058+-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4059+-- termination value of rdtsc in the beginning and then in each iteration just
4060+-- get rdtsc and check if we should terminate.
4061+--
4062+data TakeByTime st s
4063+ = TakeByTimeInit st
4064+ | TakeByTimeCheck st s
4065+ | TakeByTimeYield st s
4066+
40514067{-# INLINE_NORMAL takeByTime #-}
40524068takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4053−takeByTime lim' (Stream step state) = Stream step' (state, Nothing)
4054− where
4055− lim = toRelTime64 lim'
4056− {-# INLINE_LATE step' #-}
4057− step' _ (s, Nothing) = do
4069+takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4070+ where
4071+
4072+ lim = toRelTime64 duration
4073+
4074+ {-# INLINE_LATE step #-}
4075+ step _ (TakeByTimeInit _) | lim == 0 = return Stop
4076+ step _ (TakeByTimeInit st) = do
4077+ t0 <- liftIO $ getTime Monotonic
4078+ return $ Skip (TakeByTimeYield st t0)
4079+ step _ (TakeByTimeCheck st t0) = do
40584080 t <- liftIO $ getTime Monotonic
4059− return $ Skip (s, Just t)
4060− step' gst (s, jt@(Just t)) = do
4061− ns <- step gst s
4062− t' <- liftIO $ getTime Monotonic
4063− if diffAbsTime64 t' t > lim
4064− then return Stop
4065− else return $ case ns of
4066− Yield a s' -> Yield a (s', jt)
4067− Skip s' -> Skip (s', jt)
4068− Stop -> Stop
4069−
4070−-- XXX This does not fuse. Removing the if function results in fusion though.
4071−-- XXX This function will fuse with the use of fusion-plugin.
4072−-- XXX This function also fuses if compiled with -funfolding-use-threshold=500
4081+ return $
4082+ if diffAbsTime64 t t0 > lim
4083+ then Stop
4084+ else Skip (TakeByTimeYield st t0)
4085+ step gst (TakeByTimeYield st t0) = do
4086+ r <- step1 gst st
4087+ return $ case r of
4088+ Yield x s -> Yield x (TakeByTimeCheck s t0)
4089+ Skip s -> Skip (TakeByTimeCheck s t0)
4090+ Stop -> Stop
4091+
4092+data DropByTime st s x
4093+ = DropByTimeInit st
4094+ | DropByTimeGen st s
4095+ | DropByTimeCheck st s x
4096+ | DropByTimeYield st
4097+
40734098{-# INLINE_NORMAL dropByTime #-}
40744099dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4075−dropByTime lim' (Stream step state) = Stream step' (state, Nothing)
4076− where
4077− lim = toRelTime64 lim'
4078− {-# INLINE_LATE step' #-}
4079− step' _ (s, Nothing) = do
4100+dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4101+ where
4102+
4103+ lim = toRelTime64 duration
4104+
4105+ {-# INLINE_LATE step #-}
4106+ step _ (DropByTimeInit st) = do
4107+ t0 <- liftIO $ getTime Monotonic
4108+ return $ Skip (DropByTimeGen st t0)
4109+ step gst (DropByTimeGen st t0) = do
4110+ r <- step1 gst st
4111+ return $ case r of
4112+ Yield x s -> Skip (DropByTimeCheck s t0 x)
4113+ Skip s -> Skip (DropByTimeGen s t0)
4114+ Stop -> Stop
4115+ step _ (DropByTimeCheck st t0 x) = do
40804116 t <- liftIO $ getTime Monotonic
4081− return $ Skip (s, Just t)
4082− step' gst (s, jt@(Just t)) = do
4083− t' <- liftIO $ getTime Monotonic
4084− ns <- step gst s
4085− if diffAbsTime64 t' t <= lim
4086− then return $
4087− case ns of
4088− Yield _ s' -> Skip (s', jt)
4089− Skip s' -> Skip (s', jt)
4090− Stop -> Stop
4091− else return $
4092− case ns of
4093− Yield a s' -> Yield a (s', jt)
4094− Skip s' -> Skip (s', jt)
4095− Stop -> Stop
4117+ if diffAbsTime64 t t0 <= lim
4118+ then return $ Skip $ DropByTimeGen st t0
4119+ else return $ Yield x $ DropByTimeYield st
4120+ step gst (DropByTimeYield st) = do
4121+ r <- step1 gst st
4122+ return $ case r of
4123+ Yield x s -> Yield x (DropByTimeYield s)
4124+ Skip s -> Skip (DropByTimeYield s)
4125+ Stop -> Stop