Skip to main content

yo_resp/dispatch/
blocking.rs

1//! The six list commands that wait, and the machinery that lets a client wait.
2//!
3//! `BLPOP` is `LPOP` with one difference: when there is nothing to pop, the
4//! client waits instead of being told no. Everything here is about that wait,
5//! and nothing here knows anything about lists that [`super::lists`] does not
6//! already know.
7//!
8//! # The command is kept, not the client
9//!
10//! A parked client is a [`Waiter`]: the keys it named, what it wanted to do with
11//! them, and when to give up. It is not a suspended stack and it is not a task.
12//! Answering it later is running the same attempt again against a database that
13//! has changed since, which is why [`Want::attempt`] is the whole of both paths.
14//! The command handler calls it once to see whether the client has to wait at
15//! all, and the retry calls it again each time something might have arrived.
16//!
17//! That is also why the six commands cost nothing when they do not block. A
18//! `BLPOP` on a list with something in it runs the same three lines `LPOP` runs
19//! and never touches the waiter list.
20//!
21//! # A slot is reused and a client id is not
22//!
23//! A waiter remembers both. The slot is where that connection's reply buffer
24//! is, and the client id is what says the connection sitting on that slot is
25//! still the one that blocked. The engine takes a waiter off the list when its
26//! connection closes, so the check should never fail, and it is there because
27//! the cost of being wrong about it is a reply written into somebody else's
28//! socket.
29//!
30//! # What wakes a waiter
31//!
32//! Any command at all, which is more than is needed and is not the cost it
33//! sounds like: the engine looks at whether anybody is parked before it looks at
34//! anything else, so a server with no blocked clients pays one load and one
35//! branch per command and nothing more. Narrowing it to writes would save
36//! nothing measurable and would need a rule about which commands can put a list
37//! under a key, which `RENAME`, `COPY` and `RESTORE` all make longer than it
38//! looks.
39//!
40//! What is left is that the waiter list is walked rather than indexed by key, so
41//! a server with a thousand parked workers walks a thousand entries per command.
42//! The fix when that matters is an index from key to waiter, not a different
43//! rule about when to look.
44
45use std::sync::atomic::Ordering::Relaxed;
46use yo_common::lock::Held;
47use yo_common::{Code, Error, Result, num};
48use yo_kv::{Db, End, Entry, Member, Movem, ZEnd};
49
50use super::args::{self, Args, NOT_AN_INT};
51use super::lists::{BAD_MPOP_COUNT, BAD_NUMKEYS, end_of, movem_options};
52use super::streams;
53use super::table::Spec;
54use super::zsets;
55use super::{Flow, Server, Session};
56use crate::reply::Out;
57
58/// What Redis says about a timeout it cannot read as a number.
59const NOT_A_FLOAT: &str = "timeout is not a float or out of range";
60/// What it says about one it can read and will not take.
61const NEGATIVE: &str = "timeout is negative";
62/// And about one so far away that milliseconds do not fit in an `i64`.
63const OUT_OF_RANGE: &str = "timeout is out of range";
64/// `WAIT` and `WAITAOF` take their timeout in whole milliseconds rather than in
65/// seconds, so a timeout they cannot read is a different complaint again.
66const TIMEOUT_NOT_AN_INT: &str = "timeout is not an integer or out of range";
67/// What `WAITAOF` says about a `numlocal` that is neither of the two it takes.
68const NOT_ZERO_OR_ONE: &str = "value is out of range, value must between 0 and 1";
69/// And about a negative `numreplicas`.
70const NOT_POSITIVE: &str = "value is out of range, must be positive";
71/// And what it says when asked to wait for a file the server does not keep. The
72/// full stop at the end is Redis's and is the one message in the group that has
73/// one, which is why it is worth writing down rather than tidying up.
74const NO_AOF: &str = "WAITAOF cannot be used when numlocal is set but appendonly is disabled.";
75
76/// Run one blocking command.
77///
78/// `Flow::Block` means nothing was written and the client is on the waiter
79/// list. The engine is what knows which socket that client is on, so it is the
80/// engine that finishes the registration and the engine that stops reading
81/// commands from a connection that is now waiting for one.
82///
83/// # Errors
84///
85/// A timeout that is not a timeout, a direction that is not a direction, and a
86/// key holding something that is not a list.
87pub(super) fn execute(
88    server: &Server,
89    session: &Session,
90    spec: &Spec,
91    args: Args<'_>,
92    out: &mut Out,
93) -> Result<Flow> {
94    // The two that wait on replication rather than on a key. They are here
95    // because they carry the blocking flag and that flag is what routes a
96    // command to this file, and they leave immediately because there is nothing
97    // for them to wait for yet. See [`replication`] for what they answer.
98    if spec.name == "wait" || spec.name == "waitaof" {
99        return replication(spec.name, args, out).map(|()| Flow::Continue);
100    }
101    let now = server.now_ms();
102    // The two stream reads, which are here for the same reason the list six
103    // are and leave through a different door. `BLOCK` is optional on both, so
104    // where `BLPOP` always has a timeout to read, `XREAD` may have been told to
105    // answer now and take nothing for an answer. That is the difference between
106    // parking and writing the null, and it cannot be said with a deadline of
107    // `None`, which already means wait for as long as it takes.
108    if spec.name == "xread" || spec.name == "xreadgroup" {
109        let db = session.db();
110        let want = streams::parse_read(spec.name, args, server.striped(db), now)?;
111        let block = Block::xread(want.keys, want.reads);
112        if block.now(server.striped(db), now, out)? {
113            return Ok(Flow::Continue);
114        }
115        // A script is the other way there is nothing to wait for. It cannot
116        // park, because the thing it would be waiting for is a command from
117        // another client and the script is what that client is queued behind,
118        // so the wait would never end. Timing out at once is what a real server
119        // does and it answers the same null a full timeout would have.
120        let Some(deadline) = want.wait.filter(|_| !session.scripted()) else {
121            // No `BLOCK` at all, so nothing arriving is the answer and not a
122            // reason to wait for it. A null array on both protocols, which is
123            // also what a `BLOCK` that runs out sends.
124            out.nil_array();
125            return Ok(Flow::Continue);
126        };
127        server.park(session.id(), db, deadline, block);
128        return Ok(Flow::Block);
129    }
130    let last = args.len() - 1;
131    let (deadline, block) = match spec.name {
132        // The keys are everything between the name and the timeout, so `BLPOP a
133        // b c 0` waits on three keys and answers with whichever one arrives
134        // first rather than with the first one named.
135        "blpop" | "brpop" => {
136            let end = if spec.name == "blpop" {
137                End::Left
138            } else {
139                End::Right
140            };
141            let deadline = timeout(args.get(last), now)?;
142            (deadline, Block::pop((1..last).map(|i| args.get(i)), end))
143        }
144        // The directions before the timeout, which is the order Redis checks
145        // them in, so `BLMOVE a b UP DOWN nonsense` is a syntax error and not a
146        // complaint about the timeout.
147        "blmove" => {
148            let (from, to) = (end_of(args.get(3))?, end_of(args.get(4))?);
149            let deadline = timeout(args.get(5), now)?;
150            (deadline, Block::moved(args.get(1), args.get(2), from, to))
151        }
152        // The same order again with one more thing to read: ends, then timeout,
153        // then the options behind it. `BLMOVEM s d UP DOWN abc` complains about
154        // the directions and `BLMOVEM s d LEFT RIGHT abc COUNT abc BULK` about
155        // the timeout, both measured against 8.10.1 rather than assumed, because
156        // a line wrong in two places has exactly one right answer.
157        "blmovem" => {
158            let (from, to) = (end_of(args.get(3))?, end_of(args.get(4))?);
159            let deadline = timeout(args.get(5), now)?;
160            let mv = movem_options(args, 6, from, to)?;
161            (deadline, Block::movem(args.get(1), args.get(2), mv))
162        }
163        "brpoplpush" => {
164            let deadline = timeout(args.get(3), now)?;
165            (
166                deadline,
167                Block::moved(args.get(1), args.get(2), End::Right, End::Left),
168            )
169        }
170        "blmpop" => mpop(args, now)?,
171        // The sorted set three, which are the same three shapes again with a
172        // different collection under them. `BZPOPMIN` reads its keys up to the
173        // timeout the way `BLPOP` does, and `BZMPOP` counts them the way
174        // `BLMPOP` does.
175        "bzpopmin" | "bzpopmax" => {
176            let end = zsets::end_of_name(spec.name);
177            let deadline = timeout(args.get(last), now)?;
178            (deadline, Block::zpop((1..last).map(|i| args.get(i)), end))
179        }
180        "bzmpop" => {
181            let deadline = timeout(args.get(1), now)?;
182            let (end, from, to, count) = zsets::parse_mpop(args, 2)?;
183            (
184                deadline,
185                Block::zmpop((from..to).map(|i| args.get(i)), end, count),
186            )
187        }
188        // The table and this match are checked against each other by
189        // `cargo xtask check`, so a name reaching here is a table row without a
190        // handler and there is nothing sensible to answer.
191        _ => return Err(args::syntax()),
192    };
193
194    let db = session.db();
195    if block.now(server.striped(db), now, out)? {
196        return Ok(Flow::Continue);
197    }
198    // Called from a script, so there is nobody left to deliver what it is
199    // waiting for. The same null a timeout writes, for the reason the stream
200    // reads above give.
201    if session.scripted() {
202        out.nil_array();
203        return Ok(Flow::Continue);
204    }
205    server.park(session.id(), db, deadline, block);
206    Ok(Flow::Block)
207}
208
209/// `BLMPOP timeout numkeys key [key ...] LEFT|RIGHT [COUNT count]`.
210///
211/// The same parse as `LMPOP` shifted along by one, including the check that the
212/// key count leaves room for the direction behind it. `BLMPOP 0 2 k LEFT` names
213/// two keys and only gives one, so the word that should have been the direction
214/// is a key and there is no direction left, which Redis calls a syntax error
215/// rather than anything about counts.
216fn mpop(args: Args<'_>, now: u64) -> Result<(Option<u64>, Block)> {
217    let deadline = timeout(args.get(1), now)?;
218    let numkeys = match args.int(2) {
219        Ok(n) if n > 0 => usize::try_from(n).unwrap_or(usize::MAX),
220        _ => return Err(Error::new(Code::Invalid, BAD_NUMKEYS)),
221    };
222    if numkeys >= args.len() - 3 {
223        return Err(args::syntax());
224    }
225    let at = 3 + numkeys;
226    let end = end_of(args.get(at))?;
227    let mut want = 1usize;
228    if at + 1 < args.len() {
229        if args.len() != at + 3 || !args::is(args.get(at + 1), b"count") {
230            return Err(args::syntax());
231        }
232        want = match args.int(at + 2) {
233            Ok(n) if n > 0 => usize::try_from(n).unwrap_or(usize::MAX),
234            _ => return Err(Error::new(Code::Invalid, BAD_MPOP_COUNT)),
235        };
236    }
237    Ok((
238        deadline,
239        Block::mpop((3..at).map(|i| args.get(i)), end, want),
240    ))
241}
242
243/// `WAIT numreplicas timeout` and `WAITAOF numlocal numreplicas timeout`.
244///
245/// Both of them ask the same question, which is whether this connection's writes
246/// have got somewhere durable, and both of them answer zero here. There are no
247/// replicas because there is no replication, and there is no append only file
248/// because `appendonly` is fixed at `no`, so nothing can ever move either count
249/// off zero and there is nothing to wait for. Redis in the same state gives the
250/// same numbers, it just takes the timeout to do it, and that is registered as
251/// D-25.
252///
253/// What is not a formality is the argument checking, because that is what a
254/// client sees when it gets something wrong, and the three numbers are read by
255/// three different Redis helpers with three different complaints. `numlocal` is
256/// a range and says so. `numreplicas` is a positive number for `WAITAOF` and any
257/// number at all for `WAIT`, where a negative one is accepted and satisfied on
258/// the spot because zero replicas is already more than it asked for. The timeout
259/// is milliseconds here and not the seconds the list commands take, so it does
260/// not go through [`timeout`] above, and a negative one is refused with its own
261/// message rather than the range one.
262fn replication(name: &str, args: Args<'_>, out: &mut Out) -> Result<()> {
263    let aof = name == "waitaof";
264    // `WAITAOF` has one number in front of the two `WAIT` has, and everything
265    // after it is in the same place, so the offset is the whole difference.
266    let at = usize::from(aof);
267    let mut wants_local = false;
268    if aof {
269        let local = whole(args.get(1))?;
270        if !(0..=1).contains(&local) {
271            return Err(Error::new(Code::Invalid, NOT_ZERO_OR_ONE));
272        }
273        wants_local = local == 1;
274    }
275    let replicas = whole(args.get(at + 1))?;
276    if aof && replicas < 0 {
277        return Err(Error::new(Code::Invalid, NOT_POSITIVE));
278    }
279    let ms = whole(args.get(at + 2)).map_err(|_| Error::new(Code::Invalid, TIMEOUT_NOT_AN_INT))?;
280    if ms < 0 {
281        return Err(Error::new(Code::Invalid, NEGATIVE));
282    }
283    // The one complaint here that is about the server rather than about the
284    // arguments, and the reason it comes last is that Redis reads all three
285    // arguments before it looks at itself. `appendonly` is `no` here and cannot
286    // be set, so asking to wait for a local copy is asking for something that
287    // cannot happen rather than something that has not happened yet.
288    if wants_local {
289        return Err(Error::new(Code::Invalid, NO_AOF));
290    }
291    if aof {
292        // Two integers and not a map, whichever protocol is in use. The local
293        // count is first and it is zero for the same reason the other one is:
294        // this server has no append only file to be behind.
295        out.array(2);
296        out.int(0);
297        out.int(0);
298    } else {
299        out.int(0);
300    }
301    Ok(())
302}
303
304/// A whole number argument, with the message Redis gives when it is not one.
305fn whole(arg: &[u8]) -> Result<i64> {
306    num::parse_i64(arg).ok_or_else(|| Error::new(Code::Invalid, NOT_AN_INT))
307}
308
309/// The moment to give up at, or `None` for a wait with no end to it.
310///
311/// Seconds as a float on the wire and a millisecond deadline here. Redis reads
312/// it as a long double, refuses a negative one, multiplies by a thousand and
313/// refuses what will not fit in an `i64`, and treats a timeout of exactly zero
314/// as no timeout at all. All four of those are visible from a client:
315///
316/// - `-0.0` is not negative, so it is accepted, and it is zero, so it waits
317///   forever. `-0.1` is refused.
318/// - `1e400` and `inf` parse, so they are not the not-a-float error, and both
319///   are further away than an `i64` of milliseconds reaches, so they are the out
320///   of range one.
321/// - `0.0000001` is a real timeout however small, so it expires on the next turn
322///   of the loop rather than waiting for anything.
323fn timeout(arg: &[u8], now: u64) -> Result<Option<u64>> {
324    let Some(secs) = num::parse_f64(arg) else {
325        return Err(Error::new(Code::Invalid, NOT_A_FLOAT));
326    };
327    if secs < 0.0 {
328        return Err(Error::new(Code::Invalid, NEGATIVE));
329    }
330    let ms = secs * 1000.0;
331    // `>` rather than a negated `<=`, and the two are not the same: an infinite
332    // timeout is greater than the bound and lands here, while a NaN would be
333    // neither, which is why the parse refuses one before this line is reached.
334    if ms > i64::MAX as f64 {
335        return Err(Error::new(Code::Invalid, OUT_OF_RANGE));
336    }
337    if ms <= 0.0 {
338        return Ok(None);
339    }
340    Ok(Some(now.saturating_add(ms as u64)))
341}
342
343/// What a parked client is still trying to do.
344enum Want {
345    /// `BLPOP` and `BRPOP`: one element off the first key that has one, with the
346    /// reply saying which key that turned out to be.
347    Pop { end: End },
348    /// `BLMOVE` and `BRPOPLPUSH`: one element, onto an end of another list.
349    Move { dst: Vec<u8>, from: End, to: End },
350    /// `BLMOVEM`: a block of them, onto an end of another list.
351    ///
352    /// The only want in this file where how many elements are there decides
353    /// whether the client is ready, rather than just whether any are. `COUNT`
354    /// takes what has arrived and so wakes on the first push, and `EXACTLY`
355    /// waits until the source actually holds the whole block.
356    MoveM { dst: Vec<u8>, mv: Movem },
357    /// `BLMPOP`: up to `count` elements off the first key that has any.
358    Mpop { end: End, count: usize },
359    /// `BZPOPMIN` and `BZPOPMAX`: one member and its score off the first sorted
360    /// set that has one, with the reply saying which key that turned out to be.
361    ZPop { end: ZEnd },
362    /// `BZMPOP`: up to `count` members off the first sorted set that has any.
363    ZMpop { end: ZEnd, count: usize },
364    /// `XREAD BLOCK` and `XREADGROUP BLOCK`: whatever has arrived on any of the
365    /// streams since the ID this asked from.
366    ///
367    /// Unlike the other six this takes nothing away, so several clients parked
368    /// on one stream all get the same entry rather than one of them getting it.
369    /// That is the whole point of a stream over a list, and it costs nothing
370    /// here because the attempt is a read.
371    XRead(streams::Reads),
372}
373
374impl Want {
375    /// Try to do it now.
376    ///
377    /// `Ok(true)` means a reply was written and the client is finished with.
378    /// `Ok(false)` means there was nothing to take and nothing was written.
379    ///
380    /// `strict` is the difference between the two callers. The command handler
381    /// passes `true`, so `BLPOP string 0` is a `WRONGTYPE` on the spot the way
382    /// `LPOP string` is. The retry passes `false`, so a key somebody has since
383    /// made into a set is skipped rather than turned into an error on a command
384    /// that was accepted seconds ago. That is what a running Redis does: a
385    /// `SADD` to a key a client is blocked on leaves it blocked, and it times
386    /// out in its own time.
387    ///
388    /// # Errors
389    ///
390    /// Whatever the keyspace says, which under `strict` includes a key of
391    /// another type.
392    fn attempt(
393        &self,
394        keys: &[Vec<u8>],
395        db: &Db,
396        now: u64,
397        out: &mut Out,
398        strict: bool,
399    ) -> Result<bool> {
400        match self {
401            // The one arm that needs to know what time it is, because a group
402            // read records when each entry was handed out. The other six take
403            // an element off a collection and the clock does not come into it.
404            Want::XRead(r) => streams::read(db, keys, r, now, strict, out),
405            Want::Pop { end } => {
406                for key in keys {
407                    if !ready(db, key, strict)? {
408                        continue;
409                    }
410                    out.array(2);
411                    out.bulk(key);
412                    db.hold(key).pop_into(key, *end, 1, |e| element(out, e))?;
413                    return Ok(true);
414                }
415                Ok(false)
416            }
417            Want::Mpop { end, count } => {
418                for key in keys {
419                    if !ready(db, key, strict)? {
420                        continue;
421                    }
422                    out.array(2);
423                    out.bulk(key);
424                    let mark = out.len();
425                    let n = db
426                        .hold(key)
427                        .pop_into(key, *end, *count, |e| element(out, e))?;
428                    out.close_array(mark, n);
429                    return Ok(true);
430                }
431                Ok(false)
432            }
433            // Three elements and not two, because `BZPOPMIN` puts the key, the
434            // member and the score side by side rather than pairing the last
435            // two. That is Redis's shape and it is not the shape `ZPOPMIN` has.
436            Want::ZPop { end } => {
437                for key in keys {
438                    if !zready(db, key, strict)? {
439                        continue;
440                    }
441                    out.array(3);
442                    out.bulk(key);
443                    db.hold(key).zpop(key, *end, 1, |m, sc| {
444                        member(out, m);
445                        out.double(sc);
446                    })?;
447                    return Ok(true);
448                }
449                Ok(false)
450            }
451            Want::ZMpop { end, count } => {
452                for key in keys {
453                    if !zready(db, key, strict)? {
454                        continue;
455                    }
456                    out.array(2);
457                    out.bulk(key);
458                    let mark = out.len();
459                    let n = db.hold(key).zpop(key, *end, *count, |m, sc| {
460                        out.array(2);
461                        member(out, m);
462                        out.double(sc);
463                    })?;
464                    out.close_array(mark, n);
465                    return Ok(true);
466                }
467                Ok(false)
468            }
469            // The source's length first, so that an empty source never reaches
470            // the destination's type check. `BLMOVE empty string LEFT RIGHT 0.1`
471            // times out on a running Redis rather than answering `WRONGTYPE`,
472            // because the destination is only looked at once there is something
473            // to put in it, and this order gives that answer.
474            Want::Move { dst, from, to } => {
475                let src = &keys[0];
476                if !ready(db, src, strict)? {
477                    return Ok(false);
478                }
479                match db.lmove(src, dst, *from, *to, |v| out.bulk(v)) {
480                    Ok(true) => Ok(true),
481                    // The source had something in it a line ago and this is the
482                    // only thread that could have taken it.
483                    Ok(false) => Ok(false),
484                    Err(e) if strict => Err(e),
485                    // The destination is not a list any more. Nothing was taken,
486                    // because `lmove` checks the destination before it pops, so
487                    // the client goes back to waiting with the queue as it was.
488                    Err(_) => Ok(false),
489                }
490            }
491            // The same shape as `Move` with a different question about the
492            // source. `ready` asks whether there is anything and that is not
493            // enough here, because an `EXACTLY` client is not ready until the
494            // whole block has arrived, and asking it any earlier would take
495            // nothing and answer nothing while looking like it had tried.
496            Want::MoveM { dst, mv } => {
497                let src = &keys[0];
498                let have = match db.hold(src).llen(src) {
499                    Ok(n) => n,
500                    Err(e) if strict => return Err(e),
501                    Err(_) => return Ok(false),
502                };
503                // Not ready is not the same as nothing to do, so the
504                // destination is never looked at from here. `BLMOVEM empty
505                // string LEFT RIGHT 0.1` times out on a running 8.10.1 rather
506                // than answering `WRONGTYPE`, and so does an `EXACTLY` whose
507                // source is short, both of which were measured.
508                if have == 0 || (mv.exactly && have < mv.count) {
509                    return Ok(false);
510                }
511                let mark = out.len();
512                let mut n = 0;
513                match db.lmovem(src, dst, *mv, |v| {
514                    out.bulk(v);
515                    n += 1;
516                }) {
517                    Ok(_) => {}
518                    Err(e) if strict => return Err(e),
519                    // As `Move`: the destination stopped being a list while
520                    // this client waited, and nothing was taken.
521                    Err(_) => {
522                        out.truncate(mark);
523                        return Ok(false);
524                    }
525                }
526                out.close_array(mark, n);
527                Ok(true)
528            }
529        }
530    }
531}
532
533/// Whether this key is a list with something in it.
534///
535/// A key of the wrong type is an error to the command handler and not one to the
536/// retry, which is the whole of what `strict` decides.
537fn ready(db: &Db, key: &[u8], strict: bool) -> Result<bool> {
538    match db.hold(key).llen(key) {
539        Ok(n) => Ok(n > 0),
540        Err(e) if strict => Err(e),
541        Err(_) => Ok(false),
542    }
543}
544
545/// The same for a sorted set, which has its own emptiness to ask about.
546fn zready(db: &Db, key: &[u8], strict: bool) -> Result<bool> {
547    match db.hold(key).zcard(key) {
548        Ok(n) => Ok(n > 0),
549        Err(e) if strict => Err(e),
550        Err(_) => Ok(false),
551    }
552}
553
554/// One element as the client sees it, the same as [`super::lists`] writes it.
555#[inline]
556fn element(out: &mut Out, e: Entry<'_>) {
557    match e {
558        Entry::Int(n) => out.bulk_int(n),
559        Entry::Str(s) => out.bulk(s),
560    }
561}
562
563/// One member as the client sees it, the same as [`super::zsets`] writes it.
564#[inline]
565fn member(out: &mut Out, m: Member<'_>) {
566    match m {
567        Member::Int(n) => out.bulk_int(n),
568        Member::Str(s) => out.bulk(s),
569    }
570}
571
572/// A parsed blocking command, ready to be tried or to be parked.
573pub struct Block {
574    /// The keys, already copied out of the connection's read buffer.
575    ///
576    /// This is the allocation blocking costs and it is once per block rather
577    /// than once per attempt. The arguments are slices of a buffer that is
578    /// reused as soon as the batch is over, and a waiter outlives the batch.
579    keys: Vec<Vec<u8>>,
580    want: Want,
581}
582
583impl Block {
584    /// `BLPOP` and `BRPOP`.
585    fn pop<'a>(keys: impl Iterator<Item = &'a [u8]>, end: End) -> Block {
586        Block {
587            keys: owned(keys),
588            want: Want::Pop { end },
589        }
590    }
591
592    /// `BLMPOP`.
593    fn mpop<'a>(keys: impl Iterator<Item = &'a [u8]>, end: End, count: usize) -> Block {
594        Block {
595            keys: owned(keys),
596            want: Want::Mpop { end, count },
597        }
598    }
599
600    /// `BZPOPMIN` and `BZPOPMAX`.
601    fn zpop<'a>(keys: impl Iterator<Item = &'a [u8]>, end: ZEnd) -> Block {
602        Block {
603            keys: owned(keys),
604            want: Want::ZPop { end },
605        }
606    }
607
608    /// `BZMPOP`.
609    fn zmpop<'a>(keys: impl Iterator<Item = &'a [u8]>, end: ZEnd, count: usize) -> Block {
610        Block {
611            keys: owned(keys),
612            want: Want::ZMpop { end, count },
613        }
614    }
615
616    /// `BLMOVE` and `BRPOPLPUSH`.
617    fn moved(src: &[u8], dst: &[u8], from: End, to: End) -> Block {
618        yo_alloc::allow(|| Block {
619            keys: vec![src.to_vec()],
620            want: Want::Move {
621                dst: dst.to_vec(),
622                from,
623                to,
624            },
625        })
626    }
627
628    /// `BLMOVEM`.
629    fn movem(src: &[u8], dst: &[u8], mv: Movem) -> Block {
630        yo_alloc::allow(|| Block {
631            keys: vec![src.to_vec()],
632            want: Want::MoveM {
633                dst: dst.to_vec(),
634                mv,
635            },
636        })
637    }
638
639    /// Do it now if it can be done now.
640    ///
641    /// # Errors
642    ///
643    /// A key of another type, which is an error rather than a wait.
644    fn now(&self, db: &Db, now: u64, out: &mut Out) -> Result<bool> {
645        self.want.attempt(&self.keys, db, now, out, true)
646    }
647
648    /// `XREAD BLOCK` and `XREADGROUP BLOCK`, whose keys and IDs were read
649    /// together by [`streams::parse_read`] because neither makes sense alone.
650    fn xread(keys: Vec<Vec<u8>>, reads: streams::Reads) -> Block {
651        Block {
652            keys,
653            want: Want::XRead(reads),
654        }
655    }
656}
657
658/// The keys a blocking command named, copied so they outlive the read buffer.
659fn owned<'a>(keys: impl Iterator<Item = &'a [u8]>) -> Vec<Vec<u8>> {
660    yo_alloc::allow(|| keys.map(<[u8]>::to_vec).collect())
661}
662
663/// One parked client.
664struct Waiter {
665    /// The client id, which is never reused.
666    client: u64,
667    /// The slot its reply buffer is on, which is.
668    conn: u32,
669    /// The thread holding that slot.
670    ///
671    /// A reply goes into a buffer the accepting thread owns, so a waiter can
672    /// only be answered by the thread it blocked on, and a list every thread
673    /// walks has to say which entries are whose. The slot number alone will not
674    /// do it: two threads number their connections from zero.
675    thread: usize,
676    /// The database it was on when it blocked. A push into another database is
677    /// not this client's push, even when the key has the same name.
678    db: usize,
679    /// The millisecond to give up at, or `None` for `BLPOP key 0`, which waits
680    /// for as long as the connection is open.
681    deadline: Option<u64>,
682    keys: Vec<Vec<u8>>,
683    want: Want,
684}
685
686/// Every parked client, oldest first.
687///
688/// The order is the order they blocked in and it is the order they are served
689/// in, which is what makes a queue with several workers on it fair: two clients
690/// blocked on the same key take the two elements a `RPUSH q first second` adds
691/// in the order they arrived. A `Vec` is the right structure for that while the
692/// list is short, and it is short, because a waiter is a client doing nothing.
693#[derive(Default)]
694pub struct Waiters {
695    list: Vec<Waiter>,
696}
697
698/// Where the reply to a parked client has to go.
699#[derive(Debug, Clone, Copy)]
700pub struct Parked {
701    /// The slot holding its reply buffer.
702    pub conn: u32,
703    /// The client that was on that slot when it blocked.
704    pub client: u64,
705}
706
707impl Waiters {
708    /// Whether anybody is waiting.
709    #[must_use]
710    #[inline]
711    pub fn is_empty(&self) -> bool {
712        self.list.is_empty()
713    }
714
715    /// How many clients are parked, which is what `INFO clients` calls
716    /// `blocked_clients`.
717    #[must_use]
718    #[inline]
719    pub fn len(&self) -> usize {
720        self.list.len()
721    }
722
723    /// Copy out the waiters `thread` has to answer, oldest first.
724    ///
725    /// The caller works from the copy rather than from the list, because
726    /// answering a waiter needs a connection's reply buffer and the list is
727    /// behind a lock that another thread is waiting on. It brings its own
728    /// vector, which after the first parked client is a vector it already has
729    /// the room in.
730    pub fn mine(&self, thread: usize, into: &mut Vec<Parked>) {
731        into.clear();
732        yo_alloc::allow(|| {
733            for w in self.list.iter().filter(|w| w.thread == thread) {
734                into.push(Parked {
735                    conn: w.conn,
736                    client: w.client,
737                });
738            }
739        });
740    }
741
742    /// Where a parked client sits in the list.
743    fn find(&self, client: u64) -> Option<usize> {
744        self.list.iter().position(|w| w.client == client)
745    }
746
747    /// The database the waiter at `at` blocked on.
748    fn db_of(&self, at: usize) -> usize {
749        self.list[at].db
750    }
751
752    /// Take off every waiter belonging to a client that has gone, and say how
753    /// many that was.
754    ///
755    /// Called when a connection closes rather than left for the deadline sweep
756    /// to find, because a `BLPOP key 0` on a connection nobody will ever write
757    /// to again has no deadline to be found by. The count goes back because the
758    /// caller keeps its own tally of what its thread has waiting, and a client
759    /// that was never parked has to leave that tally alone.
760    fn forget(&mut self, client: u64) -> usize {
761        let before = self.list.len();
762        self.list.retain(|w| w.client != client);
763        before - self.list.len()
764    }
765
766    /// Say which slot the waiter this client just registered is answered on.
767    ///
768    /// The command layer knows which client blocked and the engine knows which
769    /// slot that client is on, so the slot is filled in afterwards by the half
770    /// that has it. A client can only be parked once, since it is not reading
771    /// commands while it waits, so the search finds the one that was just added.
772    fn bind(&mut self, client: u64, conn: u32) {
773        if let Some(w) = self.list.iter_mut().rev().find(|w| w.client == client) {
774            w.conn = conn;
775        }
776    }
777
778    /// Park a client that could not be answered.
779    ///
780    /// The slot is filled in by [`Waiters::bind`] once the engine has it, so
781    /// this leaves it at zero rather than pretending to know.
782    fn park(&mut self, client: u64, thread: usize, db: usize, deadline: Option<u64>, block: Block) {
783        yo_alloc::allow(|| {
784            self.list.push(Waiter {
785                client,
786                conn: 0,
787                thread,
788                db,
789                deadline,
790                keys: block.keys,
791                want: block.want,
792            });
793        });
794    }
795
796    /// Try to answer the waiter at `at`, and say whether it is finished with.
797    ///
798    /// `true` means a reply is in `out` and the caller should take the waiter
799    /// off the list, which covers both a client that got what it asked for and
800    /// one that ran out of time.
801    ///
802    /// The attempt comes before the deadline, so a push that landed in the same
803    /// millisecond the client gave up in serves it rather than racing it.
804    fn try_serve(&self, at: usize, dbs: &[Db], now: u64, out: &mut Out) -> bool {
805        let w = &self.list[at];
806        let mark = out.len();
807        match w.want.attempt(&w.keys, &dbs[w.db], now, out, false) {
808            Ok(true) => return true,
809            Ok(false) => {}
810            // `strict` is off, so nothing in there returns an error today.
811            // Putting the buffer back is what makes it safe to be wrong about
812            // that later.
813            Err(_) => out.truncate(mark),
814        }
815        if w.deadline.is_some_and(|d| now >= d) {
816            // A null array for all six, `BLMOVE` and `BRPOPLPUSH` included,
817            // even though what they send when they succeed is a single element.
818            // That is Redis's and it is not what reading the reply schema would
819            // suggest: a RESP2 client sees `*-1` and not `$-1`.
820            out.nil_array();
821            return true;
822        }
823        false
824    }
825}
826
827impl Server {
828    /// The clock reading this batch is working against.
829    #[must_use]
830    pub fn now_ms(&self) -> u64 {
831        self.clock.now_ms()
832    }
833
834    /// Who is parked, for the engine walking the list.
835    ///
836    /// Takes the lock for as long as the answer is held, so a caller that only
837    /// wants to know whether anybody is waiting asks [`Server::parked`] instead
838    /// and does not take it at all.
839    #[must_use]
840    pub fn waiters(&self) -> Held<'_, Waiters> {
841        self.waiters.lock()
842    }
843
844    /// How many clients are parked, without taking the lock.
845    ///
846    /// What `INFO clients` calls `blocked_clients`, and what every command asks
847    /// before it goes looking for somebody to wake.
848    #[must_use]
849    #[inline]
850    pub fn parked(&self) -> usize {
851        self.parked.load(Relaxed)
852    }
853
854    /// How many of the calling thread's clients are parked.
855    ///
856    /// The number to branch on before reaching for the waiter list, because a
857    /// thread can only answer the waiters it parked itself. [`Server::parked`]
858    /// counts the whole server, so branching on that puts every thread through
859    /// the shared lock as soon as one client blocks anywhere.
860    #[must_use]
861    #[inline]
862    pub fn parked_here(&self) -> usize {
863        self.mine().parked.load(Relaxed)
864    }
865
866    /// Park a client on a command that could not be answered yet.
867    ///
868    /// Filed under the calling thread, which is the thread that will answer it,
869    /// because a command runs on the thread that read it and a reply goes back
870    /// into that thread's buffer for the connection.
871    pub(super) fn park(&self, client: u64, db: usize, deadline: Option<u64>, block: Block) {
872        let thread = self.my_slot();
873        let mut list = self.waiters.lock();
874        list.park(client, thread, db, deadline, block);
875        self.note(&list);
876        self.mine().blocked(1);
877    }
878
879    /// Take off every waiter belonging to a client that has gone.
880    ///
881    /// Called on the thread that parked it, which is the only thread that can
882    /// have parked it, so the count of what this thread has waiting comes down
883    /// by however many the list actually held. A thread with nothing waiting
884    /// does not take the lock, which is what keeps a server with one blocked
885    /// client from paying for it on every disconnect on every other thread.
886    pub fn forget_waiters(&self, client: u64) {
887        if self.parked_here() == 0 {
888            return;
889        }
890        let mut list = self.waiters.lock();
891        let gone = list.forget(client);
892        self.note(&list);
893        self.mine().woke(gone);
894    }
895
896    /// Say which slot the waiter this client just registered is answered on.
897    pub fn bind_waiter(&self, client: u64, conn: u32) {
898        self.waiters.lock().bind(client, conn);
899    }
900
901    /// Publish how long the list is now.
902    ///
903    /// Called with the list held and by whoever changed it, which is what keeps
904    /// the number and the list from disagreeing about anything except a change
905    /// that has not finished.
906    fn note(&self, list: &Waiters) {
907        self.parked.store(list.len(), Relaxed);
908    }
909
910    /// Try to answer a parked client, writing into the buffer the engine found
911    /// for it, and say whether it is finished with.
912    ///
913    /// The engine cannot reach the databases and this cannot reach the
914    /// connections, so the two meet here: the caller hands in one connection's
915    /// reply buffer and gets back whether to unpark the client behind it.
916    ///
917    /// By client and not by position, because the caller let go of the list
918    /// between finding the client and asking about it, and in that gap another
919    /// thread can take one of its own waiters off and move everything behind it
920    /// up one. A client that is no longer parked answers `false`, which is the
921    /// same answer as one that is parked and has nothing waiting for it.
922    pub fn serve_waiter(&self, client: u64, now: u64, out: &mut Out) -> bool {
923        let list = self.waiters.lock();
924        let Some(at) = list.find(client) else {
925            return false;
926        };
927        // Serving a waiter pops an element, which makes garbage, and it happens
928        // outside `execute` so nothing else has marked the database for the
929        // maintenance turn.
930        self.mine().mark(1u64 << list.db_of(at));
931        list.try_serve(at, &self.dbs, now, out)
932    }
933}