#![warn(missing_docs)]
#![doc = include_str!("../README.md")]
mod cipher;
mod error;
mod segment;
#[cfg(feature = "encryption")]
pub use cipher::AesGcmCipher;
pub use cipher::{CipherError, SegmentCipher};
pub use error::{Result, SegmentError};
use std::fs;
use std::path::PathBuf;
use std::time::Instant;
use parking_lot::Mutex;
use serde::de::DeserializeOwned;
use serde::Serialize;
use tracing::{debug, info};
use segment::SegmentRange;
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum FlushPolicy {
Batch(usize),
Interval(std::time::Duration),
BatchOrInterval {
batch_size: usize,
interval: std::time::Duration,
},
Manual,
}
impl Default for FlushPolicy {
fn default() -> Self {
FlushPolicy::BatchOrInterval {
batch_size: 256,
interval: std::time::Duration::from_secs(5),
}
}
}
impl FlushPolicy {
fn should_flush(&self, pending_len: usize, time_since_last_flush: std::time::Duration) -> bool {
match self {
FlushPolicy::Batch(n) => pending_len >= *n,
FlushPolicy::Interval(d) => time_since_last_flush >= *d,
FlushPolicy::BatchOrInterval {
batch_size,
interval,
} => pending_len >= *batch_size || time_since_last_flush >= *interval,
FlushPolicy::Manual => false,
}
}
}
#[non_exhaustive]
pub struct SegmentConfig {
pub flush_policy: FlushPolicy,
pub max_size_bytes: u64,
pub compression_level: i32,
pub cipher: Option<Box<dyn SegmentCipher>>,
}
impl std::fmt::Debug for SegmentConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SegmentConfig")
.field("flush_policy", &self.flush_policy)
.field("max_size_bytes", &self.max_size_bytes)
.field("compression_level", &self.compression_level)
.field("cipher", &self.cipher.as_ref().map(|_| "[set]"))
.finish()
}
}
impl Default for SegmentConfig {
fn default() -> Self {
Self {
flush_policy: FlushPolicy::default(),
max_size_bytes: 10 * 1024 * 1024 * 1024,
compression_level: 3,
cipher: None,
}
}
}
#[derive(Debug)]
pub struct SegmentConfigBuilder {
inner: SegmentConfig,
}
impl SegmentConfigBuilder {
pub fn flush_policy(mut self, policy: FlushPolicy) -> Self {
self.inner.flush_policy = policy;
self
}
pub fn flush_at_batch_size(self, batch_size: usize) -> Self {
self.flush_policy(FlushPolicy::Batch(batch_size))
}
pub fn flush_at_interval(self, interval: std::time::Duration) -> Self {
self.flush_policy(FlushPolicy::Interval(interval))
}
pub fn flush_at_batch_or_interval(
self,
batch_size: usize,
interval: std::time::Duration,
) -> Self {
self.flush_policy(FlushPolicy::BatchOrInterval {
batch_size,
interval,
})
}
pub fn flush_manually(self) -> Self {
self.flush_policy(FlushPolicy::Manual)
}
pub fn max_size_bytes(mut self, max_size_bytes: u64) -> Self {
self.inner.max_size_bytes = max_size_bytes;
self
}
pub fn compression_level(mut self, compression_level: i32) -> Self {
self.inner.compression_level = compression_level;
self
}
pub fn cipher(mut self, cipher: Box<dyn SegmentCipher>) -> Self {
self.inner.cipher = Some(cipher);
self
}
pub fn build(self) -> SegmentConfig {
self.inner
}
}
impl SegmentConfig {
#[must_use = "the builder is meaningless if discarded"]
pub fn builder() -> SegmentConfigBuilder {
SegmentConfigBuilder {
inner: SegmentConfig::default(),
}
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct BufferStats {
pub pending_count: u64,
pub latest_sequence: u64,
pub head_sequence: u64,
pub next_sequence: u64,
pub approx_disk_bytes: u64,
pub max_size_bytes: u64,
pub store_pressure: f32,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct RecoveryReport {
pub segment_count: usize,
pub head_seq: u64,
pub next_seq: u64,
pub disk_bytes: u64,
pub removed_tmp_files: usize,
}
struct BufferInner<T> {
unflushed: Vec<T>,
next_seq: u64,
head_seq: u64,
last_flush: Instant,
}
pub struct SegmentBuffer<T> {
dir: PathBuf,
config: SegmentConfig,
inner: Mutex<BufferInner<T>>,
approx_disk_bytes: std::sync::atomic::AtomicU64,
scan_cache: Mutex<Option<Vec<segment::SegmentRange>>>,
}
impl<T> std::fmt::Debug for SegmentBuffer<T>
where
T: Serialize + DeserializeOwned + Clone + Send + 'static,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let stats = self.stats();
f.debug_struct("SegmentBuffer")
.field("dir", &self.dir)
.field("pending_count", &stats.pending_count)
.field("latest_sequence", &stats.latest_sequence)
.field("head_sequence", &stats.head_sequence)
.field("next_sequence", &stats.next_sequence)
.field("approx_disk_bytes", &stats.approx_disk_bytes)
.field("max_size_bytes", &stats.max_size_bytes)
.field("store_pressure", &stats.store_pressure)
.finish()
}
}
impl<T> SegmentBuffer<T>
where
T: Serialize + DeserializeOwned + Clone + Send + 'static,
{
pub fn open(dir: impl Into<PathBuf>, config: SegmentConfig) -> Result<Self> {
let (buffer, _report) = Self::open_with_report(dir, config)?;
Ok(buffer)
}
pub fn open_with_report(
dir: impl Into<PathBuf>,
config: SegmentConfig,
) -> Result<(Self, RecoveryReport)> {
let dir = dir.into();
fs::create_dir_all(&dir)?;
let buffer = Self {
dir,
config,
inner: Mutex::new(BufferInner {
unflushed: Vec::new(),
next_seq: 0,
head_seq: 0,
last_flush: Instant::now(),
}),
approx_disk_bytes: std::sync::atomic::AtomicU64::new(0),
scan_cache: Mutex::new(None),
};
let report = buffer.recover()?;
Ok((buffer, report))
}
pub fn append(&self, event: T) -> Result<u64> {
let (should_flush, seq) = {
let mut inner = self.inner.lock();
inner.unflushed.push(event);
inner.next_seq += 1;
let seq = inner.next_seq - 1;
let should_flush = self
.config
.flush_policy
.should_flush(inner.unflushed.len(), inner.last_flush.elapsed());
(should_flush, seq)
};
if should_flush {
self.flush()?;
}
Ok(seq)
}
pub fn flush(&self) -> Result<()> {
let (events, start_seq, end_seq) = {
let mut inner = self.inner.lock();
inner.last_flush = Instant::now();
if inner.unflushed.is_empty() {
return Ok(());
}
let events = std::mem::take(&mut inner.unflushed);
let count = events.len() as u64;
let end_seq = inner.next_seq - 1;
let start_seq = end_seq + 1 - count;
(events, start_seq, end_seq)
};
let compressed_len = self.write_segment(start_seq, end_seq, &events)?;
self.approx_disk_bytes
.fetch_add(compressed_len, std::sync::atomic::Ordering::Relaxed);
self.invalidate_scan_cache();
debug!(
path = self.segment_path(start_seq, end_seq).display().to_string(),
seq = start_seq,
end_seq,
count = events.len(),
bytes = compressed_len,
"Flushed segment"
);
Ok(())
}
pub fn read_from(&self, start_seq: u64, limit: usize) -> Result<Vec<T>> {
if limit == 0 {
return Ok(Vec::new());
}
let mut result: Vec<T> = Vec::with_capacity(limit.min(1024));
let segments = self.scan_segments()?;
for seg in &segments {
if result.len() >= limit {
break;
}
if seg.end < start_seq {
continue;
}
let events = self.read_segment(*seg)?;
let skip = if seg.start < start_seq {
(start_seq - seg.start) as usize
} else {
0
};
for event in events.into_iter().skip(skip) {
if result.len() >= limit {
break;
}
result.push(event);
}
}
if result.len() < limit {
let inner = self.inner.lock();
let pending_start = inner.next_seq.saturating_sub(inner.unflushed.len() as u64);
for (i, event) in inner.unflushed.iter().enumerate() {
let seq = pending_start + i as u64;
if seq < start_seq {
continue;
}
if result.len() >= limit {
break;
}
result.push(event.clone());
}
}
Ok(result)
}
pub fn for_each_from<F>(&self, start_seq: u64, limit: usize, mut f: F) -> Result<usize>
where
F: FnMut(u64, &T),
{
if limit == 0 {
return Ok(0);
}
let mut visited = 0usize;
let segments = self.scan_segments()?;
for seg in &segments {
if visited >= limit {
break;
}
if seg.end < start_seq {
continue;
}
let events = self.read_segment(*seg)?;
let skip = if seg.start < start_seq {
(start_seq - seg.start) as usize
} else {
0
};
for (offset, event) in events.iter().enumerate().skip(skip) {
if visited >= limit {
break;
}
let seq = seg.start + offset as u64;
f(seq, event);
visited += 1;
}
}
if visited < limit {
let inner = self.inner.lock();
let pending_start = inner.next_seq.saturating_sub(inner.unflushed.len() as u64);
for (i, event) in inner.unflushed.iter().enumerate() {
if visited >= limit {
break;
}
let seq = pending_start + i as u64;
if seq < start_seq {
continue;
}
f(seq, event);
visited += 1;
}
}
Ok(visited)
}
pub fn delete_acked(&self, acked_seq: u64) -> Result<usize> {
let segments = self.scan_segments()?;
let mut deleted = 0;
let mut freed_bytes: u64 = 0;
let mut new_head = None;
for seg in &segments {
if seg.end <= acked_seq {
let path = self.segment_path(seg.start, seg.end);
let file_bytes = fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
freed_bytes += file_bytes;
match fs::remove_file(&path) {
Ok(()) => {
deleted += 1;
debug!(
path = path.display().to_string(),
seq = seg.start,
end_seq = seg.end,
bytes = file_bytes,
"Deleted acked segment"
);
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
Err(e) => return Err(e.into()),
}
} else if new_head.is_none() {
new_head = Some(seg.start);
}
}
self.approx_disk_bytes
.fetch_sub(freed_bytes, std::sync::atomic::Ordering::Relaxed);
self.invalidate_scan_cache();
{
let mut inner = self.inner.lock();
let pending_start = inner.next_seq.saturating_sub(inner.unflushed.len() as u64);
inner.head_seq = new_head.unwrap_or(inner.next_seq).min(pending_start);
}
if deleted > 0 {
info!(
path = self.dir.display().to_string(),
deleted,
bytes = freed_bytes,
seq = acked_seq,
"Deleted acked segments"
);
}
Ok(deleted)
}
#[must_use = "the sequence number is meaningless if discarded"]
pub fn latest_sequence(&self) -> u64 {
let inner = self.inner.lock();
if inner.next_seq == 0 {
0
} else {
inner.next_seq - 1
}
}
#[must_use = "the backlog size is meaningless if discarded"]
pub fn pending_count(&self) -> u64 {
let inner = self.inner.lock();
inner.next_seq.saturating_sub(inner.head_seq)
}
#[must_use = "the backlog size is meaningless if discarded"]
pub fn len(&self) -> u64 {
self.pending_count()
}
#[must_use = "the emptiness flag is meaningless if discarded"]
pub fn is_empty(&self) -> bool {
self.pending_count() == 0
}
#[must_use = "the pressure value is meaningless if discarded"]
pub fn store_pressure(&self) -> f32 {
if self.config.max_size_bytes == 0 {
return 0.0;
}
let bytes = self
.approx_disk_bytes
.load(std::sync::atomic::Ordering::Relaxed);
(bytes as f32 / self.config.max_size_bytes as f32).min(1.0)
}
#[must_use = "the overload flag is meaningless if discarded"]
pub fn is_overloaded(&self) -> bool {
self.store_pressure() > 0.9
}
#[must_use = "the snapshot is meaningless if discarded"]
pub fn stats(&self) -> BufferStats {
let inner = self.inner.lock();
let pending_count = inner.next_seq.saturating_sub(inner.head_seq);
let latest_sequence = if inner.next_seq == 0 {
0
} else {
inner.next_seq - 1
};
let approx_disk_bytes = self
.approx_disk_bytes
.load(std::sync::atomic::Ordering::Relaxed);
let store_pressure = if self.config.max_size_bytes == 0 {
0.0
} else {
(approx_disk_bytes as f32 / self.config.max_size_bytes as f32).min(1.0)
};
BufferStats {
pending_count,
latest_sequence,
head_sequence: inner.head_seq,
next_sequence: inner.next_seq,
approx_disk_bytes,
max_size_bytes: self.config.max_size_bytes,
store_pressure,
}
}
fn recover(&self) -> Result<RecoveryReport> {
let removed_tmp_files = segment::clean_tmp(&self.dir)?;
let segments = self.scan_segments()?;
let total_bytes: u64 = segments
.iter()
.filter_map(|s| fs::metadata(self.segment_path(s.start, s.end)).ok())
.map(|m| m.len())
.sum();
let (head_seq, next_seq) = match (segments.first(), segments.last()) {
(Some(first), Some(last)) => (first.start, last.end + 1),
_ => (0, 0),
};
let segment_count = segments.len();
{
let mut inner = self.inner.lock();
inner.head_seq = head_seq;
inner.next_seq = next_seq;
}
self.approx_disk_bytes
.store(total_bytes, std::sync::atomic::Ordering::Relaxed);
*self.scan_cache.lock() = Some(segments.clone());
info!(
path = self.dir.display().to_string(),
segments = segment_count,
seq = head_seq,
end_seq = next_seq,
bytes = total_bytes,
removed_tmp = removed_tmp_files,
"Segment buffer recovered"
);
Ok(RecoveryReport {
segment_count,
head_seq,
next_seq,
disk_bytes: total_bytes,
removed_tmp_files,
})
}
fn write_segment(&self, start: u64, end: u64, events: &[T]) -> Result<u64> {
let path = self.segment_path(start, end);
segment::write(
&self.dir,
self.config.cipher.as_deref(),
self.config.compression_level,
SegmentRange::new(start, end),
events,
)
.map_err(|e| e.with_path(&path))
}
fn read_segment(&self, seg: SegmentRange) -> Result<Vec<T>> {
let path = self.segment_path(seg.start, seg.end);
segment::read(&self.dir, self.config.cipher.as_deref(), seg).map_err(|e| e.with_path(&path))
}
fn scan_segments(&self) -> Result<Vec<SegmentRange>> {
{
let cache = self.scan_cache.lock();
if let Some(ref segments) = *cache {
return Ok(segments.clone());
}
}
let segments = segment::scan(&self.dir).map_err(|e| e.with_path(&self.dir))?;
let mut cache = self.scan_cache.lock();
*cache = Some(segments.clone());
Ok(segments)
}
fn invalidate_scan_cache(&self) {
let mut cache = self.scan_cache.lock();
*cache = None;
}
fn segment_path(&self, start: u64, end: u64) -> PathBuf {
self.dir.join(segment::filename(start, end))
}
}
const _: () = {
const fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<SegmentBuffer<()>>();
};
#[cfg(test)]
mod tests;
#[cfg(test)]
mod property_tests;