From f4d91ac6498f8e5b7d1797a5dbe42bf7c29c529c Mon Sep 17 00:00:00 2001 From: Alex Snaps Date: Wed, 15 May 2024 11:10:46 -0400 Subject: [PATCH] Add distributed_storage feature --- limitador-server/src/config.rs | 2 + limitador-server/src/main.rs | 83 ++++++++++++++++++---------------- limitador/Cargo.toml | 1 + limitador/src/storage/mod.rs | 4 +- 4 files changed, 50 insertions(+), 40 deletions(-) diff --git a/limitador-server/src/config.rs b/limitador-server/src/config.rs index 8309b6fe..cada0032 100644 --- a/limitador-server/src/config.rs +++ b/limitador-server/src/config.rs @@ -141,6 +141,7 @@ pub enum StorageConfiguration { InMemory(InMemoryStorageConfiguration), Disk(DiskStorageConfiguration), Redis(RedisStorageConfiguration), + #[cfg(feature = "distributed_storage")] Distributed(DistributedStorageConfiguration), #[cfg(feature = "infinispan")] Infinispan(InfinispanStorageConfiguration), @@ -152,6 +153,7 @@ pub struct InMemoryStorageConfiguration { } #[derive(PartialEq, Eq, Debug)] +#[cfg(feature = "distributed_storage")] pub struct DistributedStorageConfiguration { pub name: String, pub cache_size: Option, diff --git a/limitador-server/src/main.rs b/limitador-server/src/main.rs index ffceb9d9..b4bdeac5 100644 --- a/limitador-server/src/main.rs +++ b/limitador-server/src/main.rs @@ -5,12 +5,13 @@ extern crate log; extern crate clap; +#[cfg(feature = "distributed_storage")] +use crate::config::DistributedStorageConfiguration; #[cfg(feature = "infinispan")] use crate::config::InfinispanStorageConfiguration; use crate::config::{ - Configuration, DiskStorageConfiguration, DistributedStorageConfiguration, - InMemoryStorageConfiguration, RedisStorageCacheConfiguration, RedisStorageConfiguration, - StorageConfiguration, + Configuration, DiskStorageConfiguration, InMemoryStorageConfiguration, + RedisStorageCacheConfiguration, RedisStorageConfiguration, StorageConfiguration, }; use crate::envoy_rls::server::{run_envoy_rls_server, RateLimitHeaders}; use crate::http_api::server::run_http_server; @@ -28,6 +29,7 @@ use limitador::storage::redis::{ AsyncRedisStorage, CachedRedisStorage, CachedRedisStorageBuilder, DEFAULT_BATCH_SIZE, DEFAULT_FLUSHING_PERIOD_SEC, DEFAULT_MAX_CACHED_COUNTERS, DEFAULT_RESPONSE_TIMEOUT_MS, }; +#[cfg(feature = "distributed_storage")] use limitador::storage::DistributedInMemoryStorage; use limitador::storage::{AsyncCounterStorage, AsyncStorage, Storage}; use limitador::{ @@ -95,6 +97,7 @@ impl Limiter { #[cfg(feature = "infinispan")] StorageConfiguration::Infinispan(cfg) => Self::infinispan_limiter(cfg).await, StorageConfiguration::InMemory(cfg) => Self::in_memory_limiter(cfg), + #[cfg(feature = "distributed_storage")] StorageConfiguration::Distributed(cfg) => Self::distributed_limiter(cfg), StorageConfiguration::Disk(cfg) => Self::disk_limiter(cfg), }; @@ -210,6 +213,7 @@ impl Limiter { Self::Blocking(rate_limiter_builder.build()) } + #[cfg(feature = "distributed_storage")] fn distributed_limiter(cfg: DistributedStorageConfiguration) -> Self { let storage = DistributedInMemoryStorage::new( cfg.name, @@ -647,43 +651,45 @@ fn create_config() -> (Configuration, &'static str) { .display_order(6) .help("Timeout for Redis commands in milliseconds"), ), - ) - .subcommand( - Command::new("distributed") - .about("Replicates CRDT-based counters across multiple Limitador servers") - .display_order(5) - .arg( - Arg::new("NAME") - .action(ArgAction::Set) - .required(true) - .display_order(2) - .help("Unique name to identify this Limitador instance"), - ) - .arg( - Arg::new("LOCAL") - .action(ArgAction::Set) - .required(true) - .display_order(2) - .help("Local IP:PORT to send datagrams from"), - ) - .arg( - Arg::new("BROADCAST") - .action(ArgAction::Set) - .required(true) - .display_order(3) - .help("Broadcast IP:PORT to send datagrams to"), - ) - .arg( - Arg::new("CACHE_SIZE") - .long("cache") - .short('c') - .action(ArgAction::Set) - .value_parser(value_parser!(u64)) - .display_order(4) - .help("Sets the size of the cache for 'qualified counters'"), - ), ); + #[cfg(feature = "distributed_storage")] + let cmdline = cmdline.subcommand( + Command::new("distributed") + .about("Replicates CRDT-based counters across multiple Limitador servers") + .display_order(5) + .arg( + Arg::new("NAME") + .action(ArgAction::Set) + .required(true) + .display_order(2) + .help("Unique name to identify this Limitador instance"), + ) + .arg( + Arg::new("LOCAL") + .action(ArgAction::Set) + .required(true) + .display_order(2) + .help("Local IP:PORT to send datagrams from"), + ) + .arg( + Arg::new("BROADCAST") + .action(ArgAction::Set) + .required(true) + .display_order(3) + .help("Broadcast IP:PORT to send datagrams to"), + ) + .arg( + Arg::new("CACHE_SIZE") + .long("cache") + .short('c') + .action(ArgAction::Set) + .value_parser(value_parser!(u64)) + .display_order(4) + .help("Sets the size of the cache for 'qualified counters'"), + ), + ); + #[cfg(feature = "infinispan")] let cmdline = cmdline.subcommand( Command::new("infinispan") @@ -801,6 +807,7 @@ fn create_config() -> (Configuration, &'static str) { Some(("memory", sub)) => StorageConfiguration::InMemory(InMemoryStorageConfiguration { cache_size: sub.get_one::("CACHE_SIZE").copied(), }), + #[cfg(feature = "distributed_storage")] Some(("distributed", sub)) => { StorageConfiguration::Distributed(DistributedStorageConfiguration { name: sub.get_one::("NAME").unwrap().to_owned(), diff --git a/limitador/Cargo.toml b/limitador/Cargo.toml index 0e300804..a6fbcc9d 100644 --- a/limitador/Cargo.toml +++ b/limitador/Cargo.toml @@ -16,6 +16,7 @@ edition = "2021" [features] default = ["disk_storage", "redis_storage"] disk_storage = ["rocksdb"] +distributed_storage = [] redis_storage = ["redis", "r2d2", "tokio"] infinispan_storage = ["infinispan", "reqwest", "base64", "tokio"] lenient_conditions = [] diff --git a/limitador/src/storage/mod.rs b/limitador/src/storage/mod.rs index 70aab7c9..0926f642 100644 --- a/limitador/src/storage/mod.rs +++ b/limitador/src/storage/mod.rs @@ -8,12 +8,12 @@ use thiserror::Error; #[cfg(feature = "disk_storage")] pub mod disk; -#[cfg(not(target_arch = "wasm32"))] +#[cfg(feature = "distributed_storage")] pub mod distributed; pub mod in_memory; pub mod wasm; -#[cfg(not(target_arch = "wasm32"))] +#[cfg(feature = "distributed_storage")] pub use crate::storage::distributed::CrInMemoryStorage as DistributedInMemoryStorage; #[cfg(feature = "redis_storage")]