Add distributed_storage feature

Kuadrant · May 15, 2024 · 13259fc · 13259fc
1 parent 0afcfd1
commit 13259fc
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 40 deletions.
diff --git a/limitador-server/Cargo.toml b/limitador-server/Cargo.toml
@@ -14,6 +14,7 @@ edition = "2021"
 
 [features]
 infinispan = ["limitador/infinispan_storage"]
+distributed_storage = ["limitador/distributed_storage"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 

diff --git 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<u64>,

diff --git 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::<u64>("CACHE_SIZE").copied(),
         }),
+        #[cfg(feature = "distributed_storage")]
         Some(("distributed", sub)) => {
             StorageConfiguration::Distributed(DistributedStorageConfiguration {
                 name: sub.get_one::<String>("NAME").unwrap().to_owned(),

diff --git 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
@@ -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")]