Support pending writes within CachedCounterValue

limitador/Cargo.toml

@@ -22,6 +22,7 @@ lenient_conditions = []
 
 [dependencies]
 moka = { version = "0.12", features = ["sync"] }
+dashmap = "5.5.3"
 getrandom = { version = "0.2", features = ["js"] }
 serde = { version = "1", features = ["derive"] }
 postcard = { version = "1.0.4", features = ["use-std"] }

limitador/src/storage/redis/counters_cache.rs

@@ -4,26 +4,132 @@
     DEFAULT_MAX_CACHED_COUNTERS, DEFAULT_MAX_TTL_CACHED_COUNTERS_SEC,
     DEFAULT_TTL_RATIO_CACHED_COUNTERS,
 };
+use dashmap::DashMap;
 use moka::sync::Cache;
+use std::collections::HashMap;
+use std::future::Future;
+use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
+use tokio::select;
+use tokio::sync::Notify;
+use tokio::time::interval;
 
 pub struct CachedCounterValue {
     value: AtomicExpiringValue,
+    initial_value: AtomicI64,
     expiry: AtomicExpiryTime,
 }
 
+pub struct Batcher {
+    updates: DashMap<Counter, Arc<CachedCounterValue>>,
+    notifier: Notify,
+    interval: Duration,
+    priority_flush: AtomicBool,
+}
+
+impl Batcher {
+    fn new(period: Duration) -> Self {
+        Self {
+            updates: Default::default(),
+            notifier: Default::default(),
+            interval: period,
+            priority_flush: AtomicBool::new(false),
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.updates.is_empty()
+    }
+
+    pub async fn consume<F, Fut, O>(&self, min: usize, consumer: F) -> O
+    where
+        F: FnOnce(HashMap<Counter, Arc<CachedCounterValue>>) -> Fut,
+        Fut: Future<Output = O>,
+    {
+        let mut interval = interval(self.interval);
+        let mut ready = self.updates.len() >= min;
+        loop {
+            if ready {
+                let mut batch = Vec::with_capacity(min);
+                let mut probably_fake = Vec::with_capacity(min);
+                for entry in &self.updates {
+                    if entry.value().value.ttl() < self.interval {
+                        batch.push(entry.key().clone());
+                        if batch.len() == min {
+                            break;
+                        }
+                    }
+                    if entry.value().expiry.duration() == Duration::from_secs(entry.key().seconds())
+                    {
+                        probably_fake.push(entry.key().clone());
+                        if probably_fake.len() == min {
+                            break;
+                        }
+                    }
+                }
+                if let Some(remaining) = min.checked_sub(batch.len()) {
+                    let take = probably_fake.into_iter().take(remaining);
+                    batch.append(&mut take.collect());
+                }
+                if let Some(remaining) = min.checked_sub(batch.len()) {
+                    let take = self.updates.iter().take(remaining);
+                    batch.append(&mut take.map(|e| e.key().clone()).collect());
+                }
+                let mut result = HashMap::new();
+                for counter in &batch {
+                    let value = self.updates.get(counter).unwrap().clone();
+                    result.insert(counter.clone(), value);
+                }
+                let result = consumer(result).await;
+                for counter in &batch {
+                    self.updates
+                        .remove_if(counter, |_, v| v.no_pending_writes());
+                }
+                return result;
+            } else {
+                ready = select! {
+                    _ = self.notifier.notified() => {
+                        self.updates.len() >= min ||
+                        self.priority_flush
+                            .compare_exchange(true, false, Ordering::Release, Ordering::Acquire)
+                            .is_ok()
+                    },
+                    _ = interval.tick() => true,
+                }
+            }
+        }
+    }
+
+    pub fn add(&self, counter: Counter, value: Arc<CachedCounterValue>, priority: bool) {
+        let priority = priority || value.value.ttl() < self.interval;
+        self.updates.entry(counter).or_insert(value);
+        if priority {
+            self.priority_flush.store(true, Ordering::Release);
+        }
+        self.notifier.notify_one();
+    }
+}
+
+impl Default for Batcher {
+    fn default() -> Self {
+        Self::new(Duration::from_millis(100))
+    }
+}
+
 pub struct CountersCache {
     max_ttl_cached_counters: Duration,
     pub ttl_ratio_cached_counters: u64,
     cache: Cache<Counter, Arc<CachedCounterValue>>,
+    batcher: Batcher,
 }
 
 impl CachedCounterValue {
     pub fn from(counter: &Counter, value: i64, ttl: Duration) -> Self {
         let now = SystemTime::now();
         Self {
             value: AtomicExpiringValue::new(value, now + Duration::from_secs(counter.seconds())),
+            initial_value: AtomicI64::new(value),
             expiry: AtomicExpiryTime::from_now(ttl),
         }
     }
@@ -34,13 +140,57 @@
 
     pub fn set_from_authority(&self, counter: &Counter, value: i64, expiry: Duration) {
         let time_window = Duration::from_secs(counter.seconds());
+        self.initial_value.store(value, Ordering::SeqCst);
         self.value.set(value, time_window);
         self.expiry.update(expiry);
     }
 
     pub fn delta(&self, counter: &Counter, delta: i64) -> i64 {
-        self.value
-            .update(delta, counter.seconds(), SystemTime::now())
+        let value = self
+            .value
+            .update(delta, counter.seconds(), SystemTime::now());
+        if value == delta {
+            // new window, invalidate initial value
+            self.initial_value.store(0, Ordering::SeqCst);
+        }
+        value
+    }
+
+    pub fn pending_writes(&self) -> Result<i64, ()> {
+        let start = self.initial_value.load(Ordering::SeqCst);
+        let value = self.value.value_at(SystemTime::now());
+        let offset = if start == 0 {
+            value
+        } else {
+            let writes = value - start;
+            if writes > 0 {
+                writes
+            } else {
+                value
+            }
+        };
+        match self
+            .initial_value
+            .compare_exchange(start, value, Ordering::SeqCst, Ordering::SeqCst)
+        {
+            Ok(_) => Ok(offset),
+            Err(newer) => {
+                if newer == 0 {
+                    // We got expired in the meantime, this fresh value can wait the next iteration
+                    Ok(0)
+                } else {
+                    // Concurrent call to this method?
+                    // We could support that with a CAS loop in the future if needed
+                    Err(())
+                }
+            }
+        }
+    }
+
+    fn no_pending_writes(&self) -> bool {
+        let start = self.initial_value.load(Ordering::SeqCst);
+        let value = self.value.value_at(SystemTime::now());
+        value - start == 0
     }
 
     pub fn hits(&self, _: &Counter) -> i64 {
@@ -90,18 +240,31 @@
         self
     }
 
-    pub fn build(&self) -> CountersCache {
+    pub fn build(&self, period: Duration) -> CountersCache {
         CountersCache {
             max_ttl_cached_counters: self.max_ttl_cached_counters,
             ttl_ratio_cached_counters: self.ttl_ratio_cached_counters,
             cache: Cache::new(self.max_cached_counters as u64),
+            batcher: Batcher::new(period),
         }
     }
 }
 
 impl CountersCache {
     pub fn get(&self, counter: &Counter) -> Option<Arc<CachedCounterValue>> {
-        self.cache.get(counter)
+        let option = self.cache.get(counter);
+        if option.is_none() {
+            let from_queue = self.batcher.updates.get(counter);
+            if let Some(entry) = from_queue {
+                self.cache.insert(counter.clone(), entry.value().clone());
+                return Some(entry.value().clone())
+            }
+        }
+        option
+    }
+
+    pub fn batcher(&self) -> &Batcher {
+        &self.batcher
     }
 
     pub fn insert(
@@ -122,7 +285,11 @@
         if let Some(ttl) = cache_ttl.checked_sub(ttl_margin) {
             if ttl > Duration::ZERO {
                 let previous = self.cache.get_with(counter.clone(), || {
-                    Arc::new(CachedCounterValue::from(&counter, counter_val, cache_ttl))
+                    if let Some(entry) = self.batcher.updates.get(&counter) {
+                        entry.value().clone()
+                    } else {
+                        Arc::new(CachedCounterValue::from(&counter, counter_val, cache_ttl))
+                    }
                 });
                 if previous.expired_at(now) || previous.value.value() < counter_val {
                     previous.set_from_authority(&counter, counter_val, cache_ttl);
@@ -138,9 +305,21 @@
     }
 
     pub fn increase_by(&self, counter: &Counter, delta: i64) {
-        if let Some(val) = self.cache.get(counter) {
-            val.delta(counter, delta);
-        };
+        let mut priority = false;
+        let val = self.cache.get_with_by_ref(counter, || {
+            if let Some(entry) = self.batcher.updates.get(&counter) {
+                entry.value().clone()
+            } else {
+                priority = true;
+                Arc::new(
+                    // this TTL is wrong, it needs to be the cache's TTL, not the time window of our limit
+                    // todo fix when introducing the Batcher type!
+                    CachedCounterValue::from(counter, 0, Duration::from_secs(counter.seconds())),
+                )
+            }
+        });
+        val.delta(counter, delta);
+        self.batcher.add(counter.clone(), val.clone(), priority);
     }
 
     fn ttl_from_redis_ttl(
@@ -209,7 +388,7 @@
             values,
         );
 
-        let cache = CountersCacheBuilder::new().build();
+        let cache = CountersCacheBuilder::new().build(Duration::default());
         cache.insert(
             counter.clone(),
             Some(10),
@@ -236,7 +415,7 @@
             values,
         );
 
-        let cache = CountersCacheBuilder::new().build();
+        let cache = CountersCacheBuilder::new().build(Duration::default());
 
         assert!(cache.get(&counter).is_none());
     }
@@ -258,7 +437,7 @@
             values,
         );
 
-        let cache = CountersCacheBuilder::new().build();
+        let cache = CountersCacheBuilder::new().build(Duration::default());
         cache.insert(
             counter.clone(),
             Some(current_value),
@@ -289,7 +468,7 @@
             values,
         );
 
-        let cache = CountersCacheBuilder::new().build();
+        let cache = CountersCacheBuilder::new().build(Duration::default());
         cache.insert(
             counter.clone(),
             None,
@@ -318,7 +497,7 @@
             values,
         );
 
-        let cache = CountersCacheBuilder::new().build();
+        let cache = CountersCacheBuilder::new().build(Duration::default());
         cache.insert(
             counter.clone(),
             Some(current_val),