Skip to main content

yo_resp/
front.rs

1//! What one thread owns: the connections, their buffers and the framing.
2//!
3//! A server on several threads is two halves that have to be told apart before
4//! either of them can move. One half is per connection and belongs to whichever
5//! thread accepted it: the read buffer, the decoder holding a half read command,
6//! the session, the reply buffer and the queue of commands framed and not yet
7//! run. The other half is the keyspace, which every thread reaches and which is
8//! behind the stripes. This module is the first half, and the line is drawn by
9//! the compiler rather than by a comment: nothing in this file can name a
10//! [`Server`], because it does not import one.
11//!
12//! [`Wire`] is where the two meet. Everything that needs both, which is running
13//! a command, answering a blocked client and forgetting a client that has gone,
14//! is a method there and calls into here for the connection half. Everything
15//! that needs only the connections is a method here, which is why framing can be
16//! tested against a [`Front`] with no database anywhere in the test.
17//!
18//! [`Server`]: crate::dispatch::Server
19//! [`Wire`]: crate::engine::Wire
20
21use std::collections::VecDeque;
22
23use yo_reactor::BATCH_MAX;
24
25use crate::dispatch::table::lookup_index;
26use crate::dispatch::{Args, Session};
27use crate::engine::{ConnId, Sink};
28use crate::error::ProtocolError;
29use crate::proto::{Limits, Proto};
30use crate::reply::Out;
31use crate::request::{Argv, Step};
32
33/// The read buffer a connection starts with.
34///
35/// Redis's query buffer starts at sixteen kilobytes for the same reason: it is
36/// larger than every command a client actually sends, so the buffer grows once
37/// at accept time and then never again.
38const READ_BUF: usize = 16 * 1024;
39
40/// The reply buffer a connection starts with.
41const OUT_BUF: usize = 16 * 1024;
42
43/// How many arguments a decoder has room for before it grows.
44const ARGV_HINT: usize = 8;
45
46/// One framed command, waiting to run.
47///
48/// Names the bytes rather than holding them, so the reactor can queue a batch
49/// of these while the front keeps ownership of every buffer they point into.
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub struct Cmd {
52    pub(crate) conn: ConnId,
53    pub(crate) slot: u32,
54    pub(crate) base: usize,
55    /// Which command this is, as a position in the command table.
56    ///
57    /// Resolved once, here, because the name is otherwise looked up twice more
58    /// on the way to running it: once to work out which key to prefetch and once
59    /// to dispatch. A position rather than a reference because this struct is
60    /// queued by the thousand and two bytes is what it costs.
61    ///
62    /// Past the end of the table for a name that is no command, which needs no
63    /// flag of its own and no `Option`, because that is what the lookup already
64    /// answers and what the dispatcher already has a reply for.
65    pub(crate) spec: u16,
66}
67
68impl Cmd {
69    /// The connection this command arrived on.
70    #[must_use]
71    pub const fn conn(&self) -> ConnId {
72        self.conn
73    }
74}
75
76/// What one connection's replies did on their way to the socket.
77pub(crate) enum Wrote {
78    /// The socket took less than was offered, so what is left is held for the
79    /// next flush and the connection stays on the dirty list.
80    Owed,
81    /// Everything went out and the connection is still open.
82    Done,
83    /// Everything went out and the connection ended with it. The client id is
84    /// the one thing the server has to hear about, because a waiter is found by
85    /// it and the slot is about to belong to somebody else.
86    Ended(u64),
87}
88
89/// One connection's state.
90struct Conn {
91    live: bool,
92    session: Session,
93    out: Out,
94    /// What has arrived and not yet been framed away.
95    buf: Vec<u8>,
96    /// How much of `buf` the framing has consumed.
97    head: usize,
98    /// The decoder holding a command that has not all arrived.
99    partial: Option<u32>,
100    /// Commands framed out of this buffer and not yet run.
101    pending: u32,
102    /// This connection is on its way out, once what is buffered has gone.
103    closing: bool,
104    /// A protocol error waiting for the commands in front of it to answer.
105    ///
106    /// The framing finds the error before any of the batch it was framed with
107    /// has run, and writing the error there would put it in front of replies
108    /// the client is still owed. Redis answers in order, so this waits until
109    /// nothing is pending and goes out last.
110    deferred: Option<ProtocolError>,
111    /// Everything still queued for this connection is thrown away unanswered.
112    ///
113    /// `QUIT` sets this and a protocol error does not, which is the difference
114    /// between the two ways a connection ends. A client that pipelines `QUIT`
115    /// and then `SET` has said goodbye and then said something after it, and
116    /// Redis answers the goodbye and drops the rest. A client that sends two
117    /// good commands and then a malformed one gets both good ones answered,
118    /// because they were complete and correct before the stream went wrong.
119    skip: bool,
120    /// The peer is gone, so there is nothing to answer and nothing to write.
121    gone: bool,
122    /// Already on the dirty list.
123    dirty: bool,
124    /// This client is parked on a blocking command.
125    ///
126    /// While it is set, framing stops: whatever the client pipelined behind its
127    /// `BLPOP` stays in the read buffer unread, which is what a client waiting
128    /// for an answer means and is what Redis does with the same bytes.
129    blocked: bool,
130    /// Commands framed before it blocked and not run yet.
131    ///
132    /// A batch is framed before any of it runs, so a `BLPOP` can be the first of
133    /// sixty four commands and the other sixty three are already on their way to
134    /// the reactor when it parks. They come back here and go to the front of the
135    /// queue when the client wakes up, in the order they arrived.
136    ///
137    /// They are still counted in `pending`, which is what stops the read buffer
138    /// being compacted under the offsets they hold.
139    parked: Vec<Cmd>,
140    /// What the two buffers were holding the last time anybody counted.
141    ///
142    /// The connection's share of `INFO memory`, kept here so that reporting it
143    /// is a subtraction against this rather than a walk over every connection.
144    held: usize,
145}
146
147impl Conn {
148    fn new(id: u64) -> Conn {
149        // Accept time, which is the one moment a connection is allowed to cost
150        // an allocation. Everything after this reuses these two buffers.
151        yo_alloc::allow(|| Conn {
152            live: true,
153            session: Session::new(id),
154            out: Out::with_capacity(Proto::Resp2, OUT_BUF),
155            buf: Vec::with_capacity(READ_BUF),
156            head: 0,
157            partial: None,
158            pending: 0,
159            closing: false,
160            deferred: None,
161            skip: false,
162            gone: false,
163            dirty: false,
164            blocked: false,
165            parked: Vec::new(),
166            held: 0,
167        })
168    }
169
170    /// What the two buffers cost the process, which is the room they are
171    /// holding and not the bytes in use: both keep their capacity between
172    /// batches on purpose.
173    fn size(&self) -> usize {
174        self.buf.capacity() + self.out.capacity()
175    }
176
177    /// Back to how it was at accept time, buffers kept.
178    fn reset(&mut self, id: u64) {
179        self.live = true;
180        self.session = Session::new(id);
181        self.out.clear();
182        // The protocol lives in the reply buffer and the reply buffer is kept,
183        // so it has to be put back by hand. Without this a client that opened a
184        // connection into a slot the last client had spoken RESP3 on would be
185        // answered in RESP3 without ever sending `HELLO`, which is a nil it
186        // cannot parse on the first `GET` that misses.
187        self.out.set_proto(Proto::Resp2);
188        self.buf.clear();
189        self.head = 0;
190        self.partial = None;
191        self.pending = 0;
192        self.closing = false;
193        self.deferred = None;
194        self.skip = false;
195        self.gone = false;
196        self.dirty = false;
197        self.blocked = false;
198        // The room it took stays, the way the two buffers' does.
199        self.parked.clear();
200    }
201
202    /// Drop what the framing has already read, when nothing points into it.
203    ///
204    /// A framed command's arguments are offsets from the front of this buffer,
205    /// so this waits for the batch to run. After a batch is where a pipelining
206    /// connection spends most of its life, so that is not much of a wait.
207    ///
208    /// A half read command is not in the way. Its decoder was handed
209    /// `buf[head..]` and every offset it kept is from the front of that slice,
210    /// and `head` does not move until the command is complete, so the bytes it
211    /// is waiting on are exactly the bytes this keeps. They arrive at the front
212    /// instead of at `head` and the decoder cannot tell the difference.
213    ///
214    /// Waiting for it anyway is what made a read buffer grow to everything the
215    /// connection had ever sent. The framing loop only ever stops on an
216    /// incomplete command, and a buffer that ends on a command boundary gives
217    /// one of those on the next turn round: an empty slice, nothing decoded,
218    /// `Step::Incomplete`. So a connection that is exactly up to date always had
219    /// a decoder parked on it, this always returned early, and `head` walked
220    /// forward with the bytes behind it kept forever. Measured on server3, four
221    /// connections sending 100000 sets each held 16 MiB of read buffer apiece,
222    /// and fifty connections sending 8000 each held 1 MiB apiece: in both cases
223    /// every byte the connection had ever sent.
224    fn compact(&mut self) {
225        if self.pending > 0 || self.head == 0 {
226            return;
227        }
228        if self.head == self.buf.len() {
229            self.buf.clear();
230        } else {
231            self.buf.drain(..self.head);
232        }
233        self.head = 0;
234    }
235}
236
237/// The connection side of the server, and all of it belongs to one thread.
238///
239/// Connections, their buffers, the decoder pool, the framing and the queue of
240/// work it produces. There is one of these per I/O thread and they share
241/// nothing, which is why none of it is behind a lock and none of it is atomic.
242pub(crate) struct Front<S> {
243    sink: S,
244    conns: Vec<Conn>,
245    /// Connection slots that closed and can be handed out again.
246    free: Vec<ConnId>,
247    /// The decoder pool.
248    argvs: Vec<Argv>,
249    spare: Vec<u32>,
250    /// Framed and not yet handed to the reactor.
251    ready: VecDeque<Cmd>,
252    /// Connections this batch wrote to.
253    dirty: Vec<ConnId>,
254    /// Where a protocol error line is built before it is copied into a reply.
255    scratch: Vec<u8>,
256    limits: Limits,
257    /// How much the buffers have grown or shrunk since anybody last asked.
258    ///
259    /// `INFO memory` reports what every connection is holding and that total
260    /// lives on the server, which this side cannot reach. So the change is kept
261    /// here and taken by [`Wire`] at the end of whatever call made it, which is
262    /// as timely as reporting it on the spot and does not put the server on the
263    /// other end of a framing call.
264    ///
265    /// [`Wire`]: crate::engine::Wire
266    moved: isize,
267}
268
269impl<S: Sink> Front<S> {
270    /// A front with no connections and nothing pooled.
271    pub(crate) fn new(sink: S) -> Front<S> {
272        Front {
273            sink,
274            conns: Vec::new(),
275            free: Vec::new(),
276            argvs: Vec::new(),
277            spare: Vec::new(),
278            ready: VecDeque::with_capacity(BATCH_MAX),
279            dirty: Vec::with_capacity(16),
280            scratch: Vec::with_capacity(128),
281            limits: Limits::default(),
282            moved: 0,
283        }
284    }
285
286    /// Where the replies went.
287    pub(crate) const fn sink(&self) -> &S {
288        &self.sink
289    }
290
291    /// The same, mutably.
292    pub(crate) const fn sink_mut(&mut self) -> &mut S {
293        &mut self.sink
294    }
295
296    /// Change the protocol limits, which is `proto-max-bulk-len` and friends.
297    pub(crate) fn set_limits(&mut self, limits: Limits) {
298        self.limits = limits;
299    }
300
301    /// Open a connection under the given client id and give back its slot.
302    ///
303    /// The id comes from the caller because CLIENT LIST and CLIENT KILL name a
304    /// client by it across the whole server, so two fronts handing out the same
305    /// number would be two clients answering to one name. A front has no way to
306    /// reach the other fronts, so the one thing they share mints it.
307    ///
308    /// Reuses a closed connection's slot and its two buffers when there is one,
309    /// so a server with a churning client population allocates for the high
310    /// water mark and not for the total.
311    pub(crate) fn open(&mut self, id: u64) -> ConnId {
312        let at = match self.free.pop() {
313            Some(at) => {
314                // A reused slot keeps its buffers, so what it holds is already
315                // counted and this only puts the id back in service.
316                self.conns[at as usize].reset(id);
317                at
318            }
319            None => {
320                let conn = Conn::new(id);
321                yo_alloc::allow(|| self.conns.push(conn));
322                (self.conns.len() - 1) as ConnId
323            }
324        };
325        // The session carries the slot from here, because the slot is what a
326        // subscription on the server names and the front is the only place that
327        // knows it. Both arms above make a fresh session, so this is the one
328        // place it has to be said.
329        self.conns[at as usize].session.set_conn(at);
330        self.note_size(at);
331        at
332    }
333
334    /// Take bytes off a connection and frame whatever commands they complete.
335    ///
336    /// Anything left over stays in the connection's buffer, half a command
337    /// included, so the caller hands over whatever the socket gave it without
338    /// looking at it.
339    pub(crate) fn feed(&mut self, conn: ConnId, bytes: &[u8]) {
340        {
341            let c = &mut self.conns[conn as usize];
342            if !c.live || c.closing {
343                return;
344            }
345            // The buffer is sized for a command at accept time, so this only
346            // grows for a client sending a bulk larger than that, which is a
347            // real allocation for a real reason.
348            yo_alloc::allow(|| c.buf.extend_from_slice(bytes));
349            c.session.read_bytes(bytes.len());
350        }
351        self.frame(conn);
352        self.note_size(conn);
353    }
354
355    /// Note what this connection's buffers are holding now, if it has changed
356    /// since the last time anybody asked.
357    ///
358    /// Once per read and once per flush, which is where a buffer can grow, and
359    /// two loads and a compare when nothing has moved. The alternative is a
360    /// walk over every connection on a turn of the loop, which puts the cost of
361    /// a report nobody has asked for on the command path.
362    fn note_size(&mut self, conn: ConnId) {
363        let c = &mut self.conns[conn as usize];
364        // What `CLIENT INFO` reports about the two buffers, taken here because
365        // this already runs at the two moments they can change and because a
366        // command has no way to reach a connection's buffers.
367        c.session.note_buffers(
368            c.buf.len() - c.head,
369            c.buf.capacity() - c.buf.len(),
370            c.out.len(),
371            c.out.capacity(),
372        );
373        let now = c.size();
374        if now == c.held {
375            return;
376        }
377        let delta = now as isize - c.held as isize;
378        c.held = now;
379        self.moved += delta;
380    }
381
382    /// How much the buffers have moved since this was last called.
383    pub(crate) fn buffer_delta(&mut self) -> isize {
384        core::mem::take(&mut self.moved)
385    }
386
387    /// Move as many complete commands as possible out of the read buffer.
388    ///
389    /// Nothing at all while the client is parked. The bytes stay where they are
390    /// and `head` does not move, so a client that pipelines `BLPOP` and then
391    /// `PING` gets the `PING` answered when the `BLPOP` is, and in that order.
392    fn frame(&mut self, conn: ConnId) {
393        if self.conns[conn as usize].blocked {
394            return;
395        }
396        loop {
397            let base = self.conns[conn as usize].head;
398            let slot = match self.conns[conn as usize].partial.take() {
399                Some(slot) => slot,
400                None => self.take_decoder(),
401            };
402
403            let step = {
404                let c = &self.conns[conn as usize];
405                self.argvs[slot as usize].decode(&c.buf[base..], &self.limits)
406            };
407
408            match step {
409                Ok(Step::Command { consumed }) => {
410                    self.conns[conn as usize].head += consumed;
411                    if self.argvs[slot as usize].is_empty() {
412                        // `*0` and a blank inline line: consumed, not answered.
413                        self.spare.push(slot);
414                    } else {
415                        if self.ready.len() == self.ready.capacity() {
416                            yo_alloc::allow(|| self.ready.reserve(BATCH_MAX));
417                        }
418                        // Here and not later, because the name is in front of
419                        // the argument list that was just decoded and this is
420                        // the last place that holds both it and nothing else to
421                        // do. Everything downstream takes the number.
422                        let spec = {
423                            let c = &self.conns[conn as usize];
424                            let args = Args::new(&self.argvs[slot as usize], &c.buf[base..]);
425                            lookup_index(args.name())
426                        };
427                        self.ready.push_back(Cmd {
428                            conn,
429                            slot,
430                            base,
431                            spec,
432                        });
433                        self.conns[conn as usize].pending += 1;
434                    }
435                }
436                Ok(Step::Incomplete) => {
437                    // Hold the decoder so the rest of this command resumes
438                    // where it stopped instead of being read again from the
439                    // front every time more of it arrives.
440                    self.conns[conn as usize].partial = Some(slot);
441                    break;
442                }
443                Err(e) => {
444                    self.spare.push(slot);
445                    let c = &mut self.conns[conn as usize];
446                    // Held rather than written, so it lands behind the replies
447                    // to the commands that were framed in front of it out of
448                    // the same read.
449                    c.deferred = Some(e);
450                    // Redis closes after a protocol error and so do we: the two
451                    // ends no longer agree on where the next command starts.
452                    c.closing = true;
453                    self.soil(conn);
454                    break;
455                }
456            }
457        }
458        self.conns[conn as usize].compact();
459    }
460
461    /// A decoder from the pool, or a new one the first time round.
462    ///
463    /// The one from the pool is reset before it goes out, because a decoder can
464    /// come back to the pool part way through a command: a protocol error stops
465    /// framing where it is, and a connection that hangs up with half a command
466    /// in its buffer hands its decoder back too. Either one leaves a resume
467    /// point behind, and a resume point is an offset into a buffer that is
468    /// about to stop being the same buffer. A decoder taken here is always
469    /// starting a command, never continuing one, since a continuation comes off
470    /// the connection's own `partial` and never off the pool.
471    fn take_decoder(&mut self) -> u32 {
472        match self.spare.pop() {
473            Some(slot) => {
474                self.argvs[slot as usize].reset();
475                slot
476            }
477            None => yo_alloc::allow(|| {
478                self.argvs.push(Argv::with_capacity(ARGV_HINT));
479                // Every slot handed out here comes back to `spare` exactly
480                // once, so `spare` never holds more than `argvs` has slots.
481                // Sizing it here means the pushes that give a slot back never
482                // touch the allocator, and those are on the command path while
483                // this is not: a decoder is made once per depth of pipelining
484                // the connection has ever reached. `spare` is empty right now,
485                // which is why we are down here at all.
486                self.spare.reserve(self.argvs.len());
487                (self.argvs.len() - 1) as u32
488            }),
489        }
490    }
491
492    /// Note that this connection has something to write.
493    pub(crate) fn soil(&mut self, conn: ConnId) {
494        let c = &mut self.conns[conn as usize];
495        if !c.dirty {
496            c.dirty = true;
497            if self.dirty.len() == self.dirty.capacity() {
498                yo_alloc::allow(|| self.dirty.reserve(16));
499            }
500            self.dirty.push(conn);
501        }
502    }
503
504    /// The session on a connection, for the server side of it going away.
505    ///
506    /// `None` for a slot that is already free, so that closing twice is not two
507    /// chances to hand back the same watches.
508    pub(crate) fn session_mut(&mut self, conn: ConnId) -> Option<&mut Session> {
509        let c = &mut self.conns[conn as usize];
510        c.live.then_some(&mut c.session)
511    }
512
513    /// Hand the slot and its buffers back, and say which client has gone.
514    ///
515    /// `None` for a slot that was already closed. The id is what the server
516    /// finds a waiter by, and the caller forgets it before anything else runs,
517    /// because this slot is on the free list from here and the next accept
518    /// hands it to somebody else.
519    pub(crate) fn close(&mut self, conn: ConnId) -> Option<u64> {
520        {
521            let c = &mut self.conns[conn as usize];
522            if !c.live {
523                return None;
524            }
525            if let Some(slot) = c.partial.take() {
526                self.spare.push(slot);
527            }
528            c.live = false;
529            c.dirty = false;
530            c.blocked = false;
531            c.out.clear();
532            c.buf.clear();
533            c.head = 0;
534        }
535        let client = self.conns[conn as usize].session.id();
536        self.sink.closed(conn);
537        yo_alloc::allow(|| self.free.push(conn));
538        Some(client)
539    }
540
541    /// Move up to `max` framed commands into `into`.
542    ///
543    /// The reactor wants a batch it owns, and the front keeps the buffers, so
544    /// what crosses between them is this: numbers, no borrows.
545    pub(crate) fn take_ready(&mut self, into: &mut Vec<Cmd>, max: usize) -> usize {
546        let n = max.min(self.ready.len());
547        into.extend(self.ready.drain(..n));
548        n
549    }
550
551    /// Offer one connection's replies to the sink.
552    pub(crate) fn write_out(&mut self, conn: ConnId) -> Wrote {
553        {
554            let c = &self.conns[conn as usize];
555            if !c.live {
556                return Wrote::Done;
557            }
558        }
559        // A protocol error goes out once everything in front of it has.
560        if self.conns[conn as usize].pending == 0
561            && let Some(e) = self.conns[conn as usize].deferred.take()
562        {
563            self.scratch.clear();
564            e.write_reply(&mut self.scratch);
565            self.conns[conn as usize].out.raw(&self.scratch);
566        }
567
568        let taken = {
569            let c = &self.conns[conn as usize];
570            if c.out.is_empty() {
571                0
572            } else {
573                // One write for the whole batch's replies, never one per reply.
574                self.sink.write(conn, c.out.as_slice())
575            }
576        };
577
578        let c = &mut self.conns[conn as usize];
579        c.session.wrote_bytes(taken);
580        if taken >= c.out.len() {
581            c.out.clear();
582        } else {
583            c.out.consume(taken);
584        }
585
586        if !c.out.is_empty() {
587            return Wrote::Owed;
588        }
589        c.dirty = false;
590        let ending = c.closing && c.pending == 0;
591        if ending {
592            if let Some(client) = self.close(conn) {
593                return Wrote::Ended(client);
594            }
595        } else {
596            c.compact();
597            self.note_size(conn);
598        }
599        Wrote::Done
600    }
601
602    /// The dirty list, taken so the caller can walk it and reach the rest of
603    /// the front at the same time. The capacity comes back with it, so this is
604    /// not an allocation.
605    pub(crate) fn take_dirty(&mut self) -> Vec<ConnId> {
606        core::mem::take(&mut self.dirty)
607    }
608
609    /// The dirty list, given back with whatever is still owed on it.
610    pub(crate) fn give_dirty(&mut self, dirty: Vec<ConnId>) {
611        self.dirty = dirty;
612    }
613
614    /// How many connections are open.
615    pub(crate) fn clients(&self) -> usize {
616        self.conns.iter().filter(|c| c.live).count()
617    }
618
619    /// Commands framed and waiting for the reactor.
620    pub(crate) fn ready(&self) -> usize {
621        self.ready.len()
622    }
623
624    /// Connections with a reply that has not gone out yet.
625    pub(crate) fn owed(&self) -> usize {
626        self.dirty.len()
627    }
628
629    /// Decoders in the pool, which is the high water mark of one batch.
630    pub(crate) fn decoders(&self) -> usize {
631        self.argvs.len()
632    }
633
634    /// What every connection's read and reply buffers are holding.
635    ///
636    /// The walk is fine here because this is a test and a report, and the
637    /// number the running server uses is the one kept by `note_size`.
638    pub(crate) fn buffer_bytes(&self) -> usize {
639        self.conns.iter().map(Conn::size).sum()
640    }
641
642    /// Whether the slot is open.
643    pub(crate) fn live(&self, conn: ConnId) -> bool {
644        self.conns[conn as usize].live
645    }
646
647    /// Whether the peer has gone.
648    pub(crate) fn gone(&self, conn: ConnId) -> bool {
649        self.conns[conn as usize].gone
650    }
651
652    /// Commands framed out of this connection's buffer and not yet run.
653    pub(crate) fn pending(&self, conn: ConnId) -> u32 {
654        self.conns[conn as usize].pending
655    }
656
657    /// Whether this client is parked on a blocking command.
658    pub(crate) fn blocked(&self, conn: ConnId) -> bool {
659        self.conns[conn as usize].blocked
660    }
661
662    /// The client id, which is what the server knows a connection by.
663    pub(crate) fn client(&self, conn: ConnId) -> u64 {
664        self.conns[conn as usize].session.id()
665    }
666
667    /// The database this connection has selected.
668    pub(crate) fn db(&self, conn: ConnId) -> usize {
669        self.conns[conn as usize].session.db()
670    }
671
672    /// Whether this slot is still the client the server thinks it is.
673    ///
674    /// A slot is reused and a client id is not, so a waiter that named a client
675    /// is only about this connection while both agree.
676    pub(crate) fn answers(&self, conn: ConnId, client: u64) -> bool {
677        let c = &self.conns[conn as usize];
678        c.live && c.session.id() == client
679    }
680
681    /// Where a reply for this connection goes.
682    pub(crate) fn out(&mut self, conn: ConnId) -> &mut Out {
683        &mut self.conns[conn as usize].out
684    }
685
686    /// The peer went away.
687    pub(crate) fn mark_gone(&mut self, conn: ConnId) {
688        let c = &mut self.conns[conn as usize];
689        c.gone = true;
690        c.closing = true;
691    }
692
693    /// The client said goodbye.
694    ///
695    /// Anything it pipelined behind the `QUIT` was sent before it knew the
696    /// answer, and running it would be acting on a connection that has already
697    /// been said goodbye to.
698    pub(crate) fn quit(&mut self, conn: ConnId) {
699        let c = &mut self.conns[conn as usize];
700        c.closing = true;
701        c.skip = true;
702    }
703
704    /// The client is waiting on a blocking command.
705    pub(crate) fn block(&mut self, conn: ConnId) {
706        self.conns[conn as usize].blocked = true;
707    }
708
709    /// Hold a command that was framed with the batch that blocked.
710    pub(crate) fn park(&mut self, conn: ConnId, cmd: Cmd) {
711        yo_alloc::allow(|| self.conns[conn as usize].parked.push(cmd));
712    }
713
714    /// Take a command back off the connection that was about to run it.
715    ///
716    /// The pause is the one thing that can stop a command after it has been
717    /// handed over, so this is `park` with the two things `start` already did
718    /// undone: the command goes back into `pending`, which is what stops the
719    /// buffer its arguments point into being compacted underneath it, and the
720    /// caller keeps its decoder rather than handing it back.
721    pub(crate) fn hold(&mut self, conn: ConnId, cmd: Cmd) {
722        let c = &mut self.conns[conn as usize];
723        c.pending += 1;
724        c.blocked = true;
725        yo_alloc::allow(|| c.parked.push(cmd));
726    }
727
728    /// The client is not waiting any more: give it back its commands.
729    ///
730    /// The ones it had already sent go to the front of the queue in the order
731    /// they arrived, ahead of anything any other connection has waiting, because
732    /// they were framed before any of that was. Then framing starts again on
733    /// whatever arrived while it was parked.
734    pub(crate) fn unpark(&mut self, conn: ConnId) {
735        let mut parked = {
736            let c = &mut self.conns[conn as usize];
737            c.blocked = false;
738            core::mem::take(&mut c.parked)
739        };
740        // Back to front, since each one goes on the front.
741        while let Some(cmd) = parked.pop() {
742            if self.ready.len() == self.ready.capacity() {
743                yo_alloc::allow(|| self.ready.reserve(BATCH_MAX));
744            }
745            self.ready.push_front(cmd);
746        }
747        // Empty now, and back where it lives so its room is not paid for twice.
748        self.conns[conn as usize].parked = parked;
749        if !self.conns[conn as usize].closing {
750            self.frame(conn);
751        }
752    }
753
754    /// A command is off the queue: take it out of the count, and say whether it
755    /// should run at all.
756    ///
757    /// It should not when the peer has gone or has said goodbye, and the answer
758    /// is `false` rather than an early return because the decoder still has to
759    /// come back and the slot still has to be released.
760    pub(crate) fn start(&mut self, cmd: &Cmd) -> bool {
761        let c = &mut self.conns[cmd.conn as usize];
762        c.pending -= 1;
763        !(c.gone || c.skip)
764    }
765
766    /// The three things running a command needs from this side: the arguments,
767    /// the session they run against, and where the reply goes.
768    pub(crate) fn parts(&mut self, cmd: &Cmd) -> (Args<'_>, &mut Session, &mut Out) {
769        let c = &mut self.conns[cmd.conn as usize];
770        let args = Args::new(&self.argvs[cmd.slot as usize], &c.buf[cmd.base..]);
771        (args, &mut c.session, &mut c.out)
772    }
773
774    /// The arguments alone, for a caller that is only reading them.
775    pub(crate) fn args(&self, cmd: &Cmd) -> Args<'_> {
776        let c = &self.conns[cmd.conn as usize];
777        Args::new(&self.argvs[cmd.slot as usize], &c.buf[cmd.base..])
778    }
779
780    /// The command is finished with its decoder.
781    pub(crate) fn done(&mut self, cmd: &Cmd) {
782        self.spare.push(cmd.slot);
783    }
784}
785
786#[cfg(test)]
787mod tests {
788    use super::*;
789    use crate::engine::Recorder;
790
791    /// The wire bytes for a command, built the way a client would.
792    fn wire(args: &[&[u8]]) -> Vec<u8> {
793        let mut b = format!("*{}\r\n", args.len()).into_bytes();
794        for a in args {
795            b.extend_from_slice(format!("${}\r\n", a.len()).as_bytes());
796            b.extend_from_slice(a);
797            b.extend_from_slice(b"\r\n");
798        }
799        b
800    }
801
802    /// A front and one connection on it. No server anywhere, which is the
803    /// point: framing is this side's work alone.
804    fn front() -> (Front<Recorder>, ConnId) {
805        let mut f = Front::new(Recorder::new());
806        let conn = f.open(1);
807        (f, conn)
808    }
809
810    #[test]
811    fn a_pipelined_read_frames_every_command_in_it() {
812        let (mut f, conn) = front();
813        let mut bytes = wire(&[b"SET", b"k", b"v"]);
814        bytes.extend_from_slice(&wire(&[b"GET", b"k"]));
815        f.feed(conn, &bytes);
816
817        let mut batch = Vec::new();
818        assert_eq!(f.take_ready(&mut batch, 64), 2);
819        assert_eq!(f.args(&batch[0]).name(), b"SET");
820        assert_eq!(f.args(&batch[1]).name(), b"GET");
821        assert_eq!(f.pending(conn), 2);
822    }
823
824    #[test]
825    fn a_command_split_across_reads_is_framed_once_it_is_whole() {
826        let (mut f, conn) = front();
827        let bytes = wire(&[b"SET", b"k", b"v"]);
828        let (head, tail) = bytes.split_at(9);
829
830        f.feed(conn, head);
831        let mut batch = Vec::new();
832        assert_eq!(f.take_ready(&mut batch, 64), 0);
833
834        f.feed(conn, tail);
835        assert_eq!(f.take_ready(&mut batch, 64), 1);
836        assert_eq!(f.args(&batch[0]).name(), b"SET");
837    }
838
839    #[test]
840    fn a_protocol_error_stops_the_framing_and_closes_the_connection() {
841        let (mut f, conn) = front();
842        f.feed(conn, b"*x\r\n");
843        assert_eq!(f.take_ready(&mut Vec::new(), 64), 0);
844        assert_eq!(f.owed(), 1);
845
846        // Nothing is owed to the client afterwards and the slot has gone back,
847        // which is what a closed connection means on this side.
848        assert!(matches!(f.write_out(conn), Wrote::Ended(_)));
849        assert!(!f.live(conn));
850        assert!(f.sink().sent(conn).starts_with(b"-ERR"));
851    }
852
853    #[test]
854    fn a_closed_slot_is_handed_out_again_with_its_buffers() {
855        let (mut f, conn) = front();
856        f.feed(conn, &wire(&[b"PING"]));
857        let held = f.buffer_bytes();
858        assert_eq!(f.close(conn), Some(1));
859
860        let next = f.open(2);
861        assert_eq!(next, conn, "the slot comes back");
862        assert_eq!(f.client(next), 2, "the client id does not");
863        assert_eq!(f.buffer_bytes(), held, "and neither buffer was given up");
864    }
865
866    #[test]
867    fn the_buffers_are_reported_as_they_move_and_only_once() {
868        let (mut f, conn) = front();
869        assert!(f.buffer_delta() > 0, "accept made two buffers");
870        assert_eq!(f.buffer_delta(), 0, "and nobody is told about them twice");
871
872        f.feed(conn, &wire(&[b"PING"]));
873        assert_eq!(f.buffer_delta(), 0, "a command that fits moves nothing");
874    }
875}