use std::{
sync::{
Arc,
atomic::{AtomicI64, Ordering},
},
thread,
time::Duration,
};
use parking_lot::Mutex;
use super::{
aof_address::AofAddress,
aof_backpressure::AofBackpressure,
aof_entry_type::AofEntryType,
aof_header::{AofHeader, AofHeaderType, AofShardedHeader},
sharded_log::ShardedLogLockMap,
};
use crate::config::runtime_server_options::RuntimeServerOptions;
#[derive(Debug, Clone)]
pub struct LogRecord {
pub address: i64,
pub payload: Vec<u8>,
}
pub trait SublogBackend: Send + Sync {
fn enqueue(&self, sublog_idx: usize, payload: &[u8]) -> i64;
fn tail_address(&self, sublog_idx: usize) -> i64;
fn begin_address(&self, sublog_idx: usize) -> i64;
fn committed_until_address(&self, sublog_idx: usize) -> i64;
fn commit(&self, sublog_idx: usize, until_address: i64);
fn flushed_until_address(&self, sublog_idx: usize) -> i64;
fn scan(&self, sublog_idx: usize, begin_address: i64, end_address: i64) -> Vec<LogRecord>;
fn shift_begin_address(&self, sublog_idx: usize, new_begin: i64);
fn log_page_size_bits(&self, sublog_idx: usize) -> i32;
fn memory_size_bytes(&self, sublog_idx: usize) -> i64;
fn reset(&self, sublog_idx: usize);
}
#[derive(Default)]
pub struct InMemorySublog {
records: parking_lot::Mutex<Vec<LogRecord>>,
begin: AtomicI64,
committed_until: AtomicI64,
}
impl InMemorySublog {
pub fn new() -> Self {
Self {
records: parking_lot::Mutex::new(Vec::new()),
begin: AtomicI64::new(1),
committed_until: AtomicI64::new(1),
}
}
}
impl SublogBackend for InMemorySublog {
fn enqueue(&self, _sublog_idx: usize, payload: &[u8]) -> i64 {
let mut records = self.records.lock();
let address = records
.last()
.map_or(self.begin.load(Ordering::Relaxed), |last| {
last.address + last.payload.len() as i64
});
records.push(LogRecord {
address,
payload: payload.to_vec(),
});
address
}
fn tail_address(&self, _sublog_idx: usize) -> i64 {
self.records.lock().last().map_or_else(
|| self.begin.load(Ordering::Relaxed),
|last| last.address + last.payload.len() as i64,
)
}
fn begin_address(&self, _sublog_idx: usize) -> i64 {
self.begin.load(Ordering::Relaxed)
}
fn committed_until_address(&self, _sublog_idx: usize) -> i64 {
self.committed_until.load(Ordering::Relaxed)
}
fn commit(&self, _sublog_idx: usize, until_address: i64) {
self
.committed_until
.fetch_max(until_address, Ordering::Release);
}
fn flushed_until_address(&self, _sublog_idx: usize) -> i64 {
self.committed_until.load(Ordering::Relaxed)
}
fn scan(&self, _sublog_idx: usize, begin_address: i64, end_address: i64) -> Vec<LogRecord> {
self
.records
.lock()
.iter()
.filter(|r| r.address >= begin_address && r.address < end_address)
.cloned()
.collect()
}
fn shift_begin_address(&self, _sublog_idx: usize, new_begin: i64) {
self.begin.store(new_begin, Ordering::Release);
}
fn log_page_size_bits(&self, _sublog_idx: usize) -> i32 {
22
}
fn memory_size_bytes(&self, _sublog_idx: usize) -> i64 {
self
.records
.lock()
.iter()
.map(|r| r.payload.len() as i64)
.sum()
}
fn reset(&self, _sublog_idx: usize) {
self.records.lock().clear();
self.begin.store(1, Ordering::Release);
self.committed_until.store(1, Ordering::Release);
}
}
pub struct GarnetLog {
single_log: Option<Arc<dyn SublogBackend>>,
sharded_log: Vec<Arc<dyn SublogBackend>>,
physical_sublog_count: usize,
replay_task_count: usize,
lock_map: ShardedLogLockMap,
backpressure: Option<Arc<AofBackpressure>>,
shift_tail_callback: Mutex<Option<ShiftTailCallback>>,
tail_witness: AtomicI64,
}
impl GarnetLog {
pub fn new(server_options: &RuntimeServerOptions, backends: Vec<Arc<dyn SublogBackend>>) -> Self {
let physical_sublog_count = server_options.aof_physical_sublog_count.max(1) as usize;
let replay_task_count = server_options.aof_replay_task_count.max(1) as usize;
let using_single = physical_sublog_count == 1 && replay_task_count == 1;
let mut iter = backends.into_iter();
let single_log = using_single.then(|| iter.next().expect("单日志拓扑需 1 个后端"));
let sharded_log = if using_single {
Vec::new()
} else {
iter.collect()
};
Self {
single_log,
sharded_log,
physical_sublog_count,
replay_task_count,
lock_map: ShardedLogLockMap::new(),
backpressure: Some(Arc::new(AofBackpressure::new(
physical_sublog_count,
server_options.aof_sync_max_lag_bytes,
))),
shift_tail_callback: Mutex::new(None),
tail_witness: AtomicI64::new(0),
}
}
pub fn get_sequence_number_from_cookie(cookie: &[u8]) -> i64 {
cookie
.first_chunk::<8>()
.map_or(0, |b| i64::from_le_bytes(*b))
}
pub fn hash(key: &[u8]) -> i64 {
gxhash::gxhash64(key, 0) as i64
}
#[inline]
pub fn get_physical_sublog_idx(&self, hash: i64) -> usize {
((hash as u64) % (self.physical_sublog_count as u64)) as usize
}
#[inline]
pub fn get_replay_task_idx(&self, hash: i64) -> usize {
(((hash as u64) / (self.physical_sublog_count as u64)) % (self.replay_task_count as u64))
as usize
}
#[inline]
pub fn get_virtual_sublog_idx(&self, hash: i64) -> usize {
self.get_physical_sublog_idx(hash) * self.replay_task_count + self.get_replay_task_idx(hash)
}
pub fn recover_latest_sequence_number(&self, recover_until: i64, cookie: &[u8]) -> i64 {
let latest = Self::get_sequence_number_from_cookie(cookie);
if recover_until == -1 {
latest
} else {
recover_until.min(latest)
}
}
pub fn size(&self) -> usize {
if self.single_log.is_some() {
1
} else {
self.sharded_log.len()
}
}
pub fn all_logs_bitmask(&self) -> u64 {
(1u64 << self.size()) - 1
}
pub fn lock_sublogs(&self, log_access_bitmap: u64) {
self.lock_map.lock_sublogs(log_access_bitmap);
}
pub fn unlock_sublogs(&self, log_access_bitmap: u64) {
self.lock_map.unlock_sublogs(log_access_bitmap);
}
fn sublog(&self, sublog_idx: usize) -> &Arc<dyn SublogBackend> {
if let Some(single) = &self.single_log {
single
} else {
&self.sharded_log[sublog_idx]
}
}
pub fn get_sub_log(&self, sublog_idx: usize) -> &Arc<dyn SublogBackend> {
self.sublog(sublog_idx)
}
pub fn get_begin_address(&self) -> AofAddress {
let len = self.size() as i32;
let mut result = AofAddress::create(len, 0);
for i in 0..len as usize {
result.set(i, self.sublog(i).begin_address(i));
}
result
}
pub fn get_tail_address_vector(&self) -> AofAddress {
let len = self.size() as i32;
let mut result = AofAddress::create(len, 0);
for i in 0..len as usize {
result.set(i, self.sublog(i).tail_address(i));
}
result
}
pub fn get_tail_address(&self, sublog_idx: usize) -> i64 {
self.sublog(sublog_idx).tail_address(sublog_idx)
}
pub fn set_log_shift_tail_callback(&self, callback: Box<dyn Fn(i64) + Send + Sync>) {
*self.shift_tail_callback.lock() = Some(callback);
}
pub fn scan_single(
&self,
sublog_idx: usize,
begin_address: i64,
end_address: i64,
) -> Vec<LogRecord> {
self
.sublog(sublog_idx)
.scan(sublog_idx, begin_address, end_address)
}
pub fn unsafe_get_log_page_size_bits(&self) -> i32 {
self.sublog(0).log_page_size_bits(0)
}
pub fn unsafe_get_read_only_address_above(&self, sublog_idx: usize, address: i64) -> i64 {
self.get_tail_address(sublog_idx).max(address)
}
pub fn unsafe_shift_begin_address(&self, sublog_idx: usize, new_begin: i64) {
self
.sublog(sublog_idx)
.shift_begin_address(sublog_idx, new_begin);
if let Some(callback) = self.shift_tail_callback.lock().as_ref() {
callback(new_begin);
}
}
pub fn truncate_until(&self, until: &AofAddress) {
for i in 0..self.size() {
self
.sublog(i)
.shift_begin_address(i, until.get(i).unwrap_or(0));
}
}
pub fn safe_initialize(&self) -> bool {
true
}
pub fn initialize_if(&self, condition: bool) -> bool {
condition && self.safe_initialize()
}
pub fn wait_for_commit(&self, sublog_idx: usize, address: i64) {
while self.sublog(sublog_idx).committed_until_address(sublog_idx) < address {
thread::sleep(Duration::from_micros(50));
}
}
pub fn commit(&self, sublog_idx: usize) -> i64 {
let tail = self.sublog(sublog_idx).tail_address(sublog_idx);
self.sublog(sublog_idx).commit(sublog_idx, tail);
self.tail_witness.store(tail, Ordering::Release);
tail
}
pub fn wait_for_commit_async(&self, sublog_idx: usize, address: i64) {
self.wait_for_commit(sublog_idx, address);
}
fn backpressure_wait_key(&self, key: &[u8]) {
let Some(backpressure) = &self.backpressure else {
return;
};
let sublog_idx = self.get_physical_sublog_idx(Self::hash(key));
backpressure.wait(sublog_idx, self.get_tail_address(sublog_idx));
}
fn backpressure_wait_key_hash(&self, key_hash: i64) {
let Some(backpressure) = &self.backpressure else {
return;
};
let sublog_idx = self.get_physical_sublog_idx(key_hash);
backpressure.wait(sublog_idx, self.get_tail_address(sublog_idx));
}
fn backpressure_wait_vector(&self, mut physical_sublog_access_vector: u64) {
let Some(backpressure) = &self.backpressure else {
return;
};
while physical_sublog_access_vector > 0 {
let sublog_idx = physical_sublog_access_vector.trailing_zeros() as usize;
backpressure.wait(sublog_idx, self.get_tail_address(sublog_idx));
physical_sublog_access_vector &= physical_sublog_access_vector - 1;
}
}
pub fn is_chunkable(key_len: usize, value_len: usize, input_serialized_length: usize) -> bool {
(key_len + value_len + input_serialized_length) as i64 > MIN_PARTIAL_ALLOC_SIZE
}
fn enqueue_with_header(&self, record: &RecordShape<'_>) -> i64 {
let RecordShape {
op_type,
version,
session_id,
key,
value,
input,
database_id,
} = *record;
let using_single_physical_log = self.single_log.is_some();
let physical_sublog_idx = if using_single_physical_log {
0
} else {
self.get_physical_sublog_idx(Self::hash(key))
};
let mut payload =
Vec::with_capacity(AofHeader::TOTAL_SIZE + key.len() + value.len() + input.len());
if using_single_physical_log {
let mut header = AofHeader::new();
header.set_header_type(AofHeaderType::BasicHeader);
header.op_type = op_type as u8;
header.store_version = version;
header.session_id = session_id;
header.database_id = database_id;
payload.extend_from_slice(&header.to_bytes());
} else {
let mut header = AofHeader::new();
header.set_header_type(AofHeaderType::ShardedHeader);
header.op_type = op_type as u8;
header.store_version = version;
header.session_id = session_id;
header.database_id = database_id;
let sequence_number = self.tail_witness.load(Ordering::Relaxed);
let sharded = AofShardedHeader {
basic: header,
sequence_number,
};
payload.extend_from_slice(&sharded.basic.to_bytes());
payload.extend_from_slice(&sequence_number.to_le_bytes());
}
payload.extend_from_slice(key);
payload.extend_from_slice(value);
payload.extend_from_slice(input);
self
.sublog(physical_sublog_idx)
.enqueue(physical_sublog_idx, &payload)
}
pub fn enqueue(&self, record: &RecordShape<'_>) -> i64 {
self.backpressure_wait_key(record.key);
self.enqueue_with_header(record)
}
pub fn enqueue_span_chunked(&self, chunk: &ChunkedShape<'_>) -> i64 {
let ChunkedShape {
record:
RecordShape {
op_type,
version,
session_id,
key,
value,
input,
database_id,
},
write_value,
write_input,
} = *chunk;
self.backpressure_wait_key_hash(Self::hash(key));
let chunk_header = super::aof_header::AofChunkHeader {
overflow_key_length: key.len() as u32,
overflow_value_length: if write_value { value.len() as u32 } else { 0 },
input_length: if write_input { input.len() as u32 } else { 0 },
object_id: 0,
key_hash: Self::hash(key),
};
let using_single_physical_log = self.single_log.is_some();
let header_size = if using_single_physical_log {
AofHeader::TOTAL_SIZE
} else {
AofShardedHeader::TOTAL_SIZE
};
let page_payload = (1usize << self.unsafe_get_log_page_size_bits() as u32)
- header_size
- super::aof_header::AofChunkHeader::TOTAL_SIZE;
let mut chunks: Vec<Vec<u8>> = Vec::new();
let mut emit = |mut chunk: Vec<u8>, key: &[u8], value: &[u8]| {
chunk.extend_from_slice(key);
chunk.extend_from_slice(value);
chunks.push(chunk);
};
let mut first = Vec::with_capacity(header_size + super::aof_header::AofChunkHeader::TOTAL_SIZE);
let mut header = AofHeader::new();
header.set_header_type(if using_single_physical_log {
AofHeaderType::BasicChunkHeader
} else {
AofHeaderType::ShardedChunkHeader
});
header.op_type = op_type as u8;
header.store_version = version;
header.session_id = session_id;
header.database_id = database_id;
first.extend_from_slice(&header.to_bytes());
if !using_single_physical_log {
first.extend_from_slice(&0i64.to_le_bytes());
}
first.extend_from_slice(&chunk_header.overflow_key_length.to_le_bytes());
first.extend_from_slice(&chunk_header.overflow_value_length.to_le_bytes());
first.extend_from_slice(&chunk_header.input_length.to_le_bytes());
first.extend_from_slice(&chunk_header.object_id.to_le_bytes());
first.extend_from_slice(&chunk_header.key_hash.to_le_bytes());
emit(first, key, &[]);
let remaining = if write_value { value } else { &[][..] };
for piece in remaining.chunks(page_payload.max(1)) {
let piece_chunk = Vec::with_capacity(piece.len());
emit(piece_chunk, &[], piece);
}
if write_input {
let input_chunk = Vec::with_capacity(input.len());
emit(input_chunk, &[], input);
}
let mut address = 0;
for chunk in chunks {
let physical_sublog_idx = if using_single_physical_log {
0
} else {
self.get_physical_sublog_idx(chunk_header.key_hash)
};
address = self
.sublog(physical_sublog_idx)
.enqueue(physical_sublog_idx, &chunk);
}
address
}
pub fn enqueue_object_chunked(&self, chunk: &ChunkedShape<'_>) -> i64 {
self.enqueue_span_chunked(&ChunkedShape {
write_value: true,
..chunk.clone()
})
}
pub fn chunk_buffer_size(
key_len: usize,
value_len: usize,
input_len: usize,
chunk_count: usize,
) -> usize {
key_len + value_len + input_len + chunk_count * super::aof_header::AofChunkHeader::TOTAL_SIZE
}
pub fn enqueue_stored_proc(
&self,
op_type: AofEntryType,
version: i64,
session_id: i32,
procedure_id: u8,
body: &[u8],
physical_sublog_access_vector: u64,
) -> i64 {
self.backpressure_wait_vector(physical_sublog_access_vector);
self.lock_sublogs(physical_sublog_access_vector);
let result = self.enqueue_with_header(&RecordShape {
op_type,
version,
session_id,
key: &[],
value: body,
input: &[],
database_id: procedure_id,
});
self.unlock_sublogs(physical_sublog_access_vector);
result
}
pub fn enqueue_txn(
&self,
op_type: AofEntryType,
version: i64,
session_id: i32,
body: &[u8],
physical_sublog_access_vector: u64,
) -> i64 {
self.enqueue_stored_proc(
op_type,
version,
session_id,
0,
body,
physical_sublog_access_vector,
)
}
pub fn enqueue_broadcast_entry(&self, op_type: AofEntryType, version: i64, body: &[u8]) -> i64 {
self.enqueue_with_header(&RecordShape {
op_type,
version,
session_id: 0,
key: &[],
value: body,
input: &[],
database_id: 0,
})
}
pub fn enqueue_database_commit(
&self,
op_type: AofEntryType,
version: i64,
database_id: u8,
unsafe_truncate: bool,
) -> i64 {
let mut header = AofHeader::new();
header.set_header_type(AofHeaderType::BasicHeader);
header.op_type = op_type as u8;
header.store_version = version;
header.database_id = database_id;
if unsafe_truncate {
header.set_unsafe_truncate_log(true);
}
self.sublog(0).enqueue(0, &header.to_bytes())
}
pub fn enqueue_safe_flush_aof(&self, version: i64) -> i64 {
self.enqueue_database_commit(AofEntryType::FlushAll, version, 0, false)
}
}
#[derive(Clone)]
pub struct ChunkedShape<'a> {
pub record: RecordShape<'a>,
pub write_value: bool,
pub write_input: bool,
}
#[derive(Clone)]
pub struct RecordShape<'a> {
pub op_type: AofEntryType,
pub version: i64,
pub session_id: i32,
pub key: &'a [u8],
pub value: &'a [u8],
pub input: &'a [u8],
pub database_id: u8,
}
type ShiftTailCallback = Box<dyn Fn(i64) + Send + Sync>;
pub const MIN_PARTIAL_ALLOC_SIZE: i64 = 8 * 1024 * 1024;
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::{ChunkedShape, GarnetLog, InMemorySublog, RecordShape, SublogBackend};
use crate::{
aof::{aof_address::AofAddress, aof_entry_type::AofEntryType},
config::runtime_server_options::RuntimeServerOptions,
};
fn log_with(sublogs: usize, replay_tasks: i32) -> GarnetLog {
let options = RuntimeServerOptions {
aof_physical_sublog_count: sublogs as i32,
aof_replay_task_count: replay_tasks,
..RuntimeServerOptions::default()
};
let backends: Vec<Arc<dyn SublogBackend>> = (0..sublogs.max(1))
.map(|_| Arc::new(InMemorySublog::new()) as Arc<dyn SublogBackend>)
.collect();
GarnetLog::new(&options, backends)
}
#[test]
fn sharding_routes_deterministically() {
let log = log_with(4, 2);
let hash = GarnetLog::hash(b"key");
let physical = log.get_physical_sublog_idx(hash);
assert!(physical < 4);
assert!(log.get_replay_task_idx(hash) < 2);
assert_eq!(
log.get_virtual_sublog_idx(hash),
physical * 2 + log.get_replay_task_idx(hash)
);
}
#[test]
fn enqueue_scan_roundtrip() {
let log = log_with(1, 1);
let address = log.enqueue(&RecordShape {
op_type: AofEntryType::StoreUpsert,
version: 1,
session_id: 7,
key: b"key1",
value: b"value1",
input: &[],
database_id: 0,
});
assert!(address > 0);
let records = log.scan_single(0, 1, i64::MAX);
assert_eq!(records.len(), 1);
let payload = &records[0].payload;
assert_eq!(&payload[16..20], b"key1");
assert_eq!(&payload[20..26], b"value1");
}
#[test]
fn commit_and_bitmask() {
let log = log_with(2, 1);
assert_eq!(log.all_logs_bitmask(), 0b11);
let address = log.enqueue(&RecordShape {
op_type: AofEntryType::StoreUpsert,
version: 1,
session_id: 1,
key: b"k",
value: b"v",
input: &[],
database_id: 0,
});
let physical = log.get_physical_sublog_idx(GarnetLog::hash(b"k"));
let tail = log.commit(physical);
assert!(tail > address);
log.wait_for_commit(physical, tail);
let begins = log.get_begin_address();
assert_eq!(begins.length(), 2);
}
#[test]
fn chunked_write_reassembles() {
let log = log_with(1, 1);
let value = vec![b'x'; 200];
let address = log.enqueue_object_chunked(&ChunkedShape {
record: RecordShape {
op_type: AofEntryType::ObjectStoreUpsert,
version: 3,
session_id: 9,
key: b"big",
value: &value,
input: &[],
database_id: 0,
},
write_value: true,
write_input: false,
});
assert!(address > 0);
let records = log.scan_single(0, 1, i64::MAX);
assert!(records.len() >= 2);
let chunk = super::super::aof_header::AofChunkHeader::parse(&records[0].payload[16..]).unwrap();
assert_eq!(chunk.overflow_value_length, 200);
assert_eq!(chunk.key_hash, GarnetLog::hash(b"big"));
let data: Vec<u8> = records[1..]
.iter()
.flat_map(|r| r.payload.clone())
.collect();
assert_eq!(data, value);
}
#[test]
fn sequence_number_from_cookie() {
let cookie = 123456789i64.to_le_bytes();
assert_eq!(
GarnetLog::get_sequence_number_from_cookie(&cookie),
123456789
);
}
#[test]
fn lock_bitmap_and_truncate() {
let log = log_with(2, 1);
log.lock_sublogs(0b11);
log.unlock_sublogs(0b11);
let until = AofAddress::create(2, 5);
log.truncate_until(&until);
assert_eq!(log.get_begin_address().get(0), Some(5));
}
}