Skip to content

Commit

Permalink
Use binary encoded keys with disk storage only
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexsnaps
alexsnaps committed Apr 21, 2023
1 parent 3db3c09 commit e3ddc24
Show file tree
Hide file tree
Showing 6 changed files with 182 additions and 153 deletions.
2 changes: 1 addition & 1 deletion limitador/src/storage/disk/sled_storage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use crate::counter::Counter;
use crate::limit::Limit;
use crate::storage::disk::expiring_value::ExpiringValue;
use crate::storage::disk::OptimizeFor;
use crate::storage::keys::{
use crate::storage::keys::bin::{
key_for_counter, partial_counter_from_counter_key, prefix_for_namespace,
};
use crate::storage::{Authorization, CounterStorage, StorageErr};
Expand Down
21 changes: 1 addition & 20 deletions limitador/src/storage/infinispan/infinispan_storage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use crate::storage::infinispan::response::response_to_string;
use crate::storage::infinispan::{
counters, sets, DEFAULT_INFINISPAN_CONSISTENCY, DEFAULT_INFINISPAN_LIMITS_CACHE_NAME,
};
use crate::storage::keys;
use crate::storage::keys::*;
use crate::storage::{AsyncCounterStorage, Authorization, StorageErr};
use async_trait::async_trait;
use base64::{engine::general_purpose::STANDARD_NO_PAD as base64engine, Engine};
Expand Down Expand Up @@ -292,22 +292,3 @@ impl InfinispanStorageBuilder {
.await
}
}

fn key_for_counter(counter: &Counter) -> String {
to_base64(keys::key_for_counter(counter))
}

fn key_for_counters_of_limit(limit: &Limit) -> String {
to_base64(keys::key_for_counters_of_limit(limit))
}

fn counter_from_counter_key(key: &str, limit: &Limit) -> Counter {
keys::counter_from_counter_key(from_base64(key.to_string()).as_slice(), limit)
}
fn to_base64(data: Vec<u8>) -> String {
base64engine.encode(data)
}

fn from_base64(data: String) -> Vec<u8> {
base64engine.decode(data).unwrap()
}
290 changes: 169 additions & 121 deletions limitador/src/storage/keys.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,80 +12,31 @@
// reusing this module for other storage implementations make sure that using
// "{}" for sharding applies.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use crate::counter::Counter;
use crate::limit::Limit;

#[derive(PartialEq, Debug, Serialize, Deserialize)]
struct CounterKey<'a> {
ns: &'a str,
seconds: u64,
conditions: Vec<String>,
variables: Vec<(&'a str, &'a str)>,
}

impl<'a> From<&'a Counter> for CounterKey<'a> {
fn from(counter: &'a Counter) -> Self {
let set = counter.limit().conditions();
let mut conditions = Vec::with_capacity(set.len());
for cond in &set {
conditions.push(cond.clone());
}
conditions.sort();

CounterKey {
ns: counter.namespace().as_ref(),
seconds: counter.seconds(),
conditions,
variables: counter.variables_for_key(),
}
}
}

#[derive(PartialEq, Debug, Serialize, Deserialize)]
struct LimitCountersKey<'a> {
ns: &'a str,
seconds: u64,
conditions: Vec<String>,
variables: Vec<&'a str>,
pub fn key_for_counter(counter: &Counter) -> String {
format!(
"{},counter:{}",
prefix_for_namespace(counter.namespace().as_ref()),
serde_json::to_string(counter).unwrap()
)
}

impl<'a> From<&'a Limit> for LimitCountersKey<'a> {
fn from(limit: &'a Limit) -> Self {
let set = limit.conditions();
let mut conditions = Vec::with_capacity(set.len());
for cond in &set {
conditions.push(cond.clone());
}
conditions.sort();

LimitCountersKey {
ns: limit.namespace().as_ref(),
seconds: limit.seconds(),
conditions,
variables: limit.variables_for_key(),
}
}
pub fn key_for_counters_of_limit(limit: &Limit) -> String {
format!(
"namespace:{{{}}},counters_of_limit:{}",
limit.namespace().as_ref(),
serde_json::to_string(limit).unwrap()
)
}

