Merge pull request #195 from Kuadrant/memleak-fix

Store qualified counters in dedicated capacity limited cache

limitador-server/Cargo.toml

@@ -36,6 +36,7 @@ notify = "6.0.1"
 const_format = "0.2.31"
 lazy_static = "1.4.0"
 clap = "4.3"
+sysinfo = "0.29.7"
 
 [build-dependencies]
 tonic-build = "0.9.2"

limitador-server/src/config.rs

@@ -111,7 +111,9 @@ impl Default for Configuration {
     fn default() -> Self {
         Configuration {
             limits_file: "".to_string(),
-            storage: StorageConfiguration::InMemory,
+            storage: StorageConfiguration::InMemory(InMemoryStorageConfiguration {
+                cache_size: Some(10_000),
+            }),
             rls_host: "".to_string(),
             rls_port: 0,
             http_host: "".to_string(),
@@ -125,13 +127,18 @@ impl Default for Configuration {
 
 #[derive(PartialEq, Eq, Debug)]
 pub enum StorageConfiguration {
-    InMemory,
+    InMemory(InMemoryStorageConfiguration),
     Disk(DiskStorageConfiguration),
     Redis(RedisStorageConfiguration),
     #[cfg(feature = "infinispan")]
     Infinispan(InfinispanStorageConfiguration),
 }
 
+#[derive(PartialEq, Eq, Debug)]
+pub struct InMemoryStorageConfiguration {
+    pub cache_size: Option<u64>,
+}
+
 #[derive(PartialEq, Eq, Debug)]
 pub struct DiskStorageConfiguration {
     pub path: String,

limitador-server/src/envoy_rls/server.rs

@@ -246,7 +246,7 @@ mod tests {
             vec!["app_id"],
         );
 
-        let limiter = RateLimiter::default();
+        let limiter = RateLimiter::new(10_000);
         limiter.add_limit(limit);
 
         let rate_limiter = MyRateLimiter::new(
@@ -366,7 +366,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_takes_into_account_all_the_descriptors() {
-        let limiter = RateLimiter::default();
+        let limiter = RateLimiter::new(10_000);
 
         let namespace = "test_namespace";
 
@@ -434,7 +434,7 @@ mod tests {
         let namespace = "test_namespace";
         let limit = Limit::new(namespace, 10, 60, vec!["x == '1'"], vec!["y"]);
 
-        let limiter = RateLimiter::default();
+        let limiter = RateLimiter::new(10_000);
         limiter.add_limit(limit);
 
         let rate_limiter = MyRateLimiter::new(
@@ -499,7 +499,7 @@ mod tests {
         let namespace = "test_namespace";
         let limit = Limit::new(namespace, 1, 60, vec!["x == '1'"], vec!["y"]);
 
-        let limiter = RateLimiter::default();
+        let limiter = RateLimiter::new(10_000);
         limiter.add_limit(limit);
 
         let rate_limiter = MyRateLimiter::new(

limitador-server/src/main.rs

@@ -7,14 +7,15 @@ extern crate clap;
 #[cfg(feature = "infinispan")]
 use crate::config::InfinispanStorageConfiguration;
 use crate::config::{
-    Configuration, DiskStorageConfiguration, 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;
 use clap::{value_parser, Arg, ArgAction, Command};
 use const_format::formatcp;
 use env_logger::Builder;
+use limitador::counter::Counter;
 use limitador::errors::LimitadorError;
 use limitador::limit::Limit;
 use limitador::storage::disk::DiskStorage;
@@ -38,6 +39,7 @@ use std::path::Path;
 use std::sync::Arc;
 use std::time::Duration;
 use std::{env, process};
+use sysinfo::{RefreshKind, System, SystemExt};
 use thiserror::Error;
 use tokio::runtime::Handle;
 
@@ -82,7 +84,9 @@ impl Limiter {
             StorageConfiguration::Infinispan(cfg) => {
                 Self::infinispan_limiter(cfg, config.limit_name_in_labels).await
             }
-            StorageConfiguration::InMemory => Self::in_memory_limiter(config),
+            StorageConfiguration::InMemory(cfg) => {
+                Self::in_memory_limiter(cfg, config.limit_name_in_labels)
+            }
             StorageConfiguration::Disk(cfg) => Self::disk_limiter(cfg, config.limit_name_in_labels),
         };
 
@@ -210,7 +214,7 @@ impl Limiter {
             }
         };
         let mut rate_limiter_builder =
-            RateLimiterBuilder::new().storage(Storage::with_counter_storage(Box::new(storage)));
+            RateLimiterBuilder::with_storage(Storage::with_counter_storage(Box::new(storage)));
 
         if limit_name_in_labels {
             rate_limiter_builder = rate_limiter_builder.with_prometheus_limit_name_labels()
@@ -219,10 +223,11 @@ impl Limiter {
         Self::Blocking(rate_limiter_builder.build())
     }
 
-    fn in_memory_limiter(cfg: Configuration) -> Self {
-        let mut rate_limiter_builder = RateLimiterBuilder::new();
+    fn in_memory_limiter(cfg: InMemoryStorageConfiguration, limit_name_in_labels: bool) -> Self {
+        let mut rate_limiter_builder =
+            RateLimiterBuilder::new(cfg.cache_size.or_else(guess_cache_size).unwrap());
 
-        if cfg.limit_name_in_labels {
+        if limit_name_in_labels {
             rate_limiter_builder = rate_limiter_builder.with_prometheus_limit_name_labels()
         }
 
@@ -513,7 +518,16 @@ fn create_config() -> (Configuration, &'static str) {
         .subcommand(
             Command::new("memory")
                 .display_order(1)
-                .about("Counters are held in Limitador (ephemeral)"),
+                .about("Counters are held in Limitador (ephemeral)")
+                .arg(
+                    Arg::new("CACHE_SIZE")
+                        .long("cache")
+                        .short('c')
+                        .action(ArgAction::Set)
+                        .value_parser(value_parser!(u64))
+                        .display_order(1)
+                        .help("Sets the size of the cache for 'qualified counters'"),
+                ),
         )
         .subcommand(
             Command::new("disk")
@@ -698,7 +712,9 @@ fn create_config() -> (Configuration, &'static str) {
                 consistency: Some(sub.get_one::<String>("consistency").unwrap().to_string()),
             })
         }
-        Some(("memory", _sub)) => StorageConfiguration::InMemory,
+        Some(("memory", sub)) => StorageConfiguration::InMemory(InMemoryStorageConfiguration {
+            cache_size: sub.get_one::<u64>("CACHE_SIZE").copied(),
+        }),
         None => match storage_config_from_env() {
             Ok(storage_cfg) => storage_cfg,
             Err(_) => {
@@ -785,10 +801,26 @@ fn storage_config_from_env() -> Result<StorageConfiguration, ()> {
                 consistency: env::var("INFINISPAN_COUNTERS_CONSISTENCY").ok(),
             },
         )),
-        _ => Ok(StorageConfiguration::InMemory),
+        _ => Ok(StorageConfiguration::InMemory(
+            InMemoryStorageConfiguration { cache_size: None },
+        )),
     }
 }
 
+fn guess_cache_size() -> Option<u64> {
+    let sys = System::new_with_specifics(RefreshKind::new().with_memory());
+    let free_mem = sys.available_memory();
+    let memory = free_mem as f64 * 0.7;
+    let size = (memory
+        / (std::mem::size_of::<Counter>() + 16/* size_of::<AtomicExpiringValue>() */) as f64)
+        as u64;
+    warn!(
+        "No cache size provided, aiming at 70% of {}MB, i.e. {size} entries",
+        free_mem / 1024 / 1024
+    );
+    Some(size)
+}
+
 fn env_option_is_enabled(env_name: &str) -> bool {
     match env::var(env_name) {
         Ok(value) => value == "1",

limitador/Cargo.toml

@@ -21,6 +21,8 @@ infinispan_storage = ["infinispan", "reqwest", "base64"]
 lenient_conditions = []
 
 [dependencies]
+moka = "0.11.2"
+getrandom = { version = "0.2", features = ["js"] }
 ttl_cache = "0.5"
 serde = { version = "1", features = ["derive"] }
 postcard = { version = "1.0.4", features = ["use-std"] }

limitador/benches/bench.rs

@@ -103,7 +103,7 @@ fn bench_in_mem(c: &mut Criterion) {
 #[cfg(feature = "disk_storage")]
 fn bench_disk(c: &mut Criterion) {
     let mut group = c.benchmark_group("Disk");
-    for (index, scenario) in TEST_SCENARIOS.iter().enumerate() {
+    for scenario in TEST_SCENARIOS.iter() {
         group.bench_with_input(
             BenchmarkId::new("is_rate_limited", scenario),
             scenario,

limitador/src/counter.rs

@@ -87,6 +87,10 @@ impl Counter {
         self.expires_in = Some(duration)
     }
 
+    pub fn is_qualified(&self) -> bool {
+        !self.set_variables.is_empty()
+    }
+
     #[cfg(feature = "disk_storage")]
     pub(crate) fn variables_for_key(&self) -> Vec<(&str, &str)> {
         let mut variables = Vec::with_capacity(self.set_variables.len());

limitador/src/lib.rs

@@ -7,10 +7,12 @@
 //! Limitador. Storing the limits in Redis is slower, but they can be shared
 //! between instances.
 //!
-//! By default, the rate limiter is configured to store the counters in memory:
+//! By default, the rate limiter is configured to store the counters in memory.
+//! It'll store only a limited amount of "qualified counters", specified as a
+//! `u64` value in the constructor.
 //! ```
 //! use limitador::RateLimiter;
-//! let rate_limiter = RateLimiter::default();
+//! let rate_limiter = RateLimiter::new(1000);
 //! ```
 //!
 //! To use Redis:
@@ -72,7 +74,7 @@
 //!      vec!["req.method == 'GET'"],
 //!      vec!["user_id"],
 //! );
-//! let mut rate_limiter = RateLimiter::default();
+//! let mut rate_limiter = RateLimiter::new(1000);
 //!
 //! // Add a limit
 //! rate_limiter.add_limit(limit.clone());
@@ -95,7 +97,7 @@
 //! use limitador::limit::Limit;
 //! use std::collections::HashMap;
 //!
-//! let mut rate_limiter = RateLimiter::default();
+//! let mut rate_limiter = RateLimiter::new(1000);
 //!
 //! let limit = Limit::new(
 //!     "my_namespace",
@@ -236,9 +238,16 @@ impl From<CheckResult> for bool {
 }
 
 impl RateLimiterBuilder {
-    pub fn new() -> Self {
+    pub fn with_storage(storage: Storage) -> Self {
         Self {
-            storage: Storage::new(),
+            storage,
+            prometheus_limit_name_labels_enabled: false,
+        }
+    }
+
+    pub fn new(cache_size: u64) -> Self {
+        Self {
+            storage: Storage::new(cache_size),
             prometheus_limit_name_labels_enabled: false,
         }
     }
@@ -267,12 +276,6 @@ impl RateLimiterBuilder {
     }
 }
 
-impl Default for RateLimiterBuilder {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 pub struct AsyncRateLimiterBuilder {
     storage: AsyncStorage,
     prometheus_limit_name_labels_enabled: bool,
@@ -306,9 +309,9 @@ impl AsyncRateLimiterBuilder {
 }
 
 impl RateLimiter {
-    pub fn new() -> Self {
+    pub fn new(cache_size: u64) -> Self {
         Self {
-            storage: Storage::new(),
+            storage: Storage::new(cache_size),
             prometheus_metrics: PrometheusMetrics::new(),
         }
     }
@@ -492,12 +495,6 @@ impl RateLimiter {
     }
 }
 
-impl Default for RateLimiter {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 // TODO: the code of this implementation is almost identical to the blocking
 // one. The only exception is that the functions defined are "async" and all the
 // calls to the storage need to include ".await". We'll need to think about how

limitador/src/storage/atomic_expiring_value.rs

@@ -141,4 +141,11 @@ mod tests {
         });
         assert!([2i64, 3i64].contains(&atomic_expiring_value.value.load(Ordering::SeqCst)));
     }
+
+    #[test]
+    fn size_of_struct() {
+        // This is ugly, but we don't have access to `AtomicExpiringValue` in the server,
+        // so this is hardcoded in main.rs
+        assert_eq!(16, std::mem::size_of::<AtomicExpiringValue>());
+    }
 }