@@ -174,7 +174,6 @@ import Streamly.Internal.Data.Stream.IsStream.Type
174174import Streamly.Internal.Data.Time.Units
175175 ( AbsTime, MilliSecond64(..), addToAbsTime, toRelTime
176176 , toAbsTime)
177−import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
178177
179178import qualified Data.Heap as H
180179import qualified Data.Map.Strict as Map
@@ -1173,6 +1172,8 @@ data SessionState t m k a b = SessionState
11731172 , sessionOutputStream :: t (m :: Type -> Type) (k, b) -- ^ Completed sessions
11741173 }
11751174
1175+data SessionEntry a b = LiveSession !a !b | ZombieSession
1176+
11761177-- | @classifySessionsBy tick keepalive predicate timeout fold stream@
11771178-- classifies an input event @stream@ consisting of @(timestamp, (key,
11781179-- value))@ into sessions based on the @key@, folding all the values
@@ -1252,15 +1253,34 @@ classifySessionsBy tick reset ejectPred tmout
12521253 -- XXX we should use a heap in pinned memory to scale it to a large
12531254 -- size
12541255 --
1255− -- To detect session inactivity we keep a timestamp of the latest event
1256− -- in the Map along with the fold result. When we purge the session
1257− -- from the heap we match the timestamp in the heap with the timestamp
1258− -- in the Map, if the latest timestamp is newer and has not expired we
1259− -- reinsert the key in the heap.
1260− --
12611256 -- XXX if the key is an Int, we can also use an IntMap for slightly
12621257 -- better performance.
12631258 --
1259+ -- When we insert a key in the Map we insert an entry into the heap as
1260+ -- well with the session expiry as the sort key. The Map entry
1261+ -- consists of the fold result, and the expiry time of the session. If
1262+ -- "reset" is True the expiry time is readjusted whenever a new event
1263+ -- is processed. If the fold terminates and a new session is started
1264+ -- for the same key the expiry time is set to the first timestamp of
1265+ -- the new session.
1266+ --
1267+ -- The heap must have at most one entry for any given key. The heap is
1268+ -- processed periodically to remove the expired entries. We pick up an
1269+ -- expired entry from the top of the heap and if the session has
1270+ -- expired based on the expiry time in the Map entry then we remove the
1271+ -- session from the Map and yield its fold output. Otherwise, we
1272+ -- reinsert the entry into the heap based on the current expiry in the
1273+ -- Map entry.
1274+ --
1275+ -- If an entry is removed from the Map and not removed from the heap
1276+ -- and in the meantime it is inserted again in the Map (using the same
1277+ -- key) then how do we avoid inserting multiple entries in the heap?
1278+ -- For this reason we maintain the invariant that the Map entry is
1279+ -- removed only when the heap entry is removed. Even if the fold has
1280+ -- finished we still keep a dummy Map entry (ZombieSession) until the
1281+ -- heap entry is removed. That way if we have a Map entry we do not
1282+ -- insert a heap entry because we know it is already there.
1283+ --
12641284 let curTime = max sessionEventTime timestamp
12651285 mOld = Map.lookup key sessionKeyValueMap
12661286 let done fb = do
@@ -1271,9 +1291,11 @@ classifySessionsBy tick reset ejectPred tmout
12711291 -- should not be expensive.
12721292 --
12731293 let (mp, cnt) = case mOld of
1274− Nothing -> (sessionKeyValueMap, sessionCount)
1275− Just _ -> (Map.delete key sessionKeyValueMap
1276− , sessionCount - 1)
1294+ Just (LiveSession _ _) ->
1295+ ( Map.insert key ZombieSession sessionKeyValueMap
1296+ , sessionCount - 1
1297+ )
1298+ _ -> (sessionKeyValueMap, sessionCount)
12771299 return $ session
12781300 { sessionCurTime = curTime
12791301 , sessionEventTime = curTime
@@ -1282,14 +1304,14 @@ classifySessionsBy tick reset ejectPred tmout
12821304 , sessionOutputStream = fromPure (key, fb)
12831305 }
12841306 partial fs1 = do
1285− let acc = Tuple' timestamp fs1
1307+ let expiry = addToAbsTime timestamp timeoutMs
12861308 (hp1, mp1, out1, cnt1) <- do
12871309 let vars = (sessionTimerHeap, sessionKeyValueMap,
12881310 IsStream.nil, sessionCount)
12891311 case mOld of
12901312 -- inserting new entry
12911313 Nothing -> do
1292− -- Eject a session from heap and map is needed
1314+ -- Eject a session from heap and map if needed
12931315 eject <- ejectPred sessionCount
12941316 (hp, mp, out, cnt) <-
12951317 if eject
@@ -1297,13 +1319,13 @@ classifySessionsBy tick reset ejectPred tmout
12971319 else return vars
12981320
12991321 -- Insert the new session in heap
1300− let expiry = addToAbsTime timestamp timeoutMs
1301− hp' = H.insert (Entry expiry key) hp
1322+ let hp' = H.insert (Entry expiry key) hp
13021323 in return (hp', mp, out, cnt + 1)
13031324 -- updating old entry
13041325 Just _ -> return vars
13051326
1306− let mp2 = Map.insert key acc mp1
1327+ let acc = LiveSession expiry fs1
1328+ mp2 = Map.insert key acc mp1
13071329 return $ SessionState
13081330 { sessionCurTime = curTime
13091331 , sessionEventTime = curTime
@@ -1314,10 +1336,12 @@ classifySessionsBy tick reset ejectPred tmout
13141336 }
13151337 res0 <- do
13161338 case mOld of
1317− Nothing -> initial
1318− Just (Tuple' _ acc) -> return $ FL.Partial acc
1339+ Just (LiveSession _ acc) -> return $ FL.Partial acc
1340+ _ -> initial
13191341 case res0 of
1320− FL.Done fb -> done fb
1342+ FL.Done _ ->
1343+ error $ "classifySessionsBy: "
1344+ ++ "The supplied fold must consume at least one input"
13211345 FL.Partial fs -> do
13221346 res <- step fs value
13231347 case res of
@@ -1357,7 +1381,10 @@ classifySessionsBy tick reset ejectPred tmout
13571381 Just (Entry _ key, hp1) -> do
13581382 r <- case Map.lookup key mp of
13591383 Nothing -> return (hp1, mp, out, cnt)
1360− Just (Tuple' _ acc) -> ejectEntry hp1 mp out cnt acc key
1384+ Just ZombieSession ->
1385+ return (hp1, Map.delete key mp, out, cnt)
1386+ Just (LiveSession _ acc) ->
1387+ ejectEntry hp1 mp out cnt acc key
13611388 ejectAll r
13621389 Nothing -> do
13631390 assert (Map.null mp) (return ())
@@ -1369,8 +1396,9 @@ classifySessionsBy tick reset ejectPred tmout
13691396 Just (Entry expiry key, hp1) ->
13701397 case Map.lookup key mp of
13711398 Nothing -> ejectOne (hp1, mp, out, cnt)
1372− Just (Tuple' latestTS acc) -> do
1373− let expiry1 = addToAbsTime latestTS timeoutMs
1399+ Just ZombieSession ->
1400+ ejectOne (hp1, Map.delete key mp, out, cnt)
1401+ Just (LiveSession expiry1 acc) -> do
13741402 if not reset || expiry1 <= expiry
13751403 then ejectEntry hp1 mp out cnt acc key
13761404 else
@@ -1409,8 +1437,9 @@ classifySessionsBy tick reset ejectPred tmout
14091437 then
14101438 case Map.lookup key mp of
14111439 Nothing -> ejectLoop hp1 mp out cnt
1412− Just (Tuple' latestTS acc) -> do
1413− let expiry1 = addToAbsTime latestTS timeoutMs
1440+ Just ZombieSession ->
1441+ ejectLoop hp1 (Map.delete key mp) out cnt
1442+ Just (LiveSession expiry1 acc) -> do
14141443 if expiry1 <= curTime || not reset || force
14151444 then do
14161445 (hp2,mp1,out1,cnt1) <-