ref(quotas): Return consumed quota from Redis script (#5390)

Author: Dav1dde

CHANGELOG.md

@@ -7,6 +7,10 @@
 - Support comparing release versions without build code. ([#5376](https://github.com/getsentry/relay/pull/5376))
 - Support uploading attachments directly to objectstore. ([#5367](https://github.com/getsentry/relay/pull/5367))
 
+**Internal**:
+
+- Derive the rate limiting decision in Relay from consumed quota. ([#5390](https://github.com/getsentry/relay/pull/5390))
+
 ## 25.11.0
 
 **Breaking Changes**:

relay-quotas/src/redis.rs

@@ -2,7 +2,7 @@ use std::fmt::{self, Debug};
 
 use relay_common::time::UnixTimestamp;
 use relay_log::protocol::value;
-use relay_redis::redis::Script;
+use relay_redis::redis::{self, FromRedisValue, Script};
 use relay_redis::{AsyncRedisClient, RedisError, RedisScripts};
 use thiserror::Error;
 
@@ -346,13 +346,13 @@ impl<T: GlobalLimiter> RedisRateLimiter<T> {
         // client across await points, otherwise it might be held for too long, and we will run out
         // of connections.
         let mut connection = self.client.get_connection().await?;
-        let rejections: Vec<bool> = invocation
+        let results: ScriptResult = invocation
             .invoke_async(&mut connection)
             .await
             .map_err(RedisError::Redis)?;
 
-        for (quota, is_rejected) in tracked_quotas.iter().zip(rejections) {
-            if is_rejected {
+        for (quota, state) in tracked_quotas.iter().zip(results.0) {
+            if state.is_rejected {
                 let retry_after = self.retry_after((quota.expiry() - timestamp).as_secs());
                 rate_limits.add(RateLimit::from_quota(quota, *item_scoping, retry_after));
             }
@@ -373,6 +373,51 @@ impl<T: GlobalLimiter> RedisRateLimiter<T> {
     }
 }
 
+/// The result returned from the rate limiting Redis script.
+#[derive(Debug)]
+struct ScriptResult(Vec<QuotaState>);
+
+impl FromRedisValue for ScriptResult {
+    fn from_redis_value(v: &redis::Value) -> redis::RedisResult<Self> {
+        let Some(seq) = v.as_sequence() else {
+            return Err(redis::RedisError::from((
+                redis::ErrorKind::TypeError,
+                "Expected a sequence from the rate limiting script",
+                format!("{v:?}"),
+            )));
+        };
+
+        let (chunks, rem) = seq.as_chunks();
+        if !rem.is_empty() {
+            return Err(redis::RedisError::from((
+                redis::ErrorKind::TypeError,
+                "Expected an even number of values from the rate limiting script",
+                format!("{v:?}"),
+            )));
+        }
+
+        let mut result = Vec::with_capacity(chunks.len());
+        for [is_rejected, consumed] in chunks {
+            result.push(QuotaState {
+                is_rejected: bool::from_redis_value(is_rejected)?,
+                consumed: i64::from_redis_value(consumed)?,
+            });
+        }
+
+        Ok(Self(result))
+    }
+}
+
+/// The state returned from the rate limiting script for a single quota.
+#[derive(Debug)]
+struct QuotaState {
+    /// Whether the quota rejects the request.
+    is_rejected: bool,
+    /// How much of the quota has already been consumed, before adding the requested quantity.
+    #[expect(unused, reason = "not yet used")]
+    consumed: i64,
+}
+
 #[cfg(test)]
 mod tests {
     use std::time::{SystemTime, UNIX_EPOCH};
@@ -1003,7 +1048,6 @@ mod tests {
     }
 
     #[tokio::test]
-    #[allow(clippy::disallowed_names, clippy::let_unit_value)]
     async fn test_is_rate_limited_script() {
         let now = SystemTime::now()
             .duration_since(UNIX_EPOCH)
@@ -1024,6 +1068,17 @@ mod tests {
 
         let script = RedisScripts::load_is_rate_limited();
 
+        macro_rules! assert_invocation {
+            ($invocation:expr, $($tt:tt)*) => {{
+                let result = $invocation
+                    .invoke_async::<ScriptResult>(&mut conn)
+                    .await
+                    .unwrap();
+
+                insta::assert_debug_snapshot!(result, $($tt)*);
+            }};
+        }
+
         let mut invocation = script.prepare_invoke();
         invocation
             .key(&foo) // key
@@ -1039,42 +1094,124 @@ mod tests {
             .arg(1) // quantity
             .arg(false); // over accept once
 
-        // The item should not be rate limited by either key.
-        assert_eq!(
-            invocation
-                .invoke_async::<Vec<bool>>(&mut conn)
-                .await
-                .unwrap(),
-            vec![false, false]
+        // Craft a new invocation similar to the previous one, but it only applies to the quota
+        // with a higher limit (2).
+        let mut invocation2 = script.prepare_invoke();
+        invocation2
+            .key(&bar) // key
+            .key(&r_bar) // refund key
+            .arg(2) // limit
+            .arg(now + 120) // expiry
+            .arg(1) // quantity
+            .arg(false); // over accept once
+
+        // Current usage is 0. But current values are now incremented by 1 (quantity).
+        assert_invocation!(invocation, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: false,
+                    consumed: 0,
+                },
+                QuotaState {
+                    is_rejected: false,
+                    consumed: 0,
+                },
+            ],
+        )
+        "
         );
 
-        // The item should be rate limited by the first key (1).
-        assert_eq!(
-            invocation
-                .invoke_async::<Vec<bool>>(&mut conn)
-                .await
-                .unwrap(),
-            vec![true, false]
+        // The usage was incremented in the last invocation, this invocation fails the rate limit
+        // on the first quota. -> No changes are made to the counters, the next invocation still
+        // needs to be `[1, 1]`.
+        assert_invocation!(invocation, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: true,
+                    consumed: 1,
+                },
+                QuotaState {
+                    is_rejected: false,
+                    consumed: 1,
+                },
+            ],
+        )
+        "
         );
 
         // The item should still be rate limited by the first key (1), but *not*
         // rate limited by the second key (2) even though this is the third time
         // we've checked the quotas. This ensures items that are rejected by a lower
         // quota don't affect unrelated items that share a parent quota.
-        assert_eq!(
-            invocation
-                .invoke_async::<Vec<bool>>(&mut conn)
-                .await
-                .unwrap(),
-            vec![true, false]
+        assert_invocation!(invocation, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: true,
+                    consumed: 1,
+                },
+                QuotaState {
+                    is_rejected: false,
+                    consumed: 1,
+                },
+            ],
+        )
+        "
+        );
+
+        // Using the second invocation which only considers a quota with a higher limit, this
+        // should still yield the current value of `1` and the next invocation should yield `2`.
+        assert_invocation!(invocation2, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: false,
+                    consumed: 1,
+                },
+            ],
+        )
+        "
+        );
+
+        // This now yields `2`. This is also the invocation at the limit, which means it should no
+        // longer increment the counter.
+        assert_invocation!(invocation2, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: true,
+                    consumed: 2,
+                },
+            ],
+        )
+        "
+        );
+
+        // Check again with the original invocation, this now yields `[1, 2]`.
+        assert_invocation!(invocation, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: true,
+                    consumed: 1,
+                },
+                QuotaState {
+                    is_rejected: true,
+                    consumed: 2,
+                },
+            ],
+        )
+        "
         );
 
         assert_eq!(conn.get::<_, String>(&foo).await.unwrap(), "1");
         let ttl: u64 = conn.ttl(&foo).await.unwrap();
         assert!(ttl >= 59);
         assert!(ttl <= 60);
 
