use std::mem::size_of;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum AofHeaderType {
BasicHeader = 0,
ShardedHeader = 1,
SingleLogTransactionHeader = 2,
ShardedLogTransactionHeader = 3,
BasicChunkHeader = 4,
ShardedChunkHeader = 5,
}
impl AofHeaderType {
pub const ALL: [AofHeaderType; 6] = [
Self::BasicHeader,
Self::ShardedHeader,
Self::SingleLogTransactionHeader,
Self::ShardedLogTransactionHeader,
Self::BasicChunkHeader,
Self::ShardedChunkHeader,
];
pub fn total_size(self) -> usize {
match self {
Self::BasicHeader => AofHeader::TOTAL_SIZE,
Self::ShardedHeader => AofShardedHeader::TOTAL_SIZE,
Self::SingleLogTransactionHeader => AofSingleLogTransactionHeader::TOTAL_SIZE,
Self::ShardedLogTransactionHeader => AofShardedLogTransactionHeader::TOTAL_SIZE,
Self::BasicChunkHeader => AofHeader::TOTAL_SIZE + AofChunkHeader::TOTAL_SIZE,
Self::ShardedChunkHeader => AofShardedHeader::TOTAL_SIZE + AofChunkHeader::TOTAL_SIZE,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AofHeader {
pub aof_header_version: u8,
pub flags: u8,
pub op_type: u8,
pub procedure_id: u8,
pub database_id: u8,
pub store_version: i64,
pub session_id: i32,
}
impl AofHeader {
pub const TOTAL_SIZE: usize = 16;
pub const AOF_HEADER_VERSION: u8 = 5;
pub const MAX_SUPPORTED_AOF_HEADER_VERSION: u8 = Self::AOF_HEADER_VERSION;
pub const AOF_HEADER_TYPE_MASK: u8 = 0b0111;
pub const CHUNKED_RECORD_FLAG: u8 = 0b0100;
pub const UNSAFE_TRUNCATE_LOG_FLAG: u8 = 0b1000;
pub fn new() -> Self {
Self {
aof_header_version: Self::AOF_HEADER_VERSION,
flags: 0,
op_type: 0,
procedure_id: 0,
database_id: 0,
store_version: 0,
session_id: 0,
}
}
pub fn unsafe_truncate_log(&self) -> bool {
(self.flags & Self::UNSAFE_TRUNCATE_LOG_FLAG) != 0
}
pub fn set_unsafe_truncate_log(&mut self, value: bool) {
if value {
self.flags |= Self::UNSAFE_TRUNCATE_LOG_FLAG;
} else {
self.flags &= !Self::UNSAFE_TRUNCATE_LOG_FLAG;
}
}
pub fn header_type(&self) -> Option<AofHeaderType> {
let raw = self.flags & Self::AOF_HEADER_TYPE_MASK;
AofHeaderType::ALL.iter().copied().find(|t| *t as u8 == raw)
}
pub fn set_header_type(&mut self, value: AofHeaderType) {
debug_assert!((value as u8) <= Self::AOF_HEADER_TYPE_MASK);
self.flags = (self.flags & !Self::AOF_HEADER_TYPE_MASK) | value as u8;
}
pub fn is_chunked(&self) -> bool {
(self.flags & Self::CHUNKED_RECORD_FLAG) != 0
}
pub fn parse(entry: &[u8]) -> Option<Self> {
if entry.len() < Self::TOTAL_SIZE {
return None;
}
Some(Self {
aof_header_version: entry[0],
flags: entry[1],
op_type: entry[2],
procedure_id: entry[3],
database_id: entry[3],
store_version: i64::from_le_bytes(entry[4..12].try_into().expect("长度恰为 8")),
session_id: i32::from_le_bytes(entry[12..16].try_into().expect("长度恰为 4")),
})
}
pub fn to_bytes(&self) -> [u8; Self::TOTAL_SIZE] {
let mut out = [0u8; Self::TOTAL_SIZE];
out[0] = self.aof_header_version;
out[1] = self.flags;
out[2] = self.op_type;
out[3] = self.procedure_id;
out[4..12].copy_from_slice(&self.store_version.to_le_bytes());
out[12..16].copy_from_slice(&self.session_id.to_le_bytes());
out
}
pub fn skip_header(entry: &[u8]) -> Option<usize> {
let header = AofHeader::parse(entry)?;
AofHeaderType::ALL
.iter()
.copied()
.find(|t| *t as u8 == (header.flags & Self::AOF_HEADER_TYPE_MASK))
.map(AofHeaderType::total_size)
}
pub fn get_chunked_header_ref(entry: &[u8]) -> Option<(usize, AofChunkHeader)> {
let header = AofHeader::parse(entry)?;
let offset = match header.header_type()? {
AofHeaderType::BasicChunkHeader => AofHeader::TOTAL_SIZE,
AofHeaderType::ShardedChunkHeader => AofShardedHeader::TOTAL_SIZE,
_ => return None,
};
Some((offset, AofChunkHeader::parse(&entry[offset..])?))
}
}
impl Default for AofHeader {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AofShardedHeader {
pub basic: AofHeader,
pub sequence_number: i64,
}
impl AofShardedHeader {
pub const TOTAL_SIZE: usize = AofHeader::TOTAL_SIZE + 8;
pub fn parse(entry: &[u8]) -> Option<Self> {
if entry.len() < Self::TOTAL_SIZE {
return None;
}
Some(Self {
basic: AofHeader::parse(entry)?,
sequence_number: i64::from_le_bytes(
entry[AofHeader::TOTAL_SIZE..Self::TOTAL_SIZE]
.try_into()
.expect("长度恰为 8"),
),
})
}
}
pub const REPLAY_TASK_ACCESS_VECTOR_BYTES: usize = 32;
#[derive(Debug, Clone, Copy)]
pub struct AofSingleLogTransactionHeader {
pub basic: AofHeader,
pub participant_count: i16,
pub replay_task_access_vector: [u8; REPLAY_TASK_ACCESS_VECTOR_BYTES],
}
impl AofSingleLogTransactionHeader {
pub const TOTAL_SIZE: usize = AofHeader::TOTAL_SIZE + 2 + REPLAY_TASK_ACCESS_VECTOR_BYTES;
pub fn parse(entry: &[u8]) -> Option<Self> {
if entry.len() < Self::TOTAL_SIZE {
return None;
}
let mut vector = [0u8; REPLAY_TASK_ACCESS_VECTOR_BYTES];
vector.copy_from_slice(&entry[AofHeader::TOTAL_SIZE + 2..Self::TOTAL_SIZE]);
Some(Self {
basic: AofHeader::parse(entry)?,
participant_count: i16::from_le_bytes(
entry[AofHeader::TOTAL_SIZE..AofHeader::TOTAL_SIZE + 2]
.try_into()
.expect("长度恰为 2"),
),
replay_task_access_vector: vector,
})
}
}
#[derive(Debug, Clone, Copy)]
pub struct AofShardedLogTransactionHeader {
pub sharded: AofShardedHeader,
pub participant_count: i16,
pub replay_task_access_vector: [u8; REPLAY_TASK_ACCESS_VECTOR_BYTES],
}
impl AofShardedLogTransactionHeader {
pub const TOTAL_SIZE: usize = AofShardedHeader::TOTAL_SIZE + 2 + REPLAY_TASK_ACCESS_VECTOR_BYTES;
pub fn parse_sharded(entry: &[u8]) -> Option<Self> {
if entry.len() < Self::TOTAL_SIZE {
return None;
}
let sharded = AofShardedHeader::parse(entry)?;
let mut vector = [0u8; REPLAY_TASK_ACCESS_VECTOR_BYTES];
vector.copy_from_slice(&entry[AofShardedHeader::TOTAL_SIZE + 2..Self::TOTAL_SIZE]);
Some(Self {
sharded,
participant_count: i16::from_le_bytes(
entry[AofShardedHeader::TOTAL_SIZE..AofShardedHeader::TOTAL_SIZE + 2]
.try_into()
.expect("长度恰为 2"),
),
replay_task_access_vector: vector,
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AofChunkHeader {
pub overflow_key_length: u32,
pub overflow_value_length: u32,
pub input_length: u32,
pub object_id: u64,
pub key_hash: i64,
}
impl AofChunkHeader {
pub const TOTAL_SIZE: usize = 3 * size_of::<u32>() + size_of::<u64>() + size_of::<i64>();
pub const OBJECT_ID_OFFSET: usize = 3 * size_of::<u32>();
pub fn parse(entry: &[u8]) -> Option<Self> {
if entry.len() < Self::TOTAL_SIZE {
return None;
}
Some(Self {
overflow_key_length: u32::from_le_bytes(entry[0..4].try_into().expect("长度恰为 4")),
overflow_value_length: u32::from_le_bytes(entry[4..8].try_into().expect("长度恰为 4")),
input_length: u32::from_le_bytes(entry[8..12].try_into().expect("长度恰为 4")),
object_id: u64::from_le_bytes(
entry[Self::OBJECT_ID_OFFSET..Self::OBJECT_ID_OFFSET + 8]
.try_into()
.expect("长度恰为 8"),
),
key_hash: i64::from_le_bytes(
entry[Self::OBJECT_ID_OFFSET + 8..Self::TOTAL_SIZE]
.try_into()
.expect("长度恰为 8"),
),
})
}
}
#[cfg(test)]
mod tests {
use super::{AofChunkHeader, AofHeader, AofHeaderType};
#[test]
fn header_roundtrip_and_flags() {
let mut h = AofHeader::new();
h.set_header_type(AofHeaderType::BasicHeader);
h.op_type = 0x00;
h.store_version = 42;
h.session_id = -7;
let bytes = h.to_bytes();
let parsed = AofHeader::parse(&bytes).unwrap();
assert_eq!(parsed, h);
assert_eq!(parsed.header_type(), Some(AofHeaderType::BasicHeader));
assert!(!parsed.is_chunked());
assert!(!parsed.unsafe_truncate_log());
h.set_unsafe_truncate_log(true);
h.set_header_type(AofHeaderType::ShardedChunkHeader);
assert!(h.unsafe_truncate_log());
assert!(h.is_chunked());
assert_eq!(h.header_type(), Some(AofHeaderType::ShardedChunkHeader));
}
#[test]
fn skip_header_offsets() {
for (t, size) in [
(AofHeaderType::BasicHeader, 16),
(AofHeaderType::ShardedHeader, 24),
(AofHeaderType::SingleLogTransactionHeader, 50),
(AofHeaderType::ShardedLogTransactionHeader, 58),
] {
let mut h = AofHeader::new();
h.set_header_type(t);
assert_eq!(AofHeader::skip_header(&h.to_bytes()), Some(size));
}
}
#[test]
fn chunk_header_ref() {
let mut h = AofHeader::new();
h.set_header_type(AofHeaderType::BasicChunkHeader);
let mut entry = h.to_bytes().to_vec();
let chunk = AofChunkHeader {
overflow_key_length: 8,
overflow_value_length: 0,
input_length: 4,
object_id: 7,
key_hash: -1,
};
let mut chunk_bytes = Vec::new();
chunk_bytes.extend_from_slice(&chunk.overflow_key_length.to_le_bytes());
chunk_bytes.extend_from_slice(&chunk.overflow_value_length.to_le_bytes());
chunk_bytes.extend_from_slice(&chunk.input_length.to_le_bytes());
chunk_bytes.extend_from_slice(&chunk.object_id.to_le_bytes());
chunk_bytes.extend_from_slice(&chunk.key_hash.to_le_bytes());
entry.extend_from_slice(&chunk_bytes);
let (offset, parsed) = AofHeader::get_chunked_header_ref(&entry).unwrap();
assert_eq!(offset, 16);
assert_eq!(parsed, chunk);
let mut plain = AofHeader::new();
plain.set_header_type(AofHeaderType::BasicHeader);
assert!(AofHeader::get_chunked_header_ref(&plain.to_bytes()).is_none());
}
}