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

kevy_embedded/
ops_blocking.rs
