Support providing an optional id to limits/counters

* add a key_for_counters_v2 function that uses the id as the key if set, otherwise uses the previous key encoding strategy. * updated the distributed store to use key_for_counters_v2. Since we can’t decode a partial counter from id based keys, we now also keep in memory the Counter in a counter field of the limits map. Signed-off-by: Hiram Chirino <[email protected]>
Kuadrant · Jun 12, 2024 · f3d4442 · f3d4442
1 parent 72c2b12
commit f3d4442
Show file tree

Hide file tree

Showing 16 changed files with 384 additions and 171 deletions.
diff --git a/limitador-server/src/envoy_rls/server.rs b/limitador-server/src/envoy_rls/server.rs
@@ -250,6 +250,7 @@ mod tests {
     async fn test_returns_ok_and_overlimit_correctly() {
         let namespace = "test_namespace";
         let limit = Limit::new(
+            None,
             namespace,
             1,
             60,
@@ -394,8 +395,15 @@ mod tests {
         let namespace = "test_namespace";
 
         vec![
-            Limit::new(namespace, 10, 60, vec!["x == '1'"], vec!["z"]),
-            Limit::new(namespace, 0, 60, vec!["x == '1'", "y == '2'"], vec!["z"]),
+            Limit::new(None, namespace, 10, 60, vec!["x == '1'"], vec!["z"]),
+            Limit::new(
+                None,
+                namespace,
+                0,
+                60,
+                vec!["x == '1'", "y == '2'"],
+                vec!["z"],
+            ),
         ]
         .into_iter()
         .for_each(|limit| {
@@ -459,7 +467,7 @@ mod tests {
     #[tokio::test]
     async fn test_takes_into_account_the_hits_addend_param() {
         let namespace = "test_namespace";
-        let limit = Limit::new(namespace, 10, 60, vec!["x == '1'"], vec!["y"]);
+        let limit = Limit::new(None, namespace, 10, 60, vec!["x == '1'"], vec!["y"]);
 
         let limiter = RateLimiter::new(10_000);
         limiter.add_limit(limit);
@@ -528,7 +536,7 @@ mod tests {
         // "hits_addend" is optional according to the spec, and should default
         // to 1, However, with the autogenerated structs it defaults to 0.
         let namespace = "test_namespace";
-        let limit = Limit::new(namespace, 1, 60, vec!["x == '1'"], vec!["y"]);
+        let limit = Limit::new(None, namespace, 1, 60, vec!["x == '1'"], vec!["y"]);
 
         let limiter = RateLimiter::new(10_000);
         limiter.add_limit(limit);

diff --git a/limitador-server/src/http_api/request_types.rs b/limitador-server/src/http_api/request_types.rs
@@ -18,6 +18,7 @@ pub struct CheckAndReportInfo {
 
 #[derive(Debug, Eq, PartialEq, Serialize, Deserialize, Apiv2Schema)]
 pub struct Limit {
+    id: Option<String>,
     namespace: String,
     max_value: u64,
     seconds: u64,
@@ -29,6 +30,7 @@ pub struct Limit {
 impl From<&LimitadorLimit> for Limit {
     fn from(ll: &LimitadorLimit) -> Self {
         Self {
+            id: ll.id().clone(),
             namespace: ll.namespace().as_ref().to_string(),
             max_value: ll.max_value(),
             seconds: ll.seconds(),
@@ -42,6 +44,7 @@ impl From<&LimitadorLimit> for Limit {
 impl From<Limit> for LimitadorLimit {
     fn from(limit: Limit) -> Self {
         let mut limitador_limit = Self::new(
+            limit.id,
             limit.namespace,
             limit.max_value,
             limit.seconds,

diff --git a/limitador-server/src/http_api/server.rs b/limitador-server/src/http_api/server.rs
@@ -550,6 +550,7 @@ mod tests {
     async fn create_test_limit(limiter: &Limiter, namespace: &str, max: u64) -> LimitadorLimit {
         // Create a limit
         let limit = LimitadorLimit::new(
+            None,
             namespace,
             max,
             60,

diff --git a/limitador/Cargo.toml b/limitador/Cargo.toml
@@ -23,7 +23,7 @@ lenient_conditions = []
 moka = { version = "0.12", features = ["sync"] }
 dashmap = "5.5.3"
 getrandom = { version = "0.2", features = ["js"] }
-serde = { version = "1", features = ["derive"] }
+serde = { version = "1", features = ["derive", "rc"] }
 postcard = { version = "1.0.4", features = ["use-std"] }
 serde_json = "1"
 rmp-serde = "1.1.0"

diff --git a/limitador/benches/bench.rs b/limitador/benches/bench.rs
@@ -548,6 +548,7 @@ fn generate_test_limits(scenario: &TestScenario) -> (Vec<Limit>, Vec<TestCallPar
 
         for limit_idx in 0..scenario.n_limits_per_ns {
             test_limits.push(Limit::new(
+                None,
                 namespace.clone(),
                 u64::MAX,
                 ((limit_idx * 60) + 10) as u64,

diff --git a/limitador/src/counter.rs b/limitador/src/counter.rs
@@ -15,6 +15,7 @@ pub struct Counter {
     set_variables: HashMap<String, String>,
 
     remaining: Option<u64>,
+
     expires_in: Option<Duration>,
 }
 
@@ -72,6 +73,10 @@ impl Counter {
         Duration::from_secs(self.limit.seconds())
     }
 
+    pub fn id(&self) -> &Option<String> {
+        self.limit.id()
+    }
+
     pub fn namespace(&self) -> &Namespace {
         self.limit.namespace()
     }

diff --git a/limitador/src/lib.rs b/limitador/src/lib.rs
@@ -51,6 +51,7 @@
 //! ```
 //! use limitador::limit::Limit;
 //! let limit = Limit::new(
+//!      None,
 //!     "my_namespace",
 //!      10,
 //!      60,
@@ -68,6 +69,7 @@
 //! use limitador::RateLimiter;
 //! use limitador::limit::Limit;
 //! let limit = Limit::new(
+//!      None,
 //!     "my_namespace",
 //!      10,
 //!      60,
@@ -100,6 +102,7 @@
 //! let mut rate_limiter = RateLimiter::new(1000);
 //!
 //! let limit = Limit::new(
+//!      None,
 //!     "my_namespace",
 //!      2,
 //!      60,
@@ -164,6 +167,7 @@
 //! use limitador::limit::Limit;
 //! use limitador::storage::redis::AsyncRedisStorage;
 //! let limit = Limit::new(
+//!      None,
 //!      "my_namespace",
 //!      10,
 //!      60,

diff --git a/limitador/src/limit.rs b/limitador/src/limit.rs
@@ -1,11 +1,16 @@
-use crate::limit::conditions::{ErrorType, Literal, SyntaxError, Token, TokenType};
-use serde::{Deserialize, Serialize, Serializer};
 use std::cmp::Ordering;
 use std::collections::{BTreeSet, HashMap, HashSet};
 use std::error::Error;
 use std::fmt::{Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
 
+use serde::{Deserialize, Serialize, Serializer};
+
+#[cfg(feature = "lenient_conditions")]
+pub use deprecated::check_deprecated_syntax_usages_and_reset;
+
+use crate::limit::conditions::{ErrorType, Literal, SyntaxError, Token, TokenType};
+
 #[cfg(feature = "lenient_conditions")]
 mod deprecated {
     use std::sync::atomic::{AtomicBool, Ordering};
@@ -25,9 +30,6 @@ mod deprecated {
     }
 }
 
-#[cfg(feature = "lenient_conditions")]
-pub use deprecated::check_deprecated_syntax_usages_and_reset;
-
 #[derive(Debug, Hash, Eq, PartialEq, Clone, Serialize, Deserialize)]
 pub struct Namespace(String);
 
@@ -51,6 +53,8 @@ impl From<String> for Namespace {
 
 #[derive(Eq, Debug, Clone, Serialize, Deserialize)]
 pub struct Limit {
+    #[serde(skip_serializing, default)]
+    id: Option<String>,
     namespace: Namespace,
     #[serde(skip_serializing, default)]
     max_value: u64,
@@ -307,6 +311,7 @@ where
 
 impl Limit {
     pub fn new<N: Into<Namespace>, T: TryInto<Condition>>(
+        id: Option<String>,
         namespace: N,
         max_value: u64,
         seconds: u64,
@@ -319,6 +324,7 @@ impl Limit {
     {
         // the above where-clause is needed in order to call unwrap().
         Self {
+            id,
             namespace: namespace.into(),
             max_value,
             seconds,
@@ -335,6 +341,10 @@ impl Limit {
         &self.namespace
     }
 
+    pub fn id(&self) -> &Option<String> {
+        &self.id
+    }
+
     pub fn max_value(&self) -> u64 {
         self.max_value
     }
@@ -821,7 +831,14 @@ mod tests {
 
     #[test]
     fn limit_can_have_an_optional_name() {
-        let mut limit = Limit::new("test_namespace", 10, 60, vec!["x == \"5\""], vec!["y"]);
+        let mut limit = Limit::new(
+            None,
+            "test_namespace",
+            10,
+            60,
+            vec!["x == \"5\""],
+            vec!["y"],
+        );
         assert!(limit.name.is_none());
 
         let name = "Test Limit";
@@ -831,7 +848,14 @@ mod tests {
 
     #[test]
     fn limit_applies() {
-        let limit = Limit::new("test_namespace", 10, 60, vec!["x == \"5\""], vec!["y"]);
+        let limit = Limit::new(
+            None,
+            "test_namespace",
+            10,
+            60,
+            vec!["x == \"5\""],
+            vec!["y"],
+        );
 
         let mut values: HashMap<String, String> = HashMap::new();
         values.insert("x".into(), "5".into());
@@ -842,7 +866,14 @@ mod tests {
 
     #[test]
     fn limit_does_not_apply_when_cond_is_false() {
-        let limit = Limit::new("test_namespace", 10, 60, vec!["x == \"5\""], vec!["y"]);
+        let limit = Limit::new(
+            None,
+            "test_namespace",
+            10,
+            60,
+            vec!["x == \"5\""],
+            vec!["y"],
+        );
 
         let mut values: HashMap<String, String> = HashMap::new();
         values.insert("x".into(), "1".into());
@@ -854,7 +885,7 @@ mod tests {
     #[test]
     #[cfg(feature = "lenient_conditions")]
     fn limit_does_not_apply_when_cond_is_false_deprecated_style() {
-        let limit = Limit::new("test_namespace", 10, 60, vec!["x == 5"], vec!["y"]);
+        let limit = Limit::new(None, "test_namespace", 10, 60, vec!["x == 5"], vec!["y"]);
 
         let mut values: HashMap<String, String> = HashMap::new();
         values.insert("x".into(), "1".into());
@@ -864,7 +895,14 @@ mod tests {
         assert!(check_deprecated_syntax_usages_and_reset());
         assert!(!check_deprecated_syntax_usages_and_reset());
 
-        let limit = Limit::new("test_namespace", 10, 60, vec!["x == foobar"], vec!["y"]);
+        let limit = Limit::new(
+            None,
+            "test_namespace",
+            10,
+            60,
+            vec!["x == foobar"],
+            vec!["y"],
+        );
 
         let mut values: HashMap<String, String> = HashMap::new();
         values.insert("x".into(), "foobar".into());
@@ -877,7 +915,14 @@ mod tests {
 
     #[test]
     fn limit_does_not_apply_when_cond_var_is_not_set() {
-        let limit = Limit::new("test_namespace", 10, 60, vec!["x == \"5\""], vec!["y"]);
+        let limit = Limit::new(
+            None,
+            "test_namespace",
+            10,
+            60,
+            vec!["x == \"5\""],
+            vec!["y"],
+        );
 
         // Notice that "x" is not set
         let mut values: HashMap<String, String> = HashMap::new();
@@ -889,7 +934,14 @@ mod tests {
 
     #[test]
     fn limit_does_not_apply_when_var_not_set() {
-        let limit = Limit::new("test_namespace", 10, 60, vec!["x == \"5\""], vec!["y"]);
+        let limit = Limit::new(
+            None,
+            "test_namespace",
+            10,
+            60,
+            vec!["x == \"5\""],
+            vec!["y"],
+        );
 
         // Notice that "y" is not set
         let mut values: HashMap<String, String> = HashMap::new();
@@ -901,6 +953,7 @@ mod tests {
     #[test]
     fn limit_applies_when_all_its_conditions_apply() {
         let limit = Limit::new(
+            None,
             "test_namespace",
             10,
             60,
@@ -919,6 +972,7 @@ mod tests {
     #[test]
     fn limit_does_not_apply_if_one_cond_doesnt() {
         let limit = Limit::new(
+            None,
             "test_namespace",
             10,
             60,
@@ -998,4 +1052,18 @@ mod tests {
         let result = serde_json::to_string(&condition).expect("Should serialize");
         assert_eq!(result, r#""foobar == \"ok\"""#.to_string());
     }
+
+    #[test]
+    fn limit_id() {
+        let limit = Limit::new(
+            Some("test_id".to_string()),
+            "test_namespace",
+            10,
+            60,
+            vec!["req.method == 'GET'"],
+            vec!["app_id"],
+        );
+
+        assert_eq!(limit.id().clone(), Some("test_id".to_string()))
+    }
 }
diff --git a/limitador/src/storage/disk/rocksdb_storage.rs b/limitador/src/storage/disk/rocksdb_storage.rs
@@ -242,7 +242,14 @@ mod tests {
     #[test]
     fn opens_db_on_disk() {
         let namespace = "test_namespace";
-        let limit = Limit::new(namespace, 1, 2, vec!["req.method == 'GET'"], vec!["app_id"]);
+        let limit = Limit::new(
+            None,
+            namespace,
+            1,
+            2,
+            vec!["req.method == 'GET'"],
+            vec!["app_id"],
+        );
         let counter = Counter::new(limit, HashMap::default());
 
         let tmp = TempDir::new().expect("We should have a dir!");