fix(cardinality): Fix edge case where keys have no expiry (#3058)

In the case with "perfect" cardinality, no additional values and below
the limit, it was possible that new sets were correctly created but
their expiry was not updated.

Author: Dav1dde

Cargo.lock

@@ -28,6 +28,7 @@ thiserror = { workspace = true }
 [dev-dependencies]
 criterion = { workspace = true }
 serde_json = { workspace = true }
+uuid = { workspace = true }
 
 [[bench]]
 name = "redis"

relay-cardinality/Cargo.toml

@@ -86,20 +86,23 @@ local results = {
 --
 -- Returns the total amount of values added to the 'working set'.
 local function sadd(t, offset, max)
-    local added = 0;
+    local working_set_cardinality = 0
+    local any_modifications = false
 
     for i = 1, #KEYS do
         local is_working_set = i == 1
 
         for from, to in batches(#t, 7000, offset, max) do
             local r = redis.call('SADD', KEYS[i], unpack(t, from, to))
+
+            any_modifications = any_modifications or r > 0
             if is_working_set then
-                added = added + r
+                working_set_cardinality = working_set_cardinality + r
             end
         end
     end
 
-    return added
+    return working_set_cardinality, any_modifications
 end
 
 -- Bumps to expiry of all sets by the passed expiry
@@ -116,15 +119,15 @@ local budget = math.max(0, max_cardinality - current_cardinality)
 
 -- Fast Path: we have enough budget to fit all elements
 if budget >= num_hashes then
-    local added = sadd(ARGV, HASHES_OFFSET)
+    local added, any_modifications = sadd(ARGV, HASHES_OFFSET)
     -- New current cardinality is current + amount of keys that have been added to the set
     current_cardinality = current_cardinality + added
 
     for _ = 1, num_hashes do
         table.insert(results, ACCEPTED)
     end
 
-    if added > 0 then
+    if any_modifications then
         bump_expiry()
     end
 
@@ -137,10 +140,10 @@ local offset = HASHES_OFFSET
 local needs_expiry_bumped = false
 while budget > 0 and offset < #ARGV do
     local len = math.min(#ARGV - offset, budget)
-    local added = sadd(ARGV, offset, len)
+    local added, any_modifications = sadd(ARGV, offset, len)
 
     current_cardinality = current_cardinality + added
-    needs_expiry_bumped = needs_expiry_bumped or added > 0
+    needs_expiry_bumped = needs_expiry_bumped or any_modifications
 
     for _ = 1, len do
         table.insert(results, ACCEPTED)

relay-cardinality/src/redis/cardinality.lua

@@ -320,7 +320,11 @@ mod tests {
         let url = std::env::var("RELAY_REDIS_URL")
             .unwrap_or_else(|_| "redis://127.0.0.1:6379".to_owned());
 
-        let redis = RedisPool::single(&url, RedisConfigOptions::default()).unwrap();
+        let opts = RedisConfigOptions {
+            max_connections: 1,
+            ..Default::default()
+        };
+        let redis = RedisPool::single(&url, opts).unwrap();
 
         RedisSetLimiter::new(
             RedisSetLimiterOptions {

relay-cardinality/src/redis/limiter.rs

@@ -119,3 +119,68 @@ impl CardinalityScript {
         Ok(result)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use relay_redis::{RedisConfigOptions, RedisPool};
+    use uuid::Uuid;
+
+    use super::*;
+
+    fn build_redis() -> RedisPool {
+        let url = std::env::var("RELAY_REDIS_URL")
+            .unwrap_or_else(|_| "redis://127.0.0.1:6379".to_owned());
+
+        let opts = RedisConfigOptions {
+            max_connections: 1,
+            ..Default::default()
+        };
+        RedisPool::single(&url, opts).unwrap()
+    }
+
+    fn keys(prefix: Uuid, keys: &[&str]) -> impl Iterator<Item = String> {
+        keys.iter()
+            .map(move |key| format!("{prefix}-{key}"))
+            .collect::<Vec<_>>()
+            .into_iter()
+    }
+
+    fn assert_ttls(connection: &mut Connection, prefix: Uuid) {
+        let keys = redis::cmd("KEYS")
+            .arg(format!("{prefix}-*"))
+            .query::<Vec<String>>(connection)
+            .unwrap();
+
+        for key in keys {
+            let ttl = redis::cmd("TTL")
+                .arg(&key)
+                .query::<i64>(connection)
+                .unwrap();
+
+            assert!(ttl >= 0, "Key {key} has no TTL");
+        }
+    }
+
+    #[test]
+    fn test_below_limit_perfect_cardinality_ttl() {
+        let redis = build_redis();
+        let mut client = redis.client().unwrap();
+        let mut connection = client.connection().unwrap();
+
+        let script = CardinalityScript::load();
+
+        let prefix = Uuid::new_v4();
+        let k1 = &["a", "b", "c"];
+        let k2 = &["b", "c", "d"];
+
+        script
+            .invoke(&mut connection, 50, 3600, 0..30, keys(prefix, k1))
+            .unwrap();
+
+        script
+            .invoke(&mut connection, 50, 3600, 0..30, keys(prefix, k2))
+            .unwrap();
+
+        assert_ttls(&mut connection, prefix);
+    }
+}