pub fn key_for_counter(counter: &Counter) -> Vec<u8> {
let key: CounterKey = counter.into();
postcard::to_stdvec(&key).unwrap()
pub fn prefix_for_namespace(namespace: &str) -> String {
format!("namespace:{{{namespace}}},")
}

pub fn key_for_counters_of_limit(limit: &Limit) -> Vec<u8> {
let key: LimitCountersKey = limit.into();
postcard::to_stdvec(&key).unwrap()
}

pub fn prefix_for_namespace(namespace: &str) -> Vec<u8> {
postcard::to_stdvec(namespace).unwrap()
}

pub fn counter_from_counter_key(key: &[u8], limit: &Limit) -> Counter {
let mut counter = partial_counter_from_counter_key(key);
pub fn counter_from_counter_key(key: &str, limit: &Limit) -> Counter {
let mut counter = partial_counter_from_counter_key(key, limit.namespace().as_ref());
if !counter.update_to_limit(limit) {
// this means some kind of data corruption _or_ most probably
// an out of sync `impl PartialEq for Limit` vs `pub fn key_for_counter(counter: &Counter) -> String`
Expand All @@ -98,35 +49,25 @@ pub fn counter_from_counter_key(key: &[u8], limit: &Limit) -> Counter {
counter
}

pub fn partial_counter_from_counter_key(key: &[u8]) -> Counter {
let key: CounterKey = postcard::from_bytes(key).unwrap();
let CounterKey {
ns,
seconds,
conditions,
variables,
} = key;

let map: HashMap<String, String> = variables
.into_iter()
.map(|(var, value)| (var.to_string(), value.to_string()))
.collect();
let limit = Limit::new(ns, i64::default(), seconds, conditions, map.keys());
Counter::new(limit, map)
pub fn partial_counter_from_counter_key(key: &str, namespace: &str) -> Counter {
let offset = ",counter:".len();
let start_pos_counter = prefix_for_namespace(namespace).len() + offset;

let counter: Counter = serde_json::from_str(&key[start_pos_counter..]).unwrap();
counter
}

#[cfg(test)]
mod tests {
use crate::counter::Counter;
use crate::storage::keys::{
use super::{
key_for_counter, key_for_counters_of_limit, partial_counter_from_counter_key,
prefix_for_namespace, CounterKey,
prefix_for_namespace,
};
use crate::counter::Counter;
use crate::Limit;
use std::collections::HashMap;

#[test]
#[ignore]
fn key_for_limit_format() {
let limit = Limit::new(
"example.com",
Expand All @@ -136,52 +77,159 @@ mod tests {
vec!["app_id"],
);
assert_eq!(
&key_for_counters_of_limit(&limit),
b"namespace:{example.com},counters_of_limit:{\"namespace\":\"example.com\",\"seconds\":60,\"conditions\":[\"req.method == \\\"GET\\\"\"],\"variables\":[\"app_id\"]}")
}

#[test]
fn counter_key_serializes_and_back() {
let namespace = "ns_counter:";
let limit = Limit::new(
namespace,
1,
2,
vec!["foo == 'bar'"],
vec!["app_id", "role", "wat"],
);
let mut vars = HashMap::default();
vars.insert("role".to_string(), "admin".to_string());
vars.insert("app_id".to_string(), "123".to_string());
vars.insert("wat".to_string(), "dunno".to_string());
let counter = Counter::new(limit.clone(), vars);

let raw = key_for_counter(&counter);
let key_back: CounterKey =
postcard::from_bytes(&raw).expect("This should deserialize back!");
let key: CounterKey = (&counter).into();
assert_eq!(key_back, key);
"namespace:{example.com},counters_of_limit:{\"namespace\":\"example.com\",\"seconds\":60,\"conditions\":[\"req.method == \\\"GET\\\"\"],\"variables\":[\"app_id\"]}",
key_for_counters_of_limit(&limit))
}

#[test]
fn counter_key_and_counter_are_symmetric() {
let namespace = "ns_counter:";
let limit = Limit::new(namespace, 1, 1, vec!["req.method == 'GET'"], vec!["app_id"]);
let mut variables = HashMap::default();
variables.insert("app_id".to_string(), "123".to_string());
let counter = Counter::new(limit.clone(), variables);
let counter = Counter::new(limit.clone(), HashMap::default());
let raw = key_for_counter(&counter);
assert_eq!(counter, partial_counter_from_counter_key(&raw));
assert_eq!(counter, partial_counter_from_counter_key(&raw, namespace));
let prefix = prefix_for_namespace(namespace);
assert_eq!(&raw[0..prefix.len()], &prefix);
}
}

#[test]
fn counter_key_starts_with_namespace_prefix() {
let namespace = "ns_counter:";
let limit = Limit::new(namespace, 1, 1, vec!["req.method == 'GET'"], vec!["app_id"]);
let counter = Counter::new(limit, HashMap::default());
let serialized_counter = key_for_counter(&counter);
pub mod bin {
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

let prefix = prefix_for_namespace(namespace);
assert_eq!(&serialized_counter[..prefix.len()], &prefix);
use crate::counter::Counter;
use crate::limit::Limit;

#[derive(PartialEq, Debug, Serialize, Deserialize)]
struct CounterKey<'a> {
ns: &'a str,
seconds: u64,
conditions: Vec<String>,
variables: Vec<(&'a str, &'a str)>,
}

impl<'a> From<&'a Counter> for CounterKey<'a> {
fn from(counter: &'a Counter) -> Self {
let set = counter.limit().conditions();
let mut conditions = Vec::with_capacity(set.len());
for cond in &set {
conditions.push(cond.clone());
}
conditions.sort();

CounterKey {
ns: counter.namespace().as_ref(),
seconds: counter.seconds(),
conditions,
variables: counter.variables_for_key(),
}
}
}

#[derive(PartialEq, Debug, Serialize, Deserialize)]
struct LimitCountersKey<'a> {
ns: &'a str,
seconds: u64,
conditions: Vec<String>,
variables: Vec<&'a str>,
}

impl<'a> From<&'a Limit> for LimitCountersKey<'a> {
fn from(limit: &'a Limit) -> Self {
let set = limit.conditions();
let mut conditions = Vec::with_capacity(set.len());
for cond in &set {
conditions.push(cond.clone());
}
conditions.sort();

LimitCountersKey {
ns: limit.namespace().as_ref(),
seconds: limit.seconds(),
conditions,
variables: limit.variables_for_key(),
}
}
}

pub fn key_for_counter(counter: &Counter) -> Vec<u8> {
let key: CounterKey = counter.into();
postcard::to_stdvec(&key).unwrap()
}

pub fn prefix_for_namespace(namespace: &str) -> Vec<u8> {
postcard::to_stdvec(namespace).unwrap()
}

pub fn partial_counter_from_counter_key(key: &[u8]) -> Counter {
let key: CounterKey = postcard::from_bytes(key).unwrap();
let CounterKey {
ns,
seconds,
conditions,
variables,
} = key;

let map: HashMap<String, String> = variables
.into_iter()
.map(|(var, value)| (var.to_string(), value.to_string()))
.collect();
let limit = Limit::new(ns, i64::default(), seconds, conditions, map.keys());
Counter::new(limit, map)
}

#[cfg(test)]
mod tests {
use super::{
key_for_counter, partial_counter_from_counter_key, prefix_for_namespace, CounterKey,
};
use crate::counter::Counter;
use crate::Limit;
use std::collections::HashMap;

#[test]
fn counter_key_serializes_and_back() {
let namespace = "ns_counter:";
let limit = Limit::new(
namespace,
1,
2,
vec!["foo == 'bar'"],
vec!["app_id", "role", "wat"],
);
let mut vars = HashMap::default();
vars.insert("role".to_string(), "admin".to_string());
vars.insert("app_id".to_string(), "123".to_string());
vars.insert("wat".to_string(), "dunno".to_string());
let counter = Counter::new(limit.clone(), vars);

let raw = key_for_counter(&counter);
let key_back: CounterKey =
postcard::from_bytes(&raw).expect("This should deserialize back!");
let key: CounterKey = (&counter).into();
assert_eq!(key_back, key);
}

#[test]
fn counter_key_and_counter_are_symmetric() {
let namespace = "ns_counter:";
let limit = Limit::new(namespace, 1, 1, vec!["req.method == 'GET'"], vec!["app_id"]);
let mut variables = HashMap::default();
variables.insert("app_id".to_string(), "123".to_string());
let counter = Counter::new(limit.clone(), variables);
let raw = key_for_counter(&counter);
assert_eq!(counter, partial_counter_from_counter_key(&raw));
}

#[test]
fn counter_key_starts_with_namespace_prefix() {
let namespace = "ns_counter:";
let limit = Limit::new(namespace, 1, 1, vec!["req.method == 'GET'"], vec!["app_id"]);
let counter = Counter::new(limit, HashMap::default());
let serialized_counter = key_for_counter(&counter);

let prefix = prefix_for_namespace(namespace);
assert_eq!(&serialized_counter[..prefix.len()], &prefix);
}
}
}
Loading

0 comments on commit e3ddc24

Please sign in to comment.