From d35a2665960167114308af03a065c611b470bb85 Mon Sep 17 00:00:00 2001
From: Alex Snaps <alex@wcgw.dev>
Date: Mon, 27 May 2024 09:37:25 -0400
Subject: [PATCH] Use a BTreeMap in InMemoryStorage

Signed-off-by: Alex Snaps <alex@wcgw.dev>
---
 limitador/src/limit.rs             | 32 +++++++-----------------------
 limitador/src/storage/in_memory.rs |  8 ++++----
 2 files changed, 11 insertions(+), 29 deletions(-)

diff --git a/limitador/src/limit.rs b/limitador/src/limit.rs
index fa0e124c..ebf06e0f 100644
--- a/limitador/src/limit.rs
+++ b/limitador/src/limit.rs
@@ -1,5 +1,5 @@
 use crate::limit::conditions::{ErrorType, Literal, SyntaxError, Token, TokenType};
-use serde::{Deserialize, Serialize, Serializer};
+use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
 use std::collections::{BTreeSet, HashMap, HashSet};
 use std::error::Error;
@@ -28,7 +28,7 @@ mod deprecated {
 #[cfg(feature = "lenient_conditions")]
 pub use deprecated::check_deprecated_syntax_usages_and_reset;
 
-#[derive(Debug, Hash, Eq, PartialEq, Clone, Serialize, Deserialize)]
+#[derive(Debug, Hash, Eq, PartialEq, Clone, PartialOrd, Ord, Serialize, Deserialize)]
 pub struct Namespace(String);
 
 impl From<&str> for Namespace {
@@ -49,7 +49,7 @@ impl From<String> for Namespace {
     }
 }
 
-#[derive(Eq, Debug, Clone, Serialize, Deserialize)]
+#[derive(Eq, Debug, Clone, PartialOrd, Ord, Serialize, Deserialize)]
 pub struct Limit {
     #[serde(skip_serializing, default)]
     id: Option<String>,
@@ -62,13 +62,11 @@ pub struct Limit {
 
     // Need to sort to generate the same object when using the JSON as a key or
     // value in Redis.
-    #[serde(serialize_with = "ordered_condition_set")]
-    conditions: HashSet<Condition>,
-    #[serde(serialize_with = "ordered_set")]
-    variables: HashSet<String>,
+    conditions: BTreeSet<Condition>,
+    variables: BTreeSet<String>,
 }
 
-#[derive(Deserialize, Serialize, PartialEq, Eq, Debug, Clone, Hash)]
+#[derive(Deserialize, Serialize, PartialEq, Eq, Debug, Clone, Hash, PartialOrd, Ord)]
 #[serde(try_from = "String", into = "String")]
 pub struct Condition {
     var_name: String,
@@ -267,7 +265,7 @@ impl From<Condition> for String {
     }
 }
 
-#[derive(PartialEq, Eq, Debug, Clone, Hash)]
+#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Hash)]
 pub enum Predicate {
     Equal,
     NotEqual,
@@ -291,22 +289,6 @@ impl From<Predicate> for String {
     }
 }
 
-fn ordered_condition_set<S>(value: &HashSet<Condition>, serializer: S) -> Result<S::Ok, S::Error>
-where
-    S: Serializer,
-{
-    let ordered: BTreeSet<String> = value.iter().map(|c| c.clone().into()).collect();
-    ordered.serialize(serializer)
-}
-
-fn ordered_set<S>(value: &HashSet<String>, serializer: S) -> Result<S::Ok, S::Error>
-where
-    S: Serializer,
-{
-    let ordered: BTreeSet<_> = value.iter().collect();
-    ordered.serialize(serializer)
-}
-
 impl Limit {
     pub fn new<N: Into<Namespace>, T: TryInto<Condition>>(
         namespace: N,
diff --git a/limitador/src/storage/in_memory.rs b/limitador/src/storage/in_memory.rs
index afa15d49..b01eeb88 100644
--- a/limitador/src/storage/in_memory.rs
+++ b/limitador/src/storage/in_memory.rs
@@ -3,14 +3,14 @@ use crate::limit::{Limit, Namespace};
 use crate::storage::atomic_expiring_value::AtomicExpiringValue;
 use crate::storage::{Authorization, CounterStorage, StorageErr};
 use moka::sync::Cache;
-use std::collections::hash_map::Entry;
-use std::collections::{HashMap, HashSet};
+use std::collections::btree_map::Entry;
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::ops::Deref;
 use std::sync::{Arc, RwLock};
 use std::time::{Duration, SystemTime};
 
 pub struct InMemoryStorage {
-    simple_limits: RwLock<HashMap<Limit, AtomicExpiringValue>>,
+    simple_limits: RwLock<BTreeMap<Limit, AtomicExpiringValue>>,
     qualified_counters: Cache<Counter, Arc<AtomicExpiringValue>>,
 }
 
@@ -197,7 +197,7 @@ impl CounterStorage for InMemoryStorage {
 impl InMemoryStorage {
     pub fn new(cache_size: u64) -> Self {
         Self {
-            simple_limits: RwLock::new(HashMap::new()),
+            simple_limits: RwLock::new(BTreeMap::new()),
             qualified_counters: Cache::new(cache_size),
         }
     }