8000 optimize: reduce persistent overlay memory with mmap offsets by nuskey8 · Pull Request #11 · nuskey8/lkv · GitHub
[go: up one dir, main page]
Skip to content

optimize: reduce persistent overlay memory with mmap offsets #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nuskey8 merged 1 commit into from
Aug 12, 2026
Merged
8000

optimize: reduce persistent overlay memory with mmap offsets #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 70 additions & 84 deletions lkv/src/database/mod.rs
View file 8000
Open in desktop
10000
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ use crate::format::segment::{self, EMPTY_SEGMENT_SIZE, segment_metadata_checksum
use crate::format::superblock::{self, DATA_START, Superblock};
use crate::options::DatabaseOptions;
use fs2::FileExt;
use state::{ActiveBase, MAPPED_VALUE_THRESHOLD};
use state::{ActiveBase, MappedMutation, OverlayMap};
use std::fs::{self, File, OpenOptions};
use std::io::{BufWriter, ErrorKind, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
Expand All @@ -23,7 +23,9 @@ use storage::Storage;
use transaction::RawEntries;
use view::{BaseView, ReadView};

pub use state::{BaseBytes, KeyMap, OverlayEntry, OverlayState, ValueBytes};
#[cfg(test)]
pub use state::{BaseBytes, ValueBytes};
pub use state::{KeyMap, OverlayEntry, OverlayState};
pub use transaction::{Entries, ReadTransaction, ReservedValue, Snapshot, WriteTransaction};

const LOG_WRITE_BUFFER_SIZE: usize = 64 * 1024;
Expand Down Expand Up @@ -314,41 +316,46 @@ impl Database {
.ok_or_else(|| Error::from_io(ErrorKind::InvalidInput, "batch record size overflow"))?;
self.ensure_capacity(batch_record_len)?;
let rollback_offset = self.storage.seek(SeekFrom::End(0))?;
let batch_offset = rollback_offset
.checked_sub(self.base.log_start)
.and_then(|offset| usize::try_from(offset).ok())
.ok_or_else(|| Error::other("Overlay batch offset overflow"))?;
let mutations = mapped_mutations(&staged, batch_offset)?;
if let Err(error) = append(&mut self.storage, &staged, &value_checksums) {
return self.rollback_or_poison(rollback_offset, error);
}
let mapped = if !self.storage.is_memory()
&& staged.values().any(|entry| {
matches!(entry, OverlayEntry::Put(value) if value.len() >= MAPPED_VALUE_THRESHOLD)
})
{
let record_len = match self.storage.len().and_then(|len| {
len.checked_sub(rollback_offset)
.ok_or_else(|| Error::other("storage shrank while committing a batch"))
}) {
Ok(len) => len,
Err(error) => return self.rollback_or_poison(rollback_offset, error),
};
match self.storage.load_immutable(rollback_offset, record_len) {
Ok(mapping) => Some(mapping),
Err(error) => return self.rollback_or_poison(rollback_offset, error),
}
} else {
None
};
let committed = match committed_entries(staged, mapped) {
Ok(entries) => entries,
Err(error) => return self.rollback_or_poison(rollback_offset, error),
};
failpoints::crash_process_if_requested("after_batch_write");
if let Err(error) = self.storage.sync_data() {
// Windows cannot truncate a range while it is still mapped.
drop(committed);
return self.rollback_or_poison(rollback_offset, error);
}
failpoints::crash_process_if_requested("after_batch_sync");
for (key, entry) in committed {
self.set_overlay_entry(key, entry);
let overlay_size = match self.storage.len().and_then(|len| {
len.checked_sub(self.base.log_start)
.ok_or_else(|| Error::other("Overlay range precedes the active Base"))
}) {
Ok(size) => size,
Err(error) => return self.rollback_or_poison(rollback_offset, error),
};
let overlay_size_usize = match usize::try_from(overlay_size) {
Ok(size) => size,
Err(_) => {
return self.rollback_or_poison(
rollback_offset,
Error::other("Overlay mapping size overflow"),
);
}
};
if self.overlay.needs_mapping(overlay_size_usize) {
let mapping = match self
.storage
.load_immutable(self.base.log_start, overlay_size)
{
Ok(mapping) => mapping,
Err(error) => return self.rollback_or_poison(rollback_offset, error),
};
self.overlay.install_mapped(mapping, mutations);
} else {
self.overlay.install_tail(staged, mutations);
}
Ok(())
}
Expand Down Expand Up @@ -507,68 +514,47 @@ fn load_storage_state(
let overlay_size = valid_len
.checked_sub(base.log_start)
.ok_or_else(|| Error::invalid_base("overlay range precedes active Base"))?;
let overlay_mapping = if overlay_size == 0 {
None
let overlay = if storage.is_memory() {
OverlayState::new(OverlayMap::memory())
} else {
Some(storage.load_immutable(base.log_start, overlay_size)?)
let overlay_mapping = if overlay_size == 0 {
None
} else {
Some(storage.load_immutable(base.log_start, overlay_size)?)
};
let (index, memory) = overlay_scan.into_index(overlay_mapping, base.log_start)?;
OverlayState::with_memory(index, memory)
};
let overlay_index = overlay_scan.into_index(overlay_mapping, base.log_start)?;
storage.seek(SeekFrom::End(0))?;
Ok((base, OverlayState::new(overlay_index)))
Ok((base, overlay))
}

fn committed_entries(
mut staged: KeyMap<OverlayEntry>,
mapping: Option<BaseBytes>,
) -> Result<KeyMap<OverlayEntry>> {
let Some(mapping) = mapping else {
return Ok(staged);
};
let bytes = &mapping[..];
let count_end = LOG_HEADER_SIZE + 4;
let count = bytes
.get(LOG_HEADER_SIZE..count_end)
.map(|bytes| u32::from_le_bytes(bytes.try_into().unwrap()) as usize)
.ok_or_else(|| Error::other("written batch is missing its operation count"))?;
let mut cursor = count_end;
let mut committed = KeyMap::default();
for _ in 0..count {
let header_end = cursor
.checked_add(9)
.filter(|end| *end <= bytes.len())
.ok_or_else(|| Error::other("written batch has a truncated operation"))?;
let header = &bytes[cursor..header_end];
let key_len = u32::from_le_bytes(header[1..5].try_into().unwrap()) as usize;
let value_len = u32::from_le_bytes(header[5..9].try_into().unwrap()) as usize;
let key_end = header_end
.checked_add(key_len)
.filter(|end| *end <= bytes.len())
.ok_or_else(|| Error::other("written batch has a truncated key"))?;
let value_end = key_end
.checked_add(value_len)
.filter(|end| *end <= bytes.len())
.ok_or_else(|| Error::other("written batch has a truncated value"))?;
let (key, entry) = staged
.remove_entry(&bytes[header_end..key_end])
.ok_or_else(|| Error::other("written batch key is missing from staged entries"))?;
let entry = match entry {
OverlayEntry::Put(value) if value.len() >= MAPPED_VALUE_THRESHOLD => {
OverlayEntry::Put(ValueBytes::Mapped {
bytes: mapping.clone(),
range: key_end..value_end,
})
}
entry => entry,
fn mapped_mutations(
staged: &KeyMap<OverlayEntry>,
batch_offset: usize,
) -> Result<Vec<MappedMutation>> {
let mut cursor = batch_offset
.checked_add(LOG_HEADER_SIZE + 4)
.ok_or_else(|| Error::other("Overlay batch layout overflow"))?;
let mut mutations = Vec::new();
mutations
.try_reserve_exact(staged.len())
.map_err(|_| Error::from_io(ErrorKind::OutOfMemory, "could not reserve Overlay entries"))?;
for (key, entry) in staged {
mutations.push(MappedMutation {
operation_offset: cursor,
});
let value_len = match entry {
OverlayEntry::Put(value) => value.len(),
OverlayEntry::Delete => 0,
};
committed.insert(key, entry);
cursor = value_end;
}
if cursor != bytes.len() || !staged.is_empty() {
return Err(Error::other(
"written batch layout does not match staged entries",
));
cursor = cursor
.checked_add(9)
.and_then(|cursor| cursor.checked_add(key.len()))
.and_then(|cursor| cursor.checked_add(value_len))
.ok_or_else(|| Error::other("Overlay operation range overflow"))?;
}
Ok(committed)
Ok(mutations)
}

fn open_database_file(path: &Path) -> Result<File> {
Expand Down
86 changes: 40 additions & 46 deletions lkv/src/database/recovery.rs
View file
Open in desktop
69E8
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use super::state::{BaseBytes, KeyMap, OverlayEntry, OverlayMap, ValueBytes};
use super::state::{BaseBytes, KeyMap, MAPPED_VALUE_THRESHOLD, MappedMutation, OverlayMap};
use crate::format::log::{
LOG_HEADER_CHECKSUM_OFFSET, LOG_HEADER_SIZE, MAX_BATCH_OPERATIONS, MAX_LOG_PAYLOAD_SIZE,
Marker, record_checksum, record_checksum_start,
Expand All @@ -24,43 +24,37 @@ impl OverlayScan {
self,
mapping: Option<BaseBytes>,
mapping_offset: u64,
) -> Result<OverlayMap<OverlayEntry>> {
let mut index = OverlayMap::default();
) -> Result<(OverlayMap, usize)> {
let mut index = OverlayMap::mapped(mapping);
index
.try_reserve(self.index.len())
.try_reserve_mapped(self.index.len())
.map_err(|_| Error::invalid_base("Overlay index allocation failed"))?;
let mut memory = 0usize;
for (key, entry) in self.index {
let entry = match entry {
RecoveredOverlayEntry::Put { offset, len } => {
let bytes = mapping.as_ref().ok_or_else(|| {
Error::invalid_base("recovered Overlay is missing its mapping")
})?;
let start = offset
.checked_sub(mapping_offset)
.and_then(|offset| usize::try_from(offset).ok())
.ok_or_else(|| Error::invalid_base("overlay value offset overflow"))?;
let end = start
.checked_add(len)
.filter(|end| *end <= bytes.len())
.ok_or_else(|| {
Error::invalid_base("overlay value lies outside its mapping")
})?;
OverlayEntry::Put(ValueBytes::Mapped {
bytes: bytes.clone(),
range: start..end,
})
}
RecoveredOverlayEntry::Delete => OverlayEntry::Delete,
};
index.insert(key, entry);
let operation_offset = entry
.operation_offset
.checked_sub(mapping_offset)
.and_then(|offset| usize::try_from(offset).ok())
.ok_or_else(|| Error::invalid_base("Overlay operation offset overflow"))?;
if operation_offset
.checked_add(9)
.is_none_or(|end| end > index.mapped_bytes_len())
{
return Err(Error::invalid_base(
"Overlay operation lies outside its mapping",
));
}
let mutation = MappedMutation { operation_offset };
memory = memory.saturating_add(key.len() + entry.value_memory);
index.install_recovered(mutation, xxhash_rust::xxh3::xxh3_64(&key));
}
Ok(index)
Ok((index, memory))
}
}

enum RecoveredOverlayEntry {
Put { offset: u64, len: usize },
Delete,
struct RecoveredOverlayEntry {
operation_offset: u64,
value_memory: usize,
}

pub fn scan_overlay(file: &mut (impl Read + Seek), start: u64, end: u64) -> Result<OverlayScan> {
Expand Down Expand Up @@ -202,6 +196,12 @@ fn apply_batch_payload<R: BufRead>(
.map_err(|_| Error::invalid_base("Overlay index allocation failed"))?;
}
for _ in 0..count {
let operation_offset = payload_offset
.checked_add(
u64::try_from(reader.consumed())
.map_err(|_| Error::invalid_base("overlay operation offset overflow"))?,
)
.ok_or_else(|| Error::invalid_base("overlay operation offset overflow"))?;
let operation = reader.read_operation()?;
if operation.key_len > recovery.key_buffer.capacity() {
recovery
Expand All @@ -211,20 +211,10 @@ fn apply_batch_payload<R: BufRead>(
}
recovery.key_buffer.resize(operation.key_len, 0);
reader.read_exact(recovery.key_buffer, "truncated batch key")?;
let value_offset = payload_offset
.checked_add(
u64::try_from(reader.consumed())
.map_err(|_| Error::invalid_base("overlay value offset overflow"))?,
)
.ok_or_else(|| Error::invalid_base("overlay value offset overflow"))?;
reader.skip(operation.value_len)?;
let entry = match operation.marker {
Marker::Put => RecoveredOverlayEntry::Put {
offset: value_offset,
len: operation.value_len,
},
Marker::Delete => RecoveredOverlayEntry::Delete,
_ => unreachable!("batch operation validation rejects non-mutations"),
let entry = RecoveredOverlayEntry {
operation_offset,
value_memory: operation.value_memory,
};
if let Some(previous) = recovery.index.get_mut(recovery.key_buffer.as_slice()) {
*previous = entry;
Expand All @@ -241,9 +231,9 @@ fn apply_batch_payload<R: BufRead>(
}

struct BatchOperation {
marker: Marker,
key_len: usize,
value_len: usize,
value_memory: usize,
}

struct BatchReader<R> {
Expand Down Expand Up @@ -298,9 +288,13 @@ impl<R: BufRead> BatchReader<R> {
.ok_or_else(|| Error::from_io(ErrorKind::InvalidData, "truncated batch value"))?;
debug_assert!(data_len <= self.remaining);
Ok(BatchOperation {
marker: marker.expect("validated batch marker"),
key_len,
value_len,
value_memory: if marker == Some(Marker::Put) && value_len < MAPPED_VALUE_THRESHOLD {
value_len
} else {
0
},
})
}

Expand Down
Loading
0