kevy_embedded/
ops_blocking.rs

1//! 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 crate::KevyResult;
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 { waiters: AtomicUsize::new(0), generation: Mutex::new(0), cv: Condvar::new() }
43    }
44
45    /// Writer side — called from `commit_write`. One Relaxed load when
46    /// idle; lock + notify only with live waiters.
47    #[inline]
48    pub(crate) fn wake_all(&self) {
49        if self.waiters.load(Ordering::Relaxed) == 0 {
50            return;
51        }
52        let mut g = self.generation.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
53        *g += 1;
54        self.cv.notify_all();
55    }
56
57    /// Current generation (observe BEFORE polling, wait against it).
58    fn generation(&self) -> u64 {
59        *self.generation.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
60    }
61
62    /// Park until the generation moves past `seen` or `deadline`
63    /// passes. Returns `false` on timeout.
64    fn wait_past(&self, seen: u64, deadline: Option<Instant>) -> bool {
65        self.waiters.fetch_add(1, Ordering::Relaxed);
66        let mut g = self.generation.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
67        let ok = loop {
68            if *g != seen {
69                break true;
70            }
71            match deadline {
72                None => {
73                    g = self.cv.wait(g).unwrap_or_else(std::sync::PoisonError::into_inner);
74                }
75                Some(d) => {
76                    let now = Instant::now();
77                    if now >= d {
78                        break false;
79                    }
80                    let (guard, res) = self
81                        .cv
82                        .wait_timeout(g, d - now)
83                        .unwrap_or_else(std::sync::PoisonError::into_inner);
84                    g = guard;
85                    if res.timed_out() && *g == seen {
86                        break false;
87                    }
88                }
89            }
90        };
91        drop(g);
92        self.waiters.fetch_sub(1, Ordering::Relaxed);
93        ok
94    }
95}
96
97impl Store {
98    /// `BLPOP` — block until one of `keys` has a head element (checked
99    /// in argument order each round) or `timeout` passes (`None` =
100    /// wait forever). Returns `(key, value)`.
101    pub fn blpop(
102        &self,
103        keys: &[&[u8]],
104        timeout: Option<Duration>,
105    ) -> KevyResult<Option<(Vec<u8>, Vec<u8>)>> {
106        self.block_on(keys, timeout, |s, k| {
107            Ok(s.lpop(k, 1)?.into_iter().next().map(|v| (k.to_vec(), v)))
108        })
109    }
110
111    /// `BRPOP` — tail-end counterpart of [`Self::blpop`].
112    pub fn brpop(
113        &self,
114        keys: &[&[u8]],
115        timeout: Option<Duration>,
116    ) -> KevyResult<Option<(Vec<u8>, Vec<u8>)>> {
117        self.block_on(keys, timeout, |s, k| {
118            Ok(s.rpop(k, 1)?.into_iter().next().map(|v| (k.to_vec(), v)))
119        })
120    }
121
122    /// `BZPOPMIN` — block until one of `keys` has a zset member;
123    /// returns `(key, member, score)`.
124    pub fn bzpopmin(
125        &self,
126        keys: &[&[u8]],
127        timeout: Option<Duration>,
128    ) -> KevyResult<Option<ZPopHit>> {
129        self.block_on(keys, timeout, |s, k| {
130            Ok(s.zpopmin(k, 1)?.into_iter().next().map(|(m, sc)| (k.to_vec(), m, sc)))
131        })
132    }
133
134    /// The shared park-wait loop: `try_pop` is the non-blocking probe
135    /// run against each key in order, every wake round.
136    fn block_on<T>(
137        &self,
138        keys: &[&[u8]],
139        timeout: Option<Duration>,
140        try_pop: impl Fn(&Self, &[u8]) -> KevyResult<Option<T>>,
141    ) -> KevyResult<Option<T>> {
142        let deadline = timeout.map(|t| Instant::now() + t);
143        loop {
144            let seen = self.blocker.generation();
145            for k in keys {
146                if let Some(hit) = try_pop(self, k)? {
147                    return Ok(Some(hit));
148                }
149            }
150            if !self.blocker.wait_past(seen, deadline) {
151                return Ok(None); // timed out
152            }
153        }
154    }
155}
kevy_embedded/ops_blocking.rs

kevy_embedded/
ops_blocking.rs
