use reifydb_codec::{key::encoded::EncodedKey, row::bytes::EncodedBytes};
use reifydb_value::util::cowvec::CowVec;
use reifydb_wal::{
body::Body,
error::{Result, WalError},
};
use crate::{
common::CommitVersion,
delta::{Delta, RemoveAnnounce, SetPre},
key::any::TaggedKey,
};
const OP_SET: u8 = 0;
const OP_REMOVE: u8 = 1;
const PRE_ABSENT: u8 = 0;
const PRE_PRESENT: u8 = 1;
const ANNOUNCE_SILENT: u8 = 0;
const ANNOUNCE_ANNOUNCED: u8 = 1;
const ANNOUNCE_UNOBSERVED: u8 = 2;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CommitRecord {
pub version: CommitVersion,
pub deltas: Vec<Delta>,
}
impl CommitRecord {
pub const KIND: u32 = 1;
pub fn stamp_version(encoded: &mut [u8], version: CommitVersion) {
encoded[..8].copy_from_slice(&version.0.to_le_bytes());
}
pub fn encode_refs(version: CommitVersion, deltas: &[&Delta], out: &mut Vec<u8>) {
out.extend_from_slice(&version.0.to_le_bytes());
put_len(out, deltas.len());
for delta in deltas {
match delta {
Delta::Set {
key,
bytes,
pre,
} => {
out.push(OP_SET);
put_bytes(out, key.encode().as_bytes());
put_bytes(out, bytes.as_ref());
match pre {
SetPre::Absent => out.push(PRE_ABSENT),
SetPre::Present(pre) => {
out.push(PRE_PRESENT);
put_bytes(out, pre.as_ref());
}
SetPre::Unknown => panic!(
"commit {} has an unresolved pre for key {:?}",
version.0, key
),
}
}
Delta::Remove {
key,
announce,
} => {
out.push(OP_REMOVE);
put_bytes(out, key.encode().as_bytes());
match announce {
RemoveAnnounce::Silent => out.push(ANNOUNCE_SILENT),
RemoveAnnounce::Announced {
pre,
} => {
out.push(ANNOUNCE_ANNOUNCED);
put_bytes(out, pre.as_ref());
}
RemoveAnnounce::Unobserved {
pre,
} => {
out.push(ANNOUNCE_UNOBSERVED);
put_bytes(out, pre.as_ref());
}
}
}
}
}
}
}
impl Body for CommitRecord {
fn kind(&self) -> u32 {
Self::KIND
}
fn encode(&self, out: &mut Vec<u8>) {
Self::encode_refs(self.version, &self.deltas.iter().collect::<Vec<_>>(), out);
}
fn decode(kind: u32, bytes: &[u8]) -> Result<Self> {
if kind != Self::KIND {
return Err(WalError::Body(format!("unknown commit record kind {kind}")));
}
let mut reader = Reader {
bytes,
at: 0,
};
let version = CommitVersion(reader.u64()?);
let count = reader.u32()?;
let mut deltas = Vec::new();
for _ in 0..count {
deltas.push(reader.delta()?);
}
if reader.at != bytes.len() {
return Err(WalError::Body(format!(
"{} trailing bytes after delta {count}",
bytes.len() - reader.at
)));
}
Ok(Self {
version,
deltas,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SingleRecord {
pub deltas: Vec<Delta>,
}
impl SingleRecord {
pub const KIND: u32 = 2;
pub fn encode_deltas(deltas: &[Delta], out: &mut Vec<u8>) {
put_len(out, deltas.len());
for delta in deltas {
match delta {
Delta::Set {
key,
bytes,
..
} => {
out.push(OP_SET);
put_bytes(out, key.encode().as_bytes());
put_bytes(out, bytes.as_ref());
}
Delta::Remove {
key,
..
} => {
out.push(OP_REMOVE);
put_bytes(out, key.encode().as_bytes());
}
}
}
}
}
impl Body for SingleRecord {
fn kind(&self) -> u32 {
Self::KIND
}
fn encode(&self, out: &mut Vec<u8>) {
Self::encode_deltas(&self.deltas, out);
}
fn decode(kind: u32, bytes: &[u8]) -> Result<Self> {
if kind != Self::KIND {
return Err(WalError::Body(format!("unknown single record kind {kind}")));
}
let mut reader = Reader {
bytes,
at: 0,
};
let count = reader.u32()?;
let mut deltas = Vec::new();
for _ in 0..count {
deltas.push(reader.single_delta()?);
}
if reader.at != bytes.len() {
return Err(WalError::Body(format!(
"{} trailing bytes after single delta {count}",
bytes.len() - reader.at
)));
}
Ok(Self {
deltas,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WalRecord {
Multi(CommitRecord),
Single(SingleRecord),
}
impl Body for WalRecord {
fn kind(&self) -> u32 {
match self {
Self::Multi(record) => record.kind(),
Self::Single(record) => record.kind(),
}
}
fn encode(&self, out: &mut Vec<u8>) {
match self {
Self::Multi(record) => record.encode(out),
Self::Single(record) => record.encode(out),
}
}
fn decode(kind: u32, bytes: &[u8]) -> Result<Self> {
match kind {
CommitRecord::KIND => Ok(Self::Multi(CommitRecord::decode(kind, bytes)?)),
SingleRecord::KIND => Ok(Self::Single(SingleRecord::decode(kind, bytes)?)),
kind => Err(WalError::Body(format!("unknown wal record kind {kind}"))),
}
}
}
fn put_len(out: &mut Vec<u8>, len: usize) {
let len = u32::try_from(len).unwrap_or_else(|_| panic!("commit record length {len} does not fit in u32"));
out.extend_from_slice(&len.to_le_bytes());
}
fn put_bytes(out: &mut Vec<u8>, bytes: &[u8]) {
put_len(out, bytes.len());
out.extend_from_slice(bytes);
}
struct Reader<'a> {
bytes: &'a [u8],
at: usize,
}
impl<'a> Reader<'a> {
fn take(&mut self, len: usize) -> Result<&'a [u8]> {
let end = self.at.checked_add(len).filter(|end| *end <= self.bytes.len()).ok_or_else(|| {
WalError::Body(format!(
"truncated at byte {}: needs {len} of {} remaining",
self.at,
self.bytes.len() - self.at
))
})?;
let slice = &self.bytes[self.at..end];
self.at = end;
Ok(slice)
}
fn u8(&mut self) -> Result<u8> {
Ok(self.take(1)?[0])
}
fn u32(&mut self) -> Result<u32> {
let mut buf = [0u8; 4];
buf.copy_from_slice(self.take(4)?);
Ok(u32::from_le_bytes(buf))
}
fn u64(&mut self) -> Result<u64> {
let mut buf = [0u8; 8];
buf.copy_from_slice(self.take(8)?);
Ok(u64::from_le_bytes(buf))
}
fn bytes(&mut self) -> Result<&'a [u8]> {
let len = self.u32()? as usize;
self.take(len)
}
fn encoded(&mut self) -> Result<EncodedBytes> {
Ok(EncodedBytes(CowVec::new(self.bytes()?.to_vec())))
}
fn key(&mut self) -> Result<TaggedKey> {
let at = self.at;
let key = EncodedKey::new(self.bytes()?);
TaggedKey::decode(&key).ok_or_else(|| WalError::Body(format!("undecodable key at byte {at}")))
}
fn delta(&mut self) -> Result<Delta> {
let at = self.at;
match self.u8()? {
OP_SET => {
let key = self.key()?;
let bytes = self.encoded()?;
let pre_at = self.at;
let pre = match self.u8()? {
PRE_ABSENT => SetPre::Absent,
PRE_PRESENT => SetPre::Present(self.encoded()?),
tag => {
return Err(WalError::Body(format!(
"unknown pre tag {tag} at byte {pre_at}"
)));
}
};
Ok(Delta::Set {
key,
bytes,
pre,
})
}
OP_REMOVE => {
let key = self.key()?;
let announce_at = self.at;
let announce = match self.u8()? {
ANNOUNCE_SILENT => RemoveAnnounce::Silent,
ANNOUNCE_ANNOUNCED => RemoveAnnounce::Announced {
pre: self.encoded()?,
},
ANNOUNCE_UNOBSERVED => RemoveAnnounce::Unobserved {
pre: self.encoded()?,
},
tag => {
return Err(WalError::Body(format!(
"unknown announce tag {tag} at byte {announce_at}"
)));
}
};
Ok(Delta::Remove {
key,
announce,
})
}
op => Err(WalError::Body(format!("unknown delta op {op} at byte {at}"))),
}
}
fn single_delta(&mut self) -> Result<Delta> {
let at = self.at;
match self.u8()? {
OP_SET => {
let key = self.key()?;
let bytes = self.encoded()?;
Ok(Delta::Set {
key,
bytes,
pre: SetPre::Unknown,
})
}
OP_REMOVE => Ok(Delta::remove_silent(self.key()?)),
op => Err(WalError::Body(format!("unknown single delta op {op} at byte {at}"))),
}
}
}