use std::sync::atomic::Ordering::Relaxed;
use yo_common::lock::Held;
use yo_common::{Code, Error, Result, num};
use yo_kv::{Db, End, Entry, Member, Movem, ZEnd};
use super::args::{self, Args, NOT_AN_INT};
use super::lists::{BAD_MPOP_COUNT, BAD_NUMKEYS, end_of, movem_options};
use super::streams;
use super::table::Spec;
use super::zsets;
use super::{Flow, Server, Session};
use crate::reply::Out;
const NOT_A_FLOAT: &str = "timeout is not a float or out of range";
const NEGATIVE: &str = "timeout is negative";
const OUT_OF_RANGE: &str = "timeout is out of range";
const TIMEOUT_NOT_AN_INT: &str = "timeout is not an integer or out of range";
const NOT_ZERO_OR_ONE: &str = "value is out of range, value must between 0 and 1";
const NOT_POSITIVE: &str = "value is out of range, must be positive";
const NO_AOF: &str = "WAITAOF cannot be used when numlocal is set but appendonly is disabled.";
pub(super) fn execute(
server: &Server,
session: &Session,
spec: &Spec,
args: Args<'_>,
out: &mut Out,
) -> Result<Flow> {
if spec.name == "wait" || spec.name == "waitaof" {
return replication(spec.name, args, out).map(|()| Flow::Continue);
}
let now = server.now_ms();
if spec.name == "xread" || spec.name == "xreadgroup" {
let db = session.db();
let want = streams::parse_read(spec.name, args, server.striped(db), now)?;
let block = Block::xread(want.keys, want.reads);
if block.now(server.striped(db), now, out)? {
return Ok(Flow::Continue);
}
let Some(deadline) = want.wait else {
out.nil_array();
return Ok(Flow::Continue);
};
server.park(session.id(), db, deadline, block);
return Ok(Flow::Block);
}
let last = args.len() - 1;
let (deadline, block) = match spec.name {
"blpop" | "brpop" => {
let end = if spec.name == "blpop" {
End::Left
} else {
End::Right
};
let deadline = timeout(args.get(last), now)?;
(deadline, Block::pop((1..last).map(|i| args.get(i)), end))
}
"blmove" => {
let (from, to) = (end_of(args.get(3))?, end_of(args.get(4))?);
let deadline = timeout(args.get(5), now)?;
(deadline, Block::moved(args.get(1), args.get(2), from, to))
}
"blmovem" => {
let (from, to) = (end_of(args.get(3))?, end_of(args.get(4))?);
let deadline = timeout(args.get(5), now)?;
let mv = movem_options(args, 6, from, to)?;
(deadline, Block::movem(args.get(1), args.get(2), mv))
}
"brpoplpush" => {
let deadline = timeout(args.get(3), now)?;
(
deadline,
Block::moved(args.get(1), args.get(2), End::Right, End::Left),
)
}
"blmpop" => mpop(args, now)?,
"bzpopmin" | "bzpopmax" => {
let end = zsets::end_of_name(spec.name);
let deadline = timeout(args.get(last), now)?;
(deadline, Block::zpop((1..last).map(|i| args.get(i)), end))
}
"bzmpop" => {
let deadline = timeout(args.get(1), now)?;
let (end, from, to, count) = zsets::parse_mpop(args, 2)?;
(
deadline,
Block::zmpop((from..to).map(|i| args.get(i)), end, count),
)
}
_ => return Err(args::syntax()),
};
let db = session.db();
if block.now(server.striped(db), now, out)? {
return Ok(Flow::Continue);
}
server.park(session.id(), db, deadline, block);
Ok(Flow::Block)
}
fn mpop(args: Args<'_>, now: u64) -> Result<(Option<u64>, Block)> {
let deadline = timeout(args.get(1), now)?;
let numkeys = match args.int(2) {
Ok(n) if n > 0 => usize::try_from(n).unwrap_or(usize::MAX),
_ => return Err(Error::new(Code::Invalid, BAD_NUMKEYS)),
};
if numkeys >= args.len() - 3 {
return Err(args::syntax());
}
let at = 3 + numkeys;
let end = end_of(args.get(at))?;
let mut want = 1usize;
if at + 1 < args.len() {
if args.len() != at + 3 || !args::is(args.get(at + 1), b"count") {
return Err(args::syntax());
}
want = match args.int(at + 2) {
Ok(n) if n > 0 => usize::try_from(n).unwrap_or(usize::MAX),
_ => return Err(Error::new(Code::Invalid, BAD_MPOP_COUNT)),
};
}
Ok((
deadline,
Block::mpop((3..at).map(|i| args.get(i)), end, want),
))
}
fn replication(name: &str, args: Args<'_>, out: &mut Out) -> Result<()> {
let aof = name == "waitaof";
let at = usize::from(aof);
let mut wants_local = false;
if aof {
let local = whole(args.get(1))?;
if !(0..=1).contains(&local) {
return Err(Error::new(Code::Invalid, NOT_ZERO_OR_ONE));
}
wants_local = local == 1;
}
let replicas = whole(args.get(at + 1))?;
if aof && replicas < 0 {
return Err(Error::new(Code::Invalid, NOT_POSITIVE));
}
let ms = whole(args.get(at + 2)).map_err(|_| Error::new(Code::Invalid, TIMEOUT_NOT_AN_INT))?;
if ms < 0 {
return Err(Error::new(Code::Invalid, NEGATIVE));
}
if wants_local {
return Err(Error::new(Code::Invalid, NO_AOF));
}
if aof {
out.array(2);
out.int(0);
out.int(0);
} else {
out.int(0);
}
Ok(())
}
fn whole(arg: &[u8]) -> Result<i64> {
num::parse_i64(arg).ok_or_else(|| Error::new(Code::Invalid, NOT_AN_INT))
}
fn timeout(arg: &[u8], now: u64) -> Result<Option<u64>> {
let Some(secs) = num::parse_f64(arg) else {
return Err(Error::new(Code::Invalid, NOT_A_FLOAT));
};
if secs < 0.0 {
return Err(Error::new(Code::Invalid, NEGATIVE));
}
let ms = secs * 1000.0;
if ms > i64::MAX as f64 {
return Err(Error::new(Code::Invalid, OUT_OF_RANGE));
}
if ms <= 0.0 {
return Ok(None);
}
Ok(Some(now.saturating_add(ms as u64)))
}
enum Want {
Pop { end: End },
Move { dst: Vec<u8>, from: End, to: End },
MoveM { dst: Vec<u8>, mv: Movem },
Mpop { end: End, count: usize },
ZPop { end: ZEnd },
ZMpop { end: ZEnd, count: usize },
XRead(streams::Reads),
}
impl Want {
fn attempt(
&self,
keys: &[Vec<u8>],
db: &Db,
now: u64,
out: &mut Out,
strict: bool,
) -> Result<bool> {
match self {
Want::XRead(r) => streams::read(db, keys, r, now, strict, out),
Want::Pop { end } => {
for key in keys {
if !ready(db, key, strict)? {
continue;
}
out.array(2);
out.bulk(key);
db.hold(key).pop_into(key, *end, 1, |e| element(out, e))?;
return Ok(true);
}
Ok(false)
}
Want::Mpop { end, count } => {
for key in keys {
if !ready(db, key, strict)? {
continue;
}
out.array(2);
out.bulk(key);
let mark = out.len();
let n = db
.hold(key)
.pop_into(key, *end, *count, |e| element(out, e))?;
out.close_array(mark, n);
return Ok(true);
}
Ok(false)
}
Want::ZPop { end } => {
for key in keys {
if !zready(db, key, strict)? {
continue;
}
out.array(3);
out.bulk(key);
db.hold(key).zpop(key, *end, 1, |m, sc| {
member(out, m);
out.double(sc);
})?;
return Ok(true);
}
Ok(false)
}
Want::ZMpop { end, count } => {
for key in keys {
if !zready(db, key, strict)? {
continue;
}
out.array(2);
out.bulk(key);
let mark = out.len();
let n = db.hold(key).zpop(key, *end, *count, |m, sc| {
out.array(2);
member(out, m);
out.double(sc);
})?;
out.close_array(mark, n);
return Ok(true);
}
Ok(false)
}
Want::Move { dst, from, to } => {
let src = &keys[0];
if !ready(db, src, strict)? {
return Ok(false);
}
match db.lmove(src, dst, *from, *to, |v| out.bulk(v)) {
Ok(true) => Ok(true),
Ok(false) => Ok(false),
Err(e) if strict => Err(e),
Err(_) => Ok(false),
}
}
Want::MoveM { dst, mv } => {
let src = &keys[0];
let have = match db.hold(src).llen(src) {
Ok(n) => n,
Err(e) if strict => return Err(e),
Err(_) => return Ok(false),
};
if have == 0 || (mv.exactly && have < mv.count) {
return Ok(false);
}
let mark = out.len();
let mut n = 0;
match db.lmovem(src, dst, *mv, |v| {
out.bulk(v);
n += 1;
}) {
Ok(_) => {}
Err(e) if strict => return Err(e),
Err(_) => {
out.truncate(mark);
return Ok(false);
}
}
out.close_array(mark, n);
Ok(true)
}
}
}
}
fn ready(db: &Db, key: &[u8], strict: bool) -> Result<bool> {
match db.hold(key).llen(key) {
Ok(n) => Ok(n > 0),
Err(e) if strict => Err(e),
Err(_) => Ok(false),
}
}
fn zready(db: &Db, key: &[u8], strict: bool) -> Result<bool> {
match db.hold(key).zcard(key) {
Ok(n) => Ok(n > 0),
Err(e) if strict => Err(e),
Err(_) => Ok(false),
}
}
#[inline]
fn element(out: &mut Out, e: Entry<'_>) {
match e {
Entry::Int(n) => out.bulk_int(n),
Entry::Str(s) => out.bulk(s),
}
}
#[inline]
fn member(out: &mut Out, m: Member<'_>) {
match m {
Member::Int(n) => out.bulk_int(n),
Member::Str(s) => out.bulk(s),
}
}
pub struct Block {
keys: Vec<Vec<u8>>,
want: Want,
}
impl Block {
fn pop<'a>(keys: impl Iterator<Item = &'a [u8]>, end: End) -> Block {
Block {
keys: owned(keys),
want: Want::Pop { end },
}
}
fn mpop<'a>(keys: impl Iterator<Item = &'a [u8]>, end: End, count: usize) -> Block {
Block {
keys: owned(keys),
want: Want::Mpop { end, count },
}
}
fn zpop<'a>(keys: impl Iterator<Item = &'a [u8]>, end: ZEnd) -> Block {
Block {
keys: owned(keys),
want: Want::ZPop { end },
}
}
fn zmpop<'a>(keys: impl Iterator<Item = &'a [u8]>, end: ZEnd, count: usize) -> Block {
Block {
keys: owned(keys),
want: Want::ZMpop { end, count },
}
}
fn moved(src: &[u8], dst: &[u8], from: End, to: End) -> Block {
yo_alloc::allow(|| Block {
keys: vec![src.to_vec()],
want: Want::Move {
dst: dst.to_vec(),
from,
to,
},
})
}
fn movem(src: &[u8], dst: &[u8], mv: Movem) -> Block {
yo_alloc::allow(|| Block {
keys: vec![src.to_vec()],
want: Want::MoveM {
dst: dst.to_vec(),
mv,
},
})
}
fn now(&self, db: &Db, now: u64, out: &mut Out) -> Result<bool> {
self.want.attempt(&self.keys, db, now, out, true)
}
fn xread(keys: Vec<Vec<u8>>, reads: streams::Reads) -> Block {
Block {
keys,
want: Want::XRead(reads),
}
}
}
fn owned<'a>(keys: impl Iterator<Item = &'a [u8]>) -> Vec<Vec<u8>> {
yo_alloc::allow(|| keys.map(<[u8]>::to_vec).collect())
}
struct Waiter {
client: u64,
conn: u32,
db: usize,
deadline: Option<u64>,
keys: Vec<Vec<u8>>,
want: Want,
}
#[derive(Default)]
pub struct Waiters {
list: Vec<Waiter>,
}
#[derive(Debug, Clone, Copy)]
pub struct Parked {
pub conn: u32,
pub client: u64,
}
impl Waiters {
#[must_use]
#[inline]
pub fn is_empty(&self) -> bool {
self.list.is_empty()
}
#[must_use]
#[inline]
pub fn len(&self) -> usize {
self.list.len()
}
#[must_use]
pub fn at(&self, at: usize) -> Parked {
let w = &self.list[at];
Parked {
conn: w.conn,
client: w.client,
}
}
#[must_use]
pub fn db_of(&self, at: usize) -> usize {
self.list[at].db
}
fn drop_at(&mut self, at: usize) {
self.list.remove(at);
}
fn forget(&mut self, client: u64) {
self.list.retain(|w| w.client != client);
}
fn bind(&mut self, client: u64, conn: u32) {
if let Some(w) = self.list.iter_mut().rev().find(|w| w.client == client) {
w.conn = conn;
}
}
fn park(&mut self, client: u64, db: usize, deadline: Option<u64>, block: Block) {
yo_alloc::allow(|| {
self.list.push(Waiter {
client,
conn: 0,
db,
deadline,
keys: block.keys,
want: block.want,
});
});
}
fn try_serve(&self, at: usize, dbs: &[Db], now: u64, out: &mut Out) -> bool {
let w = &self.list[at];
let mark = out.len();
match w.want.attempt(&w.keys, &dbs[w.db], now, out, false) {
Ok(true) => return true,
Ok(false) => {}
Err(_) => out.truncate(mark),
}
if w.deadline.is_some_and(|d| now >= d) {
out.nil_array();
return true;
}
false
}
}
impl Server {
#[must_use]
pub fn now_ms(&self) -> u64 {
self.clock.now_ms()
}
#[must_use]
pub fn waiters(&self) -> Held<'_, Waiters> {
self.waiters.lock()
}
#[must_use]
#[inline]
pub fn parked(&self) -> usize {
self.parked.load(Relaxed)
}
pub(super) fn park(&self, client: u64, db: usize, deadline: Option<u64>, block: Block) {
let mut list = self.waiters.lock();
list.park(client, db, deadline, block);
self.note(&list);
}
pub fn drop_waiter(&self, at: usize) {
let mut list = self.waiters.lock();
list.drop_at(at);
self.note(&list);
}
pub fn forget_waiters(&self, client: u64) {
let mut list = self.waiters.lock();
list.forget(client);
self.note(&list);
}
pub fn bind_waiter(&self, client: u64, conn: u32) {
self.waiters.lock().bind(client, conn);
}
fn note(&self, list: &Waiters) {
self.parked.store(list.len(), Relaxed);
}
pub fn serve_waiter(&self, at: usize, now: u64, out: &mut Out) -> bool {
let list = self.waiters.lock();
self.mine().mark(1u64 << list.db_of(at));
list.try_serve(at, &self.dbs, now, out)
}
}