Skip to main content

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}