featherbit/stores/
counter.rs

1//! Redis-backed fixed-window counter (`policy: redis` for limit-count and
2//! the workflow limit-count action).
3//!
4//! Increment-and-check runs as one server-side Lua script so concurrent
5//! gateway instances count atomically. Window boundaries are wall-clock
6//! aligned (`now / window`), so every instance agrees on them — unlike the
7//! local store's per-process `Instant` windows. Backend errors bump
8//! `gateway_counter_store_errors_total{store}` and surface as
9//! [`CounterError`]; the calling plugin decides fail-open vs reject.
10
11use std::sync::Arc;
12use std::time::{Duration, SystemTime, UNIX_EPOCH};
13
14use async_trait::async_trait;
15
16use crate::metrics::GatewayMetrics;
17use crate::ratelimit::{CounterError, CounterStore, WindowResult};
18
19use super::redis_store::RedisStoreClient;
20
21/// `INCR` + first-increment `PEXPIRE` + `PTTL`, atomically.
22const FIXED_WINDOW_SCRIPT: &str = r#"
23local current = redis.call('INCR', KEYS[1])
24if current == 1 then
25  redis.call('PEXPIRE', KEYS[1], ARGV[1])
26end
27local ttl = redis.call('PTTL', KEYS[1])
28return {current, ttl}
29"#;
30
31/// Wall-clock window slot: the window index (key component shared by all
32/// instances) and the milliseconds from `now_ms` to the window's end (the
33/// PEXPIRE argument). Pure, so the boundary math is unit-testable.
34fn window_slot(now_ms: u64, window_ms: u64) -> (u64, u64) {
35    let window_ms = window_ms.max(1);
36    let start = now_ms / window_ms;
37    let expire_ms = (start + 1) * window_ms - now_ms;
38    (start, expire_ms)
39}
40
41/// The key for one fixed window of one counter.
42pub(crate) fn window_key(prefix: &str, slot: u64, key: &str) -> String {
43    format!(
44        "{}:{}:{}:{}",
45        prefix,
46        super::namespaces::COUNTERS,
47        slot,
48        key
49    )
50}
51
52pub struct RedisCounterStore {
53    client: Arc<RedisStoreClient>,
54    store_name: String,
55    metrics: Option<Arc<GatewayMetrics>>,
56    script: redis::Script,
57}
58
59impl RedisCounterStore {
60    pub fn new(
61        client: Arc<RedisStoreClient>,
62        store_name: String,
63        metrics: Option<Arc<GatewayMetrics>>,
64    ) -> Self {
65        Self {
66            client,
67            store_name,
68            metrics,
69            script: redis::Script::new(FIXED_WINDOW_SCRIPT),
70        }
71    }
72
73    fn backend_err(&self, msg: String) -> CounterError {
74        if let Some(ref m) = self.metrics {
75            m.counter_store_errors
76                .with_label_values(&[&self.store_name])
77                .inc();
78        }
79        tracing::warn!(store = %self.store_name, "counter store error: {}", msg);
80        CounterError(msg)
81    }
82}
83
84#[async_trait]
85impl CounterStore for RedisCounterStore {
86    async fn incr_fixed_window(
87        &self,
88        key: &str,
89        limit: u64,
90        window: Duration,
91    ) -> Result<WindowResult, CounterError> {
92        let now_ms = SystemTime::now()
93            .duration_since(UNIX_EPOCH)
94            .unwrap_or_default()
95            .as_millis() as u64;
96        let (slot, expire_ms) = window_slot(now_ms, window.as_millis() as u64);
97        let redis_key = window_key(self.client.key_prefix(), slot, key);
98
99        let mut conn = self.client.conn().await.map_err(|e| self.backend_err(e))?;
100        let (count, pttl): (u64, i64) = self
101            .script
102            .key(redis_key.as_str())
103            .arg(expire_ms)
104            .invoke_async(&mut conn)
105            .await
106            .map_err(|e| self.backend_err(format!("fixed-window script: {}", e)))?;
107
108        Ok(WindowResult {
109            allowed: count <= limit,
110            remaining: limit.saturating_sub(count),
111            reset: Duration::from_millis(pttl.max(0) as u64),
112            limit,
113        })
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    /// Window slots are wall-clock aligned and expiry lands exactly on the
122    /// window boundary — the property that makes limits cluster-consistent.
123    #[test]
124    fn test_window_slot_alignment() {
125        // 10s window: 25_000ms is 5s into slot 2, 5s left.
126        assert_eq!(window_slot(25_000, 10_000), (2, 5_000));
127        // Exactly on a boundary: full window remains.
128        assert_eq!(window_slot(30_000, 10_000), (3, 10_000));
129        // 1ms before the boundary.
130        assert_eq!(window_slot(29_999, 10_000), (2, 1));
131        // Degenerate zero window is clamped, never divides by zero.
132        assert_eq!(window_slot(5, 0), (5, 1));
133    }
134
135    /// Live-backend atomicity test; skipped unless FEATHERBIT_TEST_REDIS_URL
136    /// is set. N concurrent tasks never over-admit past the limit.
137    #[tokio::test]
138    async fn test_concurrent_increments_never_exceed_limit_live() {
139        let Ok(url) = std::env::var("FEATHERBIT_TEST_REDIS_URL") else {
140            eprintln!(
141                "skipping test_concurrent_increments_never_exceed_limit_live: FEATHERBIT_TEST_REDIS_URL not set"
142            );
143            return;
144        };
145        let cfg: crate::config::StoreConfig = serde_yaml::from_str(&format!(
146            "name: live\ntype: redis\nurl: {url}\nkey_prefix: fbtest{}\n",
147            std::process::id()
148        ))
149        .unwrap();
150        let client = Arc::new(RedisStoreClient::build(&cfg).unwrap());
151        let store = Arc::new(RedisCounterStore::new(client, "live".to_string(), None));
152
153        let limit = 10u64;
154        let window = Duration::from_secs(60);
155        let mut handles = Vec::new();
156        for _ in 0..40 {
157            let store = store.clone();
158            handles.push(tokio::spawn(async move {
159                store
160                    .incr_fixed_window("conc-key", limit, window)
161                    .await
162                    .unwrap()
163                    .allowed
164            }));
165        }
166        let mut admitted = 0;
167        for h in handles {
168            if h.await.unwrap() {
169                admitted += 1;
170            }
171        }
172        assert_eq!(admitted as u64, limit, "exactly `limit` requests admitted");
173    }
174}
featherbit/stores/counter.rs

featherbit/stores/
counter.rs
