3e22f1b · streamly · git.iko.soy

git.iko.soy/streamly

flush sessions when classifySessionsOf ends

Earlier the sessions that were incomplete and in the hash were dropped when the stream ended.

commit3e22f1bae01e7e0aeb13e5a176a825cdb165d88c authorHarendra Kumar <harendra@composewell.com> dateFri 29 May 2020, 17:08 (6 years ago) parent07efcbd

2 files changed+81−2

compare with master

unifiedsplit

src/Streamly/Internal/Data/Stream/StreamD.hs+34 src/Streamly/Internal/Prelude.hs+47−2

src/Streamly/Internal/Data/Stream/StreamD.hs+34

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

226226

227227 -- ** By folding (scans)

228228 , scanlM'

229+ , scanlMAfter'

229230 , scanl'

230231 , scanlM

231232 , scanl

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

241242 , postscanlM

242243 , postscanl'

243244 , postscanlM'

245+ , postscanlMAfter'

244246

245247 , postscanlx'

246248 , postscanlMx'

@@ -3609,6 +3611,31 @@ postscanlM' fstep begin (Stream step state) =

36093611postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a

36103612postscanl' f = postscanlM' (\a b -> return (f a b))

36113613

3614+-- We can possibly have the "done" function as a Maybe to provide an option to

3615+-- emit or not emit the accumulator when the stream stops.

3616+--

3617+-- TBD: use a single Yield point

3618+--

3619+{-# INLINE_NORMAL postscanlMAfter' #-}

3620+postscanlMAfter' :: Monad m

3621+ => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

3622+postscanlMAfter' fstep initial done (Stream step1 state1) = do

3623+ Stream step (Just (state1, initial))

3624+

3625+ where

3626+

3627+ {-# INLINE_LATE step #-}

3628+ step gst (Just (st, acc)) = do

3629+ r <- step1 (adaptState gst) st

3630+ case r of

3631+ Yield x s -> do

3632+ old <- acc

3633+ y <- fstep old x

3634+ y `seq` return (Yield y (Just (s, return y)))

3635+ Skip s -> return $ Skip $ Just (s, acc)

3636+ Stop -> acc >>= done >>= \res -> return (Yield res Nothing)

3637+ step _ Nothing = return Stop

3638+

36123639{-# INLINE_NORMAL postscanlM #-}

36133640postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b

36143641postscanlM fstep begin (Stream step state) = Stream step' (state, begin)

@@ -3631,6 +3658,13 @@ postscanl f = postscanlM (\a b -> return (f a b))

36313658scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b

36323659scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)

36333660

3661+{-# INLINE scanlMAfter' #-}

3662+scanlMAfter' :: Monad m

3663+ => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b

3664+scanlMAfter' fstep initial done s =

3665+ (initial >>= \x -> x `seq` return x) `consM`

3666+ postscanlMAfter' fstep initial done s

3667+

36343668{-# INLINE scanl' #-}

36353669scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b

36363670scanl' f = scanlM' (\a b -> return (f a b))

src/Streamly/Internal/Prelude.hs+47−2

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

202202 -- ** Left scans

203203 , scanl'

204204 , scanlM'

205+ , scanlMAfter'

205206 , postscanl'

206207 , postscanlM'

207208 , prescanl'

@@ -2125,6 +2126,23 @@ scanx = P.scanlx'

21252126scanlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> t m b

21262127scanlM' step begin m = fromStreamD $ D.scanlM' step begin $ toStreamD m

21272128

2129+-- | @scanlMAfter' accumulate initial done stream@ is like 'scanlM'' except

2130+-- that it provides an additional @done@ function to be applied on the

2131+-- accumulator when the stream stops. The result of @done@ is also emitted in

2132+-- the stream.

2133+--

2134+-- This function can be used to allocate a resource in the beginning of the

2135+-- scan and release it when the stream ends or to flush the internal state of

2136+-- the scan at the end.

2137+--

2138+-- /Internal/

2139+--

2140+{-# INLINE scanlMAfter' #-}

2141+scanlMAfter' :: (IsStream t, Monad m)

2142+ => (b -> a -> m b) -> m b -> (b -> m b) -> t m a -> t m b

2143+scanlMAfter' step initial done stream =

2144+ fromStreamD $ D.scanlMAfter' step initial done $ toStreamD stream

2145+

21282146-- | Strict left scan. Like 'map', 'scanl'' too is a one to one transformation,

21292147-- however it adds an extra element.

21302148--

@@ -4507,8 +4525,8 @@ classifySessionsBy

45074525 -> t m (k, b) -- ^ session key, fold result

45084526classifySessionsBy tick tmout reset ejectPred

45094527 (Fold step initial extract) str =

4510− concatMap sessionOutputStream

4511− $ scanlM' sstep szero stream

4528+ concatMap sessionOutputStream $

4529+ scanlMAfter' sstep (return szero) flush stream

45124530

45134531 where

45144532

@@ -4613,6 +4631,21 @@ classifySessionsBy tick tmout reset ejectPred

46134631 let curTime = addToAbsTime sessionCurTime tickMs

46144632 in ejectExpired sessionState curTime

46154633

4634+ flush session@SessionState{..} = do

4635+ (hp', mp', out, count) <-

4636+ ejectAll

4637+ ( sessionTimerHeap

4638+ , sessionKeyValueMap

4639+ , K.nil

4640+ , sessionCount

4641+ )

4642+ return $ session

4643+ { sessionCount = count

4644+ , sessionTimerHeap = hp'

4645+ , sessionKeyValueMap = mp'

4646+ , sessionOutputStream = out

4647+ }

4648+

46164649 fromEither e =

46174650 case e of

46184651 Left x -> x

@@ -4625,6 +4658,18 @@ classifySessionsBy tick tmout reset ejectPred

46254658 let mp1 = Map.delete key mp

46264659 return (hp, mp1, out1, cnt - 1)

46274660

4661+ ejectAll (hp, mp, out, !cnt) = do

4662+ let hres = H.uncons hp

4663+ case hres of

4664+ Just (Entry _ key, hp1) -> do

4665+ r <- case Map.lookup key mp of

4666+ Nothing -> return (hp1, mp, out, cnt)

4667+ Just (Tuple' _ acc) -> ejectEntry hp1 mp out cnt acc key

4668+ ejectAll r

4669+ Nothing -> do

4670+ assert (Map.null mp) (return ())

4671+ return (hp, mp, out, cnt)

4672+

46284673 ejectOne (hp, mp, out, !cnt) = do

46294674 let hres = H.uncons hp

46304675 case hres of