use std::borrow::Cow;
use std::fmt;
use yo_common::crc::crc64;
use crate::keys::Record;
use crate::rdb::{
self, Bad, Limits, OP_AUX, OP_EOF, OP_EXPIRETIME, OP_EXPIRETIME_MS, OP_FREQ,
OP_FUNCTION_PRE_GA, OP_FUNCTION2, OP_IDLE, OP_MODULE_AUX, OP_RESIZEDB, OP_SELECTDB,
OP_SLOT_INFO, Reader,
};
const MAGIC: &[u8] = b"REDIS";
const HEADER: usize = 9;
const FOOTER: usize = 8;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Fault {
Magic,
Version(u16),
Checksum,
Truncated,
Opcode(u8),
Value,
}
impl fmt::Display for Fault {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Fault::Magic => f.write_str("the file does not start with REDIS"),
Fault::Version(v) => write!(
f,
"the file is version {v} and this server reads up to {}",
rdb::READS_UP_TO
),
Fault::Checksum => f.write_str("the checksum does not match, so the file is damaged"),
Fault::Truncated => f.write_str("the file ends in the middle of something"),
Fault::Opcode(op) => write!(
f,
"the file holds opcode {op}, which needs whatever wrote it"
),
Fault::Value => f.write_str("a value in the file is not one this server can hold"),
}
}
}
#[derive(Debug)]
#[allow(clippy::large_enum_variant)]
pub enum Item<'a> {
Aux {
name: Cow<'a, [u8]>,
value: Cow<'a, [u8]>,
},
Library(Cow<'a, [u8]>),
Key {
db: usize,
key: Cow<'a, [u8]>,
record: Record,
},
}
pub struct Load<'a> {
r: Reader<'a>,
limits: Limits<'a>,
now: u64,
version: u16,
db: usize,
expire_at: Option<u64>,
expired: usize,
done: bool,
}
impl<'a> Load<'a> {
pub fn open(file: &'a [u8], limits: Limits<'a>, now: u64) -> Result<Load<'a>, Fault> {
if file.len() < HEADER + FOOTER {
return Err(Fault::Truncated);
}
if &file[..MAGIC.len()] != MAGIC {
return Err(Fault::Magic);
}
let mut version = 0u16;
for &d in &file[MAGIC.len()..HEADER] {
if !d.is_ascii_digit() {
return Err(Fault::Magic);
}
version = version * 10 + u16::from(d - b'0');
}
if version > rdb::READS_UP_TO {
return Err(Fault::Version(version));
}
let body = &file[HEADER..file.len() - FOOTER];
let stored = u64::from_le_bytes(
file[file.len() - FOOTER..]
.try_into()
.expect("eight bytes of checksum"),
);
if stored != 0 && stored != crc64(0, &file[..file.len() - FOOTER]) && !rdb::skipping() {
return Err(Fault::Checksum);
}
Ok(Load {
r: Reader::new(body),
limits,
now,
version,
db: 0,
expire_at: None,
expired: 0,
done: false,
})
}
#[must_use]
pub const fn version(&self) -> u16 {
self.version
}
#[must_use]
pub const fn expired(&self) -> usize {
self.expired
}
fn step(&mut self) -> Result<Option<Item<'a>>, Fault> {
loop {
if self.done {
return Ok(None);
}
let op = self.r.byte().map_err(|_| Fault::Truncated)?;
match op {
OP_EOF => {
self.done = true;
if !self.r.done() {
return Err(Fault::Truncated);
}
return Ok(None);
}
OP_SELECTDB => self.db = self.r.len().map_err(|_| Fault::Truncated)?,
OP_RESIZEDB => {
self.r.len().map_err(|_| Fault::Truncated)?;
self.r.len().map_err(|_| Fault::Truncated)?;
}
OP_EXPIRETIME_MS => {
let b = self.r.take(8).map_err(|_| Fault::Truncated)?;
self.expire_at = Some(u64::from_le_bytes(b.try_into().expect("eight bytes")));
}
OP_EXPIRETIME => {
let b = self.r.take(4).map_err(|_| Fault::Truncated)?;
let secs = u32::from_le_bytes(b.try_into().expect("four bytes"));
self.expire_at = Some(u64::from(secs) * 1_000);
}
OP_IDLE => {
self.r.num().map_err(|_| Fault::Truncated)?;
}
OP_FREQ => {
self.r.byte().map_err(|_| Fault::Truncated)?;
}
OP_AUX => {
let name = self.r.str().map_err(|_| Fault::Truncated)?;
let value = self.r.str().map_err(|_| Fault::Truncated)?;
return Ok(Some(Item::Aux { name, value }));
}
OP_FUNCTION2 => {
let code = self.r.str().map_err(|_| Fault::Truncated)?;
return Ok(Some(Item::Library(code)));
}
OP_SLOT_INFO => {
for _ in 0..3 {
self.r.len().map_err(|_| Fault::Truncated)?;
}
}
OP_FUNCTION_PRE_GA | OP_MODULE_AUX => return Err(Fault::Opcode(op)),
kind => {
let key = self.r.str().map_err(|_| Fault::Truncated)?;
let expire_at = self.expire_at.take();
let body = rdb::read_object(&mut self.r, kind, self.limits, self.now).map_err(
|e| match e {
Bad::Footer => Fault::Checksum,
Bad::Format => Fault::Value,
},
)?;
if expire_at.is_some_and(|at| at <= self.now) {
self.expired += 1;
continue;
}
return Ok(Some(Item::Key {
db: self.db,
key,
record: Record::new(body, expire_at),
}));
}
}
}
}
}
impl<'a> Iterator for Load<'a> {
type Item = Result<Item<'a>, Fault>;
fn next(&mut self) -> Option<Result<Item<'a>, Fault>> {
match self.step() {
Ok(None) => None,
Ok(Some(item)) => Some(Ok(item)),
Err(fault) => {
self.done = true;
Some(Err(fault))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::clock::Clock;
use crate::db::Db;
use crate::hash;
use crate::list;
use crate::lists::End;
use crate::set;
use crate::snapshot::Snapshot;
use crate::streams::{Add, Trim};
use crate::ttl::Ask;
use crate::ttl::Cond;
use crate::value::Kind;
use crate::zset;
use crate::zsets::ZAdd;
fn bands() -> rdb::Bands {
rdb::Bands {
set: set::Limits::DEFAULT,
hash: hash::Limits::DEFAULT,
list: list::Limits::default(),
zset: zset::Limits::DEFAULT,
}
}
fn db() -> Db {
Db::with_clock(Clock::fixed(1_000), 1)
}
fn filled() -> Db {
let mut d = db();
d.at(b"str").set_plain(b"str", b"hello").unwrap();
d.at(b"long")
.set_plain(b"long", &b"ab".repeat(500))
.unwrap();
d.at(b"list")
.push(b"list", End::Right, [b"a".as_slice(), b"b"].into_iter())
.unwrap();
d.at(b"set")
.sadd(b"set", [b"x".as_slice(), b"y"].into_iter())
.unwrap();
d.at(b"ints")
.sadd(b"ints", [b"1".as_slice(), b"2"].into_iter())
.unwrap();
d.at(b"zset")
.zadd(
b"zset",
[(1.5, b"m".as_slice())].into_iter(),
ZAdd::default(),
)
.unwrap();
d.at(b"hash")
.hset(b"hash", [(b"f".as_slice(), b"v".as_slice())].into_iter())
.unwrap();
d.at(b"stream")
.xadd(
b"stream",
Add::Auto,
&[(b"f".as_slice(), b"v".as_slice())],
Trim::None,
true,
1_000,
)
.unwrap();
d
}
fn into_dbs(file: &[u8], now: u64) -> Vec<(usize, Db)> {
let bands = bands();
let mut out: Vec<(usize, Db)> = Vec::new();
for item in Load::open(file, bands.limits(), now).unwrap() {
if let Item::Key { db: n, key, record } = item.unwrap() {
if !out.iter().any(|(i, _)| *i == n) {
out.push((n, db()));
}
let at = out.iter_mut().find(|(i, _)| *i == n).unwrap();
at.1.at(&key).import(&key, record);
}
}
out
}
#[test]
fn every_type_comes_back_the_way_it_went_in() {
let mut before = filled();
let mut snap = Snapshot::new();
snap.database(0, &before);
let file = snap.finish();
let mut after = into_dbs(&file, 1_000);
assert_eq!(after.len(), 1, "one database");
let (index, ref mut back) = after[0];
assert_eq!(index, 0);
let keys: Vec<&[u8]> = vec![
b"str", b"long", b"list", b"set", b"ints", b"zset", b"hash", b"stream",
];
for key in keys {
assert_eq!(
back.at(key).dump(key),
before.at(key).dump(key),
"{}",
String::from_utf8_lossy(key)
);
}
}
#[test]
fn a_key_lands_in_the_database_the_file_put_it_in() {
let mut zero = db();
let mut nine = db();
zero.at(b"a").set_plain(b"a", b"1").unwrap();
nine.at(b"b").set_plain(b"b", b"2").unwrap();
let mut snap = Snapshot::new();
snap.database(0, &zero);
snap.database(9, &nine);
let mut after = into_dbs(&snap.finish(), 1_000);
assert_eq!(
after.iter().map(|(i, _)| *i).collect::<Vec<_>>(),
vec![0, 9]
);
assert!(after[0].1.at(b"a").exists(b"a"));
assert!(after[1].1.at(b"b").exists(b"b"));
}
#[test]
fn a_deadline_travels_with_the_key_and_a_dead_one_is_dropped() {
let mut d = db();
for key in [&b"alive"[..], b"soon", b"gone"] {
d.at(key).set_plain(key, b"1").unwrap();
}
d.at(b"soon").expire(b"soon", 9_000, Cond::Always);
d.at(b"gone").expire(b"gone", 5_000, Cond::Always);
let mut snap = Snapshot::new();
snap.database(0, &d);
let file = snap.finish();
let bands = bands();
let mut load = Load::open(&file, bands.limits(), 7_000).unwrap();
let mut found: Vec<(Vec<u8>, Option<u64>)> = Vec::new();
for item in &mut load {
if let Item::Key { key, record, .. } = item.unwrap() {
found.push((key.into_owned(), record.expire_at()));
}
}
found.sort();
assert_eq!(
found,
vec![(b"alive".to_vec(), None), (b"soon".to_vec(), Some(9_000))],
);
assert_eq!(load.expired(), 1, "the one that had gone was counted");
}
#[test]
fn the_aux_fields_come_out_in_the_order_they_were_written() {
let mut snap = Snapshot::new();
snap.aux(b"redis-ver", b"8.8.0");
snap.aux(b"aof-base", b"0");
let file = snap.finish();
let bands = bands();
let seen: Vec<(Vec<u8>, Vec<u8>)> = Load::open(&file, bands.limits(), 0)
.unwrap()
.map(|item| match item.unwrap() {
Item::Aux { name, value } => (name.into_owned(), value.into_owned()),
other => panic!("{other:?}"),
})
.collect();
assert_eq!(seen.len(), 3, "ours and the two the test wrote");
assert_eq!(seen[0].0, b"yo-ver");
assert_eq!(seen[1], (b"redis-ver".to_vec(), b"8.8.0".to_vec()));
assert_eq!(seen[2], (b"aof-base".to_vec(), b"0".to_vec()));
}
#[test]
fn a_file_with_more_keys_than_one_batch_comes_back_whole() {
let mut d = db();
let many = 1_000;
for i in 0..many {
let key = format!("k{i}");
d.at(key.as_bytes())
.set_plain(key.as_bytes(), b"v")
.unwrap();
}
let mut snap = Snapshot::new();
snap.database(0, &d);
let after = into_dbs(&snap.finish(), 1_000);
assert_eq!(after[0].1.len(), many);
}
#[test]
fn a_newer_header_is_read_up_to_the_version_and_refused_past_it() {
let mut snap = Snapshot::new();
snap.database(0, &filled());
let file = snap.finish();
let bands = bands();
let stamped = |version: &[u8]| {
let mut edited = file.clone();
edited[5..9].copy_from_slice(version);
let end = edited.len() - FOOTER;
let crc = crc64(0, &edited[..end]);
edited[end..].copy_from_slice(&crc.to_le_bytes());
edited
};
for version in [&b"0006"[..], b"0011", b"0014", b"0015"] {
let edited = stamped(version);
let load = Load::open(&edited, bands.limits(), 1_000).unwrap();
assert_eq!(
load.filter(|item| matches!(item, Ok(Item::Key { .. })))
.count(),
8,
"{version:?}"
);
}
assert_eq!(
Load::open(&stamped(b"0016"), bands.limits(), 1_000).err(),
Some(Fault::Version(16)),
"the version is looked at before the checksum, which is right"
);
}
#[test]
fn a_file_that_is_not_one_is_turned_down_before_anything_is_built() {
let bands = bands();
let short = b"REDIS0012";
assert_eq!(
Load::open(short, bands.limits(), 0).err(),
Some(Fault::Truncated),
"a header with no room for a checksum"
);
let wrong = b"NOTAREDISFILE0000000";
assert_eq!(
Load::open(wrong, bands.limits(), 0).err(),
Some(Fault::Magic)
);
let letters = b"REDISabcd0000000000";
assert_eq!(
Load::open(letters, bands.limits(), 0).err(),
Some(Fault::Magic),
"a version that is not four digits"
);
}
#[test]
fn a_damaged_file_is_refused_rather_than_half_loaded() {
let mut snap = Snapshot::new();
snap.database(0, &filled());
let mut file = snap.finish();
let middle = file.len() / 2;
file[middle] ^= 0x01;
let bands = bands();
assert_eq!(
Load::open(&file, bands.limits(), 1_000).err(),
Some(Fault::Checksum)
);
}
#[test]
fn a_file_cut_short_is_a_fault_and_not_a_shorter_dataset() {
let mut snap = Snapshot::new();
snap.database(0, &filled());
let file = snap.finish();
let mut cut = file[..file.len() - 40].to_vec();
let crc = crc64(0, &cut);
cut.extend_from_slice(&crc.to_le_bytes());
let bands = bands();
let load = Load::open(&cut, bands.limits(), 1_000).unwrap();
let faults: Vec<Fault> = load.filter_map(|item| item.err()).collect();
assert_eq!(faults.len(), 1, "the walk stops at the first fault");
assert!(
matches!(faults[0], Fault::Truncated | Fault::Value),
"{:?}",
faults[0]
);
}
#[test]
fn the_opcodes_that_are_bookkeeping_are_read_and_the_rest_are_refused() {
let mut snap = Snapshot::new();
snap.database(0, &filled());
let good = snap.finish();
let mut edited = Vec::new();
edited.extend_from_slice(&good[..HEADER]);
edited.push(OP_IDLE);
edited.push(40); edited.push(OP_FREQ);
edited.push(255);
edited.push(OP_SLOT_INFO);
edited.extend_from_slice(&[7, 3, 0]);
edited.extend_from_slice(&good[HEADER..good.len() - FOOTER]);
let crc = crc64(0, &edited);
edited.extend_from_slice(&crc.to_le_bytes());
let bands = bands();
let keys = Load::open(&edited, bands.limits(), 1_000)
.unwrap()
.filter(|item| matches!(item, Ok(Item::Key { .. })))
.count();
assert_eq!(keys, 8, "every key still arrived");
for op in [OP_FUNCTION_PRE_GA, OP_MODULE_AUX] {
let mut edited = Vec::new();
edited.extend_from_slice(&good[..HEADER]);
edited.push(op);
edited.extend_from_slice(&good[HEADER..good.len() - FOOTER]);
let crc = crc64(0, &edited);
edited.extend_from_slice(&crc.to_le_bytes());
let faults: Vec<Fault> = Load::open(&edited, bands.limits(), 1_000)
.unwrap()
.filter_map(|item| item.err())
.collect();
assert_eq!(faults, vec![Fault::Opcode(op)]);
}
}
#[test]
fn a_file_that_says_it_has_no_checksum_is_read_anyway() {
let mut snap = Snapshot::new();
snap.database(0, &filled());
let mut file = snap.finish();
let end = file.len() - FOOTER;
file[end..].fill(0);
let bands = bands();
let keys = Load::open(&file, bands.limits(), 1_000)
.unwrap()
.filter(|item| matches!(item, Ok(Item::Key { .. })))
.count();
assert_eq!(keys, 8);
}
const REAL: &str = concat!(
"524544495330303135fa0972656469732d76657206382e31302e31fa0a72656469732d62",
"697473c040fa056374696d65c248a99f6afa08757365642d6d656dc2d0941f00fa08616f",
"662d62617365c000fe00fb0a010001730568656c6c6f1001680d0d000000020081660281",
"7602ff1902687800d8c32cbb0300001d1d00000006008161020101f400d8c32cbb030000",
"0981620202010001ff0b04696e74730e020000000300000001000200030012016c010210",
"100000000300816102816202816302ff11017a0f0f0000000200816d0283312e3504ff14",
"0273740d0d0000000200817802817902fffc00d8c32cbb030000000374746c017600016e",
"c287d612001b01780110000000000000000100000000000000011d1d0000000a00010100",
"01010181660200010201000100018176020401ff01010101010000010101670101010100",
"0000000000000100000000000000017b44ad7fa0010000010101637b44ad7fa00100007b",
"44ad7fa001000001000000000000000100000000000000010040644064000000ffa2531e",
"bc6edefb1f",
);
fn unhex(s: &str) -> Vec<u8> {
(0..s.len() / 2)
.map(|i| u8::from_str_radix(&s[i * 2..i * 2 + 2], 16).unwrap())
.collect()
}
#[test]
fn a_file_a_real_redis_wrote_comes_back_whole() {
let file = unhex(REAL);
let bands = bands();
let mut load = Load::open(&file, bands.limits(), 1_000).unwrap();
assert_eq!(load.version(), 15, "the header 8.10.1 writes");
let mut d = db();
let mut aux = Vec::new();
let mut kinds: Vec<(Vec<u8>, Kind)> = Vec::new();
let mut deadline = None;
for item in &mut load {
match item.unwrap() {
Item::Aux { name, value } => aux.push((name.into_owned(), value.into_owned())),
Item::Library(_) => panic!("the file holds no functions"),
Item::Key { db: n, key, record } => {
assert_eq!(n, 0);
kinds.push((key.to_vec(), record.kind()));
if &*key == b"ttl" {
deadline = record.expire_at();
}
d.at(&key).import(&key, record);
}
}
}
assert_eq!(load.expired(), 0);
assert!(
aux.contains(&(b"redis-ver".to_vec(), b"8.10.1".to_vec())),
"{aux:?}"
);
kinds.sort_by(|a, b| a.0.cmp(&b.0));
assert_eq!(
kinds,
vec![
(b"h".to_vec(), Kind::Hash),
(b"hx".to_vec(), Kind::Hash),
(b"ints".to_vec(), Kind::Set),
(b"l".to_vec(), Kind::List),
(b"n".to_vec(), Kind::String),
(b"s".to_vec(), Kind::String),
(b"st".to_vec(), Kind::Set),
(b"ttl".to_vec(), Kind::String),
(b"x".to_vec(), Kind::Stream),
(b"z".to_vec(), Kind::Zset),
]
);
assert_eq!(deadline, Some(4_102_444_800_000));
let mut text = Vec::new();
d.at(b"s").get(b"s").unwrap().unwrap().write_to(&mut text);
assert_eq!(text, b"hello");
text.clear();
d.at(b"n").get(b"n").unwrap().unwrap().write_to(&mut text);
assert_eq!(text, b"1234567", "an integer encoded string");
let list: Vec<Vec<u8>> = d
.at(b"l")
.lrange(b"l", 0, -1)
.unwrap()
.map(|e| e.to_vec())
.collect();
assert_eq!(list, vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()]);
let mut members: Vec<Vec<u8>> = d
.at(b"st")
.smembers(b"st")
.unwrap()
.unwrap()
.map(|m| m.to_vec())
.collect();
members.sort();
assert_eq!(members, vec![b"x".to_vec(), b"y".to_vec()]);
assert_eq!(d.at(b"ints").scard(b"ints").unwrap(), 3, "an intset");
assert_eq!(
d.at(b"h")
.hget(b"h", b"f", |v| v.map(|t| t.to_vec()))
.unwrap(),
Some(b"v".to_vec())
);
assert_eq!(d.at(b"z").zscore(b"z", b"m").unwrap(), Some(1.5));
let mut asked = Vec::new();
d.at(b"hx")
.httl(b"hx", [b"a".as_slice(), b"b"].into_iter(), |ask| {
asked.push(ask)
})
.unwrap();
assert_eq!(asked, vec![Ask::At(4_102_444_800_000), Ask::NoDeadline]);
let stream = d.at(b"x").stream(b"x").unwrap().expect("a stream");
assert_eq!(stream.len(), 1);
assert_eq!(stream.groups().count(), 1, "the group came back with it");
}
}