use std::collections::HashMap;
use yo_common::{Result, parse_i64};
use yo_search::expr::Value;
use super::super::Server;
use super::super::args::{self, Args};
use super::aggregate::writes;
use super::{Built, Fail, MISSING, Pairs, Rolled, Row, Shows, found, rolls};
use crate::reply::Out;
pub(super) const LIMIT: usize = 128;
const IDLE: u64 = 300_000;
const CHUNK: usize = 1000;
const OVER: &str = "SEARCH_LIMIT_OVER INDEX_CURSOR_LIMIT of 128 has been reached for an index";
const NOT_FOUND: &str = "Cursor not found, id: ";
const BAD_ID: &str = "Bad cursor ID";
const NO_CURSOR: &str = "Cursor does not exist";
const BAD_COUNT: &str = "Bad value for COUNT: `";
const QUOTE: &str = "`";
pub(super) const MOST: i64 = u32::MAX as i64;
#[derive(Clone, Copy)]
pub(super) struct Asks {
pub(super) count: usize,
pub(super) idle: u64,
}
impl Default for Asks {
fn default() -> Asks {
Asks {
count: CHUNK,
idle: IDLE,
}
}
}
pub(super) enum Made {
Found(Vec<(Row, Option<Pairs>)>),
Rolled(Vec<(Row, Pairs)>),
Piped {
names: Vec<Box<[u8]>>,
rows: Vec<(Option<Row>, Vec<Value>)>,
sorted: Option<usize>,
warning: Option<Vec<u8>>,
},
}
impl Made {
fn len(&self) -> usize {
match self {
Made::Found(rows) => rows.len(),
Made::Rolled(rows) => rows.len(),
Made::Piped { rows, .. } => rows.len(),
}
}
}
pub(super) struct Kept {
pub(super) made: Made,
pub(super) walk: Vec<bool>,
pub(super) shows: Shows,
pub(super) total: usize,
pub(super) whole: bool,
pub(super) loader: bool,
pub(super) offset: usize,
pub(super) window: usize,
}
struct Cursor {
kept: Kept,
at: usize,
count: usize,
idle: u64,
touched: u64,
first: bool,
buffer: usize,
pulled: usize,
skip: usize,
given: usize,
}
impl Cursor {
fn counted(&mut self, deep: bool) -> usize {
if self.kept.whole {
return match self.first {
true => self.kept.total,
false => 0,
};
}
let mut walked = 0;
let mut dropped = 0;
let mut delivered = 0;
let mut header = None;
while delivered < self.count && self.given < self.kept.window {
if self.buffer == 0 {
if self.pulled >= self.kept.walk.len() {
break;
}
let fill = self.fill(delivered);
self.buffer = fill;
walked += fill;
}
self.buffer -= 1;
let alive = self.kept.walk[self.pulled];
self.pulled += 1;
if !alive {
dropped += 1;
continue;
}
if self.skip > 0 {
self.skip -= 1;
continue;
}
delivered += 1;
self.given += 1;
if !deep && header.is_none() {
header = Some(walked.saturating_sub(dropped));
}
}
header.unwrap_or(walked.saturating_sub(dropped))
}
fn fill(&self, delivered: usize) -> usize {
let left = self.kept.walk.len().saturating_sub(self.pulled);
if !self.kept.loader {
return 1;
}
let mut want = self.count.saturating_sub(delivered);
if self.kept.offset > 0 {
want = want.min(self.kept.window);
}
self.skip.saturating_add(want).min(left).max(1)
}
fn chunk(&mut self, out: &mut Out) -> bool {
let from = self.at;
let to = (from + self.count).min(self.kept.made.len());
let count = self.counted(out.proto().is_resp3());
let shows = self.kept.shows;
let window = from..to;
match &self.kept.made {
Made::Found(rows) => {
let built: Vec<Built<'_>> = rows[window]
.iter()
.map(|(row, fields)| (row, fields.as_ref().map(|fields| pairs(fields))))
.collect();
found(count, &built, shows, out);
}
Made::Rolled(rows) => {
let built: Vec<Rolled<'_>> = rows[window]
.iter()
.map(|(row, props)| (row, pairs(props)))
.collect();
rolls(count, &built, shows, out);
}
Made::Piped {
names,
rows,
sorted,
warning,
} => {
let shown: Vec<(Option<&Row>, &Vec<Value>)> = rows[window]
.iter()
.map(|(row, values)| (row.as_ref(), values))
.collect();
writes(
count,
names,
&shown,
*sorted,
shows,
warning.as_deref(),
out,
);
}
}
self.at = to;
self.first = false;
to - from == self.count
}
}
fn pairs(held: &Pairs) -> Vec<(&[u8], &[u8])> {
held.iter()
.map(|(name, value)| (&**name, &**value))
.collect()
}
#[derive(Default)]
pub(in crate::dispatch) struct Cursors {
held: HashMap<u64, Held>,
seed: u64,
}
struct Held {
index: Box<[u8]>,
cursor: Cursor,
}
impl Cursors {
fn sweep(&mut self, now: u64) {
self.held
.retain(|_, held| now.saturating_sub(held.cursor.touched) <= held.cursor.idle);
}
fn on(&self, index: &[u8]) -> usize {
self.held.iter().filter(|(_, h)| *h.index == *index).count()
}
pub(super) fn total(&self) -> usize {
self.held.len()
}
fn mint(&mut self, now: u64) -> u64 {
if self.seed == 0 {
self.seed = now | 1;
}
loop {
self.seed ^= self.seed << 13;
self.seed ^= self.seed >> 7;
self.seed ^= self.seed << 17;
let id = self.seed & 0xffff_ffff;
if id != 0 && !self.held.contains_key(&id) {
return id;
}
}
}
}
pub(super) fn stats(server: &Server, index: &[u8]) -> (u64, u64) {
let now = server.clock.now_ms();
let mut cursors = server.cursors.lock();
cursors.sweep(now);
(cursors.total() as u64, cursors.on(index) as u64)
}
pub(super) fn room<'a>(server: &Server, index: &[u8]) -> core::result::Result<(), Fail<'a>> {
match server.cursors.lock().on(index) < LIMIT {
true => Ok(()),
false => Err(Fail::plain(OVER)),
}
}
pub(super) fn open(server: &Server, index: &[u8], kept: Kept, asks: Asks, out: &mut Out) {
let now = server.clock.now_ms();
let mut cursor = Cursor {
skip: kept.offset,
kept,
at: 0,
count: asks.count,
idle: asks.idle,
touched: now,
first: true,
buffer: 0,
pulled: 0,
given: 0,
};
out.array(2);
if !cursor.chunk(out) {
out.int(0);
return;
}
let mut cursors = server.cursors.lock();
cursors.sweep(now);
let id = cursors.mint(now);
cursors.held.insert(
id,
Held {
index: index.into(),
cursor,
},
);
out.int(id as i64);
}
pub(in crate::dispatch) fn execute(server: &Server, args: Args<'_>, out: &mut Out) -> Result<()> {
let sub = args.get(1);
if args::is(sub, b"READ") {
return read(server, args, out);
}
if args::is(sub, b"DEL") {
return del(server, args, out);
}
if args::is(sub, b"GC") {
return gc(server, args, out);
}
Err(args::unknown_subcommand(sub, "FT.CURSOR"))
}
fn read(server: &Server, args: Args<'_>, out: &mut Out) -> Result<()> {
if args.len() < 4 {
return Err(args::wrong_arity_sub("FT.CURSOR", "READ"));
}
let Some(id) = names(server, args, out)? else {
return Ok(());
};
let mut count = None;
if let Some(word) = args.opt(4)
&& args::is(word, b"COUNT")
&& let Some(value) = args.opt(5)
{
let Some(asked) = parse_i64(value) else {
out.error_about(BAD_COUNT.as_bytes(), value, QUOTE.as_bytes());
return Ok(());
};
if asked > 0 {
count = usize::try_from(asked).ok();
}
}
let name = args.get(2);
let now = server.clock.now_ms();
let mut cursors = server.cursors.lock();
cursors.sweep(now);
let key = u64::try_from(id).unwrap_or(0);
let Some(held) = cursors
.held
.get_mut(&key)
.filter(|held| *held.index == *name)
else {
return missing(id, out);
};
if let Some(count) = count {
held.cursor.count = count;
}
held.cursor.touched = now;
out.array(2);
match held.cursor.chunk(out) {
true => out.int(id),
false => {
cursors.held.remove(&key);
out.int(0);
}
}
Ok(())
}
fn del(server: &Server, args: Args<'_>, out: &mut Out) -> Result<()> {
if args.len() < 4 {
return Err(args::wrong_arity_sub("FT.CURSOR", "DEL"));
}
let Some(id) = names(server, args, out)? else {
return Ok(());
};
let key = u64::try_from(id).unwrap_or(0);
match server.cursors.lock().held.remove(&key) {
Some(_) => out.ok(),
None => out.error(NO_CURSOR.as_bytes()),
}
Ok(())
}
fn gc(server: &Server, args: Args<'_>, out: &mut Out) -> Result<()> {
if args.len() < 4 {
return Err(args::wrong_arity_sub("FT.CURSOR", "GC"));
}
let name = args.get(2);
if server.search.lock().named(name).is_none() {
Fail::naming(MISSING, name).write(out);
return Ok(());
}
let now = server.clock.now_ms();
server.cursors.lock().sweep(now);
out.int(0);
Ok(())
}
fn names(server: &Server, args: Args<'_>, out: &mut Out) -> Result<Option<i64>> {
let name = args.get(2);
if server.search.lock().named(name).is_none() {
Fail::naming(MISSING, name).write(out);
return Ok(None);
}
let Some(id) = parse_i64(args.get(3)) else {
out.error(BAD_ID.as_bytes());
return Ok(None);
};
Ok(Some(id))
}
fn missing(id: i64, out: &mut Out) -> Result<()> {
let mut line = Vec::with_capacity(NOT_FOUND.len() + 20);
line.extend_from_slice(NOT_FOUND.as_bytes());
line.extend_from_slice(id.to_string().as_bytes());
out.error(&line);
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn over(walk: Vec<bool>, loader: bool, offset: usize, window: usize, count: usize) -> Cursor {
Cursor {
kept: Kept {
made: Made::Rolled(Vec::new()),
walk,
shows: Shows::default(),
total: 7,
whole: false,
loader,
offset,
window,
},
at: 0,
count,
idle: IDLE,
touched: 0,
first: true,
buffer: 0,
pulled: 0,
skip: offset,
given: 0,
}
}
fn heads(mut cursor: Cursor, deep: bool, chunks: usize) -> Vec<usize> {
(0..chunks)
.map(|_| {
let count = cursor.counted(deep);
cursor.first = false;
count
})
.collect()
}
#[test]
fn a_pipeline_that_reads_keys_pulls_a_chunk_of_them_at_a_time() {
let cursor = over(vec![true; 7], true, 0, usize::MAX, 2);
assert_eq!(heads(cursor, false, 4), vec![2, 2, 2, 1]);
}
#[test]
fn a_pipeline_that_reads_nothing_pulls_one_document_at_a_time() {
let cursor = over(vec![true; 7], false, 0, usize::MAX, 2);
assert_eq!(heads(cursor, false, 4), vec![1, 1, 1, 1]);
}
#[test]
fn a_window_with_an_offset_holds_the_pull_back_and_one_without_does_not() {
assert_eq!(
heads(over(vec![true; 7], true, 0, 3, 2), false, 2),
vec![2, 2]
);
assert_eq!(
heads(over(vec![true; 7], true, 1, 3, 2), false, 2),
vec![3, 2]
);
assert_eq!(heads(over(vec![true; 7], true, 2, 2, 3), false, 1), vec![4]);
}
#[test]
fn a_filter_takes_its_rows_off_the_count_and_the_protocol_says_when() {
let walk = vec![false, true, true, false, true, true, false];
let shallow = over(walk.clone(), true, 0, usize::MAX, 2);
assert_eq!(heads(shallow, false, 3), vec![1, 1, 0]);
let deep = over(walk, true, 0, usize::MAX, 2);
assert_eq!(heads(deep, true, 3), vec![2, 2, 0]);
}
#[test]
fn a_settled_answer_says_the_total_once_and_nought_after_that() {
let mut cursor = over(vec![true; 7], true, 0, usize::MAX, 2);
cursor.kept.whole = true;
assert_eq!(heads(cursor, false, 3), vec![7, 0, 0]);
}
#[test]
fn two_numbers_out_of_the_mint_are_not_the_same_number() {
let mut cursors = Cursors::default();
let one = cursors.mint(12);
let two = cursors.mint(12);
assert_ne!(one, two);
assert!(one > 0 && one <= u64::from(u32::MAX));
assert!(two > 0 && two <= u64::from(u32::MAX));
}
#[test]
fn a_cursor_nobody_has_read_from_for_long_enough_is_swept_away() {
let mut cursors = Cursors::default();
cursors.held.insert(
7,
Held {
index: b"i".as_slice().into(),
cursor: Cursor {
kept: Kept {
made: Made::Rolled(Vec::new()),
walk: Vec::new(),
shows: Shows::default(),
total: 0,
whole: true,
loader: false,
offset: 0,
window: 0,
},
at: 0,
count: 1,
idle: 100,
touched: 1_000,
first: true,
buffer: 0,
pulled: 0,
skip: 0,
given: 0,
},
},
);
cursors.sweep(1_100);
assert_eq!(cursors.total(), 1);
cursors.sweep(1_101);
assert_eq!(cursors.total(), 0);
}
}