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