Skip to main content

plecto_host/
backend.rs

1//! Host-held state backend (ADR 000004 / 000011).
2//!
3//! A stateless filter (Fork 4) keeps its *mutable* business state here: raw KV bytes,
4//! atomic counters, and token-bucket rate limiters. `KvBackend` is the **seam** ADR
5//! 000011 asks for: the host-API impls and the lifecycle never name a concrete store,
6//! so swapping in-memory ↔ redb is local, and when wasmtime 46 makes host calls async
7//! only the redb impl moves behind a blocking pool — callers stay put.
8//!
9//! Sync today (wasmtime 45 sync path). Locks are **non-poisoning** (`parking_lot`): a
10//! panicking filter must not cascade a poisoned lock across every later request.
11//!
12//! Keys arrive already namespaced by filter identity + primitive tag (done in
13//! `HostState`, ADR 000011); a backend treats them as opaque bytes and never inspects
14//! the namespace.
15
16use std::collections::HashMap;
17use std::sync::atomic::{AtomicU64, Ordering};
18use std::sync::mpsc;
19
20use parking_lot::Mutex;
21use redb::{Database, ReadableDatabase, ReadableTable, TableDefinition};
22
23/// A token-bucket specification (mirrors the WIT `host-ratelimit.bucket`).
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct Bucket {
26    pub capacity: u64,
27    pub refill_tokens: u64,
28    pub refill_interval_ms: u64,
29}
30
31/// The outcome of a token-bucket acquire (mirrors the WIT `host-ratelimit.acquire`).
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct Acquire {
34    pub allowed: bool,
35    pub remaining: u64,
36    pub retry_after_ms: u64,
37}
38
39/// The place a stateless filter's mutable state lives. Object-safe so the host can hold
40/// `Arc<dyn KvBackend>` and pick the backend at construction. Every method is internally
41/// synchronized and infallible from the filter's view — a backend error is logged and
42/// resolved **fail-closed** (reads vanish, rate limits deny), never a panic on the data
43/// plane (bp-rust).
44pub trait KvBackend: Send + Sync {
45    fn get(&self, key: &[u8]) -> Option<Vec<u8>>;
46    fn set(&self, key: &[u8], value: Vec<u8>);
47    fn delete(&self, key: &[u8]);
48    /// Atomic add-and-get. An unset counter starts at 0; `delta` is signed.
49    fn increment(&self, key: &[u8], delta: i64) -> i64;
50    /// Atomic token-bucket acquire against the `now_ms` request-clock snapshot. The
51    /// refill + counting stay host-native (ADR 000005) — they never cross the WASM
52    /// boundary; the filter only decided to consult the limiter.
53    fn try_acquire(&self, key: &[u8], cost: u64, spec: Bucket, now_ms: u64) -> Acquire;
54}
55
56// --- pure token-bucket math (host-native, deterministic against `now_ms`) ---
57
58/// Refill then consume. State is `(tokens, last_refill_ms)`; the host advances `last`
59/// by whole intervals only, so no fractional tokens are lost between calls. Returns the
60/// new state to persist and the acquire outcome.
61///
62/// Pure and storage-agnostic: it owns no state and does no I/O, so the same math drives both
63/// the per-filter `host-ratelimit` capability (this module's backends) and the fast path's
64/// native per-route rate limiter (ADR 000033, `plecto-control`). A zero state `(0, 0)` refills
65/// from epoch and therefore reads as a full bucket on first use — a caller backing the state
66/// with a zero-initialised table needs no separate "first sight" sentinel.
67pub fn apply_bucket(
68    state: Option<(u64, u64)>,
69    cost: u64,
70    spec: Bucket,
71    now_ms: u64,
72) -> ((u64, u64), Acquire) {
73    let no_refill = spec.refill_interval_ms == 0 || spec.refill_tokens == 0;
74    let (tokens, last_refill) = match state {
75        // first sight of this bucket: start full as of now
76        None => (spec.capacity, now_ms),
77        Some((tokens, last)) if no_refill => (tokens.min(spec.capacity), last),
78        Some((tokens, last)) => {
79            let intervals = now_ms.saturating_sub(last) / spec.refill_interval_ms;
80            let refilled = tokens
81                .saturating_add(intervals.saturating_mul(spec.refill_tokens))
82                .min(spec.capacity);
83            let advanced = last.saturating_add(intervals.saturating_mul(spec.refill_interval_ms));
84            (refilled, advanced)
85        }
86    };
87
88    if tokens >= cost {
89        let remaining = tokens - cost;
90        (
91            (remaining, last_refill),
92            Acquire {
93                allowed: true,
94                remaining,
95                retry_after_ms: 0,
96            },
97        )
98    } else {
99        // Retry-After is a deliberate over-estimate: it counts whole refill intervals and
100        // ignores the fraction of the current interval already elapsed, so it can be late by up
101        // to one interval. That is the conservative side for an advisory hint — it never invites
102        // a retry that is too early. Tightening it would mean persisting sub-interval phase,
103        // not worth it for an advisory value.
104        let retry_after_ms = if no_refill {
105            u64::MAX
106        } else {
107            let needed = cost - tokens;
108            needed
109                .div_ceil(spec.refill_tokens)
110                .saturating_mul(spec.refill_interval_ms)
111        };
112        (
113            (tokens, last_refill),
114            Acquire {
115                allowed: false,
116                remaining: tokens,
117                retry_after_ms,
118            },
119        )
120    }
121}
122
123/// Decode a stored counter (tolerant of short/empty values: missing == 0).
124#[allow(clippy::indexing_slicing)] // n = bytes.len().min(8), so buf[..n]/bytes[..n] are always in-bounds
125fn decode_i64(bytes: &[u8]) -> i64 {
126    let mut buf = [0u8; 8];
127    let n = bytes.len().min(8);
128    buf[..n].copy_from_slice(&bytes[..n]);
129    i64::from_le_bytes(buf)
130}
131
132/// Decode bucket state `(tokens, last_refill_ms)` from 16 LE bytes (None if malformed).
133#[allow(clippy::indexing_slicing)] // length is checked (== 16) three lines above
134fn decode_bucket(bytes: &[u8]) -> Option<(u64, u64)> {
135    if bytes.len() != 16 {
136        return None;
137    }
138    let tokens = u64::from_le_bytes(bytes[0..8].try_into().ok()?);
139    let last = u64::from_le_bytes(bytes[8..16].try_into().ok()?);
140    Some((tokens, last))
141}
142
143fn encode_bucket(state: (u64, u64)) -> Vec<u8> {
144    let mut out = Vec::with_capacity(16);
145    out.extend_from_slice(&state.0.to_le_bytes());
146    out.extend_from_slice(&state.1.to_le_bytes());
147    out
148}
149
150/// Map stored bucket bytes to the state `apply_bucket` consumes, **fail-closed on corruption**.
151/// `None` (key absent) is a legitimate first sight → start full. Present-but-malformed bytes must
152/// NOT decode to a full bucket (that is fail-OPEN, inconsistent with the limiter's fail-closed
153/// stance); treat corruption as an empty bucket so the call is denied and the limiter self-heals
154/// via refill.
155fn bucket_input(raw: Option<&[u8]>, now_ms: u64) -> Option<(u64, u64)> {
156    raw.map(|bytes| decode_bucket(bytes).unwrap_or((0, now_ms)))
157}
158
159// --- in-memory backend (default; tests and single-process runs) ---
160
161/// Process-lifetime, in-memory backend. The default until a durable store is configured.
162#[derive(Default)]
163pub struct MemoryBackend {
164    map: Mutex<HashMap<Vec<u8>, Vec<u8>>>,
165}
166
167impl KvBackend for MemoryBackend {
168    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
169        self.map.lock().get(key).cloned()
170    }
171
172    fn set(&self, key: &[u8], value: Vec<u8>) {
173        self.map.lock().insert(key.to_vec(), value);
174    }
175
176    fn delete(&self, key: &[u8]) {
177        self.map.lock().remove(key);
178    }
179
180    fn increment(&self, key: &[u8], delta: i64) -> i64 {
181        let mut map = self.map.lock();
182        let cur = decode_i64(map.get(key).map(Vec::as_slice).unwrap_or(&[]));
183        // delta == 0 is a counter read (host-counter.get): return the current value without
184        // creating or rewriting the key (mirrors the redb read-txn branch).
185        if delta == 0 {
186            return cur;
187        }
188        let next = cur.saturating_add(delta);
189        map.insert(key.to_vec(), next.to_le_bytes().to_vec());
190        next
191    }
192
193    fn try_acquire(&self, key: &[u8], cost: u64, spec: Bucket, now_ms: u64) -> Acquire {
194        let mut map = self.map.lock();
195        let prev = bucket_input(map.get(key).map(Vec::as_slice), now_ms);
196        let (next, result) = apply_bucket(prev, cost, spec, now_ms);
197        map.insert(key.to_vec(), encode_bucket(next));
198        result
199    }
200}
201
202// --- redb backend (durable; ADR 000004) ---
203
204const STATE_TABLE: TableDefinition<'_, &[u8], &[u8]> = TableDefinition::new("plecto_state");
205
206/// Cap on redb's lazily-filled page cache. The library default (1 GiB) is sized for a
207/// process that IS the database; embedded stores conventionally cap far lower (SQLite ~2 MiB,
208/// RocksDB 32 MiB, Kafka Streams 50 MiB per store). The cache only evicts at the cap, so a
209/// long-lived proxy would otherwise grow toward the full gigabyte; 64 MiB still dwarfs the
210/// working set of 8–16-byte counters and bucket states.
211const REDB_CACHE_BYTES: usize = 64 << 20;
212
213/// Cap on consecutive combining rounds one caller performs before releasing the lock, so
214/// sustained write pressure cannot conscript a single caller as combiner indefinitely.
215/// 16 rounds bounds the conscription to roughly a millisecond of commits while keeping
216/// timeout-based re-election (the slow path) rare.
217const MAX_COMBINE_ROUNDS: u32 = 16;
218
219/// Cap on ops drained into one write transaction. Together with `MAX_COMBINE_ROUNDS` this
220/// bounds combiner conscription by *work*, not only by round count.
221const MAX_BATCH_OPS: usize = 64;
222
223/// Every `DURABLE_FLUSH_EVERY`th hot-path **op** is upgraded from `Durability::None` to
224/// `Immediate`. redb frees pages only at a durable commit, so an unbounded None-only run
225/// (a counter/ratelimit-heavy workload with no kv writes) would grow the file without bound.
226/// 1024 amortises the fsync to noise on the hot path while bounding both the file growth
227/// between durable commits and the window of updates a crash can lose. The run counts ops,
228/// not commits, so combining many ops into one txn does not stretch that bound.
229const DURABLE_FLUSH_EVERY: u64 = 1024;
230
231/// redb-backed durable state. Atomicity comes from redb's single-writer write
232/// transaction: each op's read-modify-write is applied inside one (ADR 000004). redb is
233/// fully synchronous; ADR 000011's async-aware seam is this `KvBackend` impl — callers
234/// never see the store or the batching behind it.
235///
236/// Writes amortize begin/commit (and fsync when durable) across concurrent callers — the
237/// classical **group-commit** idea (DeWitt et al., SIGMOD'84: share durable I/O across a
238/// commit group) — shaped as **flat combining** (Hendler et al., SPAA'10): a caller queues
239/// its op, then competes for the combiner lock; the winner drains up to `MAX_BATCH_OPS`
240/// queued ops and applies them inside ONE write transaction. redb serializes writers
241/// globally (`begin_write` blocks), so per-op transactions would serialize every filter
242/// and route on N× the begin/commit cost; combining keeps the single writer but pays
243/// that cost once per batch. The batch is self-clocking — no timer, no resident thread:
244/// an uncontended caller drains only its own op and runs it inline (zero thread handoff),
245/// while contended callers batch what accumulated while the previous combiner held the lock.
246pub struct RedbBackend {
247    db: Database,
248    flush_every: u64,
249    /// Hot-path (non-durable) ops committed since the last durable commit. Mutated only
250    /// while holding the `combine` lock, and only after a successful commit.
251    non_durable_run: AtomicU64,
252    /// Durable commits forced by the cadence (observability for the flush tests).
253    forced_flushes: AtomicU64,
254    /// Write queue, open for the backend's whole life (the paired receiver lives inside
255    /// `combine`, so the channel never disconnects while a caller can reach it).
256    jobs: mpsc::Sender<WriteJob>,
257    /// The combiner election: holding this lock IS being the combiner — it grants
258    /// exclusive drain access to the queue. Non-poisoning (parking_lot): if a combiner
259    /// died mid-batch its drained jobs' reply channels disconnect, so every waiter
260    /// resolves fail-closed instead of inheriting a poisoned lock.
261    combine: Mutex<mpsc::Receiver<WriteJob>>,
262    /// Test seams for combiner/durability contracts (production `open` leaves them inert).
263    #[cfg(test)]
264    test: CombineTestHooks,
265}
266
267/// Observability + fault injection for combiner tests. Not used on the production `open` path
268/// beyond default values that match the constants above.
269#[cfg(test)]
270struct CombineTestHooks {
271    combine_rounds: u32,
272    max_batch_ops: usize,
273    fail_next_commits: AtomicU64,
274    /// High-water mark of `batch.len()` seen by `apply_batch` (asserts the per-round cap).
275    max_batch_seen: AtomicU64,
276}
277
278/// One queued write, applied by the combiner inside the next batch's transaction.
279enum WriteOp {
280    Set {
281        key: Vec<u8>,
282        value: Vec<u8>,
283    },
284    Delete {
285        key: Vec<u8>,
286    },
287    Increment {
288        key: Vec<u8>,
289        delta: i64,
290    },
291    TryAcquire {
292        key: Vec<u8>,
293        cost: u64,
294        spec: Bucket,
295        now_ms: u64,
296    },
297}
298
299impl WriteOp {
300    /// Durable KV (`set`/`delete`) keeps its per-call `Immediate` guarantee: one in a batch
301    /// upgrades the whole commit, and the co-batched hot ops become durable for free.
302    fn needs_durable_commit(&self) -> bool {
303        matches!(self, WriteOp::Set { .. } | WriteOp::Delete { .. })
304    }
305}
306
307/// A completed op's result, sent back over the job's reply channel only after its batch
308/// commits — an outcome must never be observable before it is committed.
309enum WriteOutcome {
310    Done,
311    Counter(i64),
312    Acquired(Acquire),
313}
314
315struct WriteJob {
316    op: WriteOp,
317    /// One-shot reply: the caller blocks on the paired receiver until the batch commits.
318    /// In-flight jobs are bounded by calling threads; each combine round additionally caps
319    /// how many are drained (`MAX_BATCH_OPS` / test override).
320    reply: mpsc::Sender<anyhow::Result<WriteOutcome>>,
321}
322
323impl RedbBackend {
324    /// Open (or create) the redb database at `path`.
325    pub fn open(path: impl AsRef<std::path::Path>) -> anyhow::Result<Self> {
326        Self::open_inner(path, DURABLE_FLUSH_EVERY)
327    }
328
329    #[cfg(test)]
330    fn open_with_flush_every(
331        path: impl AsRef<std::path::Path>,
332        flush_every: u64,
333    ) -> anyhow::Result<Self> {
334        Self::open_with_combine_limits(path, flush_every, MAX_COMBINE_ROUNDS, MAX_BATCH_OPS)
335    }
336
337    /// Test constructor: override combiner round / per-round batch caps.
338    #[cfg(test)]
339    fn open_with_combine_limits(
340        path: impl AsRef<std::path::Path>,
341        flush_every: u64,
342        combine_rounds: u32,
343        max_batch_ops: usize,
344    ) -> anyhow::Result<Self> {
345        let mut b = Self::open_inner(path, flush_every)?;
346        b.test.combine_rounds = combine_rounds.max(1);
347        b.test.max_batch_ops = max_batch_ops.max(1);
348        Ok(b)
349    }
350
351    fn open_inner(path: impl AsRef<std::path::Path>, flush_every: u64) -> anyhow::Result<Self> {
352        let db = redb::Builder::new()
353            .set_cache_size(REDB_CACHE_BYTES)
354            .create(path)?;
355        let (jobs, queue) = mpsc::channel();
356        Ok(Self {
357            db,
358            flush_every,
359            non_durable_run: AtomicU64::new(0),
360            forced_flushes: AtomicU64::new(0),
361            jobs,
362            combine: Mutex::new(queue),
363            #[cfg(test)]
364            test: CombineTestHooks {
365                combine_rounds: MAX_COMBINE_ROUNDS,
366                max_batch_ops: MAX_BATCH_OPS,
367                fail_next_commits: AtomicU64::new(0),
368                max_batch_seen: AtomicU64::new(0),
369            },
370        })
371    }
372
373    #[cfg(test)]
374    fn forced_flushes(&self) -> u64 {
375        self.forced_flushes.load(Ordering::Relaxed)
376    }
377
378    #[cfg(test)]
379    fn non_durable_run(&self) -> u64 {
380        self.non_durable_run.load(Ordering::Relaxed)
381    }
382
383    #[cfg(test)]
384    fn max_batch_seen(&self) -> u64 {
385        self.test.max_batch_seen.load(Ordering::Relaxed)
386    }
387
388    #[cfg(test)]
389    fn fail_next_n_commits(&self, n: u64) {
390        self.test.fail_next_commits.store(n, Ordering::Relaxed);
391    }
392
393    /// Test seam: run `ops` through `apply_batch` as one transaction (no queue).
394    #[cfg(test)]
395    fn apply_ops_as_batch(&self, ops: Vec<WriteOp>) -> anyhow::Result<Vec<WriteOutcome>> {
396        let batch: Vec<_> = ops
397            .into_iter()
398            .map(|op| {
399                let (reply, _rx) = mpsc::channel();
400                WriteJob { op, reply }
401            })
402            .collect();
403        self.apply_batch(&batch)
404    }
405
406    fn get_inner(&self, key: &[u8]) -> anyhow::Result<Option<Vec<u8>>> {
407        let rtxn = self.db.begin_read()?;
408        let table = match rtxn.open_table(STATE_TABLE) {
409            Ok(t) => t,
410            // no writer has created the table yet → empty
411            Err(redb::TableError::TableDoesNotExist(_)) => return Ok(None),
412            Err(e) => return Err(e.into()),
413        };
414        Ok(table.get(key)?.map(|g| g.value().to_vec()))
415    }
416
417    /// Queue one write, then either combine (apply everything queued, ours included) or
418    /// wait for a concurrent combiner to commit ours. An error is the batch's transaction
419    /// failing; it resolves fail-closed at the trait impl (reads vanish, limiters deny).
420    ///
421    /// Waiters park on their OWN reply channel, never on the combiner lock: replying
422    /// wakes every serviced caller in parallel, where queueing on the lock would wake
423    /// them one serial handoff at a time and hand the throughput right back.
424    fn submit(&self, op: WriteOp) -> anyhow::Result<WriteOutcome> {
425        let (reply, outcome) = mpsc::channel();
426        // Unreachable: `self.combine` owns the receiver, so the channel outlives every
427        // caller — but the data plane never panics on it (bp-rust).
428        self.jobs
429            .send(WriteJob { op, reply })
430            .map_err(|_| anyhow::anyhow!("kv write queue closed"))?;
431        loop {
432            match outcome.try_recv() {
433                Ok(result) => return result,
434                Err(mpsc::TryRecvError::Empty) => {}
435                // A combiner died mid-batch and dropped our job — fail closed.
436                Err(mpsc::TryRecvError::Disconnected) => {
437                    anyhow::bail!("kv combiner dropped a queued write")
438                }
439            }
440            let Some(queue) = self.combine.try_lock() else {
441                // Another caller is combining, and its next drain will pick our job up.
442                // The timeout covers one race — the combiner drained empty and released
443                // without seeing our job — by re-entering the election above.
444                match outcome.recv_timeout(std::time::Duration::from_micros(100)) {
445                    Ok(result) => return result,
446                    Err(mpsc::RecvTimeoutError::Timeout) => continue,
447                    Err(mpsc::RecvTimeoutError::Disconnected) => {
448                        anyhow::bail!("kv combiner dropped a queued write")
449                    }
450                }
451            };
452            // We are the combiner. Drain-apply-reply until a drain comes up empty: a job
453            // queued during our commit parks on its reply channel (not the lock), so
454            // leaving it behind would strand it for a full timeout round. Cap the rounds
455            // so one caller cannot be conscripted indefinitely under sustained load —
456            // past the cap the residual jobs' owners re-elect via the timeout above.
457            #[cfg(test)]
458            let rounds = self.test.combine_rounds;
459            #[cfg(not(test))]
460            let rounds = MAX_COMBINE_ROUNDS;
461            #[cfg(test)]
462            let max_batch = self.test.max_batch_ops;
463            #[cfg(not(test))]
464            let max_batch = MAX_BATCH_OPS;
465            for _ in 0..rounds {
466                let mut batch = Vec::new();
467                while batch.len() < max_batch {
468                    match queue.try_recv() {
469                        Ok(job) => batch.push(job),
470                        Err(_) => break,
471                    }
472                }
473                if batch.is_empty() {
474                    break;
475                }
476                match self.apply_batch(&batch) {
477                    Ok(outcomes) => {
478                        for (job, out) in batch.iter().zip(outcomes) {
479                            let _ = job.reply.send(Ok(out));
480                        }
481                    }
482                    Err(e) => {
483                        // The whole batch rode one transaction, so all of it fails
484                        // together; each caller resolves its own op fail-closed.
485                        // Stop combining on storage failure — further rounds would only
486                        // amplify conscription while the underlying error persists.
487                        tracing::error!(error = %e, ops = batch.len(), "redb batch commit failed");
488                        for job in &batch {
489                            let _ = job
490                                .reply
491                                .send(Err(anyhow::anyhow!("batch commit failed: {e:#}")));
492                        }
493                        break;
494                    }
495                }
496            }
497            // Our own op rode the first round (we enqueued before winning the election),
498            // so the next `try_recv` returns it.
499        }
500    }
501
502    /// Choose durability for one combined commit without mutating the flush cadence.
503    /// Cadence side effects run only in `record_successful_commit` after `commit` succeeds.
504    fn durability_for_batch(&self, batch: &[WriteJob]) -> redb::Durability {
505        if batch.iter().any(|j| j.op.needs_durable_commit()) {
506            return redb::Durability::Immediate;
507        }
508        let ops = batch.len() as u64;
509        let run = self.non_durable_run.load(Ordering::Relaxed) + ops;
510        if run >= self.flush_every {
511            redb::Durability::Immediate
512        } else {
513            redb::Durability::None
514        }
515    }
516
517    /// Advance / reset the hot-path flush cadence after a successful commit.
518    fn record_successful_commit(&self, batch: &[WriteJob], durability: redb::Durability) {
519        if batch.iter().any(|j| j.op.needs_durable_commit()) {
520            // Matches main's set_inner/delete_inner: reset only after Immediate lands.
521            self.non_durable_run.store(0, Ordering::Relaxed);
522            return;
523        }
524        let ops = batch.len() as u64;
525        match durability {
526            redb::Durability::Immediate => {
527                self.non_durable_run.store(0, Ordering::Relaxed);
528                self.forced_flushes.fetch_add(1, Ordering::Relaxed);
529            }
530            redb::Durability::None => {
531                self.non_durable_run.fetch_add(ops, Ordering::Relaxed);
532            }
533            // Unused variants: prefer an extra fsync over an unbounded None run.
534            _ => {
535                self.non_durable_run.store(0, Ordering::Relaxed);
536                self.forced_flushes.fetch_add(1, Ordering::Relaxed);
537            }
538        }
539    }
540
541    /// Apply one batch inside a single write transaction. Outcomes are returned, not sent —
542    /// replies must not leave before the commit they depend on.
543    fn apply_batch(&self, batch: &[WriteJob]) -> anyhow::Result<Vec<WriteOutcome>> {
544        #[cfg(test)]
545        {
546            let len = batch.len() as u64;
547            self.test.max_batch_seen.fetch_max(len, Ordering::Relaxed);
548        }
549        let durability = self.durability_for_batch(batch);
550        let mut wtxn = self.db.begin_write()?;
551        // set_durability only errors if a persistent savepoint changed in this txn; we never
552        // use savepoints.
553        wtxn.set_durability(durability)?;
554        let mut outcomes = Vec::with_capacity(batch.len());
555        {
556            let mut table = wtxn.open_table(STATE_TABLE)?;
557            for job in batch {
558                outcomes.push(apply_op(&mut table, &job.op)?);
559            }
560        }
561        #[cfg(test)]
562        if self
563            .test
564            .fail_next_commits
565            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |n| n.checked_sub(1))
566            .is_ok()
567        {
568            drop(wtxn);
569            anyhow::bail!("injected commit failure");
570        }
571        wtxn.commit()?;
572        self.record_successful_commit(batch, durability);
573        Ok(outcomes)
574    }
575}
576
577fn apply_op(
578    table: &mut redb::Table<'_, &'static [u8], &'static [u8]>,
579    op: &WriteOp,
580) -> anyhow::Result<WriteOutcome> {
581    match op {
582        WriteOp::Set { key, value } => {
583            table.insert(key.as_slice(), value.as_slice())?;
584            Ok(WriteOutcome::Done)
585        }
586        WriteOp::Delete { key } => {
587            table.remove(key.as_slice())?;
588            Ok(WriteOutcome::Done)
589        }
590        WriteOp::Increment { key, delta } => {
591            let cur = table
592                .get(key.as_slice())?
593                .map(|g| decode_i64(g.value()))
594                .unwrap_or(0);
595            let next = cur.saturating_add(*delta);
596            table.insert(key.as_slice(), next.to_le_bytes().as_slice())?;
597            Ok(WriteOutcome::Counter(next))
598        }
599        WriteOp::TryAcquire {
600            key,
601            cost,
602            spec,
603            now_ms,
604        } => {
605            let prev = {
606                let guard = table.get(key.as_slice())?;
607                bucket_input(guard.as_ref().map(|g| g.value()), *now_ms)
608            };
609            let (next, result) = apply_bucket(prev, *cost, *spec, *now_ms);
610            table.insert(key.as_slice(), encode_bucket(next).as_slice())?;
611            Ok(WriteOutcome::Acquired(result))
612        }
613    }
614}
615
616impl KvBackend for RedbBackend {
617    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
618        match self.get_inner(key) {
619            Ok(v) => v,
620            Err(e) => {
621                tracing::error!(error = %e, "redb get failed; treating key as absent");
622                None
623            }
624        }
625    }
626
627    fn set(&self, key: &[u8], value: Vec<u8>) {
628        if let Err(e) = self.submit(WriteOp::Set {
629            key: key.to_vec(),
630            value,
631        }) {
632            tracing::error!(error = %e, "redb set failed; value dropped");
633        }
634    }
635
636    fn delete(&self, key: &[u8]) {
637        if let Err(e) = self.submit(WriteOp::Delete { key: key.to_vec() }) {
638            tracing::error!(error = %e, "redb delete failed");
639        }
640    }
641
642    fn increment(&self, key: &[u8], delta: i64) -> i64 {
643        // delta == 0 is the canonical counter READ (host-counter.get). Serve it from an MVCC
644        // read txn so it never queues behind the combiner or pays a commit (ADR 000004 /
645        // 000005 hot path). Only a real mutation takes the write path.
646        if delta == 0 {
647            return match self.get_inner(key) {
648                Ok(opt) => decode_i64(opt.as_deref().unwrap_or_default()),
649                Err(e) => {
650                    tracing::error!(error = %e, "redb increment failed; returning 0");
651                    0
652                }
653            };
654        }
655        match self.submit(WriteOp::Increment {
656            key: key.to_vec(),
657            delta,
658        }) {
659            Ok(WriteOutcome::Counter(v)) => v,
660            Ok(WriteOutcome::Done | WriteOutcome::Acquired(_)) => {
661                tracing::error!("redb increment got a mismatched outcome; returning 0");
662                0
663            }
664            Err(e) => {
665                tracing::error!(error = %e, "redb increment failed; returning 0");
666                0
667            }
668        }
669    }
670
671    fn try_acquire(&self, key: &[u8], cost: u64, spec: Bucket, now_ms: u64) -> Acquire {
672        // fail-closed: a limiter that cannot read its state denies (ADR 000004).
673        let denied = Acquire {
674            allowed: false,
675            remaining: 0,
676            retry_after_ms: spec.refill_interval_ms,
677        };
678        match self.submit(WriteOp::TryAcquire {
679            key: key.to_vec(),
680            cost,
681            spec,
682            now_ms,
683        }) {
684            Ok(WriteOutcome::Acquired(r)) => r,
685            Ok(WriteOutcome::Done | WriteOutcome::Counter(_)) => {
686                tracing::error!("redb try_acquire got a mismatched outcome; denying");
687                denied
688            }
689            Err(e) => {
690                tracing::error!(error = %e, "redb try_acquire failed; denying");
691                denied
692            }
693        }
694    }
695}
696
697#[cfg(test)]
698mod tests {
699    use std::sync::Arc;
700
701    use super::*;
702
703    /// The same behaviour suite must hold for every backend (the seam, ADR 000011).
704    fn kv_roundtrip(backend: &dyn KvBackend) {
705        assert_eq!(backend.get(b"k"), None);
706        backend.set(b"k", b"v".to_vec());
707        assert_eq!(backend.get(b"k"), Some(b"v".to_vec()));
708        backend.delete(b"k");
709        assert_eq!(backend.get(b"k"), None);
710    }
711
712    fn counter_is_atomic_add_and_get(backend: &dyn KvBackend) {
713        assert_eq!(backend.increment(b"c", 1), 1);
714        assert_eq!(backend.increment(b"c", 4), 5);
715        assert_eq!(backend.increment(b"c", -2), 3);
716        // get reads the same encoding the counter wrote
717        assert_eq!(decode_i64(&backend.get(b"c").unwrap()), 3);
718    }
719
720    fn token_bucket_drains_then_refills(backend: &dyn KvBackend) {
721        let spec = Bucket {
722            capacity: 2,
723            refill_tokens: 1,
724            refill_interval_ms: 1000,
725        };
726        // capacity 2 → two acquires allowed, third denied (no time passed)
727        assert!(backend.try_acquire(b"rl", 1, spec, 0).allowed);
728        assert!(backend.try_acquire(b"rl", 1, spec, 0).allowed);
729        let denied = backend.try_acquire(b"rl", 1, spec, 0);
730        assert!(!denied.allowed);
731        assert_eq!(denied.remaining, 0);
732        assert_eq!(denied.retry_after_ms, 1000, "1 token needs one interval");
733        // after one interval, one token refills → allowed again
734        assert!(backend.try_acquire(b"rl", 1, spec, 1000).allowed);
735    }
736
737    fn counter_read_via_zero_delta_does_not_create_key(backend: &dyn KvBackend) {
738        // host-counter.get maps to increment(key, 0); a pure read must not create the key — it
739        // takes the redb read-txn / memory read branch, off the single-writer write path.
740        assert_eq!(backend.increment(b"zc", 0), 0, "unset counter reads as 0");
741        assert_eq!(
742            backend.get(b"zc"),
743            None,
744            "a zero-delta read must not create the counter key"
745        );
746        assert_eq!(backend.increment(b"zc", 5), 5);
747        assert_eq!(
748            backend.increment(b"zc", 0),
749            5,
750            "zero-delta still reads the live value"
751        );
752    }
753
754    fn token_bucket_corrupt_state_fails_closed(backend: &dyn KvBackend) {
755        // A malformed stored bucket must DENY (fail-closed), never reset to full (fail-open).
756        let spec = Bucket {
757            capacity: 5,
758            refill_tokens: 1,
759            refill_interval_ms: 1000,
760        };
761        backend.set(b"cb", vec![0xff; 3]); // not 16 bytes → corrupt
762        assert!(
763            !backend.try_acquire(b"cb", 1, spec, 0).allowed,
764            "corrupt bucket must fail closed, not start full"
765        );
766        // and it self-heals: after one interval a refilled token is granted
767        assert!(backend.try_acquire(b"cb", 1, spec, 1000).allowed);
768    }
769
770    #[test]
771    fn memory_backend_behaviour() {
772        let b = MemoryBackend::default();
773        kv_roundtrip(&b);
774        counter_is_atomic_add_and_get(&b);
775        counter_read_via_zero_delta_does_not_create_key(&b);
776        token_bucket_drains_then_refills(&b);
777        token_bucket_corrupt_state_fails_closed(&b);
778    }
779
780    #[test]
781    fn redb_backend_behaviour() {
782        let dir = tempfile::tempdir().unwrap();
783        let b = RedbBackend::open(dir.path().join("state.redb")).unwrap();
784        kv_roundtrip(&b);
785        counter_is_atomic_add_and_get(&b);
786        counter_read_via_zero_delta_does_not_create_key(&b);
787        token_bucket_drains_then_refills(&b);
788        token_bucket_corrupt_state_fails_closed(&b);
789    }
790
791    #[test]
792    fn redb_periodic_durable_flush_caps_a_non_durable_run() {
793        // Hot-path ops (increment / try_acquire) skip the per-commit fsync
794        // (Durability::None), but redb frees pages only at a durable commit — an unbounded
795        // None-only run (a counter/ratelimit-heavy workload with no kv writes) would grow the
796        // file without bound. Every `flush_every`th hot-path **op** must be durable.
797        let dir = tempfile::tempdir().unwrap();
798        let b = RedbBackend::open_with_flush_every(dir.path().join("state.redb"), 4).unwrap();
799
800        let spec = Bucket {
801            capacity: 1000,
802            refill_tokens: 0,
803            refill_interval_ms: 0,
804        };
805        b.increment(b"c", 1);
806        b.increment(b"c", 1);
807        b.increment(b"c", 0); // a read (delta 0) commits nothing and must not advance the run
808        b.try_acquire(b"rl", 1, spec, 0);
809        assert_eq!(b.forced_flushes(), 0, "three hot-path ops stay non-durable");
810
811        b.try_acquire(b"rl", 1, spec, 0);
812        assert_eq!(
813            b.forced_flushes(),
814            1,
815            "the 4th op in a run is upgraded to durable"
816        );
817
818        for _ in 0..4 {
819            b.increment(b"c", 1);
820        }
821        assert_eq!(b.forced_flushes(), 2, "the cadence repeats");
822    }
823
824    #[test]
825    fn redb_durable_kv_write_resets_the_flush_cadence() {
826        // set/delete commit with Durability::Immediate, which already makes every earlier
827        // non-durable commit durable (and frees its pages) — the flush cadence restarts
828        // instead of forcing a redundant flush shortly after.
829        let dir = tempfile::tempdir().unwrap();
830        let b = RedbBackend::open_with_flush_every(dir.path().join("state.redb"), 4).unwrap();
831
832        b.increment(b"c", 1);
833        b.increment(b"c", 1);
834        b.increment(b"c", 1);
835        b.set(b"k", b"v".to_vec()); // durable; the run restarts
836        b.increment(b"c", 1);
837        b.increment(b"c", 1);
838        b.increment(b"c", 1);
839        assert_eq!(
840            b.forced_flushes(),
841            0,
842            "a durable set resets the non-durable run"
843        );
844        b.increment(b"c", 1);
845        assert_eq!(b.forced_flushes(), 1);
846    }
847
848    #[test]
849    fn redb_concurrent_writes_are_exact_under_combining() {
850        // 8 threads hammer one counter and one bucket through the combiner; group commit
851        // must not lose or double any read-modify-write. The counter must total exactly,
852        // and a capacity-100 no-refill bucket must admit exactly 100 of the 400 acquires.
853        let dir = tempfile::tempdir().unwrap();
854        let b = Arc::new(RedbBackend::open(dir.path().join("state.redb")).unwrap());
855        let spec = Bucket {
856            capacity: 100,
857            refill_tokens: 0,
858            refill_interval_ms: 0,
859        };
860        let admitted = Arc::new(AtomicU64::new(0));
861        let workers: Vec<_> = (0..8)
862            .map(|_| {
863                let b = Arc::clone(&b);
864                let admitted = Arc::clone(&admitted);
865                std::thread::spawn(move || {
866                    for _ in 0..50 {
867                        b.increment(b"c", 1);
868                        if b.try_acquire(b"rl", 1, spec, 0).allowed {
869                            admitted.fetch_add(1, Ordering::Relaxed);
870                        }
871                    }
872                })
873            })
874            .collect();
875        for w in workers {
876            w.join().unwrap();
877        }
878        assert_eq!(b.increment(b"c", 0), 400, "no increment lost or doubled");
879        assert_eq!(
880            admitted.load(Ordering::Relaxed),
881            100,
882            "the bucket admits exactly its capacity"
883        );
884    }
885
886    #[test]
887    fn redb_failed_durable_commit_does_not_reset_flush_cadence() {
888        // A failed Immediate (set/delete) must not clear non_durable_run — otherwise the
889        // next None-only stretch can exceed flush_every ops since the last successful
890        // durable commit (ADR 000093; main reset set/delete cadence only post-commit).
891        let dir = tempfile::tempdir().unwrap();
892        let b = RedbBackend::open_with_flush_every(dir.path().join("state.redb"), 4).unwrap();
893        b.increment(b"c", 1);
894        b.increment(b"c", 1);
895        b.increment(b"c", 1);
896        assert_eq!(b.non_durable_run(), 3);
897        b.fail_next_n_commits(1);
898        b.set(b"k", b"v".to_vec()); // Immediate attempt aborts before commit
899        assert!(b.get(b"k").is_none(), "failed set must not be observable");
900        assert_eq!(
901            b.non_durable_run(),
902            3,
903            "failed durable commit must leave the hot-path run intact"
904        );
905        b.increment(b"c", 1);
906        assert_eq!(
907            b.forced_flushes(),
908            1,
909            "the next hot op must still trip the cadence at flush_every"
910        );
911    }
912
913    #[test]
914    fn redb_flush_cadence_counts_ops_not_commits() {
915        // One write txn carrying N hot ops must advance the run by N, not by 1.
916        let dir = tempfile::tempdir().unwrap();
917        let b = RedbBackend::open_with_flush_every(dir.path().join("state.redb"), 4).unwrap();
918        let ops = (0..4)
919            .map(|_| WriteOp::Increment {
920                key: b"c".to_vec(),
921                delta: 1,
922            })
923            .collect();
924        b.apply_ops_as_batch(ops).unwrap();
925        assert_eq!(
926            b.forced_flushes(),
927            1,
928            "4 ops in one commit trip flush_every=4"
929        );
930        assert_eq!(b.increment(b"c", 0), 4);
931        let ops = (0..3)
932            .map(|_| WriteOp::Increment {
933                key: b"c".to_vec(),
934                delta: 1,
935            })
936            .collect();
937        b.apply_ops_as_batch(ops).unwrap();
938        assert_eq!(
939            b.forced_flushes(),
940            1,
941            "3 more ops stay under the next threshold"
942        );
943        b.increment(b"c", 1);
944        assert_eq!(
945            b.forced_flushes(),
946            2,
947            "the 4th op of the new run forces a flush"
948        );
949    }
950
951    #[test]
952    fn redb_batch_commit_failure_fails_closed_for_every_op() {
953        let dir = tempfile::tempdir().unwrap();
954        let b = RedbBackend::open(dir.path().join("state.redb")).unwrap();
955        b.fail_next_n_commits(1);
956        let ops = vec![
957            WriteOp::Increment {
958                key: b"c".to_vec(),
959                delta: 1,
960            },
961            WriteOp::Increment {
962                key: b"c".to_vec(),
963                delta: 1,
964            },
965            WriteOp::Set {
966                key: b"k".to_vec(),
967                value: b"v".to_vec(),
968            },
969        ];
970        assert!(b.apply_ops_as_batch(ops).is_err());
971        assert_eq!(b.increment(b"c", 0), 0, "no partial counter apply");
972        assert!(b.get(b"k").is_none(), "no partial kv apply");
973    }
974
975    #[test]
976    fn redb_mixed_durable_and_hot_batch_resets_cadence_without_forced_flush() {
977        // set in a batch upgrades the commit to Immediate; co-batched hot ops must not
978        // bump forced_flushes, and the hot-path run restarts after success.
979        let dir = tempfile::tempdir().unwrap();
980        let b = RedbBackend::open_with_flush_every(dir.path().join("state.redb"), 4).unwrap();
981        b.increment(b"c", 1);
982        b.increment(b"c", 1);
983        b.increment(b"c", 1);
984        b.apply_ops_as_batch(vec![
985            WriteOp::Increment {
986                key: b"c".to_vec(),
987                delta: 1,
988            },
989            WriteOp::Set {
990                key: b"k".to_vec(),
991                value: b"v".to_vec(),
992            },
993        ])
994        .unwrap();
995        assert_eq!(b.get(b"k"), Some(b"v".to_vec()));
996        assert_eq!(b.forced_flushes(), 0);
997        assert_eq!(b.non_durable_run(), 0);
998        for _ in 0..3 {
999            b.increment(b"c", 1);
1000        }
1001        assert_eq!(b.forced_flushes(), 0);
1002        b.increment(b"c", 1);
1003        assert_eq!(b.forced_flushes(), 1);
1004    }
1005
1006    #[test]
1007    fn redb_per_round_batch_cap_bounds_combiner_work() {
1008        // Under contention, a combiner round must not drain more than max_batch_ops.
1009        let dir = tempfile::tempdir().unwrap();
1010        let b = Arc::new(
1011            RedbBackend::open_with_combine_limits(dir.path().join("state.redb"), 1024, 16, 2)
1012                .unwrap(),
1013        );
1014        let workers: Vec<_> = (0..8)
1015            .map(|_| {
1016                let b = Arc::clone(&b);
1017                std::thread::spawn(move || {
1018                    for _ in 0..40 {
1019                        b.increment(b"c", 1);
1020                    }
1021                })
1022            })
1023            .collect();
1024        for w in workers {
1025            w.join().unwrap();
1026        }
1027        assert_eq!(b.increment(b"c", 0), 320);
1028        assert!(
1029            b.max_batch_seen() <= 2,
1030            "max_batch_seen={} exceeded cap 2",
1031            b.max_batch_seen()
1032        );
1033    }
1034
1035    #[test]
1036    fn redb_residual_jobs_after_round_cap_still_commit() {
1037        // combine_rounds=1 and max_batch_ops=1 force frequent re-election; every op must land.
1038        let dir = tempfile::tempdir().unwrap();
1039        let b = Arc::new(
1040            RedbBackend::open_with_combine_limits(dir.path().join("state.redb"), 1024, 1, 1)
1041                .unwrap(),
1042        );
1043        let workers: Vec<_> = (0..8)
1044            .map(|_| {
1045                let b = Arc::clone(&b);
1046                std::thread::spawn(move || {
1047                    for _ in 0..25 {
1048                        b.increment(b"c", 1);
1049                    }
1050                })
1051            })
1052            .collect();
1053        for w in workers {
1054            w.join().unwrap();
1055        }
1056        assert_eq!(b.increment(b"c", 0), 200);
1057    }
1058
1059    #[test]
1060    fn redb_concurrent_set_delete_are_exact_under_combining() {
1061        let dir = tempfile::tempdir().unwrap();
1062        let b = Arc::new(RedbBackend::open(dir.path().join("state.redb")).unwrap());
1063        let workers: Vec<_> = (0..8)
1064            .map(|i| {
1065                let b = Arc::clone(&b);
1066                std::thread::spawn(move || {
1067                    let key = format!("k{i}").into_bytes();
1068                    for n in 0..20u8 {
1069                        b.set(&key, vec![n]);
1070                    }
1071                    b.delete(&key);
1072                })
1073            })
1074            .collect();
1075        for w in workers {
1076            w.join().unwrap();
1077        }
1078        for i in 0..8 {
1079            assert_eq!(b.get(format!("k{i}").as_bytes()), None);
1080        }
1081    }
1082
1083    /// Same partitioning the `kv_backend` bench must use: total work == Criterion `iters`.
1084    fn split_iters_across_threads(threads: u64, iters: u64) -> Vec<u64> {
1085        let threads = threads.max(1);
1086        let base = iters / threads;
1087        let rem = iters % threads;
1088        (0..threads).map(|t| base + u64::from(t < rem)).collect()
1089    }
1090
1091    #[test]
1092    fn split_iters_across_threads_sums_exactly_to_iters() {
1093        for threads in [1u64, 2, 8] {
1094            for iters in [0u64, 1, 3, 7, 8, 100] {
1095                let parts = split_iters_across_threads(threads, iters);
1096                assert_eq!(parts.len(), threads as usize);
1097                assert_eq!(
1098                    parts.iter().copied().sum::<u64>(),
1099                    iters,
1100                    "threads={threads} iters={iters} parts={parts:?}"
1101                );
1102            }
1103        }
1104    }
1105
1106    #[test]
1107    fn redb_persists_across_reopen() {
1108        let dir = tempfile::tempdir().unwrap();
1109        let path = dir.path().join("state.redb");
1110        {
1111            let b = RedbBackend::open(&path).unwrap();
1112            assert_eq!(b.increment(b"hits", 3), 3);
1113        }
1114        // reopening the same file recovers durable state (ADR 000004 durability)
1115        let b = RedbBackend::open(&path).unwrap();
1116        assert_eq!(b.increment(b"hits", 1), 4);
1117    }
1118
1119    #[test]
1120    fn token_bucket_cost_zero_always_allowed() {
1121        let spec = Bucket {
1122            capacity: 1,
1123            refill_tokens: 1,
1124            refill_interval_ms: 1000,
1125        };
1126        let (_state, r) = apply_bucket(Some((0, 0)), 0, spec, 0);
1127        assert!(r.allowed, "a zero-cost acquire never blocks");
1128    }
1129}