use std::collections::VecDeque;
use yo_reactor::BATCH_MAX;
use crate::dispatch::table::lookup_index;
use crate::dispatch::{Args, Session};
use crate::engine::{ConnId, Sink};
use crate::error::ProtocolError;
use crate::proto::{Limits, Proto};
use crate::reply::Out;
use crate::request::{Argv, Step};
const READ_BUF: usize = 16 * 1024;
const OUT_BUF: usize = 16 * 1024;
const ARGV_HINT: usize = 8;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Cmd {
pub(crate) conn: ConnId,
pub(crate) slot: u32,
pub(crate) base: usize,
pub(crate) spec: u16,
}
impl Cmd {
#[must_use]
pub const fn conn(&self) -> ConnId {
self.conn
}
}
pub(crate) enum Wrote {
Owed,
Done,
Ended(u64),
}
struct Conn {
live: bool,
session: Session,
out: Out,
buf: Vec<u8>,
head: usize,
partial: Option<u32>,
pending: u32,
closing: bool,
deferred: Option<ProtocolError>,
skip: bool,
gone: bool,
dirty: bool,
blocked: bool,
parked: Vec<Cmd>,
held: usize,
}
impl Conn {
fn new(id: u64) -> Conn {
yo_alloc::allow(|| Conn {
live: true,
session: Session::new(id),
out: Out::with_capacity(Proto::Resp2, OUT_BUF),
buf: Vec::with_capacity(READ_BUF),
head: 0,
partial: None,
pending: 0,
closing: false,
deferred: None,
skip: false,
gone: false,
dirty: false,
blocked: false,
parked: Vec::new(),
held: 0,
})
}
fn size(&self) -> usize {
self.buf.capacity() + self.out.capacity()
}
fn reset(&mut self, id: u64) {
self.live = true;
self.session = Session::new(id);
self.out.clear();
self.out.set_proto(Proto::Resp2);
self.buf.clear();
self.head = 0;
self.partial = None;
self.pending = 0;
self.closing = false;
self.deferred = None;
self.skip = false;
self.gone = false;
self.dirty = false;
self.blocked = false;
self.parked.clear();
}
fn compact(&mut self) {
if self.pending > 0 || self.head == 0 {
return;
}
if self.head == self.buf.len() {
self.buf.clear();
} else {
self.buf.drain(..self.head);
}
self.head = 0;
}
}
pub(crate) struct Front<S> {
sink: S,
conns: Vec<Conn>,
free: Vec<ConnId>,
argvs: Vec<Argv>,
spare: Vec<u32>,
ready: VecDeque<Cmd>,
dirty: Vec<ConnId>,
scratch: Vec<u8>,
limits: Limits,
moved: isize,
}
impl<S: Sink> Front<S> {
pub(crate) fn new(sink: S) -> Front<S> {
Front {
sink,
conns: Vec::new(),
free: Vec::new(),
argvs: Vec::new(),
spare: Vec::new(),
ready: VecDeque::with_capacity(BATCH_MAX),
dirty: Vec::with_capacity(16),
scratch: Vec::with_capacity(128),
limits: Limits::default(),
moved: 0,
}
}
pub(crate) const fn sink(&self) -> &S {
&self.sink
}
pub(crate) const fn sink_mut(&mut self) -> &mut S {
&mut self.sink
}
pub(crate) fn set_limits(&mut self, limits: Limits) {
self.limits = limits;
}
pub(crate) fn open(&mut self, id: u64) -> ConnId {
let at = match self.free.pop() {
Some(at) => {
self.conns[at as usize].reset(id);
at
}
None => {
let conn = Conn::new(id);
yo_alloc::allow(|| self.conns.push(conn));
(self.conns.len() - 1) as ConnId
}
};
self.conns[at as usize].session.set_conn(at);
self.note_size(at);
at
}
pub(crate) fn feed(&mut self, conn: ConnId, bytes: &[u8]) {
{
let c = &mut self.conns[conn as usize];
if !c.live || c.closing {
return;
}
yo_alloc::allow(|| c.buf.extend_from_slice(bytes));
}
self.frame(conn);
self.note_size(conn);
}
fn note_size(&mut self, conn: ConnId) {
let c = &mut self.conns[conn as usize];
let now = c.size();
if now == c.held {
return;
}
let delta = now as isize - c.held as isize;
c.held = now;
self.moved += delta;
}
pub(crate) fn buffer_delta(&mut self) -> isize {
core::mem::take(&mut self.moved)
}
fn frame(&mut self, conn: ConnId) {
if self.conns[conn as usize].blocked {
return;
}
loop {
let base = self.conns[conn as usize].head;
let slot = match self.conns[conn as usize].partial.take() {
Some(slot) => slot,
None => self.take_decoder(),
};
let step = {
let c = &self.conns[conn as usize];
self.argvs[slot as usize].decode(&c.buf[base..], &self.limits)
};
match step {
Ok(Step::Command { consumed }) => {
self.conns[conn as usize].head += consumed;
if self.argvs[slot as usize].is_empty() {
self.spare.push(slot);
} else {
if self.ready.len() == self.ready.capacity() {
yo_alloc::allow(|| self.ready.reserve(BATCH_MAX));
}
let spec = {
let c = &self.conns[conn as usize];
let args = Args::new(&self.argvs[slot as usize], &c.buf[base..]);
lookup_index(args.name())
};
self.ready.push_back(Cmd {
conn,
slot,
base,
spec,
});
self.conns[conn as usize].pending += 1;
}
}
Ok(Step::Incomplete) => {
self.conns[conn as usize].partial = Some(slot);
break;
}
Err(e) => {
self.spare.push(slot);
let c = &mut self.conns[conn as usize];
c.deferred = Some(e);
c.closing = true;
self.soil(conn);
break;
}
}
}
self.conns[conn as usize].compact();
}
fn take_decoder(&mut self) -> u32 {
match self.spare.pop() {
Some(slot) => {
self.argvs[slot as usize].reset();
slot
}
None => yo_alloc::allow(|| {
self.argvs.push(Argv::with_capacity(ARGV_HINT));
self.spare.reserve(self.argvs.len());
(self.argvs.len() - 1) as u32
}),
}
}
pub(crate) fn soil(&mut self, conn: ConnId) {
let c = &mut self.conns[conn as usize];
if !c.dirty {
c.dirty = true;
if self.dirty.len() == self.dirty.capacity() {
yo_alloc::allow(|| self.dirty.reserve(16));
}
self.dirty.push(conn);
}
}
pub(crate) fn session_mut(&mut self, conn: ConnId) -> Option<&mut Session> {
let c = &mut self.conns[conn as usize];
c.live.then_some(&mut c.session)
}
pub(crate) fn close(&mut self, conn: ConnId) -> Option<u64> {
{
let c = &mut self.conns[conn as usize];
if !c.live {
return None;
}
if let Some(slot) = c.partial.take() {
self.spare.push(slot);
}
c.live = false;
c.dirty = false;
c.blocked = false;
c.out.clear();
c.buf.clear();
c.head = 0;
}
let client = self.conns[conn as usize].session.id();
self.sink.closed(conn);
yo_alloc::allow(|| self.free.push(conn));
Some(client)
}
pub(crate) fn take_ready(&mut self, into: &mut Vec<Cmd>, max: usize) -> usize {
let n = max.min(self.ready.len());
into.extend(self.ready.drain(..n));
n
}
pub(crate) fn write_out(&mut self, conn: ConnId) -> Wrote {
{
let c = &self.conns[conn as usize];
if !c.live {
return Wrote::Done;
}
}
if self.conns[conn as usize].pending == 0
&& let Some(e) = self.conns[conn as usize].deferred.take()
{
self.scratch.clear();
e.write_reply(&mut self.scratch);
self.conns[conn as usize].out.raw(&self.scratch);
}
let taken = {
let c = &self.conns[conn as usize];
if c.out.is_empty() {
0
} else {
self.sink.write(conn, c.out.as_slice())
}
};
let c = &mut self.conns[conn as usize];
if taken >= c.out.len() {
c.out.clear();
} else {
c.out.consume(taken);
}
if !c.out.is_empty() {
return Wrote::Owed;
}
c.dirty = false;
let ending = c.closing && c.pending == 0;
if ending {
if let Some(client) = self.close(conn) {
return Wrote::Ended(client);
}
} else {
c.compact();
self.note_size(conn);
}
Wrote::Done
}
pub(crate) fn take_dirty(&mut self) -> Vec<ConnId> {
core::mem::take(&mut self.dirty)
}
pub(crate) fn give_dirty(&mut self, dirty: Vec<ConnId>) {
self.dirty = dirty;
}
pub(crate) fn clients(&self) -> usize {
self.conns.iter().filter(|c| c.live).count()
}
pub(crate) fn ready(&self) -> usize {
self.ready.len()
}
pub(crate) fn owed(&self) -> usize {
self.dirty.len()
}
pub(crate) fn decoders(&self) -> usize {
self.argvs.len()
}
pub(crate) fn buffer_bytes(&self) -> usize {
self.conns.iter().map(Conn::size).sum()
}
pub(crate) fn live(&self, conn: ConnId) -> bool {
self.conns[conn as usize].live
}
pub(crate) fn gone(&self, conn: ConnId) -> bool {
self.conns[conn as usize].gone
}
pub(crate) fn pending(&self, conn: ConnId) -> u32 {
self.conns[conn as usize].pending
}
pub(crate) fn blocked(&self, conn: ConnId) -> bool {
self.conns[conn as usize].blocked
}
pub(crate) fn client(&self, conn: ConnId) -> u64 {
self.conns[conn as usize].session.id()
}
pub(crate) fn db(&self, conn: ConnId) -> usize {
self.conns[conn as usize].session.db()
}
pub(crate) fn answers(&self, conn: ConnId, client: u64) -> bool {
let c = &self.conns[conn as usize];
c.live && c.session.id() == client
}
pub(crate) fn out(&mut self, conn: ConnId) -> &mut Out {
&mut self.conns[conn as usize].out
}
pub(crate) fn mark_gone(&mut self, conn: ConnId) {
let c = &mut self.conns[conn as usize];
c.gone = true;
c.closing = true;
}
pub(crate) fn quit(&mut self, conn: ConnId) {
let c = &mut self.conns[conn as usize];
c.closing = true;
c.skip = true;
}
pub(crate) fn block(&mut self, conn: ConnId) {
self.conns[conn as usize].blocked = true;
}
pub(crate) fn park(&mut self, conn: ConnId, cmd: Cmd) {
yo_alloc::allow(|| self.conns[conn as usize].parked.push(cmd));
}
pub(crate) fn unpark(&mut self, conn: ConnId) {
let mut parked = {
let c = &mut self.conns[conn as usize];
c.blocked = false;
core::mem::take(&mut c.parked)
};
while let Some(cmd) = parked.pop() {
if self.ready.len() == self.ready.capacity() {
yo_alloc::allow(|| self.ready.reserve(BATCH_MAX));
}
self.ready.push_front(cmd);
}
self.conns[conn as usize].parked = parked;
if !self.conns[conn as usize].closing {
self.frame(conn);
}
}
pub(crate) fn start(&mut self, cmd: &Cmd) -> bool {
let c = &mut self.conns[cmd.conn as usize];
c.pending -= 1;
!(c.gone || c.skip)
}
pub(crate) fn parts(&mut self, cmd: &Cmd) -> (Args<'_>, &mut Session, &mut Out) {
let c = &mut self.conns[cmd.conn as usize];
let args = Args::new(&self.argvs[cmd.slot as usize], &c.buf[cmd.base..]);
(args, &mut c.session, &mut c.out)
}
pub(crate) fn args(&self, cmd: &Cmd) -> Args<'_> {
let c = &self.conns[cmd.conn as usize];
Args::new(&self.argvs[cmd.slot as usize], &c.buf[cmd.base..])
}
pub(crate) fn done(&mut self, cmd: &Cmd) {
self.spare.push(cmd.slot);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::engine::Recorder;
fn wire(args: &[&[u8]]) -> Vec<u8> {
let mut b = format!("*{}\r\n", args.len()).into_bytes();
for a in args {
b.extend_from_slice(format!("${}\r\n", a.len()).as_bytes());
b.extend_from_slice(a);
b.extend_from_slice(b"\r\n");
}
b
}
fn front() -> (Front<Recorder>, ConnId) {
let mut f = Front::new(Recorder::new());
let conn = f.open(1);
(f, conn)
}
#[test]
fn a_pipelined_read_frames_every_command_in_it() {
let (mut f, conn) = front();
let mut bytes = wire(&[b"SET", b"k", b"v"]);
bytes.extend_from_slice(&wire(&[b"GET", b"k"]));
f.feed(conn, &bytes);
let mut batch = Vec::new();
assert_eq!(f.take_ready(&mut batch, 64), 2);
assert_eq!(f.args(&batch[0]).name(), b"SET");
assert_eq!(f.args(&batch[1]).name(), b"GET");
assert_eq!(f.pending(conn), 2);
}
#[test]
fn a_command_split_across_reads_is_framed_once_it_is_whole() {
let (mut f, conn) = front();
let bytes = wire(&[b"SET", b"k", b"v"]);
let (head, tail) = bytes.split_at(9);
f.feed(conn, head);
let mut batch = Vec::new();
assert_eq!(f.take_ready(&mut batch, 64), 0);
f.feed(conn, tail);
assert_eq!(f.take_ready(&mut batch, 64), 1);
assert_eq!(f.args(&batch[0]).name(), b"SET");
}
#[test]
fn a_protocol_error_stops_the_framing_and_closes_the_connection() {
let (mut f, conn) = front();
f.feed(conn, b"*x\r\n");
assert_eq!(f.take_ready(&mut Vec::new(), 64), 0);
assert_eq!(f.owed(), 1);
assert!(matches!(f.write_out(conn), Wrote::Ended(_)));
assert!(!f.live(conn));
assert!(f.sink().sent(conn).starts_with(b"-ERR"));
}
#[test]
fn a_closed_slot_is_handed_out_again_with_its_buffers() {
let (mut f, conn) = front();
f.feed(conn, &wire(&[b"PING"]));
let held = f.buffer_bytes();
assert_eq!(f.close(conn), Some(1));
let next = f.open(2);
assert_eq!(next, conn, "the slot comes back");
assert_eq!(f.client(next), 2, "the client id does not");
assert_eq!(f.buffer_bytes(), held, "and neither buffer was given up");
}
#[test]
fn the_buffers_are_reported_as_they_move_and_only_once() {
let (mut f, conn) = front();
assert!(f.buffer_delta() > 0, "accept made two buffers");
assert_eq!(f.buffer_delta(), 0, "and nobody is told about them twice");
f.feed(conn, &wire(&[b"PING"]));
assert_eq!(f.buffer_delta(), 0, "a command that fits moves nothing");
}
}