api_rate_limiter/cache/
in_memory.rs

1use std::time::{Duration, Instant};
2use dashmap::DashMap;
3use crate::limiter::CacheBackend;
4
5#[derive(Debug)]
6struct CacheEntry {
7    value: u32,
8    expires_at: Instant,
9}
10
11/// An in-memory cache implementation of the `CacheBackend` trait.
12/// It uses a concurrent DashMap to store keys with their expiration.
13pub struct InMemoryCache {
14    store: DashMap<String, CacheEntry>,
15}
16
17impl InMemoryCache {
18    /// Creates a new in-memory cache instance.
19    pub fn new() -> Self {
20        InMemoryCache {
21            store: DashMap::new(),
22        }
23    }
24}
25
26impl CacheBackend for InMemoryCache {
27    fn get(&self, key: &str) -> Option<u32> {
28        if let Some(entry) = self.store.get(key) {
29            if entry.expires_at > Instant::now() {
30                // println!("Returning the current entry");
31                return Some(entry.value);
32            } else {
33                // Expired: remove the entry.
34                // println!("removing the current entry and returning None");
35                drop(entry);
36                self.store.remove(key);
37                // println!("removed the current entry and returning None");
38                return None;
39            }
40        } else {
41            // println!("no entry found and returning None");
42            return None;
43        }
44    }
45
46    fn set(&self, key: &str, value: u32, ttl: Duration) -> Result<(), String> {
47        let expires_at = Instant::now() + ttl;
48        let entry = CacheEntry { value, expires_at };
49        self.store.insert(key.to_string(), entry);
50        Ok(())
51    }
52
53    fn incr(&self, key: &str, amount: u32) -> Result<u32, String> {
54        let now = Instant::now();
55        if let Some(mut entry) = self.store.get_mut(key) {
56            if entry.expires_at <= now {
57                // If the entry is expired, reset it.
58                entry.value = amount;
59            } else {
60                entry.value += amount;
61            }
62            Ok(entry.value)
63        } else {
64            // Insert a new entry. The TTL will be set by the caller if needed.
65            self.store.insert(key.to_string(), CacheEntry {
66                value: amount,
67                expires_at: now, // Temporary; caller should update TTL with `set`.
68            });
69            Ok(amount)
70        }
71    }
72}