1use 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#[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#[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
39pub 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 fn increment(&self, key: &[u8], delta: i64) -> i64;
50 fn try_acquire(&self, key: &[u8], cost: u64, spec: Bucket, now_ms: u64) -> Acquire;
54}
55
56pub 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 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 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#[allow(clippy::indexing_slicing)] fn 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#[allow(clippy::indexing_slicing)] fn 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
150fn 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#[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 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
202const STATE_TABLE: TableDefinition<'_, &[u8], &[u8]> = TableDefinition::new("plecto_state");
205
206const REDB_CACHE_BYTES: usize = 64 << 20;
212
213const MAX_COMBINE_ROUNDS: u32 = 16;
218
219const MAX_BATCH_OPS: usize = 64;
222
223const DURABLE_FLUSH_EVERY: u64 = 1024;
230
231pub struct RedbBackend {
247 db: Database,
248 flush_every: u64,
249 non_durable_run: AtomicU64,
252 forced_flushes: AtomicU64,
254 jobs: mpsc::Sender<WriteJob>,
257 combine: Mutex<mpsc::Receiver<WriteJob>>,
262 #[cfg(test)]
264 test: CombineTestHooks,
265}
266
267#[cfg(test)]
270struct CombineTestHooks {
271 combine_rounds: u32,
272 max_batch_ops: usize,
273 fail_next_commits: AtomicU64,
274 max_batch_seen: AtomicU64,
276}
277
278enum 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 fn needs_durable_commit(&self) -> bool {
303 matches!(self, WriteOp::Set { .. } | WriteOp::Delete { .. })
304 }
305}
306
307enum WriteOutcome {
310 Done,
311 Counter(i64),
312 Acquired(Acquire),
313}
314
315struct WriteJob {
316 op: WriteOp,
317 reply: mpsc::Sender<anyhow::Result<WriteOutcome>>,
321}
322
323impl RedbBackend {
324 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 #[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 #[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 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 fn submit(&self, op: WriteOp) -> anyhow::Result<WriteOutcome> {
425 let (reply, outcome) = mpsc::channel();
426 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 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 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 #[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 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 }
500 }
501
502 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 fn record_successful_commit(&self, batch: &[WriteJob], durability: redb::Durability) {
519 if batch.iter().any(|j| j.op.needs_durable_commit()) {
520 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 _ => {
535 self.non_durable_run.store(0, Ordering::Relaxed);
536 self.forced_flushes.fetch_add(1, Ordering::Relaxed);
537 }
538 }
539 }
540
541 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 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 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 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 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 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 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 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 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 let spec = Bucket {
757 capacity: 5,
758 refill_tokens: 1,
759 refill_interval_ms: 1000,
760 };
761 backend.set(b"cb", vec![0xff; 3]); assert!(
763 !backend.try_acquire(b"cb", 1, spec, 0).allowed,
764 "corrupt bucket must fail closed, not start full"
765 );
766 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 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); 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 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()); 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 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 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()); 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 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 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 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 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 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 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}