git.iko.soy/streamly
on master help sign in
Move the slowpath functions out of the main loop
92b80f1Harendra Kumar
src/Streamly/Internal/Data/Stream/IsStream/Reduce.hs+146−109
compare with master
unifiedsplit
@@ -1155,12 +1155,150 @@ data SessionState t m k a b = SessionState
1155
1155
1156data SessionEntry a b = LiveSession !a !b | ZombieSession
1156data SessionEntry a b = LiveSession !a !b | ZombieSession
1157
1157
1158+-- delete from map and output the fold accumulator
1159+ejectEntry :: (Monad m, IsStream t, IsMap f) =>
1160+ (acc -> m b)
1161+ -> heap
1162+ -> f entry
1163+ -> t m (Key f, b)
1164+ -> Int
1165+ -> acc
1166+ -> Key f
1167+ -> m (heap, f entry, t m (Key f, b), Int)
1168+ejectEntry extract hp mp out cnt acc key = do
1169+ sess <- extract acc
1170+ let out1 = (key, sess) `cons` out
1171+ let mp1 = IsMap.mapDelete key mp
1172+ return (hp, mp1, out1, cnt - 1)
1173+
1174+{-# NOINLINE flush #-}
1175+flush :: (Monad m, Ord k, IsStream t) =>
1176+ (s -> m b)
1177+ -> SessionState t m k (SessionEntry a s) b
1178+ -> m (SessionState t m k (SessionEntry a s) b)
1179+flush extract session@SessionState{..} = do
1180+ (hp', mp', out, count) <-
1181+ ejectAll
1182+ ( sessionTimerHeap
1183+ , sessionKeyValueMap
1184+ , IsStream.nil
1185+ , sessionCount
1186+ )
1187+ return $ session
1188+ { sessionCount = count
1189+ , sessionTimerHeap = hp'
1190+ , sessionKeyValueMap = mp'
1191+ , sessionOutputStream = out
1192+ }
1193+
1194+ where
1195+
1196+ ejectAll (hp, mp, out, !cnt) = do
1197+ let hres = H.uncons hp
1198+ case hres of
1199+ Just (Entry _ key, hp1) -> do
1200+ r <- case IsMap.mapLookup key mp of
1201+ Nothing -> return (hp1, mp, out, cnt)
1202+ Just ZombieSession ->
1203+ return (hp1, IsMap.mapDelete key mp, out, cnt)
1204+ Just (LiveSession _ acc) ->
1205+ ejectEntry extract hp1 mp out cnt acc key
1206+ ejectAll r
1207+ Nothing -> do
1208+ assert (Map.null mp) (return ())
1209+ return (hp, mp, out, cnt)
1210+
1211+{-# NOINLINE ejectOne #-}
1212+ejectOne :: (IsMap f, IsStream t, Ord k, Monad m) =>
1213+ Bool
1214+ -> (acc -> m b)
1215+ -> ( H.Heap (Entry k (Key f))
1216+ , f (SessionEntry k acc)
1217+ , t m (Key f, b)
1218+ , Int
1219+ )
1220+ -> m ( H.Heap (Entry k (Key f))
1221+ , f (SessionEntry k acc)
1222+ , t m (Key f, b), Int
1223+ )
1224+ejectOne reset extract = go
1225+
1226+ where
1227+
1228+ go (hp, mp, out, !cnt) = do
1229+ let hres = H.uncons hp
1230+ case hres of
1231+ Just (Entry expiry key, hp1) ->
1232+ case IsMap.mapLookup key mp of
1233+ Nothing -> go (hp1, mp, out, cnt)
1234+ Just ZombieSession ->
1235+ go (hp1, IsMap.mapDelete key mp, out, cnt)
1236+ Just (LiveSession expiry1 acc) -> do
1237+ if not reset || expiry1 <= expiry
1238+ then ejectEntry extract hp1 mp out cnt acc key
1239+ else
1240+ -- reset the session timeout and continue
1241+ let hp2 = H.insert (Entry expiry1 key) hp1
1242+ in go (hp2, mp, out, cnt)
1243+ Nothing -> do
1244+ assert (IsMap.mapNull mp) (return ())
1245+ return (hp, mp, out, cnt)
1246+
1247+{-# NOINLINE ejectExpired #-}
1248+ejectExpired :: (Ord k, IsStream t, Monad m) =>
1249+ Bool
1250+ -> (Int -> m Bool)
1251+ -> (acc -> m b)
1252+ -> SessionState t m k (SessionEntry AbsTime acc) b
1253+ -> AbsTime
1254+ -> m (SessionState t m k (SessionEntry AbsTime acc) b)
1255+ejectExpired reset ejectPred extract session@SessionState{..} curTime = do
1256+ (hp', mp', out, count) <-
1257+ ejectLoop
1258+ sessionTimerHeap sessionKeyValueMap IsStream.nil sessionCount
1259+ return $ session
1260+ { sessionCurTime = curTime
1261+ , sessionCount = count
1262+ , sessionTimerHeap = hp'
1263+ , sessionKeyValueMap = mp'
1264+ , sessionOutputStream = out
1265+ }
1266+
1267+ where
1268+
1269+ ejectLoop hp mp out !cnt = do
1270+ let hres = H.uncons hp
1271+ case hres of
1272+ Just (Entry expiry key, hp1) -> do
1273+ (eject, force) <-
1274+ if curTime >= expiry
1275+ then return (True, False)
1276+ else do
1277+ r <- ejectPred cnt
1278+ return (r, r)
1279+ if eject
1280+ then
1281+ case IsMap.mapLookup key mp of
1282+ Nothing -> ejectLoop hp1 mp out cnt
1283+ Just ZombieSession ->
1284+ ejectLoop hp1 (IsMap.mapDelete key mp) out cnt
1285+ Just (LiveSession expiry1 acc) -> do
1286+ if expiry1 <= curTime || not reset || force
1287+ then do
1288+ (hp2,mp1,out1,cnt1) <-
1289+ ejectEntry extract hp1 mp out cnt acc key
1290+ ejectLoop hp2 mp1 out1 cnt1
1291+ else
1292+ -- reset the session timeout and continue
1293+ let hp2 = H.insert (Entry expiry1 key) hp1
1294+ in ejectLoop hp2 mp out cnt
1295+ else return (hp, mp, out, cnt)
1296+ Nothing -> do
1297+ assert (IsMap.mapNull mp) (return ())
1298+ return (hp, mp, out, cnt)
1299+
1158-- XXX Use mutable IORef in accumulator
1300-- XXX Use mutable IORef in accumulator
1159−--
1160−-- XXX this fuses with INLINE but it may be too much work for the compiler.
1161−-- Maybe we can NOINLINE the heap processing stuff and then INLINE it to reduce
1162−-- the code bloat.
1163−{-# INLINABLE classifySessionsByGeneric #-}
1301+{-# INLINE classifySessionsByGeneric #-}
1164classifySessionsByGeneric
1302classifySessionsByGeneric
1165 :: (IsStream t, MonadAsync m, Ord (Key f))
1303 :: (IsStream t, MonadAsync m, Ord (Key f))
1166 => Proxy (f s)
1304 => Proxy (f s)
@@ -1174,7 +1312,7 @@ classifySessionsByGeneric
1174classifySessionsByGeneric _ tick reset ejectPred tmout
1312classifySessionsByGeneric _ tick reset ejectPred tmout
1175 (Fold step initial extract) input =
1313 (Fold step initial extract) input =
1176 Expand.unfoldMany (Unfold.lmap sessionOutputStream Unfold.fromStream)
1314 Expand.unfoldMany (Unfold.lmap sessionOutputStream Unfold.fromStream)
1177− $ scanlMAfter' sstep (return szero) flush
1315+ $ scanlMAfter' sstep (return szero) (flush extract)
1178 $ interjectSuffix tick (return Nothing)
1316 $ interjectSuffix tick (return Nothing)
1179 $ map Just input
1317 $ map Just input
1180
1318
@@ -1279,7 +1417,7 @@ classifySessionsByGeneric _ tick reset ejectPred tmout
1279 eject <- ejectPred sessionCount
1417 eject <- ejectPred sessionCount
1280 (hp, mp, out, cnt) <-
1418 (hp, mp, out, cnt) <-
1281 if eject
1419 if eject
1282− then ejectOne vars
1420+ then ejectOne reset extract vars
1283 else return vars
1421 else return vars
1284
1422
1285 -- Insert the new session in heap
1423 -- Insert the new session in heap
@@ -1315,108 +1453,7 @@ classifySessionsByGeneric _ tick reset ejectPred tmout
1315 -- Got a timer tick event
1453 -- Got a timer tick event
1316 sstep sessionState@SessionState{..} Nothing =
1454 sstep sessionState@SessionState{..} Nothing =
1317 let curTime = addToAbsTime sessionCurTime tickMs
1455 let curTime = addToAbsTime sessionCurTime tickMs
1318− in ejectExpired sessionState curTime
1456+ in ejectExpired reset ejectPred extract sessionState curTime
1319−
1320− flush session@SessionState{..} = do
1321− (hp', mp', out, count) <-
1322− ejectAll
1323− ( sessionTimerHeap
1324− , sessionKeyValueMap
1325− , IsStream.nil
1326− , sessionCount
1327− )
1328− return $ session
1329− { sessionCount = count
1330− , sessionTimerHeap = hp'
1331− , sessionKeyValueMap = mp'
1332− , sessionOutputStream = out
1333− }
1334−
1335− -- delete from map and output the fold accumulator
1336− ejectEntry hp mp out cnt acc key = do
1337− sess <- extract acc
1338− let out1 = (key, sess) `cons` out
1339− let mp1 = IsMap.mapDelete key mp
1340− return (hp, mp1, out1, cnt - 1)
1341−
1342− ejectAll (hp, mp, out, !cnt) = do
1343− let hres = H.uncons hp
1344− case hres of
1345− Just (Entry _ key, hp1) -> do
1346− r <- case IsMap.mapLookup key mp of
1347− Nothing -> return (hp1, mp, out, cnt)
1348− Just ZombieSession ->
1349− return (hp1, IsMap.mapDelete key mp, out, cnt)
1350− Just (LiveSession _ acc) ->
1351− ejectEntry hp1 mp out cnt acc key
1352− ejectAll r
1353− Nothing -> do
1354− assert (Map.null mp) (return ())
1355− return (hp, mp, out, cnt)
1356−
1357− ejectOne (hp, mp, out, !cnt) = do
1358− let hres = H.uncons hp
1359− case hres of
1360− Just (Entry expiry key, hp1) ->
1361− case IsMap.mapLookup key mp of
1362− Nothing -> ejectOne (hp1, mp, out, cnt)
1363− Just ZombieSession ->
1364− ejectOne (hp1, IsMap.mapDelete key mp, out, cnt)
1365− Just (LiveSession expiry1 acc) -> do
1366− if not reset || expiry1 <= expiry
1367− then ejectEntry hp1 mp out cnt acc key
1368− else
1369− -- reset the session timeout and continue
1370− let hp2 = H.insert (Entry expiry1 key) hp1
1371− in ejectOne (hp2, mp, out, cnt)
1372− Nothing -> do
1373− assert (IsMap.mapNull mp) (return ())
1374− return (hp, mp, out, cnt)
1375−
1376− ejectExpired session@SessionState{..} curTime = do
1377− (hp', mp', out, count) <-
1378− ejectLoop
1379− sessionTimerHeap sessionKeyValueMap IsStream.nil sessionCount
1380− return $ session
1381− { sessionCurTime = curTime
1382− , sessionCount = count
1383− , sessionTimerHeap = hp'
1384− , sessionKeyValueMap = mp'
1385− , sessionOutputStream = out
1386− }
1387−
1388− where
1389−
1390− ejectLoop hp mp out !cnt = do
1391− let hres = H.uncons hp
1392− case hres of
1393− Just (Entry expiry key, hp1) -> do
1394− (eject, force) <-
1395− if curTime >= expiry
1396− then return (True, False)
1397− else do
1398− r <- ejectPred cnt
1399− return (r, r)
1400− if eject
1401− then
1402− case IsMap.mapLookup key mp of
1403− Nothing -> ejectLoop hp1 mp out cnt
1404− Just ZombieSession ->
1405− ejectLoop hp1 (IsMap.mapDelete key mp) out cnt
1406− Just (LiveSession expiry1 acc) -> do
1407− if expiry1 <= curTime || not reset || force
1408− then do
1409− (hp2,mp1,out1,cnt1) <-
1410− ejectEntry hp1 mp out cnt acc key
1411− ejectLoop hp2 mp1 out1 cnt1
1412− else
1413− -- reset the session timeout and continue
1414− let hp2 = H.insert (Entry expiry1 key) hp1
1415− in ejectLoop hp2 mp out cnt
1416− else return (hp, mp, out, cnt)
1417− Nothing -> do
1418− assert (IsMap.mapNull mp) (return ())
1419− return (hp, mp, out, cnt)
1420
1457
1421-- | @classifySessionsBy tick keepalive predicate timeout fold stream@
1458-- | @classifySessionsBy tick keepalive predicate timeout fold stream@
1422-- classifies an input event @stream@ consisting of @(timestamp, (key,
1459-- classifies an input event @stream@ consisting of @(timestamp, (key,