Skip to main content

yo_resp/
engine.rs

1//! Connections, framing and buffers: the seam between the loop and the
2//! commands.
3//!
4//! `yo-reactor` knows how to run a batch and nothing about what a command is.
5//! `dispatch` knows how to run a command and nothing about where the bytes came
6//! from. This module is the piece in between, and it is the piece a server is
7//! missing until it exists: the read buffer a command's arguments point into,
8//! the framing that says where one command ends and the next begins, the reply
9//! buffer that holds an answer until the batch is done, and the state a
10//! connection keeps between the two.
11//!
12//! # What a piece of work is
13//!
14//! [`Cmd`] is three numbers: which connection, which decoder holds the
15//! arguments, and where in that connection's buffer they point. It is `Copy`
16//! and twenty four bytes, so it crosses an intake lane without touching the
17//! heap, and it carries no borrow, which is what lets the reactor hold sixty
18//! four of them while the engine owns the bytes they name.
19//!
20//! The decoders are pooled. Framing takes one out of the pool per command,
21//! `run` puts it back, and a connection with a half read command keeps hold of
22//! one so that a bulk arriving in ten reads is decoded once rather than ten
23//! times. In the steady state the pool is as large as the deepest batch and
24//! nothing here allocates at all.
25//!
26//! # One write per connection
27//!
28//! Replies accumulate in the connection's [`Out`] and go out in [`Wire::flush`],
29//! which is one call to the sink per connection touched by the batch and never
30//! one per reply. That is the syscall shape `04` section 2 asks for, and it is
31//! the one aki got wrong: its `HGETALL` profile spent 69.7 percent of its time
32//! in write syscalls.
33//!
34//! # What is not here
35//!
36//! Sockets. [`Sink`] is where the bytes go and the io_uring reactor implements
37//! it later, which keeps this module testable without a network and keeps the
38//! ring out of the crate that parses the protocol.
39//!
40//! The hash the first walk computes warms the bucket and is then thrown away,
41//! because `yo-kv`'s commands take keys rather than hashes. The prefetch is the
42//! part that is worth a cache miss; hashing a short key twice is a few
43//! nanoseconds, and removing the second one means a hashed form of every
44//! command method, which is a change to make with a benchmark rather than on
45//! the way past.
46//!
47//! ```
48//! use yo_resp::engine::{Recorder, Wire, pump};
49//! use yo_reactor::Reactor;
50//!
51//! let mut r = Reactor::inline(Wire::new(Recorder::new()));
52//! let conn = r.engine_mut().accept();
53//!
54//! r.engine_mut().feed(conn, b"*3\r\n$3\r\nSET\r\n$1\r\nk\r\n$1\r\nv\r\n*2\r\n$3\r\nGET\r\n$1\r\nk\r\n");
55//! let mut batch = Vec::new();
56//! assert_eq!(pump(&mut r, &mut batch), 2);
57//!
58//! assert_eq!(r.engine().sink().sent(conn), b"+OK\r\n$1\r\nv\r\n");
59//! ```
60
61use std::collections::VecDeque;
62
63use yo_reactor::{BATCH_MAX, Engine, Reactor};
64
65use crate::dispatch::{Args, Flow, Server, Session, execute, lookup};
66use crate::error::ProtocolError;
67use crate::proto::{Limits, Proto};
68use crate::reply::Out;
69use crate::request::{Argv, Step};
70use yo_kv::Keyspace;
71
72/// Which connection. An index, reused after a connection closes.
73pub type ConnId = u32;
74
75/// The read buffer a connection starts with.
76///
77/// Redis's query buffer starts at sixteen kilobytes for the same reason: it is
78/// larger than every command a client actually sends, so the buffer grows once
79/// at accept time and then never again.
80const READ_BUF: usize = 16 * 1024;
81
82/// The reply buffer a connection starts with.
83const OUT_BUF: usize = 16 * 1024;
84
85/// How many arguments a decoder has room for before it grows.
86const ARGV_HINT: usize = 8;
87
88/// One framed command, waiting to run.
89///
90/// Names the bytes rather than holding them, so the reactor can queue a batch
91/// of these while the engine keeps ownership of every buffer they point into.
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub struct Cmd {
94    conn: ConnId,
95    slot: u32,
96    base: usize,
97}
98
99impl Cmd {
100    /// The connection this command arrived on.
101    #[must_use]
102    pub const fn conn(&self) -> ConnId {
103        self.conn
104    }
105}
106
107/// Where replies go.
108///
109/// One call per connection per batch, with however many replies are waiting.
110/// The network reactor implements this over io_uring, a test implements it over
111/// a `Vec`, and neither this module nor `dispatch` has to know which.
112pub trait Sink {
113    /// Take up to all of `bytes` for `conn`, and say how many were taken.
114    ///
115    /// Fewer than were offered means the socket is full: what is left stays in
116    /// the connection's reply buffer and is offered again on the next flush.
117    fn write(&mut self, conn: ConnId, bytes: &[u8]) -> usize;
118
119    /// The connection is finished with and its id is about to be reused.
120    fn closed(&mut self, conn: ConnId) {
121        let _ = conn;
122    }
123}
124
125/// A sink that keeps everything, for tests and for a driver with no socket.
126#[derive(Debug, Default)]
127pub struct Recorder {
128    sent: Vec<Vec<u8>>,
129    closed: Vec<ConnId>,
130}
131
132impl Recorder {
133    /// An empty one.
134    #[must_use]
135    pub fn new() -> Recorder {
136        Recorder::default()
137    }
138
139    /// Everything written to a connection so far.
140    #[must_use]
141    pub fn sent(&self, conn: ConnId) -> &[u8] {
142        self.sent.get(conn as usize).map_or(&[], Vec::as_slice)
143    }
144
145    /// Whether a connection was closed.
146    #[must_use]
147    pub fn was_closed(&self, conn: ConnId) -> bool {
148        self.closed.contains(&conn)
149    }
150
151    /// Forget what was written, keeping the room it was written into.
152    pub fn clear(&mut self) {
153        for c in &mut self.sent {
154            c.clear();
155        }
156        self.closed.clear();
157    }
158}
159
160impl Sink for Recorder {
161    fn write(&mut self, conn: ConnId, bytes: &[u8]) -> usize {
162        // A test sink, so the growth here is not on anybody's data path.
163        yo_alloc::allow(|| {
164            if self.sent.len() <= conn as usize {
165                self.sent.resize_with(conn as usize + 1, Vec::new);
166            }
167            self.sent[conn as usize].extend_from_slice(bytes);
168        });
169        bytes.len()
170    }
171
172    fn closed(&mut self, conn: ConnId) {
173        yo_alloc::allow(|| self.closed.push(conn));
174    }
175}
176
177/// One connection's state.
178struct Conn {
179    live: bool,
180    session: Session,
181    out: Out,
182    /// What has arrived and not yet been framed away.
183    buf: Vec<u8>,
184    /// How much of `buf` the framing has consumed.
185    head: usize,
186    /// The decoder holding a command that has not all arrived.
187    partial: Option<u32>,
188    /// Commands framed out of this buffer and not yet run.
189    pending: u32,
190    /// This connection is on its way out, once what is buffered has gone.
191    closing: bool,
192    /// A protocol error waiting for the commands in front of it to answer.
193    ///
194    /// The framing finds the error before any of the batch it was framed with
195    /// has run, and writing the error there would put it in front of replies
196    /// the client is still owed. Redis answers in order, so this waits until
197    /// nothing is pending and goes out last.
198    deferred: Option<ProtocolError>,
199    /// Everything still queued for this connection is thrown away unanswered.
200    ///
201    /// `QUIT` sets this and a protocol error does not, which is the difference
202    /// between the two ways a connection ends. A client that pipelines `QUIT`
203    /// and then `SET` has said goodbye and then said something after it, and
204    /// Redis answers the goodbye and drops the rest. A client that sends two
205    /// good commands and then a malformed one gets both good ones answered,
206    /// because they were complete and correct before the stream went wrong.
207    skip: bool,
208    /// The peer is gone, so there is nothing to answer and nothing to write.
209    gone: bool,
210    /// Already on the dirty list.
211    dirty: bool,
212    /// This client is parked on a blocking command.
213    ///
214    /// While it is set, framing stops: whatever the client pipelined behind its
215    /// `BLPOP` stays in the read buffer unread, which is what a client waiting
216    /// for an answer means and is what Redis does with the same bytes.
217    blocked: bool,
218    /// Commands framed before it blocked and not run yet.
219    ///
220    /// A batch is framed before any of it runs, so a `BLPOP` can be the first of
221    /// sixty four commands and the other sixty three are already on their way to
222    /// the reactor when it parks. They come back here and go to the front of the
223    /// queue when the client wakes up, in the order they arrived.
224    ///
225    /// They are still counted in `pending`, which is what stops the read buffer
226    /// being compacted under the offsets they hold.
227    parked: Vec<Cmd>,
228    /// What the two buffers were holding the last time anybody counted.
229    ///
230    /// The connection's share of `INFO memory`, kept here so that reporting it
231    /// is a subtraction against this rather than a walk over every connection.
232    held: usize,
233}
234
235impl Conn {
236    fn new(id: u64) -> Conn {
237        // Accept time, which is the one moment a connection is allowed to cost
238        // an allocation. Everything after this reuses these two buffers.
239        yo_alloc::allow(|| Conn {
240            live: true,
241            session: Session::new(id),
242            out: Out::with_capacity(Proto::Resp2, OUT_BUF),
243            buf: Vec::with_capacity(READ_BUF),
244            head: 0,
245            partial: None,
246            pending: 0,
247            closing: false,
248            deferred: None,
249            skip: false,
250            gone: false,
251            dirty: false,
252            blocked: false,
253            parked: Vec::new(),
254            held: 0,
255        })
256    }
257
258    /// What the two buffers cost the process, which is the room they are
259    /// holding and not the bytes in use: both keep their capacity between
260    /// batches on purpose.
261    fn size(&self) -> usize {
262        self.buf.capacity() + self.out.capacity()
263    }
264
265    /// Back to how it was at accept time, buffers kept.
266    fn reset(&mut self, id: u64) {
267        self.live = true;
268        self.session = Session::new(id);
269        self.out.clear();
270        self.buf.clear();
271        self.head = 0;
272        self.partial = None;
273        self.pending = 0;
274        self.closing = false;
275        self.deferred = None;
276        self.skip = false;
277        self.gone = false;
278        self.dirty = false;
279        self.blocked = false;
280        // The room it took stays, the way the two buffers' does.
281        self.parked.clear();
282    }
283
284    /// Drop what the framing has already read, when nothing points into it.
285    ///
286    /// A framed command's arguments are offsets from the front of this buffer,
287    /// so this waits for the batch to run. After a batch is where a pipelining
288    /// connection spends most of its life, so that is not much of a wait.
289    ///
290    /// A half read command is not in the way. Its decoder was handed
291    /// `buf[head..]` and every offset it kept is from the front of that slice,
292    /// and `head` does not move until the command is complete, so the bytes it
293    /// is waiting on are exactly the bytes this keeps. They arrive at the front
294    /// instead of at `head` and the decoder cannot tell the difference.
295    ///
296    /// Waiting for it anyway is what made a read buffer grow to everything the
297    /// connection had ever sent. The framing loop only ever stops on an
298    /// incomplete command, and a buffer that ends on a command boundary gives
299    /// one of those on the next turn round: an empty slice, nothing decoded,
300    /// `Step::Incomplete`. So a connection that is exactly up to date always had
301    /// a decoder parked on it, this always returned early, and `head` walked
302    /// forward with the bytes behind it kept forever. Measured on server3, four
303    /// connections sending 100000 sets each held 16 MiB of read buffer apiece,
304    /// and fifty connections sending 8000 each held 1 MiB apiece: in both cases
305    /// every byte the connection had ever sent.
306    fn compact(&mut self) {
307        if self.pending > 0 || self.head == 0 {
308            return;
309        }
310        if self.head == self.buf.len() {
311            self.buf.clear();
312        } else {
313            self.buf.drain(..self.head);
314        }
315        self.head = 0;
316    }
317}
318
319/// The engine: connections on one side, the command layer on the other.
320///
321/// One per shard thread. Everything in it belongs to that thread, including the
322/// databases, which is what makes the whole path lock free rather than merely
323/// uncontended.
324pub struct Wire<S> {
325    server: Server,
326    sink: S,
327    conns: Vec<Conn>,
328    /// Connection slots that closed and can be handed out again.
329    free: Vec<ConnId>,
330    /// The decoder pool.
331    argvs: Vec<Argv>,
332    spare: Vec<u32>,
333    /// Framed and not yet handed to the reactor.
334    ready: VecDeque<Cmd>,
335    /// Connections this batch wrote to.
336    dirty: Vec<ConnId>,
337    /// Where a protocol error line is built before it is copied into a reply.
338    scratch: Vec<u8>,
339    limits: Limits,
340    next_id: u64,
341}
342
343impl<S: Sink> Wire<S> {
344    /// An engine with an empty server.
345    #[must_use]
346    pub fn new(sink: S) -> Wire<S> {
347        Wire::with_server(Server::new(), sink)
348    }
349
350    /// An engine over a server the caller built, which is how a test gives it a
351    /// clock it can move by hand.
352    #[must_use]
353    pub fn with_server(server: Server, sink: S) -> Wire<S> {
354        Wire {
355            server,
356            sink,
357            conns: Vec::new(),
358            free: Vec::new(),
359            argvs: Vec::new(),
360            spare: Vec::new(),
361            ready: VecDeque::with_capacity(BATCH_MAX),
362            dirty: Vec::with_capacity(16),
363            scratch: Vec::with_capacity(128),
364            limits: Limits::default(),
365            next_id: 1,
366        }
367    }
368
369    /// The databases and the numbers `INFO` reports.
370    #[must_use]
371    pub const fn server(&self) -> &Server {
372        &self.server
373    }
374
375    /// The same, for a caller that owns both ends.
376    pub const fn server_mut(&mut self) -> &mut Server {
377        &mut self.server
378    }
379
380    /// Where the replies went.
381    #[must_use]
382    pub const fn sink(&self) -> &S {
383        &self.sink
384    }
385
386    /// The same, mutably.
387    pub const fn sink_mut(&mut self) -> &mut S {
388        &mut self.sink
389    }
390
391    /// Change the protocol limits, which is `proto-max-bulk-len` and friends.
392    pub fn set_limits(&mut self, limits: Limits) {
393        self.limits = limits;
394    }
395
396    /// Open a connection and give back its id.
397    ///
398    /// Reuses a closed connection's slot and its two buffers when there is one,
399    /// so a server with a churning client population allocates for the high
400    /// water mark and not for the total.
401    pub fn accept(&mut self) -> ConnId {
402        let id = self.next_id;
403        self.next_id += 1;
404        self.server.stats.clients += 1;
405        self.server.stats.connections += 1;
406
407        let at = match self.free.pop() {
408            Some(at) => {
409                // A reused slot keeps its buffers, so what it holds is already
410                // counted and this only puts the id back in service.
411                self.conns[at as usize].reset(id);
412                at
413            }
414            None => {
415                let conn = Conn::new(id);
416                yo_alloc::allow(|| self.conns.push(conn));
417                (self.conns.len() - 1) as ConnId
418            }
419        };
420        self.note_size(at);
421        at
422    }
423
424    /// The peer went away.
425    ///
426    /// Whatever is buffered for it is dropped rather than written, and the slot
427    /// comes back as soon as the commands already framed out of its buffer have
428    /// run, because those commands' arguments still point into it.
429    pub fn hangup(&mut self, conn: ConnId) {
430        let c = &mut self.conns[conn as usize];
431        if !c.live {
432            return;
433        }
434        c.gone = true;
435        c.closing = true;
436        // A parked client holds its own commands, and those commands are what
437        // `pending` counts, so leaving it parked here would leave the slot owed
438        // to a connection that is never going to be answered. They go back to
439        // the queue and run as the no-ops a gone connection's commands are.
440        if c.blocked {
441            self.unpark(conn);
442        }
443        if self.conns[conn as usize].pending == 0 {
444            self.release(conn);
445        }
446    }
447
448    /// The client is not waiting any more: give it back its commands.
449    ///
450    /// The ones it had already sent go to the front of the queue in the order
451    /// they arrived, ahead of anything any other connection has waiting, because
452    /// they were framed before any of that was. Then framing starts again on
453    /// whatever arrived while it was parked.
454    fn unpark(&mut self, conn: ConnId) {
455        let mut parked = {
456            let c = &mut self.conns[conn as usize];
457            c.blocked = false;
458            core::mem::take(&mut c.parked)
459        };
460        // Back to front, since each one goes on the front.
461        while let Some(cmd) = parked.pop() {
462            if self.ready.len() == self.ready.capacity() {
463                yo_alloc::allow(|| self.ready.reserve(BATCH_MAX));
464            }
465            self.ready.push_front(cmd);
466        }
467        // Empty now, and back where it lives so its room is not paid for twice.
468        self.conns[conn as usize].parked = parked;
469        if !self.conns[conn as usize].closing {
470            self.frame(conn);
471        }
472    }
473
474    /// Answer everybody who can be answered, and let go of everybody whose
475    /// deadline has passed.
476    ///
477    /// The walk is over the waiter list rather than over the connections, so it
478    /// costs what blocking costs and not what the server costs. Every caller
479    /// checks that somebody is parked before calling, which is the load and the
480    /// branch a server with nobody blocked pays.
481    fn serve_waiters(&mut self) {
482        let now = self.server.now_ms();
483        let mut at = 0;
484        while at < self.server.waiters().len() {
485            let p = self.server.waiters().at(at);
486            {
487                let c = &self.conns[p.conn as usize];
488                // The slot is reused and the client id is not. `release`
489                // forgets waiters, so this should never fire; it is here
490                // because being wrong about it writes a reply into somebody
491                // else's socket rather than dropping one.
492                if !c.live || c.session.id() != p.client {
493                    self.server.waiters_mut().drop_at(at);
494                    continue;
495                }
496            }
497            // The engine cannot reach the databases and the server cannot reach
498            // the connections, so the two halves are taken apart here and the
499            // one buffer this waiter needs is handed over.
500            let served = {
501                let Wire { server, conns, .. } = self;
502                server.serve_waiter(at, now, &mut conns[p.conn as usize].out)
503            };
504            if served {
505                self.server.waiters_mut().drop_at(at);
506                self.unpark(p.conn);
507                self.soil(p.conn);
508            } else {
509                at += 1;
510            }
511        }
512    }
513
514    /// How many connections are open.
515    #[must_use]
516    pub fn clients(&self) -> usize {
517        self.conns.iter().filter(|c| c.live).count()
518    }
519
520    /// Commands framed and waiting for the reactor.
521    #[must_use]
522    pub fn ready(&self) -> usize {
523        self.ready.len()
524    }
525
526    /// Connections with a reply that has not gone out yet.
527    ///
528    /// Non zero means a socket was full and what is left is being held for a
529    /// later flush, which a driver waiting on readability needs to know: there
530    /// is work here that no incoming byte will ever wake it up for.
531    #[must_use]
532    pub fn owed(&self) -> usize {
533        self.dirty.len()
534    }
535
536    /// Decoders in the pool, which is the high water mark of one batch.
537    #[must_use]
538    pub fn decoders(&self) -> usize {
539        self.argvs.len()
540    }
541
542    /// What every connection's read and reply buffers are holding.
543    ///
544    /// The walk is fine here because this is a test and a report, and the
545    /// number the running server uses is the one kept by `note_size`.
546    #[must_use]
547    pub fn buffer_bytes(&self) -> usize {
548        self.conns.iter().map(Conn::size).sum()
549    }
550
551    /// Take bytes off a connection and frame whatever commands they complete.
552    ///
553    /// Anything left over stays in the connection's buffer, half a command
554    /// included, so the caller hands over whatever the socket gave it without
555    /// looking at it.
556    pub fn feed(&mut self, conn: ConnId, bytes: &[u8]) {
557        {
558            let c = &mut self.conns[conn as usize];
559            if !c.live || c.closing {
560                return;
561            }
562            // The buffer is sized for a command at accept time, so this only
563            // grows for a client sending a bulk larger than that, which is a
564            // real allocation for a real reason.
565            yo_alloc::allow(|| c.buf.extend_from_slice(bytes));
566        }
567        self.frame(conn);
568        self.note_size(conn);
569    }
570
571    /// Tell the server what this connection's buffers are holding now, if it
572    /// has changed since the last time anybody asked.
573    ///
574    /// Once per read and once per flush, which is where a buffer can grow, and
575    /// two loads and a compare when nothing has moved. The alternative is a
576    /// walk over every connection on a turn of the loop, which puts the cost of
577    /// a report nobody has asked for on the command path.
578    fn note_size(&mut self, conn: ConnId) {
579        let c = &mut self.conns[conn as usize];
580        let now = c.size();
581        if now == c.held {
582            return;
583        }
584        let delta = now as isize - c.held as isize;
585        c.held = now;
586        self.server.note_conn_bytes(delta);
587    }
588
589    /// Move as many complete commands as possible out of the read buffer.
590    ///
591    /// Nothing at all while the client is parked. The bytes stay where they are
592    /// and `head` does not move, so a client that pipelines `BLPOP` and then
593    /// `PING` gets the `PING` answered when the `BLPOP` is, and in that order.
594    fn frame(&mut self, conn: ConnId) {
595        if self.conns[conn as usize].blocked {
596            return;
597        }
598        loop {
599            let base = self.conns[conn as usize].head;
600            let slot = match self.conns[conn as usize].partial.take() {
601                Some(slot) => slot,
602                None => self.take_decoder(),
603            };
604
605            let step = {
606                let c = &self.conns[conn as usize];
607                self.argvs[slot as usize].decode(&c.buf[base..], &self.limits)
608            };
609
610            match step {
611                Ok(Step::Command { consumed }) => {
612                    self.conns[conn as usize].head += consumed;
613                    if self.argvs[slot as usize].is_empty() {
614                        // `*0` and a blank inline line: consumed, not answered.
615                        self.spare.push(slot);
616                    } else {
617                        if self.ready.len() == self.ready.capacity() {
618                            yo_alloc::allow(|| self.ready.reserve(BATCH_MAX));
619                        }
620                        self.ready.push_back(Cmd { conn, slot, base });
621                        self.conns[conn as usize].pending += 1;
622                    }
623                }
624                Ok(Step::Incomplete) => {
625                    // Hold the decoder so the rest of this command resumes
626                    // where it stopped instead of being read again from the
627                    // front every time more of it arrives.
628                    self.conns[conn as usize].partial = Some(slot);
629                    break;
630                }
631                Err(e) => {
632                    self.spare.push(slot);
633                    let c = &mut self.conns[conn as usize];
634                    // Held rather than written, so it lands behind the replies
635                    // to the commands that were framed in front of it out of
636                    // the same read.
637                    c.deferred = Some(e);
638                    // Redis closes after a protocol error and so do we: the two
639                    // ends no longer agree on where the next command starts.
640                    c.closing = true;
641                    self.soil(conn);
642                    break;
643                }
644            }
645        }
646        self.conns[conn as usize].compact();
647    }
648
649    /// A decoder from the pool, or a new one the first time round.
650    ///
651    /// The one from the pool is reset before it goes out, because a decoder can
652    /// come back to the pool part way through a command: a protocol error stops
653    /// framing where it is, and a connection that hangs up with half a command
654    /// in its buffer hands its decoder back too. Either one leaves a resume
655    /// point behind, and a resume point is an offset into a buffer that is
656    /// about to stop being the same buffer. A decoder taken here is always
657    /// starting a command, never continuing one, since a continuation comes off
658    /// the connection's own `partial` and never off the pool.
659    fn take_decoder(&mut self) -> u32 {
660        match self.spare.pop() {
661            Some(slot) => {
662                self.argvs[slot as usize].reset();
663                slot
664            }
665            None => yo_alloc::allow(|| {
666                self.argvs.push(Argv::with_capacity(ARGV_HINT));
667                // Every slot handed out here comes back to `spare` exactly
668                // once, so `spare` never holds more than `argvs` has slots.
669                // Sizing it here means the pushes that give a slot back never
670                // touch the allocator, and those are on the command path while
671                // this is not: a decoder is made once per depth of pipelining
672                // the connection has ever reached. `spare` is empty right now,
673                // which is why we are down here at all.
674                self.spare.reserve(self.argvs.len());
675                (self.argvs.len() - 1) as u32
676            }),
677        }
678    }
679
680    /// Note that this connection has something to write.
681    fn soil(&mut self, conn: ConnId) {
682        let c = &mut self.conns[conn as usize];
683        if !c.dirty {
684            c.dirty = true;
685            if self.dirty.len() == self.dirty.capacity() {
686                yo_alloc::allow(|| self.dirty.reserve(16));
687            }
688            self.dirty.push(conn);
689        }
690    }
691
692    /// Hand the slot and its buffers back.
693    fn release(&mut self, conn: ConnId) {
694        {
695            let c = &mut self.conns[conn as usize];
696            if !c.live {
697                return;
698            }
699            if let Some(slot) = c.partial.take() {
700                self.spare.push(slot);
701            }
702            c.live = false;
703            c.dirty = false;
704            c.blocked = false;
705            c.out.clear();
706            c.buf.clear();
707            c.head = 0;
708        }
709        // Before the slot goes back, because the slot is handed out again and a
710        // waiter on a client that has gone would then be a waiter pointing at
711        // somebody else's connection. The id is what makes it findable and the
712        // id is about to stop being this connection's.
713        let client = self.conns[conn as usize].session.id();
714        self.server.waiters_mut().forget(client);
715        self.server.stats.clients = self.server.stats.clients.saturating_sub(1);
716        self.sink.closed(conn);
717        yo_alloc::allow(|| self.free.push(conn));
718    }
719
720    /// Move up to `max` framed commands into `into`.
721    ///
722    /// The reactor wants a batch it owns, and the engine keeps the buffers, so
723    /// what crosses between them is this: numbers, no borrows.
724    pub fn take_ready(&mut self, into: &mut Vec<Cmd>, max: usize) -> usize {
725        let n = max.min(self.ready.len());
726        into.extend(self.ready.drain(..n));
727        n
728    }
729
730    /// Offer one connection's replies to the sink, and say whether it still
731    /// owes bytes afterwards.
732    fn write_out(&mut self, conn: ConnId) -> bool {
733        {
734            let c = &self.conns[conn as usize];
735            if !c.live {
736                return false;
737            }
738        }
739        // A protocol error goes out once everything in front of it has.
740        if self.conns[conn as usize].pending == 0
741            && let Some(e) = self.conns[conn as usize].deferred.take()
742        {
743            self.scratch.clear();
744            e.write_reply(&mut self.scratch);
745            self.conns[conn as usize].out.raw(&self.scratch);
746        }
747
748        let taken = {
749            let c = &self.conns[conn as usize];
750            if c.out.is_empty() {
751                0
752            } else {
753                // One write for the whole batch's replies, never one per reply.
754                self.sink.write(conn, c.out.as_slice())
755            }
756        };
757
758        let c = &mut self.conns[conn as usize];
759        if taken >= c.out.len() {
760            c.out.clear();
761        } else {
762            c.out.consume(taken);
763        }
764
765        if !c.out.is_empty() {
766            return true;
767        }
768        c.dirty = false;
769        if c.closing && c.pending == 0 {
770            self.release(conn);
771        } else {
772            c.compact();
773        }
774        self.note_size(conn);
775        false
776    }
777
778    /// Take a clock reading for the whole batch.
779    ///
780    /// `04` section 5: once per turn, never per command, so every command in a
781    /// batch compares against the same millisecond and two keys written
782    /// together expire together.
783    pub fn tick(&mut self) {
784        self.server.refresh_clock();
785    }
786
787    /// Do one batch's worth of housekeeping.
788    ///
789    /// Today that is one segment of arena compaction at most, which is what
790    /// stops a server that rewrites the same keys from holding every version of
791    /// them. It is separate from [`Wire::tick`] because the clock has to move
792    /// before a batch runs and this does not: it can wait until the replies are
793    /// out, and the driver decides when that is.
794    ///
795    /// Per batch and not per turn of the loop. A turn can carry one command or
796    /// a thousand, so a per turn call means the rate at which garbage is
797    /// collected has nothing to do with the rate at which it is made, and on a
798    /// saturated server the second one wins. That was measured: with this on
799    /// the loop's turn the server settled at seven segments for six segments'
800    /// worth of keys, which is where an unloaded process running the same
801    /// writes settled at six.
802    pub fn maintain(&mut self) -> Option<usize> {
803        self.server.compact_step()
804    }
805}
806
807impl<S: Sink> Engine for Wire<S> {
808    type Work = Cmd;
809
810    fn key_hash(&self, cmd: &Cmd) -> Option<u64> {
811        let c = &self.conns[cmd.conn as usize];
812        let args = Args::new(&self.argvs[cmd.slot as usize], &c.buf[cmd.base..]);
813        let spec = lookup(args.name())?;
814        if spec.first_key <= 0 {
815            return None;
816        }
817        // The first key only. A command with more than one, which is `MSET` and
818        // `MGET`, warms the first and takes the miss on the rest; warming all of
819        // them means a hash list per command and that is the batch's own job
820        // once multi key commands are worth measuring.
821        let key = args.opt(spec.first_key as usize)?;
822        Some(Keyspace::hash_of(key))
823    }
824
825    fn prefetch(&self, cmd: &Cmd, hash: u64) {
826        let db = self.conns[cmd.conn as usize].session.db();
827        self.server.db_ref(db).prefetch(hash);
828    }
829
830    fn run(&mut self, cmd: Cmd, _hash: Option<u64>) -> yo_reactor::Flow {
831        // Framed with the batch that blocked, so it is a command the client sent
832        // before it knew it would be waiting. It keeps its decoder and it keeps
833        // its place in `pending`, which is what stops the buffer it points into
834        // being compacted while it waits.
835        if self.conns[cmd.conn as usize].blocked {
836            yo_alloc::allow(|| self.conns[cmd.conn as usize].parked.push(cmd));
837            return yo_reactor::Flow::Next;
838        }
839
840        let flow = {
841            let c = &mut self.conns[cmd.conn as usize];
842            c.pending -= 1;
843            if c.gone || c.skip {
844                // Nobody to answer, or nobody who should be. The decoder still
845                // has to come back and the slot still has to be released, which
846                // is why this is not an early return.
847                Flow::Continue
848            } else {
849                let args = Args::new(&self.argvs[cmd.slot as usize], &c.buf[cmd.base..]);
850                execute(&mut self.server, &mut c.session, args, &mut c.out)
851            }
852        };
853
854        self.spare.push(cmd.slot);
855        let c = &self.conns[cmd.conn as usize];
856        if c.gone {
857            if c.pending == 0 {
858                self.release(cmd.conn);
859            }
860        } else {
861            match flow {
862                Flow::Close => {
863                    let c = &mut self.conns[cmd.conn as usize];
864                    c.closing = true;
865                    // Anything the client pipelined behind the `QUIT` was sent
866                    // before it knew the answer, and running it would be acting
867                    // on a connection that has already been said goodbye to.
868                    c.skip = true;
869                    self.soil(cmd.conn);
870                }
871                // Nothing was written, so there is nothing to flush and no
872                // reason to put this connection on the dirty list. The waiter
873                // carries the slot from here on, and it needs to know which one:
874                // the command layer only ever saw the client id.
875                Flow::Block => {
876                    self.conns[cmd.conn as usize].blocked = true;
877                    let client = self.conns[cmd.conn as usize].session.id();
878                    self.server.waiters_mut().bind(client, cmd.conn);
879                }
880                Flow::Continue => self.soil(cmd.conn),
881            }
882        }
883
884        // After each command and not once per batch. A client blocked on two
885        // keys and woken by `RPUSH b` then `RPUSH a` in one pipeline has to
886        // answer with `b`, because that is the push that was in front of it, and
887        // it can only do that if it was served in between the two.
888        if !self.server.waiters().is_empty() {
889            self.serve_waiters();
890        }
891        yo_reactor::Flow::Next
892    }
893
894    fn flush(&mut self) {
895        // The deadline sweep, and it is here because this is the one thing the
896        // driver calls on a turn that ran nothing at all. A client whose timeout
897        // passes while the server is idle is answered within the loop's idle
898        // wait, which is 20ms and is finer than the 10hz Redis checks its own
899        // blocked clients at.
900        if !self.server.waiters().is_empty() {
901            self.server.refresh_clock();
902            self.serve_waiters();
903        }
904
905        // Taken and put back so the loop below can reach the rest of the
906        // engine. The capacity comes back with it, so this is not an
907        // allocation.
908        let mut dirty = core::mem::take(&mut self.dirty);
909        let mut at = 0;
910        while at < dirty.len() {
911            let conn = dirty[at];
912            let owed = self.write_out(conn);
913            if owed {
914                // The socket was full. The connection stays on the list with
915                // what is left of its reply, and the next flush offers it
916                // again, which is the whole of the backpressure story here.
917                at += 1;
918            } else {
919                dirty.swap_remove(at);
920            }
921        }
922        self.dirty = dirty;
923    }
924
925    fn maintain(&mut self, budget: &mut yo_reactor::Budget) {
926        // The clock is the first thing the maintenance slice does, because
927        // everything else in it compares against a time.
928        if budget.spend(1) {
929            self.tick();
930        }
931    }
932}
933
934/// Run everything that is framed, in batches, and write the replies.
935///
936/// The inline driver: it is what a caller who is already on the shard thread
937/// uses in place of the loop, and it goes through the same two walks the loop
938/// goes through (`15` section 7). `batch` is the caller's, so a driver in a hot
939/// loop hands the same `Vec` back every time and never allocates.
940pub fn pump<S: Sink>(reactor: &mut Reactor<Wire<S>>, batch: &mut Vec<Cmd>) -> usize {
941    let mut ran = 0;
942    reactor.engine_mut().tick();
943    loop {
944        batch.clear();
945        if reactor.engine_mut().take_ready(batch, BATCH_MAX) == 0 {
946            break;
947        }
948        // The command path, and therefore the thing Y7 is about. The guard is
949        // what arms `yo-alloc`, and it covers dispatch and nothing else: framing
950        // before it and writing the replies after it are both allowed to reach
951        // for the heap, and only running the commands is not.
952        //
953        // It goes here rather than around the whole loop because `take_ready`
954        // and `flush` are on the other side of that line, and because a batch is
955        // the unit a caller can reason about. Under the default mode this is one
956        // relaxed load.
957        let armed = yo_alloc::guard();
958        ran += reactor.execute_all(batch.drain(..));
959        drop(armed);
960        reactor.engine_mut().flush();
961        // After the replies are out, so the batch that made the garbage is not
962        // the batch that waits for it to be collected.
963        reactor.engine_mut().maintain();
964    }
965    // Once more, for a connection with something to say and nothing to run: a
966    // protocol error, or a socket that was full the last time round.
967    reactor.engine_mut().flush();
968    // And once for a turn that ran nothing at all, which is where a server that
969    // has gone quiet catches up on what the last busy turn left behind.
970    reactor.engine_mut().maintain();
971    ran
972}
973
974#[cfg(test)]
975mod tests {
976    use super::*;
977
978    /// The wire bytes for a command, built the way a client would.
979    fn wire(args: &[&[u8]]) -> Vec<u8> {
980        let mut b = format!("*{}\r\n", args.len()).into_bytes();
981        for a in args {
982            b.extend_from_slice(format!("${}\r\n", a.len()).as_bytes());
983            b.extend_from_slice(a);
984            b.extend_from_slice(b"\r\n");
985        }
986        b
987    }
988
989    fn engine() -> (Reactor<Wire<Recorder>>, ConnId, Vec<Cmd>) {
990        let mut r = Reactor::inline(Wire::new(Recorder::new()));
991        let conn = r.engine_mut().accept();
992        (r, conn, Vec::new())
993    }
994
995    /// Where the fixed clock a blocking test moves by hand starts.
996    const START_MS: u64 = 1_000_000;
997
998    /// The same, on a clock the test moves rather than the system's.
999    ///
1000    /// A test about a timeout cannot wait for one: waiting a hundred
1001    /// milliseconds is a test that fails on a loaded machine and waiting a
1002    /// hundred seconds is not a test.
1003    fn timed() -> (Reactor<Wire<Recorder>>, ConnId, Vec<Cmd>) {
1004        let server = crate::dispatch::Server::with_clock(yo_kv::Clock::fixed(START_MS));
1005        let mut r = Reactor::inline(Wire::with_server(server, Recorder::new()));
1006        let conn = r.engine_mut().accept();
1007        (r, conn, Vec::new())
1008    }
1009
1010    #[test]
1011    fn a_pipelined_batch_comes_back_in_order_and_in_one_write() {
1012        let (mut r, conn, mut batch) = engine();
1013        let mut stream = wire(&[b"SET", b"k", b"v"]);
1014        stream.extend(wire(&[b"GET", b"k"]));
1015        stream.extend(wire(&[b"INCR", b"n"]));
1016
1017        r.engine_mut().feed(conn, &stream);
1018        assert_eq!(r.engine().ready(), 3);
1019        assert_eq!(pump(&mut r, &mut batch), 3);
1020
1021        assert_eq!(r.engine().sink().sent(conn), b"+OK\r\n$1\r\nv\r\n:1\r\n");
1022        assert_eq!(r.engine().ready(), 0);
1023    }
1024
1025    /// The framing has to survive a command arriving in pieces, because that is
1026    /// what a socket does.
1027    #[test]
1028    fn a_command_split_across_reads_resumes_rather_than_restarts() {
1029        let (mut r, conn, mut batch) = engine();
1030        let bytes = wire(&[b"SET", b"key", b"value"]);
1031
1032        for at in 1..bytes.len() {
1033            r.engine_mut().feed(conn, &bytes[at - 1..at]);
1034            assert_eq!(r.engine().ready(), 0, "not a command yet at {at}");
1035        }
1036        r.engine_mut().feed(conn, &bytes[bytes.len() - 1..]);
1037        assert_eq!(r.engine().ready(), 1);
1038        assert_eq!(pump(&mut r, &mut batch), 1);
1039        assert_eq!(r.engine().sink().sent(conn), b"+OK\r\n");
1040
1041        // And the value that arrived in single bytes is the value that was
1042        // stored, which is the part a naive resume gets wrong.
1043        r.engine_mut().feed(conn, &wire(&[b"GET", b"key"]));
1044        pump(&mut r, &mut batch);
1045        assert_eq!(r.engine().sink().sent(conn), b"+OK\r\n$5\r\nvalue\r\n");
1046    }
1047
1048    #[test]
1049    fn two_connections_are_two_sessions_over_one_server() {
1050        let (mut r, a, mut batch) = engine();
1051        let b = r.engine_mut().accept();
1052
1053        r.engine_mut().feed(a, &wire(&[b"SELECT", b"3"]));
1054        r.engine_mut().feed(a, &wire(&[b"SET", b"k", b"a"]));
1055        r.engine_mut().feed(b, &wire(&[b"SET", b"k", b"b"]));
1056        r.engine_mut().feed(a, &wire(&[b"GET", b"k"]));
1057        r.engine_mut().feed(b, &wire(&[b"GET", b"k"]));
1058        pump(&mut r, &mut batch);
1059
1060        assert_eq!(r.engine().sink().sent(a), b"+OK\r\n+OK\r\n$1\r\na\r\n");
1061        assert_eq!(r.engine().sink().sent(b), b"+OK\r\n$1\r\nb\r\n");
1062        assert_eq!(r.engine().clients(), 2);
1063    }
1064
1065    #[test]
1066    fn quit_is_answered_and_then_the_connection_goes() {
1067        let (mut r, conn, mut batch) = engine();
1068        r.engine_mut().feed(conn, &wire(&[b"PING"]));
1069        r.engine_mut().feed(conn, &wire(&[b"QUIT"]));
1070        pump(&mut r, &mut batch);
1071
1072        assert_eq!(r.engine().sink().sent(conn), b"+PONG\r\n+OK\r\n");
1073        assert!(r.engine().sink().was_closed(conn));
1074        assert_eq!(r.engine().clients(), 0);
1075
1076        // The slot comes back, buffers and all.
1077        let again = r.engine_mut().accept();
1078        assert_eq!(again, conn);
1079        assert_eq!(r.engine().clients(), 1);
1080    }
1081
1082    /// Redis's own unit/quit, which caught this: we answered the `QUIT` and
1083    /// then ran the `SET` behind it.
1084    #[test]
1085    fn what_a_client_pipelined_behind_quit_is_never_run() {
1086        let (mut r, conn, mut batch) = engine();
1087        let mut stream = wire(&[b"QUIT"]);
1088        stream.extend(wire(&[b"SET", b"foo", b"bar"]));
1089        r.engine_mut().feed(conn, &stream);
1090        // Both were framed, because framing happens before anything runs.
1091        assert_eq!(r.engine().ready(), 2);
1092        pump(&mut r, &mut batch);
1093
1094        // One reply and not two, and the connection is gone.
1095        assert_eq!(r.engine().sink().sent(conn), b"+OK\r\n");
1096        assert!(r.engine().sink().was_closed(conn));
1097
1098        // And the write never happened, which is the part a client can see
1099        // after it reconnects. The recorder is cleared first because the next
1100        // connection lands back in the slot this one just left, and what was
1101        // written to the slot before is still sitting in it.
1102        r.engine_mut().sink_mut().clear();
1103        let next = r.engine_mut().accept();
1104        r.engine_mut().feed(next, &wire(&[b"GET", b"foo"]));
1105        pump(&mut r, &mut batch);
1106        assert_eq!(r.engine().sink().sent(next), b"$-1\r\n");
1107    }
1108
1109    /// The other way a connection ends, which does not throw anything away.
1110    #[test]
1111    fn commands_that_arrived_before_a_protocol_error_are_still_answered() {
1112        let (mut r, conn, mut batch) = engine();
1113        let mut stream = wire(&[b"SET", b"k", b"v"]);
1114        stream.extend(wire(&[b"GET", b"k"]));
1115        stream.extend_from_slice(b"*1\r\n+notabulk\r\n");
1116        r.engine_mut().feed(conn, &stream);
1117        pump(&mut r, &mut batch);
1118
1119        // Both good commands were complete and correct before the stream went
1120        // wrong, so both are answered and the error comes after them.
1121        let sent = r.engine().sink().sent(conn);
1122        assert!(
1123            sent.starts_with(b"+OK\r\n$1\r\nv\r\n-ERR Protocol error: "),
1124            "{sent:?}"
1125        );
1126        assert!(r.engine().sink().was_closed(conn));
1127    }
1128
1129    #[test]
1130    fn a_protocol_error_is_written_and_closes_the_connection() {
1131        let (mut r, conn, mut batch) = engine();
1132        // A multibulk that says its first argument is a bulk and then does not.
1133        r.engine_mut().feed(conn, b"*1\r\n+notabulk\r\n");
1134        pump(&mut r, &mut batch);
1135
1136        let sent = r.engine().sink().sent(conn);
1137        assert!(sent.starts_with(b"-ERR Protocol error: "), "{sent:?}");
1138        assert!(r.engine().sink().was_closed(conn));
1139        assert_eq!(r.engine().clients(), 0);
1140    }
1141
1142    /// Redis's own `unit/protocol` walks a list of malformed frames, each on a
1143    /// fresh connection, which means every one of them after the first runs on
1144    /// a decoder that came back to the pool part way through a command.
1145    #[test]
1146    fn a_decoder_that_came_back_mid_command_starts_the_next_one_clean() {
1147        let (mut r, conn, mut batch) = engine();
1148        // Stops inside the third argument, on a length that is not a length.
1149        r.engine_mut()
1150            .feed(conn, b"*3\r\n$3\r\nSET\r\n$1\r\nx\r\n$blabla\r\n");
1151        pump(&mut r, &mut batch);
1152        let sent = r.engine().sink().sent(conn);
1153        assert!(
1154            sent.starts_with(b"-ERR Protocol error: invalid bulk length"),
1155            "{sent:?}"
1156        );
1157
1158        // The slot that decoder was in is now the slot the next connection
1159        // gets, and it has to be at the start of a command and not half way
1160        // through the one that went wrong.
1161        r.engine_mut().sink_mut().clear();
1162        let next = r.engine_mut().accept();
1163        r.engine_mut().feed(next, &wire(&[b"GET", b"k"]));
1164        pump(&mut r, &mut batch);
1165        assert_eq!(r.engine().sink().sent(next), b"$-1\r\n");
1166
1167        r.engine_mut().sink_mut().clear();
1168        let third = r.engine_mut().accept();
1169        r.engine_mut().feed(third, b"*1\r\n+notabulk\r\n");
1170        pump(&mut r, &mut batch);
1171        let sent = r.engine().sink().sent(third);
1172        assert!(sent.starts_with(b"-ERR Protocol error: "), "{sent:?}");
1173    }
1174
1175    /// A client that hangs up mid batch is the case that gets a server killed:
1176    /// the commands already framed still point into its buffer.
1177    #[test]
1178    fn a_hangup_with_commands_in_flight_waits_for_them() {
1179        let (mut r, conn, mut batch) = engine();
1180        r.engine_mut().feed(conn, &wire(&[b"SET", b"k", b"v"]));
1181        r.engine_mut().feed(conn, &wire(&[b"GET", b"k"]));
1182
1183        batch.clear();
1184        r.engine_mut().take_ready(&mut batch, BATCH_MAX);
1185        r.engine_mut().hangup(conn);
1186        assert_eq!(r.engine().clients(), 1, "still holding the buffer");
1187
1188        r.execute_all(batch.drain(..));
1189        r.engine_mut().flush();
1190        assert_eq!(r.engine().clients(), 0);
1191        assert!(r.engine().sink().sent(conn).is_empty(), "nobody to answer");
1192
1193        // And the slot is usable again, with the decoders both back in the
1194        // pool rather than lost with the connection.
1195        let decoders = r.engine().decoders();
1196        let again = r.engine_mut().accept();
1197        assert_eq!(again, conn);
1198        r.engine_mut().feed(again, &wire(&[b"PING"]));
1199        pump(&mut r, &mut batch);
1200        assert_eq!(r.engine().sink().sent(again), b"+PONG\r\n");
1201        assert_eq!(r.engine().decoders(), decoders);
1202    }
1203
1204    /// The claim that the steady state does not allocate, checked the only way
1205    /// a library test can check it: nothing grows.
1206    #[test]
1207    fn the_buffers_and_the_decoder_pool_stop_growing() {
1208        let (mut r, conn, mut batch) = engine();
1209        let mut stream = Vec::new();
1210        for i in 0..32 {
1211            stream.extend(wire(&[b"SET", format!("k{i}").as_bytes(), b"v"]));
1212        }
1213
1214        r.engine_mut().feed(conn, &stream);
1215        pump(&mut r, &mut batch);
1216        let decoders = r.engine().decoders();
1217        let batch_cap = batch.capacity();
1218
1219        for _ in 0..10 {
1220            r.engine_mut().feed(conn, &stream);
1221            pump(&mut r, &mut batch);
1222        }
1223        assert_eq!(r.engine().decoders(), decoders, "the pool is reused");
1224        assert_eq!(batch.capacity(), batch_cap, "the batch buffer is reused");
1225        assert!(
1226            decoders <= BATCH_MAX + 1,
1227            "{decoders} decoders for 32 commands"
1228        );
1229    }
1230
1231    /// The read buffer holds what has not been dealt with yet and nothing else.
1232    ///
1233    /// A client that pipelines sixteen commands, waits for the sixteen replies
1234    /// and goes again is what `redis-benchmark -P 16` does and what half of the
1235    /// clients in the world do. Every one of those rounds leaves the buffer
1236    /// exactly caught up, and a buffer that never drops what it has already
1237    /// dealt with grows to everything the connection has ever sent: 16 MiB
1238    /// apiece on server3 for four connections sending 100000 sets each.
1239    #[test]
1240    fn a_pipelining_client_does_not_grow_the_read_buffer() {
1241        let (mut r, conn, mut batch) = engine();
1242        let mut round = Vec::new();
1243        for i in 0..16 {
1244            round.extend(wire(&[b"SET", format!("k{i}").as_bytes(), b"v"]));
1245        }
1246
1247        r.engine_mut().feed(conn, &round);
1248        pump(&mut r, &mut batch);
1249        r.engine_mut().sink_mut().clear();
1250        let after_one = r.engine().buffer_bytes();
1251
1252        // A thousand rounds is sixteen thousand commands and about a megabyte
1253        // of wire bytes, which is a hundred times what the buffer starts with.
1254        for _ in 0..1000 {
1255            r.engine_mut().feed(conn, &round);
1256            pump(&mut r, &mut batch);
1257            r.engine_mut().sink_mut().clear();
1258        }
1259
1260        assert_eq!(
1261            r.engine().buffer_bytes(),
1262            after_one,
1263            "the buffers grew over a thousand rounds of the same sixteen commands"
1264        );
1265        assert!(
1266            r.engine().server().memory_bytes() >= after_one,
1267            "the buffers are counted in what the server reports"
1268        );
1269    }
1270
1271    /// Half a command in the buffer is the case compaction has to be careful
1272    /// about, because the decoder holding it kept offsets into those bytes.
1273    #[test]
1274    fn a_command_split_across_reads_survives_compaction() {
1275        let (mut r, conn, mut batch) = engine();
1276        let cmd = wire(&[b"SET", b"key", b"value"]);
1277        let (head, tail) = cmd.split_at(cmd.len() - 4);
1278
1279        // A complete command, so that there is something in front to drop, then
1280        // most of a second one.
1281        r.engine_mut().feed(conn, &wire(&[b"PING"]));
1282        r.engine_mut().feed(conn, head);
1283        pump(&mut r, &mut batch);
1284        assert_eq!(r.engine().sink().sent(conn), b"+PONG\r\n");
1285
1286        // The rest of it arrives after the buffer has been compacted under it.
1287        r.engine_mut().feed(conn, tail);
1288        pump(&mut r, &mut batch);
1289        assert_eq!(r.engine().sink().sent(conn), b"+PONG\r\n+OK\r\n");
1290
1291        r.engine_mut().feed(conn, &wire(&[b"GET", b"key"]));
1292        pump(&mut r, &mut batch);
1293        assert!(r.engine().sink().sent(conn).ends_with(b"$5\r\nvalue\r\n"));
1294    }
1295
1296    /// The two walks are the reactor's, not this module's, so the test is that
1297    /// the engine can be driven by them at all: same commands, same replies.
1298    #[test]
1299    fn the_batch_goes_through_the_reactors_two_walks() {
1300        let (mut r, conn, mut batch) = engine();
1301        for i in 0..100 {
1302            r.engine_mut()
1303                .feed(conn, &wire(&[b"INCR", format!("k{}", i % 7).as_bytes()]));
1304        }
1305        let ran = pump(&mut r, &mut batch);
1306
1307        assert_eq!(ran, 100);
1308        assert_eq!(r.commands(), 100);
1309        // Two batches, because a hundred commands do not fit in sixty four.
1310        assert_eq!(r.turns(), 2);
1311        // The hundredth command is the fifteenth `INCR` of `k1`.
1312        assert!(r.engine().sink().sent(conn).ends_with(b":15\r\n"));
1313    }
1314
1315    /// A sink that takes four bytes at a time, which is what a full socket
1316    /// looks like from in here.
1317    #[derive(Default)]
1318    struct Trickle {
1319        sent: Vec<u8>,
1320        writes: usize,
1321    }
1322
1323    impl Sink for Trickle {
1324        fn write(&mut self, _conn: ConnId, bytes: &[u8]) -> usize {
1325            self.writes += 1;
1326            let n = bytes.len().min(4);
1327            self.sent.extend_from_slice(&bytes[..n]);
1328            n
1329        }
1330    }
1331
1332    /// A blocking command that does not block costs nothing: no waiter, no
1333    /// allocation, the same three lines the non blocking one runs.
1334    #[test]
1335    fn a_blpop_on_a_list_with_something_in_it_never_waits() {
1336        let (mut r, conn, mut batch) = engine();
1337        r.engine_mut().feed(conn, &wire(&[b"RPUSH", b"q", b"a"]));
1338        r.engine_mut().feed(conn, &wire(&[b"BLPOP", b"q", b"0"]));
1339        pump(&mut r, &mut batch);
1340
1341        assert_eq!(
1342            r.engine().sink().sent(conn),
1343            b":1\r\n*2\r\n$1\r\nq\r\n$1\r\na\r\n"
1344        );
1345        assert_eq!(r.engine().server().waiters().len(), 0);
1346    }
1347
1348    /// The whole point: a client with nothing to pop is answered later, by
1349    /// somebody else's command.
1350    #[test]
1351    fn a_parked_client_is_answered_by_another_connections_push() {
1352        let (mut r, a, mut batch) = engine();
1353        let b = r.engine_mut().accept();
1354
1355        r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1356        pump(&mut r, &mut batch);
1357        assert!(r.engine().sink().sent(a).is_empty(), "nothing to say yet");
1358        assert_eq!(r.engine().server().waiters().len(), 1);
1359
1360        r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"one"]));
1361        pump(&mut r, &mut batch);
1362
1363        assert_eq!(r.engine().sink().sent(a), b"*2\r\n$1\r\nq\r\n$3\r\none\r\n");
1364        // The push still reports the length it made, even though the element was
1365        // gone again before the reply was written.
1366        assert_eq!(r.engine().sink().sent(b), b":1\r\n");
1367        assert_eq!(r.engine().server().waiters().len(), 0);
1368    }
1369
1370    /// A push to a key nobody named, and a key of another type on a key
1371    /// somebody did: neither is a wake up, and the client stays parked.
1372    #[test]
1373    fn only_a_list_arriving_under_a_named_key_wakes_a_waiter() {
1374        let (mut r, a, mut batch) = engine();
1375        let b = r.engine_mut().accept();
1376        r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1377        pump(&mut r, &mut batch);
1378
1379        r.engine_mut()
1380            .feed(b, &wire(&[b"RPUSH", b"elsewhere", b"x"]));
1381        r.engine_mut().feed(b, &wire(&[b"SADD", b"q", b"x"]));
1382        pump(&mut r, &mut batch);
1383
1384        assert!(r.engine().sink().sent(a).is_empty());
1385        assert_eq!(r.engine().server().waiters().len(), 1, "still waiting");
1386        // And the set is intact, so the waiter did not take anything out of it
1387        // on its way past.
1388        assert_eq!(r.engine().sink().sent(b), b":1\r\n:1\r\n");
1389    }
1390
1391    /// Two workers on one queue, which is what `BLPOP` is for. They are served
1392    /// in the order they arrived and not in whatever order the list is walked.
1393    #[test]
1394    fn two_parked_clients_are_served_in_the_order_they_arrived() {
1395        let (mut r, a, mut batch) = engine();
1396        let b = r.engine_mut().accept();
1397        let c = r.engine_mut().accept();
1398
1399        r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1400        pump(&mut r, &mut batch);
1401        r.engine_mut().feed(b, &wire(&[b"BLPOP", b"q", b"0"]));
1402        pump(&mut r, &mut batch);
1403        assert_eq!(r.engine().server().waiters().len(), 2);
1404
1405        r.engine_mut()
1406            .feed(c, &wire(&[b"RPUSH", b"q", b"first", b"second"]));
1407        pump(&mut r, &mut batch);
1408
1409        assert_eq!(
1410            r.engine().sink().sent(a),
1411            b"*2\r\n$1\r\nq\r\n$5\r\nfirst\r\n"
1412        );
1413        assert_eq!(
1414            r.engine().sink().sent(b),
1415            b"*2\r\n$1\r\nq\r\n$6\r\nsecond\r\n"
1416        );
1417        assert_eq!(r.engine().server().waiters().len(), 0);
1418    }
1419
1420    /// A client waiting for an answer is not a client that has sent another
1421    /// question, so what it pipelined behind its `BLPOP` waits for the `BLPOP`.
1422    #[test]
1423    fn what_a_client_pipelined_behind_a_block_waits_for_the_block() {
1424        let (mut r, a, mut batch) = engine();
1425        let b = r.engine_mut().accept();
1426
1427        // Framed together, so the `PING` is already on its way to the reactor
1428        // when the `BLPOP` in front of it parks.
1429        let mut stream = wire(&[b"BLPOP", b"q", b"0"]);
1430        stream.extend(wire(&[b"PING"]));
1431        r.engine_mut().feed(a, &stream);
1432        pump(&mut r, &mut batch);
1433        assert!(
1434            r.engine().sink().sent(a).is_empty(),
1435            "the PING went out in front of the answer it was sent behind"
1436        );
1437
1438        // And one that arrives while it is parked is not even framed.
1439        r.engine_mut().feed(a, &wire(&[b"ECHO", b"after"]));
1440        pump(&mut r, &mut batch);
1441        assert!(r.engine().sink().sent(a).is_empty());
1442
1443        r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"x"]));
1444        pump(&mut r, &mut batch);
1445        assert_eq!(
1446            r.engine().sink().sent(a),
1447            b"*2\r\n$1\r\nq\r\n$1\r\nx\r\n+PONG\r\n$5\r\nafter\r\n"
1448        );
1449    }
1450
1451    /// Redis serves parked clients after every command rather than once per
1452    /// turn of the loop, and a pipeline is where the difference shows: the
1453    /// waiter has to be served between the two pushes, so it answers with the
1454    /// key the first push filled and not with the one it named first.
1455    #[test]
1456    fn a_waiter_is_served_between_two_pipelined_pushes() {
1457        let (mut r, a, mut batch) = engine();
1458        let b = r.engine_mut().accept();
1459        r.engine_mut()
1460            .feed(a, &wire(&[b"BLPOP", b"p1", b"p2", b"0"]));
1461        pump(&mut r, &mut batch);
1462
1463        let mut stream = wire(&[b"RPUSH", b"p2", b"second"]);
1464        stream.extend(wire(&[b"RPUSH", b"p1", b"first"]));
1465        r.engine_mut().feed(b, &stream);
1466        pump(&mut r, &mut batch);
1467
1468        assert_eq!(
1469            r.engine().sink().sent(a),
1470            b"*2\r\n$2\r\np2\r\n$6\r\nsecond\r\n"
1471        );
1472        // Which leaves the key it named first holding what was pushed to it.
1473        r.engine_mut()
1474            .feed(b, &wire(&[b"LRANGE", b"p1", b"0", b"-1"]));
1475        pump(&mut r, &mut batch);
1476        assert!(
1477            r.engine()
1478                .sink()
1479                .sent(b)
1480                .ends_with(b"*1\r\n$5\r\nfirst\r\n")
1481        );
1482    }
1483
1484    /// A `BLMOVE` that serves itself is a push, so it wakes the client waiting
1485    /// on the key it pushed to, in the same moment and without a turn of the
1486    /// loop in between.
1487    #[test]
1488    fn a_waiter_woken_by_another_waiter() {
1489        let (mut r, a, mut batch) = engine();
1490        let b = r.engine_mut().accept();
1491        let c = r.engine_mut().accept();
1492
1493        r.engine_mut()
1494            .feed(a, &wire(&[b"BLMOVE", b"x", b"y", b"LEFT", b"RIGHT", b"0"]));
1495        pump(&mut r, &mut batch);
1496        r.engine_mut().feed(b, &wire(&[b"BLPOP", b"y", b"0"]));
1497        pump(&mut r, &mut batch);
1498        assert_eq!(r.engine().server().waiters().len(), 2);
1499
1500        r.engine_mut().feed(c, &wire(&[b"RPUSH", b"x", b"chain"]));
1501        pump(&mut r, &mut batch);
1502
1503        assert_eq!(r.engine().sink().sent(a), b"$5\r\nchain\r\n");
1504        assert_eq!(
1505            r.engine().sink().sent(b),
1506            b"*2\r\n$1\r\ny\r\n$5\r\nchain\r\n"
1507        );
1508        assert_eq!(r.engine().server().waiters().len(), 0);
1509    }
1510
1511    /// A waiter on one database is not woken by a push on another, even though
1512    /// the key has the same name.
1513    #[test]
1514    fn a_waiter_is_only_woken_on_the_database_it_blocked_on() {
1515        let (mut r, a, mut batch) = engine();
1516        let b = r.engine_mut().accept();
1517        r.engine_mut().feed(a, &wire(&[b"SELECT", b"3"]));
1518        r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1519        pump(&mut r, &mut batch);
1520        assert_eq!(r.engine().sink().sent(a), b"+OK\r\n");
1521
1522        r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"wrongdb"]));
1523        pump(&mut r, &mut batch);
1524        assert_eq!(r.engine().sink().sent(a), b"+OK\r\n", "still waiting");
1525
1526        r.engine_mut().feed(b, &wire(&[b"SELECT", b"3"]));
1527        r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"rightdb"]));
1528        pump(&mut r, &mut batch);
1529        assert!(r.engine().sink().sent(a).ends_with(b"$7\r\nrightdb\r\n"));
1530    }
1531
1532    /// The deadline sweep, which runs on a turn that has nothing else to do.
1533    #[test]
1534    fn a_client_that_waited_long_enough_gets_a_null_array() {
1535        let (mut r, conn, mut batch) = timed();
1536        r.engine_mut().feed(conn, &wire(&[b"BLPOP", b"q", b"30"]));
1537        pump(&mut r, &mut batch);
1538        assert!(r.engine().sink().sent(conn).is_empty());
1539
1540        r.engine_mut().server_mut().set_clock_ms(START_MS + 29_999);
1541        pump(&mut r, &mut batch);
1542        assert!(
1543            r.engine().sink().sent(conn).is_empty(),
1544            "a millisecond short"
1545        );
1546
1547        r.engine_mut().server_mut().set_clock_ms(START_MS + 30_000);
1548        pump(&mut r, &mut batch);
1549        // A null array and not a null string, which a RESP2 client can see.
1550        assert_eq!(r.engine().sink().sent(conn), b"*-1\r\n");
1551        assert_eq!(r.engine().server().waiters().len(), 0);
1552    }
1553
1554    /// The four that answer with something other than a two element array all
1555    /// answer a timeout the same way, which is not what the reply shape would
1556    /// suggest and is what Redis does.
1557    #[test]
1558    fn every_blocking_command_times_out_with_the_same_null_array() {
1559        for cmd in [
1560            &[b"BLPOP".as_slice(), b"q", b"0.001"][..],
1561            &[b"BRPOP", b"q", b"0.001"],
1562            &[b"BLMOVE", b"q", b"d", b"LEFT", b"RIGHT", b"0.001"],
1563            &[b"BRPOPLPUSH", b"q", b"d", b"0.001"],
1564            &[b"BLMPOP", b"0.001", b"1", b"q", b"LEFT"],
1565        ] {
1566            let (mut r, conn, mut batch) = timed();
1567            r.engine_mut().feed(conn, &wire(cmd));
1568            pump(&mut r, &mut batch);
1569            r.engine_mut().server_mut().set_clock_ms(START_MS + 1);
1570            pump(&mut r, &mut batch);
1571            assert_eq!(r.engine().sink().sent(conn), b"*-1\r\n", "for {cmd:?}");
1572        }
1573    }
1574
1575    /// A client that gave up does not go on holding a claim on the queue: the
1576    /// element that arrives after it stays where it was put.
1577    #[test]
1578    fn a_waiter_that_timed_out_does_not_eat_a_later_push() {
1579        let (mut r, a, mut batch) = timed();
1580        let b = r.engine_mut().accept();
1581        r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"1"]));
1582        pump(&mut r, &mut batch);
1583        r.engine_mut().server_mut().set_clock_ms(START_MS + 1000);
1584        pump(&mut r, &mut batch);
1585        assert_eq!(r.engine().sink().sent(a), b"*-1\r\n");
1586
1587        r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"late"]));
1588        r.engine_mut()
1589            .feed(b, &wire(&[b"LRANGE", b"q", b"0", b"-1"]));
1590        pump(&mut r, &mut batch);
1591        assert_eq!(r.engine().sink().sent(a), b"*-1\r\n", "nothing more");
1592        assert!(r.engine().sink().sent(b).ends_with(b"*1\r\n$4\r\nlate\r\n"));
1593    }
1594
1595    /// A `BLPOP key 0` has no deadline, so nothing but the connection closing
1596    /// will ever take it off the list. That makes the close path the one that
1597    /// has to be right, or a waiter outlives its client and the slot it names
1598    /// gets handed to somebody else.
1599    #[test]
1600    fn a_client_that_goes_away_while_it_waits_takes_its_waiter_with_it() {
1601        let (mut r, a, mut batch) = engine();
1602        let b = r.engine_mut().accept();
1603        r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1604        pump(&mut r, &mut batch);
1605        assert_eq!(r.engine().server().waiters().len(), 1);
1606
1607        r.engine_mut().hangup(a);
1608        pump(&mut r, &mut batch);
1609        assert_eq!(r.engine().server().waiters().len(), 0);
1610        assert_eq!(r.engine().clients(), 1);
1611
1612        // The slot is handed straight back out, which is what the waiter would
1613        // have been pointing at.
1614        let again = r.engine_mut().accept();
1615        assert_eq!(again, a);
1616        r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"x"]));
1617        r.engine_mut()
1618            .feed(again, &wire(&[b"LRANGE", b"q", b"0", b"-1"]));
1619        pump(&mut r, &mut batch);
1620        assert_eq!(r.engine().sink().sent(again), b"*1\r\n$1\r\nx\r\n");
1621    }
1622
1623    /// The same, with commands the client had already sent sitting behind the
1624    /// block. Those are what `pending` counts, so a close that forgets them is a
1625    /// connection slot that never comes back.
1626    #[test]
1627    fn a_hangup_while_parked_gives_back_the_slot_and_the_decoders() {
1628        let (mut r, a, mut batch) = engine();
1629        let mut stream = wire(&[b"BLPOP", b"q", b"0"]);
1630        stream.extend(wire(&[b"PING"]));
1631        stream.extend(wire(&[b"PING"]));
1632        r.engine_mut().feed(a, &stream);
1633        pump(&mut r, &mut batch);
1634
1635        let decoders = r.engine().decoders();
1636        r.engine_mut().hangup(a);
1637        pump(&mut r, &mut batch);
1638
1639        assert_eq!(r.engine().clients(), 0);
1640        assert!(r.engine().sink().was_closed(a));
1641        assert_eq!(r.engine().decoders(), decoders, "the pool came back whole");
1642        let again = r.engine_mut().accept();
1643        assert_eq!(again, a);
1644        r.engine_mut().feed(again, &wire(&[b"PING"]));
1645        pump(&mut r, &mut batch);
1646        assert_eq!(r.engine().sink().sent(again), b"+PONG\r\n");
1647    }
1648
1649    #[test]
1650    fn a_reply_the_socket_would_not_take_is_offered_again() {
1651        let mut r = Reactor::inline(Wire::new(Trickle::default()));
1652        let conn = r.engine_mut().accept();
1653        let mut batch = Vec::new();
1654
1655        r.engine_mut().feed(conn, &wire(&[b"PING"]));
1656        pump(&mut r, &mut batch);
1657        // Two flushes in a pump, so four bytes and then three.
1658        assert_eq!(r.engine().sink().sent, b"+PONG\r\n");
1659        assert_eq!(r.engine().sink().writes, 2);
1660    }
1661}