@@ -174,7 +174,6 @@ import Streamly.Internal.Data.Stream.IsStream.Type
174import Streamly.Internal.Data.Time.Units
174import Streamly.Internal.Data.Time.Units
175 ( AbsTime, MilliSecond64(..), addToAbsTime, toRelTime
175 ( AbsTime, MilliSecond64(..), addToAbsTime, toRelTime
176 , toAbsTime)
176 , toAbsTime)
177−import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
179import qualified Data.Heap as H
178import qualified Data.Heap as H
180import qualified Data.Map.Strict as Map
179import qualified Data.Map.Strict as Map
@@ -1173,6 +1172,8 @@ data SessionState t m k a b = SessionState
1173 , sessionOutputStream :: t (m :: Type -> Type) (k, b) -- ^ Completed sessions
1172 , sessionOutputStream :: t (m :: Type -> Type) (k, b) -- ^ Completed sessions
1175+data SessionEntry a b = LiveSession !a !b | ZombieSession
1176-- | @classifySessionsBy tick keepalive predicate timeout fold stream@
1177-- | @classifySessionsBy tick keepalive predicate timeout fold stream@
1177-- classifies an input event @stream@ consisting of @(timestamp, (key,
1178-- classifies an input event @stream@ consisting of @(timestamp, (key,
1178-- value))@ into sessions based on the @key@, folding all the values
1179-- value))@ into sessions based on the @key@, folding all the values
@@ -1252,15 +1253,34 @@ classifySessionsBy tick reset ejectPred tmout
1252 -- XXX we should use a heap in pinned memory to scale it to a large
1253 -- XXX we should use a heap in pinned memory to scale it to a large
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.
1261 -- XXX if the key is an Int, we can also use an IntMap for slightly
1256 -- XXX if the key is an Int, we can also use an IntMap for slightly
1262 -- better performance.
1257 -- better performance.
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.
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
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.
1264 let curTime = max sessionEventTime timestamp
1284 let curTime = max sessionEventTime timestamp
1265 mOld = Map.lookup key sessionKeyValueMap
1285 mOld = Map.lookup key sessionKeyValueMap
1266 let done fb = do
1286 let done fb = do
@@ -1271,9 +1291,11 @@ classifySessionsBy tick reset ejectPred tmout
1271 -- should not be expensive.
1291 -- should not be expensive.
1273 let (mp, cnt) = case mOld of
1293 let (mp, cnt) = case mOld of
1274− Nothing -> (sessionKeyValueMap, sessionCount)
1294+ Just (LiveSession _ _) ->
1275− Just _ -> (Map.delete key sessionKeyValueMap
1295+ ( Map.insert key ZombieSession sessionKeyValueMap
1276− , sessionCount - 1)
1296+ , sessionCount - 1
1298+ _ -> (sessionKeyValueMap, sessionCount)
1277 return $ session
1299 return $ session
1278 { sessionCurTime = curTime
1300 { sessionCurTime = curTime
1279 , sessionEventTime = curTime
1301 , sessionEventTime = curTime
@@ -1282,14 +1304,14 @@ classifySessionsBy tick reset ejectPred tmout
1282 , sessionOutputStream = fromPure (key, fb)
1304 , sessionOutputStream = fromPure (key, fb)
1284 partial fs1 = do
1306 partial fs1 = do
1285− let acc = Tuple' timestamp fs1
1307+ let expiry = addToAbsTime timestamp timeoutMs
1286 (hp1, mp1, out1, cnt1) <- do
1308 (hp1, mp1, out1, cnt1) <- do
1287 let vars = (sessionTimerHeap, sessionKeyValueMap,
1309 let vars = (sessionTimerHeap, sessionKeyValueMap,
1288 IsStream.nil, sessionCount)
1310 IsStream.nil, sessionCount)
1289 case mOld of
1311 case mOld of
1290 -- inserting new entry
1312 -- inserting new entry
1291 Nothing -> do
1313 Nothing -> do
1292− -- Eject a session from heap and map is needed
1314+ -- Eject a session from heap and map if needed
1293 eject <- ejectPred sessionCount
1315 eject <- ejectPred sessionCount
1294 (hp, mp, out, cnt) <-
1316 (hp, mp, out, cnt) <-
1295 if eject
1317 if eject
@@ -1297,13 +1319,13 @@ classifySessionsBy tick reset ejectPred tmout
1297 else return vars
1319 else return vars
1299 -- Insert the new session in heap
1321 -- 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
1302 in return (hp', mp, out, cnt + 1)
1323 in return (hp', mp, out, cnt + 1)
1303 -- updating old entry
1324 -- updating old entry
1304 Just _ -> return vars
1325 Just _ -> return vars
1327+ let acc = LiveSession expiry fs1
1306− let mp2 = Map.insert key acc mp1
1328+ mp2 = Map.insert key acc mp1
1307 return $ SessionState
1329 return $ SessionState
1308 { sessionCurTime = curTime
1330 { sessionCurTime = curTime
1309 , sessionEventTime = curTime
1331 , sessionEventTime = curTime
@@ -1314,10 +1336,12 @@ classifySessionsBy tick reset ejectPred tmout
1315 res0 <- do
1337 res0 <- do
1316 case mOld of
1338 case mOld of
1318− Just (Tuple' _ acc) -> return $ FL.Partial acc
1339+ Just (LiveSession _ acc) -> return $ FL.Partial acc
1319 case res0 of
1341 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"
1321 FL.Partial fs -> do
1345 FL.Partial fs -> do
1322 res <- step fs value
1346 res <- step fs value
1323 case res of
1347 case res of
@@ -1357,7 +1381,10 @@ classifySessionsBy tick reset ejectPred tmout
1357 Just (Entry _ key, hp1) -> do
1381 Just (Entry _ key, hp1) -> do
1358 r <- case Map.lookup key mp of
1382 r <- case Map.lookup key mp of
1359 Nothing -> return (hp1, mp, out, cnt)
1383 Nothing -> return (hp1, mp, out, cnt)
1384+ Just ZombieSession ->
1385+ return (hp1, Map.delete key mp, out, cnt)
1386+ Just (LiveSession _ acc) ->
1360− Just (Tuple' _ acc) -> ejectEntry hp1 mp out cnt acc key
1387+ ejectEntry hp1 mp out cnt acc key
1361 ejectAll r
1388 ejectAll r
1362 Nothing -> do
1389 Nothing -> do
1363 assert (Map.null mp) (return ())
1390 assert (Map.null mp) (return ())
@@ -1369,8 +1396,9 @@ classifySessionsBy tick reset ejectPred tmout
1369 Just (Entry expiry key, hp1) ->
1396 Just (Entry expiry key, hp1) ->
1370 case Map.lookup key mp of
1397 case Map.lookup key mp of
1371 Nothing -> ejectOne (hp1, mp, out, cnt)
1398 Nothing -> ejectOne (hp1, mp, out, cnt)
1372− Just (Tuple' latestTS acc) -> do
1399+ Just ZombieSession ->
1373− let expiry1 = addToAbsTime latestTS timeoutMs
1400+ ejectOne (hp1, Map.delete key mp, out, cnt)
1401+ Just (LiveSession expiry1 acc) -> do
1374 if not reset || expiry1 <= expiry
1402 if not reset || expiry1 <= expiry
1375 then ejectEntry hp1 mp out cnt acc key
1403 then ejectEntry hp1 mp out cnt acc key
@@ -1409,8 +1437,9 @@ classifySessionsBy tick reset ejectPred tmout
1410 case Map.lookup key mp of
1438 case Map.lookup key mp of
1411 Nothing -> ejectLoop hp1 mp out cnt
1439 Nothing -> ejectLoop hp1 mp out cnt
1412− Just (Tuple' latestTS acc) -> do
1440+ Just ZombieSession ->
1413− let expiry1 = addToAbsTime latestTS timeoutMs
1441+ ejectLoop hp1 (Map.delete key mp) out cnt
1442+ Just (LiveSession expiry1 acc) -> do
1414 if expiry1 <= curTime || not reset || force
1443 if expiry1 <= curTime || not reset || force
1416 (hp2,mp1,out1,cnt1) <-
1445 (hp2,mp1,out1,cnt1) <-