-        assert_eq!(conn.get::<_, String>(&bar).await.unwrap(), "1");
+        assert_eq!(conn.get::<_, String>(&bar).await.unwrap(), "2");
         let ttl: u64 = conn.ttl(&bar).await.unwrap();
         assert!(ttl >= 119);
         assert!(ttl <= 120);
@@ -1095,22 +1232,43 @@ mod tests {
             .arg(1) // quantity
             .arg(false);
 
-        // increment
-        assert_eq!(
-            invocation
-                .invoke_async::<Vec<bool>>(&mut conn)
-                .await
-                .unwrap(),
-            vec![false]
+        // increment, current quota is 0.
+        assert_invocation!(invocation, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: false,
+                    consumed: 0,
+                },
+            ],
+        )
+        "
         );
 
-        // test that it's rate limited without refund
-        assert_eq!(
-            invocation
-                .invoke_async::<Vec<bool>>(&mut conn)
-                .await
-                .unwrap(),
-            vec![true]
+        // test that it's rate limited without refund.
+        assert_invocation!(invocation, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: true,
+                    consumed: 1,
+                },
+            ],
+        )
+        "
+        );
+
+        // Make sure, the counter wasn't incremented.
+        assert_invocation!(invocation, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: true,
+                    consumed: 1,
+                },
+            ],
+        )
+        "
         );
 
         let mut invocation = script.prepare_invoke();
@@ -1123,12 +1281,16 @@ mod tests {
             .arg(false);
 
         // test that refund key is used
-        assert_eq!(
-            invocation
-                .invoke_async::<Vec<bool>>(&mut conn)
-                .await
-                .unwrap(),
-            vec![false]
+        assert_invocation!(invocation, @r"
+        ScriptResult(
+            [
+                QuotaState {
+                    is_rejected: false,
+                    consumed: -4,
+                },
+            ],
+        )
+        "
         );
     }
 }

relay-redis/src/scripts/is_rate_limited.lua

@@ -37,19 +37,21 @@ local all_values = redis.call('MGET', unpack(KEYS))
 local results = {}
 local failed = false
 local num_quotas = #KEYS / 2
-for i=0, num_quotas - 1 do
+for i = 0, num_quotas - 1 do
     local k = i * 2 + 1
     local v = i * 4 + 1
 
     local limit = tonumber(ARGV[v])
-    local quantity = tonumber(ARGV[v+2])
-    local over_accept_once = ARGV[v+3]
-    local rejected = false
+    local quantity = tonumber(ARGV[v + 2])
+    local over_accept_once = ARGV[v + 3]
+
+    local main_value = all_values[k] or 0
+    local refund_value = all_values[k + 1] or 0
+    local consumed = main_value - refund_value
+
+    local rejected = false;
     -- limit=-1 means "no limit"
     if limit >= 0 then
-        local main_value = all_values[k] or 0
-        local refund_value = all_values[k + 1] or 0
-        local consumed = main_value - refund_value
         -- Without over_accept_once, we never increment past the limit. if quantity is 0, check instead if we reached limit.
         -- With over_accept_once, we only reject if the previous update already reached the limit.
         -- This way, we ensure that we increment to or past the limit at some point,
@@ -63,14 +65,14 @@ for i=0, num_quotas - 1 do
         end
     end
 
-    if rejected then
-        failed = true
-    end
-    results[i + 1] = rejected
+    failed = failed or rejected
+
+    table.insert(results, rejected)
+    table.insert(results, consumed)
 end
 
 if not failed then
-    for i=0, num_quotas - 1 do
+    for i = 0, num_quotas - 1 do
         local k = i * 2 + 1
         local v = i * 4 + 1