Skip to content

Commit

Permalink
Use a BTreeMap in InMemoryStorage
Browse files Browse the repository at this point in the history
Signed-off-by: Alex Snaps <[email protected]>
  • Loading branch information
alexsnaps committed Sep 12, 2024
1 parent 0020f7a commit d35a266
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 29 deletions.
32 changes: 7 additions & 25 deletions limitador/src/limit.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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 {
Expand All @@ -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>,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
8 changes: 4 additions & 4 deletions limitador/src/storage/in_memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>>,
}

Expand Down Expand Up @@ -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),
}
}
Expand Down

0 comments on commit d35a266

Please sign in to comment.