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