kevy_embedded/ops_blocking.rs
1//! v2.4 — embedded blocking pops (`blpop` / `brpop` / `bzpopmin`).
2//!
3//! Design (the park-wait note, per ROADMAP):
4//! - One process-wide [`Blocker`] per store: a wake **generation**
5//! under a mutex + condvar, plus an atomic waiter count.
6//! - **Writers** (every `commit_write`) check `waiters` with a single
7//! Relaxed load — zero cost while nobody blocks. When waiters exist,
8//! bump the generation + `notify_all`. Any write wakes every
9//! waiter; each re-polls its own keys (spurious wakeups are cheap
10//! re-polls — correctness over per-key bookkeeping at embedded
11//! contention levels).
12//! - **Waiters** loop: non-blocking poll across their keys → if
13//! empty, wait on the condvar with the remaining deadline, keyed to
14//! the generation observed *before* the poll — the classic
15//! recheck-after-wait pattern, so a push landing between the poll
16//! and the wait is never lost.
17//! - `timeout = None` blocks indefinitely (matches `BLPOP key 0`).
18//!
19//! Fairness: multiple blocked consumers re-poll under their shard
20//! write locks; the lock queue arbitrates. No FIFO ticket order is
21//! promised (same as the server's cross-shard arbiter).
22
23use std::io;
24use std::sync::atomic::{AtomicUsize, Ordering};
25use std::sync::{Condvar, Mutex};
26use std::time::{Duration, Instant};
27
28use crate::store::Store;
29
30/// `(key, member, score)` from a zset blocking pop.
31type ZPopHit = (Vec<u8>, Vec<u8>, f64);
32
33/// Process-wide wake channel for blocking pops.
34pub(crate) struct Blocker {
35 waiters: AtomicUsize,
36 generation: Mutex<u64>,
37 cv: Condvar,
38}
39
40impl Blocker {
41 pub(crate) fn new() -> Self {
42 Self {
43 waiters: AtomicUsize::new(0),
44 generation: Mutex::new(0),
45 cv: Condvar::new(),
46 }
47 }
48
49 /// Writer side — called from `commit_write`. One Relaxed load when
50 /// idle; lock + notify only with live waiters.
51 #[inline]
52 pub(crate) fn wake_all(&self) {
53 if self.waiters.load(Ordering::Relaxed) == 0 {
54 return;
55 }
56 let mut g = self.generation.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
57 *g += 1;
58 self.cv.notify_all();
59 }
60
61 /// Current generation (observe BEFORE polling, wait against it).
62 fn generation(&self) -> u64 {
63 *self.generation.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
64 }
65
66 /// Park until the generation moves past `seen` or `deadline`
67 /// passes. Returns `false` on timeout.
68 fn wait_past(&self, seen: u64, deadline: Option<Instant>) -> bool {
69 self.waiters.fetch_add(1, Ordering::Relaxed);
70 let mut g = self.generation.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
71 let ok = loop {
72 if *g != seen {
73 break true;
74 }
75 match deadline {
76 None => {
77 g = self
78 .cv
79 .wait(g)
80 .unwrap_or_else(std::sync::PoisonError::into_inner);
81 }
82 Some(d) => {
83 let now = Instant::now();
84 if now >= d {
85 break false;
86 }
87 let (guard, res) = self
88 .cv
89 .wait_timeout(g, d - now)
90 .unwrap_or_else(std::sync::PoisonError::into_inner);
91 g = guard;
92 if res.timed_out() && *g == seen {
93 break false;
94 }
95 }
96 }
97 };
98 drop(g);
99 self.waiters.fetch_sub(1, Ordering::Relaxed);
100 ok
101 }
102}
103
104impl Store {
105 /// `BLPOP` — block until one of `keys` has a head element (checked
106 /// in argument order each round) or `timeout` passes (`None` =
107 /// wait forever). Returns `(key, value)`.
108 pub fn blpop(
109 &self,
110 keys: &[&[u8]],
111 timeout: Option<Duration>,
112 ) -> io::Result<Option<(Vec<u8>, Vec<u8>)>> {
113 self.block_on(keys, timeout, |s, k| {
114 Ok(s.lpop(k, 1)?.into_iter().next().map(|v| (k.to_vec(), v)))
115 })
116 }
117
118 /// `BRPOP` — tail-end counterpart of [`Self::blpop`].
119 pub fn brpop(
120 &self,
121 keys: &[&[u8]],
122 timeout: Option<Duration>,
123 ) -> io::Result<Option<(Vec<u8>, Vec<u8>)>> {
124 self.block_on(keys, timeout, |s, k| {
125 Ok(s.rpop(k, 1)?.into_iter().next().map(|v| (k.to_vec(), v)))
126 })
127 }
128
129 /// `BZPOPMIN` — block until one of `keys` has a zset member;
130 /// returns `(key, member, score)`.
131 pub fn bzpopmin(
132 &self,
133 keys: &[&[u8]],
134 timeout: Option<Duration>,
135 ) -> io::Result<Option<ZPopHit>> {
136 self.block_on(keys, timeout, |s, k| {
137 Ok(s
138 .zpopmin(k, 1)?
139 .into_iter()
140 .next()
141 .map(|(m, sc)| (k.to_vec(), m, sc)))
142 })
143 }
144
145 /// The shared park-wait loop: `try_pop` is the non-blocking probe
146 /// run against each key in order, every wake round.
147 fn block_on<T>(
148 &self,
149 keys: &[&[u8]],
150 timeout: Option<Duration>,
151 try_pop: impl Fn(&Self, &[u8]) -> io::Result<Option<T>>,
152 ) -> io::Result<Option<T>> {
153 let deadline = timeout.map(|t| Instant::now() + t);
154 loop {
155 let seen = self.blocker.generation();
156 for k in keys {
157 if let Some(hit) = try_pop(self, k)? {
158 return Ok(Some(hit));
159 }
160 }
161 if !self.blocker.wait_past(seen, deadline) {
162 return Ok(None); // timed out
163 }
164 }
165 }
166}