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 self.server.db_ref(db).prefetch(hash);
885 }
886
887 fn run(&mut self, cmd: Cmd, _hash: Option<u64>) -> yo_reactor::Flow {
888 // Framed with the batch that blocked, so it is a command the client sent
889 // before it knew it would be waiting. It keeps its decoder and it keeps
890 // its place in `pending`, which is what stops the buffer it points into
891 // being compacted while it waits.
892 if self.conns[cmd.conn as usize].blocked {
893 yo_alloc::allow(|| self.conns[cmd.conn as usize].parked.push(cmd));
894 return yo_reactor::Flow::Next;
895 }
896
897 let flow = {
898 let c = &mut self.conns[cmd.conn as usize];
899 c.pending -= 1;
900 if c.gone || c.skip {
901 // Nobody to answer, or nobody who should be. The decoder still
902 // has to come back and the slot still has to be released, which
903 // is why this is not an early return.
904 Flow::Continue
905 } else {
906 let args = Args::new(&self.argvs[cmd.slot as usize], &c.buf[cmd.base..]);
907 let spec = table::at(cmd.spec);
908 dispatch::resolved(&mut self.server, &mut c.session, spec, args, &mut c.out)
909 }
910 };
911
912 self.spare.push(cmd.slot);
913 let c = &self.conns[cmd.conn as usize];
914 if c.gone {
915 if c.pending == 0 {
916 self.release(cmd.conn);
917 }
918 } else {
919 match flow {
920 Flow::Close => {
921 let c = &mut self.conns[cmd.conn as usize];
922 c.closing = true;
923 // Anything the client pipelined behind the `QUIT` was sent
924 // before it knew the answer, and running it would be acting
925 // on a connection that has already been said goodbye to.
926 c.skip = true;
927 self.soil(cmd.conn);
928 }
929 // Nothing was written, so there is nothing to flush and no
930 // reason to put this connection on the dirty list. The waiter
931 // carries the slot from here on, and it needs to know which one:
932 // the command layer only ever saw the client id.
933 Flow::Block => {
934 self.conns[cmd.conn as usize].blocked = true;
935 let client = self.conns[cmd.conn as usize].session.id();
936 self.server.waiters_mut().bind(client, cmd.conn);
937 }
938 Flow::Continue => self.soil(cmd.conn),
939 }
940 }
941
942 // After each command and not once per batch. A client blocked on two
943 // keys and woken by `RPUSH b` then `RPUSH a` in one pipeline has to
944 // answer with `b`, because that is the push that was in front of it, and
945 // it can only do that if it was served in between the two.
946 if !self.server.waiters().is_empty() {
947 self.serve_waiters();
948 }
949 yo_reactor::Flow::Next
950 }
951
952 fn flush(&mut self) {
953 // The deadline sweep, and it is here because this is the one thing the
954 // driver calls on a turn that ran nothing at all. A client whose timeout
955 // passes while the server is idle is answered within the loop's idle
956 // wait, which is 20ms and is finer than the 10hz Redis checks its own
957 // blocked clients at.
958 if !self.server.waiters().is_empty() {
959 self.server.refresh_clock();
960 self.serve_waiters();
961 }
962
963 // Taken and put back so the loop below can reach the rest of the
964 // engine. The capacity comes back with it, so this is not an
965 // allocation.
966 let mut dirty = core::mem::take(&mut self.dirty);
967 let mut at = 0;
968 while at < dirty.len() {
969 let conn = dirty[at];
970 let owed = self.write_out(conn);
971 if owed {
972 // The socket was full. The connection stays on the list with
973 // what is left of its reply, and the next flush offers it
974 // again, which is the whole of the backpressure story here.
975 at += 1;
976 } else {
977 dirty.swap_remove(at);
978 }
979 }
980 self.dirty = dirty;
981 }
982
983 fn maintain(&mut self, budget: &mut yo_reactor::Budget) {
984 // The clock is the first thing the maintenance slice does, because
985 // everything else in it compares against a time.
986 if !budget.spend(1) {
987 return;
988 }
989 self.tick();
990 // Then the dead keys, which is what stops a cache that writes with a
991 // deadline and never reads back from holding every key it has ever
992 // written. One unit a key looked at, so the slice bounds the sweep the
993 // same way it bounds everything else in here, and a server where nothing
994 // has a deadline spends nothing at all.
995 let looks = budget.left() as usize;
996 let spent = self.server.expire_slice(looks);
997 budget.spend(u32::try_from(spent).unwrap_or(u32::MAX));
998 }
999}
1000
1001/// Run everything that is framed, in batches, and write the replies.
1002///
1003/// The inline driver: it is what a caller who is already on the shard thread
1004/// uses in place of the loop, and it goes through the same two walks the loop
1005/// goes through (`15` section 7). `batch` is the caller's, so a driver in a hot
1006/// loop hands the same `Vec` back every time and never allocates.
1007pub fn pump<S: Sink>(reactor: &mut Reactor<Wire<S>>, batch: &mut Vec<Cmd>) -> usize {
1008 let mut ran = 0;
1009 reactor.engine_mut().tick();
1010 loop {
1011 batch.clear();
1012 if reactor.engine_mut().take_ready(batch, BATCH_MAX) == 0 {
1013 break;
1014 }
1015 // The command path, and therefore the thing Y7 is about. The guard is
1016 // what arms `yo-alloc`, and it covers dispatch and nothing else: framing
1017 // before it and writing the replies after it are both allowed to reach
1018 // for the heap, and only running the commands is not.
1019 //
1020 // It goes here rather than around the whole loop because `take_ready`
1021 // and `flush` are on the other side of that line, and because a batch is
1022 // the unit a caller can reason about. Under the default mode this is one
1023 // relaxed load.
1024 let armed = yo_alloc::guard();
1025 ran += reactor.execute_all(batch.drain(..));
1026 drop(armed);
1027 reactor.engine_mut().flush();
1028 // After the replies are out, so the batch that made the garbage is not
1029 // the batch that waits for it to be collected.
1030 reactor.engine_mut().maintain();
1031 }
1032 // Once more, for a connection with something to say and nothing to run: a
1033 // protocol error, or a socket that was full the last time round.
1034 reactor.engine_mut().flush();
1035 // And once for a turn that ran nothing at all, which is where a server that
1036 // has gone quiet catches up on what the last busy turn left behind.
1037 reactor.engine_mut().maintain();
1038 ran
1039}
1040
1041#[cfg(test)]
1042mod tests {
1043 use super::*;
1044
1045 /// The wire bytes for a command, built the way a client would.
1046 fn wire(args: &[&[u8]]) -> Vec<u8> {
1047 let mut b = format!("*{}\r\n", args.len()).into_bytes();
1048 for a in args {
1049 b.extend_from_slice(format!("${}\r\n", a.len()).as_bytes());
1050 b.extend_from_slice(a);
1051 b.extend_from_slice(b"\r\n");
1052 }
1053 b
1054 }
1055
1056 fn engine() -> (Reactor<Wire<Recorder>>, ConnId, Vec<Cmd>) {
1057 let mut r = Reactor::inline(Wire::new(Recorder::new()));
1058 let conn = r.engine_mut().accept();
1059 (r, conn, Vec::new())
1060 }
1061
1062 /// Where the fixed clock a blocking test moves by hand starts.
1063 const START_MS: u64 = 1_000_000;
1064
1065 /// The same, on a clock the test moves rather than the system's.
1066 ///
1067 /// A test about a timeout cannot wait for one: waiting a hundred
1068 /// milliseconds is a test that fails on a loaded machine and waiting a
1069 /// hundred seconds is not a test.
1070 fn timed() -> (Reactor<Wire<Recorder>>, ConnId, Vec<Cmd>) {
1071 let server = crate::dispatch::Server::with_clock(yo_kv::Clock::fixed(START_MS));
1072 let mut r = Reactor::inline(Wire::with_server(server, Recorder::new()));
1073 let conn = r.engine_mut().accept();
1074 (r, conn, Vec::new())
1075 }
1076
1077 #[test]
1078 fn a_pipelined_batch_comes_back_in_order_and_in_one_write() {
1079 let (mut r, conn, mut batch) = engine();
1080 let mut stream = wire(&[b"SET", b"k", b"v"]);
1081 stream.extend(wire(&[b"GET", b"k"]));
1082 stream.extend(wire(&[b"INCR", b"n"]));
1083
1084 r.engine_mut().feed(conn, &stream);
1085 assert_eq!(r.engine().ready(), 3);
1086 assert_eq!(pump(&mut r, &mut batch), 3);
1087
1088 assert_eq!(r.engine().sink().sent(conn), b"+OK\r\n$1\r\nv\r\n:1\r\n");
1089 assert_eq!(r.engine().ready(), 0);
1090 }
1091
1092 /// The framing has to survive a command arriving in pieces, because that is
1093 /// what a socket does.
1094 #[test]
1095 fn a_command_split_across_reads_resumes_rather_than_restarts() {
1096 let (mut r, conn, mut batch) = engine();
1097 let bytes = wire(&[b"SET", b"key", b"value"]);
1098
1099 for at in 1..bytes.len() {
1100 r.engine_mut().feed(conn, &bytes[at - 1..at]);
1101 assert_eq!(r.engine().ready(), 0, "not a command yet at {at}");
1102 }
1103 r.engine_mut().feed(conn, &bytes[bytes.len() - 1..]);
1104 assert_eq!(r.engine().ready(), 1);
1105 assert_eq!(pump(&mut r, &mut batch), 1);
1106 assert_eq!(r.engine().sink().sent(conn), b"+OK\r\n");
1107
1108 // And the value that arrived in single bytes is the value that was
1109 // stored, which is the part a naive resume gets wrong.
1110 r.engine_mut().feed(conn, &wire(&[b"GET", b"key"]));
1111 pump(&mut r, &mut batch);
1112 assert_eq!(r.engine().sink().sent(conn), b"+OK\r\n$5\r\nvalue\r\n");
1113 }
1114
1115 #[test]
1116 fn two_connections_are_two_sessions_over_one_server() {
1117 let (mut r, a, mut batch) = engine();
1118 let b = r.engine_mut().accept();
1119
1120 r.engine_mut().feed(a, &wire(&[b"SELECT", b"3"]));
1121 r.engine_mut().feed(a, &wire(&[b"SET", b"k", b"a"]));
1122 r.engine_mut().feed(b, &wire(&[b"SET", b"k", b"b"]));
1123 r.engine_mut().feed(a, &wire(&[b"GET", b"k"]));
1124 r.engine_mut().feed(b, &wire(&[b"GET", b"k"]));
1125 pump(&mut r, &mut batch);
1126
1127 assert_eq!(r.engine().sink().sent(a), b"+OK\r\n+OK\r\n$1\r\na\r\n");
1128 assert_eq!(r.engine().sink().sent(b), b"+OK\r\n$1\r\nb\r\n");
1129 assert_eq!(r.engine().clients(), 2);
1130 }
1131
1132 #[test]
1133 fn quit_is_answered_and_then_the_connection_goes() {
1134 let (mut r, conn, mut batch) = engine();
1135 r.engine_mut().feed(conn, &wire(&[b"PING"]));
1136 r.engine_mut().feed(conn, &wire(&[b"QUIT"]));
1137 pump(&mut r, &mut batch);
1138
1139 assert_eq!(r.engine().sink().sent(conn), b"+PONG\r\n+OK\r\n");
1140 assert!(r.engine().sink().was_closed(conn));
1141 assert_eq!(r.engine().clients(), 0);
1142
1143 // The slot comes back, buffers and all.
1144 let again = r.engine_mut().accept();
1145 assert_eq!(again, conn);
1146 assert_eq!(r.engine().clients(), 1);
1147 }
1148
1149 /// Redis's own unit/quit, which caught this: we answered the `QUIT` and
1150 /// then ran the `SET` behind it.
1151 #[test]
1152 fn what_a_client_pipelined_behind_quit_is_never_run() {
1153 let (mut r, conn, mut batch) = engine();
1154 let mut stream = wire(&[b"QUIT"]);
1155 stream.extend(wire(&[b"SET", b"foo", b"bar"]));
1156 r.engine_mut().feed(conn, &stream);
1157 // Both were framed, because framing happens before anything runs.
1158 assert_eq!(r.engine().ready(), 2);
1159 pump(&mut r, &mut batch);
1160
1161 // One reply and not two, and the connection is gone.
1162 assert_eq!(r.engine().sink().sent(conn), b"+OK\r\n");
1163 assert!(r.engine().sink().was_closed(conn));
1164
1165 // And the write never happened, which is the part a client can see
1166 // after it reconnects. The recorder is cleared first because the next
1167 // connection lands back in the slot this one just left, and what was
1168 // written to the slot before is still sitting in it.
1169 r.engine_mut().sink_mut().clear();
1170 let next = r.engine_mut().accept();
1171 r.engine_mut().feed(next, &wire(&[b"GET", b"foo"]));
1172 pump(&mut r, &mut batch);
1173 assert_eq!(r.engine().sink().sent(next), b"$-1\r\n");
1174 }
1175
1176 /// A connection that never said `HELLO` is answered in RESP2, whatever the
1177 /// last client in that slot was speaking.
1178 ///
1179 /// The protocol is kept in the reply buffer and the reply buffer outlives
1180 /// the connection, so this is the one piece of connection state that a
1181 /// recycled slot used to carry over. A client got a RESP3 null back from
1182 /// the first `GET` that missed and could not parse it, which is as bad as a
1183 /// compatibility bug gets: nothing the client did caused it and nothing it
1184 /// could send would have avoided it.
1185 #[test]
1186 fn a_slot_that_last_spoke_resp3_answers_the_next_client_in_resp2() {
1187 let (mut r, conn, mut batch) = engine();
1188 r.engine_mut().feed(conn, &wire(&[b"HELLO", b"3"]));
1189 r.engine_mut().feed(conn, &wire(&[b"GET", b"nothing"]));
1190 pump(&mut r, &mut batch);
1191 assert!(r.engine().sink().sent(conn).ends_with(b"_\r\n"));
1192 r.engine_mut().feed(conn, &wire(&[b"QUIT"]));
1193 pump(&mut r, &mut batch);
1194
1195 r.engine_mut().sink_mut().clear();
1196 let next = r.engine_mut().accept();
1197 assert_eq!(next, conn, "the same slot, which is what this is about");
1198 r.engine_mut().feed(next, &wire(&[b"GET", b"nothing"]));
1199 pump(&mut r, &mut batch);
1200 assert_eq!(r.engine().sink().sent(next), b"$-1\r\n");
1201 }
1202
1203 /// The other way a connection ends, which does not throw anything away.
1204 #[test]
1205 fn commands_that_arrived_before_a_protocol_error_are_still_answered() {
1206 let (mut r, conn, mut batch) = engine();
1207 let mut stream = wire(&[b"SET", b"k", b"v"]);
1208 stream.extend(wire(&[b"GET", b"k"]));
1209 stream.extend_from_slice(b"*1\r\n+notabulk\r\n");
1210 r.engine_mut().feed(conn, &stream);
1211 pump(&mut r, &mut batch);
1212
1213 // Both good commands were complete and correct before the stream went
1214 // wrong, so both are answered and the error comes after them.
1215 let sent = r.engine().sink().sent(conn);
1216 assert!(
1217 sent.starts_with(b"+OK\r\n$1\r\nv\r\n-ERR Protocol error: "),
1218 "{sent:?}"
1219 );
1220 assert!(r.engine().sink().was_closed(conn));
1221 }
1222
1223 #[test]
1224 fn a_protocol_error_is_written_and_closes_the_connection() {
1225 let (mut r, conn, mut batch) = engine();
1226 // A multibulk that says its first argument is a bulk and then does not.
1227 r.engine_mut().feed(conn, b"*1\r\n+notabulk\r\n");
1228 pump(&mut r, &mut batch);
1229
1230 let sent = r.engine().sink().sent(conn);
1231 assert!(sent.starts_with(b"-ERR Protocol error: "), "{sent:?}");
1232 assert!(r.engine().sink().was_closed(conn));
1233 assert_eq!(r.engine().clients(), 0);
1234 }
1235
1236 /// Redis's own `unit/protocol` walks a list of malformed frames, each on a
1237 /// fresh connection, which means every one of them after the first runs on
1238 /// a decoder that came back to the pool part way through a command.
1239 #[test]
1240 fn a_decoder_that_came_back_mid_command_starts_the_next_one_clean() {
1241 let (mut r, conn, mut batch) = engine();
1242 // Stops inside the third argument, on a length that is not a length.
1243 r.engine_mut()
1244 .feed(conn, b"*3\r\n$3\r\nSET\r\n$1\r\nx\r\n$blabla\r\n");
1245 pump(&mut r, &mut batch);
1246 let sent = r.engine().sink().sent(conn);
1247 assert!(
1248 sent.starts_with(b"-ERR Protocol error: invalid bulk length"),
1249 "{sent:?}"
1250 );
1251
1252 // The slot that decoder was in is now the slot the next connection
1253 // gets, and it has to be at the start of a command and not half way
1254 // through the one that went wrong.
1255 r.engine_mut().sink_mut().clear();
1256 let next = r.engine_mut().accept();
1257 r.engine_mut().feed(next, &wire(&[b"GET", b"k"]));
1258 pump(&mut r, &mut batch);
1259 assert_eq!(r.engine().sink().sent(next), b"$-1\r\n");
1260
1261 r.engine_mut().sink_mut().clear();
1262 let third = r.engine_mut().accept();
1263 r.engine_mut().feed(third, b"*1\r\n+notabulk\r\n");
1264 pump(&mut r, &mut batch);
1265 let sent = r.engine().sink().sent(third);
1266 assert!(sent.starts_with(b"-ERR Protocol error: "), "{sent:?}");
1267 }
1268
1269 /// A client that hangs up mid batch is the case that gets a server killed:
1270 /// the commands already framed still point into its buffer.
1271 #[test]
1272 fn a_hangup_with_commands_in_flight_waits_for_them() {
1273 let (mut r, conn, mut batch) = engine();
1274 r.engine_mut().feed(conn, &wire(&[b"SET", b"k", b"v"]));
1275 r.engine_mut().feed(conn, &wire(&[b"GET", b"k"]));
1276
1277 batch.clear();
1278 r.engine_mut().take_ready(&mut batch, BATCH_MAX);
1279 r.engine_mut().hangup(conn);
1280 assert_eq!(r.engine().clients(), 1, "still holding the buffer");
1281
1282 r.execute_all(batch.drain(..));
1283 r.engine_mut().flush();
1284 assert_eq!(r.engine().clients(), 0);
1285 assert!(r.engine().sink().sent(conn).is_empty(), "nobody to answer");
1286
1287 // And the slot is usable again, with the decoders both back in the
1288 // pool rather than lost with the connection.
1289 let decoders = r.engine().decoders();
1290 let again = r.engine_mut().accept();
1291 assert_eq!(again, conn);
1292 r.engine_mut().feed(again, &wire(&[b"PING"]));
1293 pump(&mut r, &mut batch);
1294 assert_eq!(r.engine().sink().sent(again), b"+PONG\r\n");
1295 assert_eq!(r.engine().decoders(), decoders);
1296 }
1297
1298 /// The claim that the steady state does not allocate, checked the only way
1299 /// a library test can check it: nothing grows.
1300 #[test]
1301 fn the_buffers_and_the_decoder_pool_stop_growing() {
1302 let (mut r, conn, mut batch) = engine();
1303 let mut stream = Vec::new();
1304 for i in 0..32 {
1305 stream.extend(wire(&[b"SET", format!("k{i}").as_bytes(), b"v"]));
1306 }
1307
1308 r.engine_mut().feed(conn, &stream);
1309 pump(&mut r, &mut batch);
1310 let decoders = r.engine().decoders();
1311 let batch_cap = batch.capacity();
1312
1313 for _ in 0..10 {
1314 r.engine_mut().feed(conn, &stream);
1315 pump(&mut r, &mut batch);
1316 }
1317 assert_eq!(r.engine().decoders(), decoders, "the pool is reused");
1318 assert_eq!(batch.capacity(), batch_cap, "the batch buffer is reused");
1319 assert!(
1320 decoders <= BATCH_MAX + 1,
1321 "{decoders} decoders for 32 commands"
1322 );
1323 }
1324
1325 /// The read buffer holds what has not been dealt with yet and nothing else.
1326 ///
1327 /// A client that pipelines sixteen commands, waits for the sixteen replies
1328 /// and goes again is what `redis-benchmark -P 16` does and what half of the
1329 /// clients in the world do. Every one of those rounds leaves the buffer
1330 /// exactly caught up, and a buffer that never drops what it has already
1331 /// dealt with grows to everything the connection has ever sent: 16 MiB
1332 /// apiece on server3 for four connections sending 100000 sets each.
1333 #[test]
1334 fn a_pipelining_client_does_not_grow_the_read_buffer() {
1335 let (mut r, conn, mut batch) = engine();
1336 let mut round = Vec::new();
1337 for i in 0..16 {
1338 round.extend(wire(&[b"SET", format!("k{i}").as_bytes(), b"v"]));
1339 }
1340
1341 r.engine_mut().feed(conn, &round);
1342 pump(&mut r, &mut batch);
1343 r.engine_mut().sink_mut().clear();
1344 let after_one = r.engine().buffer_bytes();
1345
1346 // A thousand rounds is sixteen thousand commands and about a megabyte
1347 // of wire bytes, which is a hundred times what the buffer starts with.
1348 for _ in 0..1000 {
1349 r.engine_mut().feed(conn, &round);
1350 pump(&mut r, &mut batch);
1351 r.engine_mut().sink_mut().clear();
1352 }
1353
1354 assert_eq!(
1355 r.engine().buffer_bytes(),
1356 after_one,
1357 "the buffers grew over a thousand rounds of the same sixteen commands"
1358 );
1359 assert!(
1360 r.engine().server().memory_bytes() >= after_one,
1361 "the buffers are counted in what the server reports"
1362 );
1363 }
1364
1365 /// Half a command in the buffer is the case compaction has to be careful
1366 /// about, because the decoder holding it kept offsets into those bytes.
1367 #[test]
1368 fn a_command_split_across_reads_survives_compaction() {
1369 let (mut r, conn, mut batch) = engine();
1370 let cmd = wire(&[b"SET", b"key", b"value"]);
1371 let (head, tail) = cmd.split_at(cmd.len() - 4);
1372
1373 // A complete command, so that there is something in front to drop, then
1374 // most of a second one.
1375 r.engine_mut().feed(conn, &wire(&[b"PING"]));
1376 r.engine_mut().feed(conn, head);
1377 pump(&mut r, &mut batch);
1378 assert_eq!(r.engine().sink().sent(conn), b"+PONG\r\n");
1379
1380 // The rest of it arrives after the buffer has been compacted under it.
1381 r.engine_mut().feed(conn, tail);
1382 pump(&mut r, &mut batch);
1383 assert_eq!(r.engine().sink().sent(conn), b"+PONG\r\n+OK\r\n");
1384
1385 r.engine_mut().feed(conn, &wire(&[b"GET", b"key"]));
1386 pump(&mut r, &mut batch);
1387 assert!(r.engine().sink().sent(conn).ends_with(b"$5\r\nvalue\r\n"));
1388 }
1389
1390 /// The two walks are the reactor's, not this module's, so the test is that
1391 /// the engine can be driven by them at all: same commands, same replies.
1392 #[test]
1393 fn the_batch_goes_through_the_reactors_two_walks() {
1394 let (mut r, conn, mut batch) = engine();
1395 for i in 0..100 {
1396 r.engine_mut()
1397 .feed(conn, &wire(&[b"INCR", format!("k{}", i % 7).as_bytes()]));
1398 }
1399 let ran = pump(&mut r, &mut batch);
1400
1401 assert_eq!(ran, 100);
1402 assert_eq!(r.commands(), 100);
1403 // Two batches, because a hundred commands do not fit in sixty four.
1404 assert_eq!(r.turns(), 2);
1405 // The hundredth command is the fifteenth `INCR` of `k1`.
1406 assert!(r.engine().sink().sent(conn).ends_with(b":15\r\n"));
1407 }
1408
1409 /// A sink that takes four bytes at a time, which is what a full socket
1410 /// looks like from in here.
1411 #[derive(Default)]
1412 struct Trickle {
1413 sent: Vec<u8>,
1414 writes: usize,
1415 }
1416
1417 impl Sink for Trickle {
1418 fn write(&mut self, _conn: ConnId, bytes: &[u8]) -> usize {
1419 self.writes += 1;
1420 let n = bytes.len().min(4);
1421 self.sent.extend_from_slice(&bytes[..n]);
1422 n
1423 }
1424 }
1425
1426 /// A blocking command that does not block costs nothing: no waiter, no
1427 /// allocation, the same three lines the non blocking one runs.
1428 #[test]
1429 fn a_blpop_on_a_list_with_something_in_it_never_waits() {
1430 let (mut r, conn, mut batch) = engine();
1431 r.engine_mut().feed(conn, &wire(&[b"RPUSH", b"q", b"a"]));
1432 r.engine_mut().feed(conn, &wire(&[b"BLPOP", b"q", b"0"]));
1433 pump(&mut r, &mut batch);
1434
1435 assert_eq!(
1436 r.engine().sink().sent(conn),
1437 b":1\r\n*2\r\n$1\r\nq\r\n$1\r\na\r\n"
1438 );
1439 assert_eq!(r.engine().server().waiters().len(), 0);
1440 }
1441
1442 /// The whole point: a client with nothing to pop is answered later, by
1443 /// somebody else's command.
1444 #[test]
1445 fn a_parked_client_is_answered_by_another_connections_push() {
1446 let (mut r, a, mut batch) = engine();
1447 let b = r.engine_mut().accept();
1448
1449 r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1450 pump(&mut r, &mut batch);
1451 assert!(r.engine().sink().sent(a).is_empty(), "nothing to say yet");
1452 assert_eq!(r.engine().server().waiters().len(), 1);
1453
1454 r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"one"]));
1455 pump(&mut r, &mut batch);
1456
1457 assert_eq!(r.engine().sink().sent(a), b"*2\r\n$1\r\nq\r\n$3\r\none\r\n");
1458 // The push still reports the length it made, even though the element was
1459 // gone again before the reply was written.
1460 assert_eq!(r.engine().sink().sent(b), b":1\r\n");
1461 assert_eq!(r.engine().server().waiters().len(), 0);
1462 }
1463
1464 /// A push to a key nobody named, and a key of another type on a key
1465 /// somebody did: neither is a wake up, and the client stays parked.
1466 #[test]
1467 fn only_a_list_arriving_under_a_named_key_wakes_a_waiter() {
1468 let (mut r, a, mut batch) = engine();
1469 let b = r.engine_mut().accept();
1470 r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1471 pump(&mut r, &mut batch);
1472
1473 r.engine_mut()
1474 .feed(b, &wire(&[b"RPUSH", b"elsewhere", b"x"]));
1475 r.engine_mut().feed(b, &wire(&[b"SADD", b"q", b"x"]));
1476 pump(&mut r, &mut batch);
1477
1478 assert!(r.engine().sink().sent(a).is_empty());
1479 assert_eq!(r.engine().server().waiters().len(), 1, "still waiting");
1480 // And the set is intact, so the waiter did not take anything out of it
1481 // on its way past.
1482 assert_eq!(r.engine().sink().sent(b), b":1\r\n:1\r\n");
1483 }
1484
1485 /// Two workers on one queue, which is what `BLPOP` is for. They are served
1486 /// in the order they arrived and not in whatever order the list is walked.
1487 #[test]
1488 fn two_parked_clients_are_served_in_the_order_they_arrived() {
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().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1494 pump(&mut r, &mut batch);
1495 r.engine_mut().feed(b, &wire(&[b"BLPOP", b"q", b"0"]));
1496 pump(&mut r, &mut batch);
1497 assert_eq!(r.engine().server().waiters().len(), 2);
1498
1499 r.engine_mut()
1500 .feed(c, &wire(&[b"RPUSH", b"q", b"first", b"second"]));
1501 pump(&mut r, &mut batch);
1502
1503 assert_eq!(
1504 r.engine().sink().sent(a),
1505 b"*2\r\n$1\r\nq\r\n$5\r\nfirst\r\n"
1506 );
1507 assert_eq!(
1508 r.engine().sink().sent(b),
1509 b"*2\r\n$1\r\nq\r\n$6\r\nsecond\r\n"
1510 );
1511 assert_eq!(r.engine().server().waiters().len(), 0);
1512 }
1513
1514 /// A client waiting for an answer is not a client that has sent another
1515 /// question, so what it pipelined behind its `BLPOP` waits for the `BLPOP`.
1516 #[test]
1517 fn what_a_client_pipelined_behind_a_block_waits_for_the_block() {
1518 let (mut r, a, mut batch) = engine();
1519 let b = r.engine_mut().accept();
1520
1521 // Framed together, so the `PING` is already on its way to the reactor
1522 // when the `BLPOP` in front of it parks.
1523 let mut stream = wire(&[b"BLPOP", b"q", b"0"]);
1524 stream.extend(wire(&[b"PING"]));
1525 r.engine_mut().feed(a, &stream);
1526 pump(&mut r, &mut batch);
1527 assert!(
1528 r.engine().sink().sent(a).is_empty(),
1529 "the PING went out in front of the answer it was sent behind"
1530 );
1531
1532 // And one that arrives while it is parked is not even framed.
1533 r.engine_mut().feed(a, &wire(&[b"ECHO", b"after"]));
1534 pump(&mut r, &mut batch);
1535 assert!(r.engine().sink().sent(a).is_empty());
1536
1537 r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"x"]));
1538 pump(&mut r, &mut batch);
1539 assert_eq!(
1540 r.engine().sink().sent(a),
1541 b"*2\r\n$1\r\nq\r\n$1\r\nx\r\n+PONG\r\n$5\r\nafter\r\n"
1542 );
1543 }
1544
1545 /// Redis serves parked clients after every command rather than once per
1546 /// turn of the loop, and a pipeline is where the difference shows: the
1547 /// waiter has to be served between the two pushes, so it answers with the
1548 /// key the first push filled and not with the one it named first.
1549 #[test]
1550 fn a_waiter_is_served_between_two_pipelined_pushes() {
1551 let (mut r, a, mut batch) = engine();
1552 let b = r.engine_mut().accept();
1553 r.engine_mut()
1554 .feed(a, &wire(&[b"BLPOP", b"p1", b"p2", b"0"]));
1555 pump(&mut r, &mut batch);
1556
1557 let mut stream = wire(&[b"RPUSH", b"p2", b"second"]);
1558 stream.extend(wire(&[b"RPUSH", b"p1", b"first"]));
1559 r.engine_mut().feed(b, &stream);
1560 pump(&mut r, &mut batch);
1561
1562 assert_eq!(
1563 r.engine().sink().sent(a),
1564 b"*2\r\n$2\r\np2\r\n$6\r\nsecond\r\n"
1565 );
1566 // Which leaves the key it named first holding what was pushed to it.
1567 r.engine_mut()
1568 .feed(b, &wire(&[b"LRANGE", b"p1", b"0", b"-1"]));
1569 pump(&mut r, &mut batch);
1570 assert!(
1571 r.engine()
1572 .sink()
1573 .sent(b)
1574 .ends_with(b"*1\r\n$5\r\nfirst\r\n")
1575 );
1576 }
1577
1578 /// A `BLMOVE` that serves itself is a push, so it wakes the client waiting
1579 /// on the key it pushed to, in the same moment and without a turn of the
1580 /// loop in between.
1581 #[test]
1582 fn a_waiter_woken_by_another_waiter() {
1583 let (mut r, a, mut batch) = engine();
1584 let b = r.engine_mut().accept();
1585 let c = r.engine_mut().accept();
1586
1587 r.engine_mut()
1588 .feed(a, &wire(&[b"BLMOVE", b"x", b"y", b"LEFT", b"RIGHT", b"0"]));
1589 pump(&mut r, &mut batch);
1590 r.engine_mut().feed(b, &wire(&[b"BLPOP", b"y", b"0"]));
1591 pump(&mut r, &mut batch);
1592 assert_eq!(r.engine().server().waiters().len(), 2);
1593
1594 r.engine_mut().feed(c, &wire(&[b"RPUSH", b"x", b"chain"]));
1595 pump(&mut r, &mut batch);
1596
1597 assert_eq!(r.engine().sink().sent(a), b"$5\r\nchain\r\n");
1598 assert_eq!(
1599 r.engine().sink().sent(b),
1600 b"*2\r\n$1\r\ny\r\n$5\r\nchain\r\n"
1601 );
1602 assert_eq!(r.engine().server().waiters().len(), 0);
1603 }
1604
1605 /// A waiter on one database is not woken by a push on another, even though
1606 /// the key has the same name.
1607 #[test]
1608 fn a_waiter_is_only_woken_on_the_database_it_blocked_on() {
1609 let (mut r, a, mut batch) = engine();
1610 let b = r.engine_mut().accept();
1611 r.engine_mut().feed(a, &wire(&[b"SELECT", b"3"]));
1612 r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1613 pump(&mut r, &mut batch);
1614 assert_eq!(r.engine().sink().sent(a), b"+OK\r\n");
1615
1616 r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"wrongdb"]));
1617 pump(&mut r, &mut batch);
1618 assert_eq!(r.engine().sink().sent(a), b"+OK\r\n", "still waiting");
1619
1620 r.engine_mut().feed(b, &wire(&[b"SELECT", b"3"]));
1621 r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"rightdb"]));
1622 pump(&mut r, &mut batch);
1623 assert!(r.engine().sink().sent(a).ends_with(b"$7\r\nrightdb\r\n"));
1624 }
1625
1626 /// The deadline sweep, which runs on a turn that has nothing else to do.
1627 #[test]
1628 fn a_client_that_waited_long_enough_gets_a_null_array() {
1629 let (mut r, conn, mut batch) = timed();
1630 r.engine_mut().feed(conn, &wire(&[b"BLPOP", b"q", b"30"]));
1631 pump(&mut r, &mut batch);
1632 assert!(r.engine().sink().sent(conn).is_empty());
1633
1634 r.engine_mut().server_mut().set_clock_ms(START_MS + 29_999);
1635 pump(&mut r, &mut batch);
1636 assert!(
1637 r.engine().sink().sent(conn).is_empty(),
1638 "a millisecond short"
1639 );
1640
1641 r.engine_mut().server_mut().set_clock_ms(START_MS + 30_000);
1642 pump(&mut r, &mut batch);
1643 // A null array and not a null string, which a RESP2 client can see.
1644 assert_eq!(r.engine().sink().sent(conn), b"*-1\r\n");
1645 assert_eq!(r.engine().server().waiters().len(), 0);
1646 }
1647
1648 /// The four that answer with something other than a two element array all
1649 /// answer a timeout the same way, which is not what the reply shape would
1650 /// suggest and is what Redis does.
1651 #[test]
1652 fn every_blocking_command_times_out_with_the_same_null_array() {
1653 for cmd in [
1654 &[b"BLPOP".as_slice(), b"q", b"0.001"][..],
1655 &[b"BRPOP", b"q", b"0.001"],
1656 &[b"BLMOVE", b"q", b"d", b"LEFT", b"RIGHT", b"0.001"],
1657 &[b"BRPOPLPUSH", b"q", b"d", b"0.001"],
1658 &[b"BLMPOP", b"0.001", b"1", b"q", b"LEFT"],
1659 ] {
1660 let (mut r, conn, mut batch) = timed();
1661 r.engine_mut().feed(conn, &wire(cmd));
1662 pump(&mut r, &mut batch);
1663 r.engine_mut().server_mut().set_clock_ms(START_MS + 1);
1664 pump(&mut r, &mut batch);
1665 assert_eq!(r.engine().sink().sent(conn), b"*-1\r\n", "for {cmd:?}");
1666 }
1667 }
1668
1669 /// A client that gave up does not go on holding a claim on the queue: the
1670 /// element that arrives after it stays where it was put.
1671 #[test]
1672 fn a_waiter_that_timed_out_does_not_eat_a_later_push() {
1673 let (mut r, a, mut batch) = timed();
1674 let b = r.engine_mut().accept();
1675 r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"1"]));
1676 pump(&mut r, &mut batch);
1677 r.engine_mut().server_mut().set_clock_ms(START_MS + 1000);
1678 pump(&mut r, &mut batch);
1679 assert_eq!(r.engine().sink().sent(a), b"*-1\r\n");
1680
1681 r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"late"]));
1682 r.engine_mut()
1683 .feed(b, &wire(&[b"LRANGE", b"q", b"0", b"-1"]));
1684 pump(&mut r, &mut batch);
1685 assert_eq!(r.engine().sink().sent(a), b"*-1\r\n", "nothing more");
1686 assert!(r.engine().sink().sent(b).ends_with(b"*1\r\n$4\r\nlate\r\n"));
1687 }
1688
1689 /// A `BLPOP key 0` has no deadline, so nothing but the connection closing
1690 /// will ever take it off the list. That makes the close path the one that
1691 /// has to be right, or a waiter outlives its client and the slot it names
1692 /// gets handed to somebody else.
1693 #[test]
1694 fn a_client_that_goes_away_while_it_waits_takes_its_waiter_with_it() {
1695 let (mut r, a, mut batch) = engine();
1696 let b = r.engine_mut().accept();
1697 r.engine_mut().feed(a, &wire(&[b"BLPOP", b"q", b"0"]));
1698 pump(&mut r, &mut batch);
1699 assert_eq!(r.engine().server().waiters().len(), 1);
1700
1701 r.engine_mut().hangup(a);
1702 pump(&mut r, &mut batch);
1703 assert_eq!(r.engine().server().waiters().len(), 0);
1704 assert_eq!(r.engine().clients(), 1);
1705
1706 // The slot is handed straight back out, which is what the waiter would
1707 // have been pointing at.
1708 let again = r.engine_mut().accept();
1709 assert_eq!(again, a);
1710 r.engine_mut().feed(b, &wire(&[b"RPUSH", b"q", b"x"]));
1711 r.engine_mut()
1712 .feed(again, &wire(&[b"LRANGE", b"q", b"0", b"-1"]));
1713 pump(&mut r, &mut batch);
1714 assert_eq!(r.engine().sink().sent(again), b"*1\r\n$1\r\nx\r\n");
1715 }
1716
1717 /// The same, with commands the client had already sent sitting behind the
1718 /// block. Those are what `pending` counts, so a close that forgets them is a
1719 /// connection slot that never comes back.
1720 #[test]
1721 fn a_hangup_while_parked_gives_back_the_slot_and_the_decoders() {
1722 let (mut r, a, mut batch) = engine();
1723 let mut stream = wire(&[b"BLPOP", b"q", b"0"]);
1724 stream.extend(wire(&[b"PING"]));
1725 stream.extend(wire(&[b"PING"]));
1726 r.engine_mut().feed(a, &stream);
1727 pump(&mut r, &mut batch);
1728
1729 let decoders = r.engine().decoders();
1730 r.engine_mut().hangup(a);
1731 pump(&mut r, &mut batch);
1732
1733 assert_eq!(r.engine().clients(), 0);
1734 assert!(r.engine().sink().was_closed(a));
1735 assert_eq!(r.engine().decoders(), decoders, "the pool came back whole");
1736 let again = r.engine_mut().accept();
1737 assert_eq!(again, a);
1738 r.engine_mut().feed(again, &wire(&[b"PING"]));
1739 pump(&mut r, &mut batch);
1740 assert_eq!(r.engine().sink().sent(again), b"+PONG\r\n");
1741 }
1742
1743 #[test]
1744 fn a_reply_the_socket_would_not_take_is_offered_again() {
1745 let mut r = Reactor::inline(Wire::new(Trickle::default()));
1746 let conn = r.engine_mut().accept();
1747 let mut batch = Vec::new();
1748
1749 r.engine_mut().feed(conn, &wire(&[b"PING"]));
1750 pump(&mut r, &mut batch);
1751 // Two flushes in a pump, so four bytes and then three.
1752 assert_eq!(r.engine().sink().sent, b"+PONG\r\n");
1753 assert_eq!(r.engine().sink().writes, 2);
1754 }
1755}