cratefield_runtime_cloudflare/ports/
d1_rate_limit.rs

1//! `RateLimiter` over D1 (issue #538): a fixed window per key, so limits
2//! can vary per key where the Workers Rate Limiting binding offers one
3//! boolean per namespace. One row per key, and the whole turn is a single
4//! atomic upsert whose `RETURNING` row carries the new count — no
5//! read-before-write race, even with isolates in different regions hitting
6//! the same key inside one window.
7
8use async_trait::async_trait;
9use std::sync::Arc;
10use std::time::Duration;
11
12use cratefield_core::{Decision, Quota, RateLimitError, RateLimiter};
13use worker::D1Database as WorkerD1;
14use worker::send::IntoSendFuture;
15
16/// The forward-only DDL for the `rate_limit_counters` table the upsert
17/// below counts in. A venture wiring `.d1_rate_limiter(..)` ships this as
18/// a migration in its `migrations/` directory, next to the module SQL.
19pub const RATE_LIMIT_COUNTERS_SQL: &str = "CREATE TABLE IF NOT EXISTS rate_limit_counters (\
20key TEXT PRIMARY KEY, \
21window_start INTEGER NOT NULL, \
22count INTEGER NOT NULL);";
23
24/// One spent request, counted in the window the statement below opens.
25const UPSERT_SQL: &str = "INSERT INTO rate_limit_counters \
26(key, window_start, count) VALUES (?1, ?2, 1) \
27ON CONFLICT(key) DO UPDATE SET \
28count = CASE WHEN window_start = excluded.window_start THEN count + 1 ELSE 1 END, \
29window_start = excluded.window_start \
30RETURNING count, window_start";
31
32/// The per-key budget a [`RateLimitPolicy`] answers with: `max` requests
33/// per `period`.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct Limit {
36    /// Requests allowed in one window; the request that raises the count
37    /// to `max` is the last one allowed before the window turns.
38    pub max: u32,
39    /// How long one window runs. Windows align to whole periods since the
40    /// epoch, so every isolate agrees on the edges without talking.
41    pub period: Duration,
42}
43
44/// Maps a rate-limit key to its budget — the hook that makes limits
45/// per-key. `None` means **unlimited**: the key is allowed without a D1
46/// round trip and no row is written for it. Apps encode the plan in the
47/// key itself (`plan:pro:ip:203.0.113.7`, `plan:pro:email:n@x.co`) and
48/// match on the prefix; the `ip:`/`email:` keys of
49/// `cratefield_core::rate_limit_keys` slot in unchanged.
50pub type RateLimitPolicy = Arc<dyn Fn(&str) -> Option<Limit> + Send + Sync>;
51
52/// The answer for an unlimited key: allowed, nothing spent, nothing known.
53const UNLIMITED: Decision = Decision {
54    ok: true,
55    retry_after: None,
56    quota: None,
57};
58
59pub struct D1RateLimiter {
60    db: WorkerD1,
61    policy: RateLimitPolicy,
62}
63
64impl D1RateLimiter {
65    #[must_use]
66    pub fn new(db: WorkerD1, policy: RateLimitPolicy) -> Self {
67        Self { db, policy }
68    }
69}
70
71/// The whole period in whole milliseconds, at least one: a degenerate
72/// `period` must not divide by zero, and a window of less than a
73/// millisecond would count nothing anyway.
74fn period_ms(period: Duration) -> i64 {
75    i64::try_from(period.as_millis()).unwrap_or(i64::MAX).max(1)
76}
77
78/// The window a timestamp falls in: the whole multiple of the period the
79/// epoch last reached before it.
80fn window_start(now_ms: i64, period: Duration) -> i64 {
81    let period = period_ms(period);
82    now_ms - now_ms.rem_euclid(period)
83}
84
85/// The [`Decision`] for the count the upsert returned. `count` is what the
86/// window has now seen *including* this request, so exactly the `max`-th
87/// request passes and the next one waits for the rest of the window —
88/// which is both its `retry_after` and the quota's `reset`.
89fn decide(limit: Limit, now_ms: i64, window: i64, count: i64) -> Decision {
90    let left_ms = (window + period_ms(limit.period) - now_ms).max(0);
91    let reset = Duration::from_millis(u64::try_from(left_ms).unwrap_or(u64::MAX));
92    let remaining = i64::from(limit.max)
93        .saturating_sub(count)
94        .clamp(0, i64::from(u32::MAX));
95    let quota = Quota {
96        limit: limit.max,
97        // `unwrap_or(0)` covers the negative side (`try_from` fails); the
98        // clamp above already bounds the positive side.
99        remaining: u32::try_from(remaining).unwrap_or(0),
100        reset,
101    };
102    if count <= i64::from(limit.max) {
103        Decision {
104            ok: true,
105            retry_after: None,
106            quota: Some(quota),
107        }
108    } else {
109        Decision {
110            ok: false,
111            retry_after: Some(reset),
112            quota: Some(Quota {
113                remaining: 0,
114                ..quota
115            }),
116        }
117    }
118}
119
120fn transport(err: &worker::Error) -> RateLimitError {
121    RateLimitError::Transport(err.to_string())
122}
123
124#[async_trait]
125impl RateLimiter for D1RateLimiter {
126    async fn limit(&self, key: &str) -> Result<Decision, RateLimitError> {
127        let Some(limit) = (self.policy)(key) else {
128            return Ok(UNLIMITED);
129        };
130        let now = time::OffsetDateTime::now_utc();
131        let now_ms = now.unix_timestamp() * 1_000 + i64::from(now.millisecond());
132        let window = window_start(now_ms, limit.period);
133
134        // Bind through serde_json + serde_wasm_bindgen, the path
135        // `ports/d1.rs` proved under workerd (the `D1Type` binds fail or
136        // hang there). Both values are JSON-representable: a string and a
137        // millisecond timestamp (< 2^53).
138        let json = [
139            serde_json::Value::String(key.to_owned()),
140            serde_json::Value::from(window),
141        ];
142        let values: Vec<worker::wasm_bindgen::JsValue> = json
143            .iter()
144            .map(worker::d1::serde_wasm_bindgen::to_value)
145            .collect::<Result<Vec<_>, _>>()
146            .map_err(|err| transport(&worker::Error::SerdeWasmBindgenError(err)))?;
147        let prepared = self
148            .db
149            .prepare(UPSERT_SQL)
150            .bind(&values)
151            .map_err(|err| transport(&err))?;
152        // Writes resolve through `batch`, not `.run()`/`.all()` (see
153        // `ports/d1.rs`): plain write promises hang under local workerd.
154        // The `RETURNING` row comes back in the statement's results.
155        let results = self
156            .db
157            .batch(vec![prepared])
158            .into_send()
159            .await
160            .map_err(|err| transport(&err))?;
161        let row: serde_json::Value = results
162            .first()
163            .and_then(|result| result.results().ok())
164            .and_then(|rows| rows.into_iter().next())
165            .ok_or_else(|| {
166                RateLimitError::Transport("rate-limit upsert returned no row".to_owned())
167            })?;
168        let count = row
169            .get("count")
170            .and_then(serde_json::Value::as_i64)
171            .unwrap_or(1);
172        let returned_window = row
173            .get("window_start")
174            .and_then(serde_json::Value::as_i64)
175            .unwrap_or(window);
176        Ok(decide(limit, now_ms, returned_window, count))
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183
184    const MINUTE: Duration = Duration::from_secs(60);
185
186    #[test]
187    fn windows_align_to_whole_periods_since_the_epoch() {
188        // 12:01:59.500 falls in the window that opened at 12:01:00.
189        assert_eq!(
190            window_start(12 * 60 * 1_000 + 1_500, MINUTE),
191            12 * 60 * 1_000
192        );
193        assert_eq!(window_start(0, MINUTE), 0);
194    }
195
196    #[test]
197    fn a_degenerate_period_divides_by_one_millisecond() {
198        assert_eq!(window_start(1_500, Duration::ZERO), 1_500);
199    }
200
201    #[test]
202    fn the_nth_request_passes_and_the_next_waits_out_the_window() {
203        let limit = Limit {
204            max: 3,
205            period: MINUTE,
206        };
207        let (now, window) = (90_000, 60_000);
208
209        let third = decide(limit, now, window, 3);
210        assert!(third.ok);
211        assert_eq!(third.retry_after, None);
212        let quota = third.quota.expect("quota");
213        assert_eq!((quota.limit, quota.remaining), (3, 0));
214        assert_eq!(quota.reset, Duration::from_secs(30));
215
216        let fourth = decide(limit, now, window, 4);
217        assert!(!fourth.ok);
218        assert_eq!(fourth.retry_after, Some(Duration::from_secs(30)));
219        let quota = fourth.quota.expect("quota");
220        assert_eq!((quota.limit, quota.remaining), (3, 0));
221        assert_eq!(quota.reset, Duration::from_secs(30));
222    }
223
224    #[test]
225    fn remaining_counts_down_from_the_limit() {
226        let limit = Limit {
227            max: 5,
228            period: MINUTE,
229        };
230        let decision = decide(limit, 0, 0, 2);
231        assert!(decision.ok);
232        assert_eq!(decision.quota.expect("quota").remaining, 3);
233    }
234
235    #[test]
236    fn a_count_beyond_the_cap_clamps_to_zero_remaining() {
237        let limit = Limit {
238            max: 5,
239            period: MINUTE,
240        };
241        let decision = decide(limit, 1_000, 0, 500);
242        assert!(!decision.ok);
243        assert_eq!(decision.quota.expect("quota").remaining, 0);
244    }
245}
cratefield_runtime_cloudflare/ports/d1_rate_limit.rs

cratefield_runtime_cloudflare/ports/
d1_rate_limit.rs
