use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
use std::sync::mpsc::{sync_channel, Receiver, RecvTimeoutError, SyncSender};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
use tracing::warn;
use crate::archive::{Archive, ArchiveMut, CallerRow, Evicted, SegmentMeta, SourceMeta, WalRow};
use crate::error::{Error, Result};
use crate::segment::SegmentEncoder;
pub type StreamFilter = Box<dyn Fn(&str) -> bool + Send>;
enum Msg {
AddSource {
seed: Box<SourceMeta>,
uuid: Option<String>,
reply: SyncSender<Result<i64>>,
},
ResumeSource {
source_id: i64,
clock_anchor_wall_ns: i64,
reply: SyncSender<Result<Resumed>>,
},
Wal { ticks: Vec<(i64, Vec<WalRow>)> },
Seal { source_id: i64, batch: Vec<String> },
AdoptSegment {
source_id: i64,
segment: Box<Adopted>,
reply: SyncSender<Result<bool>>,
},
Evict {
source_id: i64,
cutoff_ts: i64,
streams: Option<StreamFilter>,
reply: SyncSender<Result<Evicted>>,
},
UpdateMetadata {
source_id: i64,
patch: BTreeMap<String, String>,
},
ClockOffset {
source_id: i64,
ts: i64,
offset_ns: i64,
},
CallerRows {
source_id: i64,
stream: String,
rows: Vec<CallerRow>,
},
Finalize {
source_id: i64,
clock_offset: (i64, i64),
},
Shutdown,
Sync(SyncSender<()>),
#[cfg(any(test, feature = "test-support"))]
Commits(SyncSender<u64>),
}
type ErrorSlot = Arc<Mutex<Option<Arc<Error>>>>;
type ShadowCounts = Arc<Mutex<BTreeMap<i64, u64>>>;
#[doc(hidden)]
pub const RECLAIM_PAGES_PER_PASS: u32 = 100;
#[doc(hidden)]
pub const RECLAIM_FREELIST_DIVISOR: u32 = 10;
pub const RECLAIM_AT_CLOSE_BUDGET: Duration = Duration::from_secs(2);
pub struct Writer {
tx: Option<SyncSender<Msg>>,
thread: Option<JoinHandle<Result<()>>>,
path: PathBuf,
err: ErrorSlot,
shadowed: ShadowCounts,
}
impl Writer {
pub fn create(path: &Path, encoder: Box<dyn SegmentEncoder + Send>) -> Result<Self> {
Self::create_checkpointing_every(path, encoder, CHECKPOINT_INTERVAL)
}
pub fn create_checkpointing_every(
path: &Path,
encoder: Box<dyn SegmentEncoder + Send>,
checkpoint_every: Duration,
) -> Result<Self> {
let db = ArchiveMut::create(path)?;
Self::spawn(db, path, encoder, checkpoint_every, true)
}
pub fn open(path: &Path, encoder: Box<dyn SegmentEncoder + Send>) -> Result<Self> {
Self::open_checkpointing_every(path, encoder, CHECKPOINT_INTERVAL)
}
pub fn open_checkpointing_every(
path: &Path,
encoder: Box<dyn SegmentEncoder + Send>,
checkpoint_every: Duration,
) -> Result<Self> {
let db = ArchiveMut::open_for_write(path)?;
Self::spawn(db, path, encoder, checkpoint_every, false)
}
#[cfg(any(test, feature = "test-support"))]
pub fn create_with_busy_timeout(
path: &Path,
encoder: Box<dyn SegmentEncoder + Send>,
checkpoint_every: Duration,
busy_timeout: Duration,
) -> Result<Self> {
let db = ArchiveMut::create(path)?;
db.set_busy_timeout(busy_timeout)?;
Self::spawn(db, path, encoder, checkpoint_every, true)
}
fn spawn(
db: ArchiveMut,
path: &Path,
encoder: Box<dyn SegmentEncoder + Send>,
checkpoint_every: Duration,
created: bool,
) -> Result<Self> {
let (tx, rx) = sync_channel(1);
let err: ErrorSlot = Arc::new(Mutex::new(None));
let thread_err = Arc::clone(&err);
let shadowed: ShadowCounts = Arc::new(Mutex::new(BTreeMap::new()));
let thread_shadowed = Arc::clone(&shadowed);
let thread = match std::thread::Builder::new()
.name("dendro-writer".to_string())
.spawn(move || {
writer_thread(
rx,
db,
thread_err,
thread_shadowed,
checkpoint_every,
encoder,
)
}) {
Ok(thread) => thread,
Err(e) => {
if created {
Archive::remove_archive(path);
}
return Err(Error::Message(format!(
"failed to spawn the archive writer thread: {e}"
)));
}
};
Ok(Self {
tx: Some(tx),
thread: Some(thread),
path: path.to_path_buf(),
err,
shadowed,
})
}
pub fn add_source(&mut self, seed: SourceMeta) -> Result<SourceWriter> {
self.add_source_with_uuid(seed, None)
}
pub fn add_source_with_uuid(
&mut self,
seed: SourceMeta,
uuid: Option<&str>,
) -> Result<SourceWriter> {
let stagger_key = crate::seal::source_stagger_key(&seed.labels);
let Some(tx) = self.tx.as_ref() else {
return Err("the archive writer thread has already been joined".into());
};
let (reply_tx, reply_rx) = sync_channel(0);
if tx
.send(Msg::AddSource {
seed: Box::new(seed),
uuid: uuid.map(str::to_string),
reply: reply_tx,
})
.is_err()
{
return Err(self.take_error());
}
let source_id = match reply_rx.recv() {
Ok(inserted) => inserted?,
Err(_) => return Err(self.take_error()),
};
Ok(SourceWriter {
tx: tx.clone(),
source_id,
stagger_key,
err: Arc::clone(&self.err),
path: self.path.clone(),
floor_ts: None,
shadowed: Arc::clone(&self.shadowed),
})
}
pub fn resume_source(
&mut self,
source_id: i64,
clock_anchor_wall_ns: i64,
) -> Result<(SourceWriter, Option<i64>)> {
let Some(tx) = self.tx.as_ref() else {
return Err(Error::WriterGone);
};
let (reply_tx, reply_rx) = sync_channel(0);
if tx
.send(Msg::ResumeSource {
source_id,
clock_anchor_wall_ns,
reply: reply_tx,
})
.is_err()
{
return Err(self.take_error());
}
let resumed = match reply_rx.recv() {
Ok(resumed) => resumed?,
Err(_) => return Err(self.take_error()),
};
Ok((
SourceWriter {
tx: tx.clone(),
source_id,
stagger_key: crate::seal::source_stagger_key(&resumed.labels),
err: Arc::clone(&self.err),
path: self.path.clone(),
floor_ts: resumed.last_ts,
shadowed: Arc::clone(&self.shadowed),
},
resumed.last_ts,
))
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn join(&mut self) -> Result<()> {
if let Some(tx) = self.tx.as_ref() {
let _ = tx.send(Msg::Shutdown);
}
self.tx = None;
match self.thread.take() {
Some(handle) => handle
.join()
.unwrap_or_else(|_| Err("the archive writer thread panicked".into())),
None => Ok(()),
}
}
fn take_error(&mut self) -> Error {
take_writer_error(&self.err)
}
pub fn wal_tick(&mut self, ticks: Vec<(i64, Vec<WalRow>)>) -> Result<()> {
let ticks: Vec<(i64, Vec<WalRow>)> = ticks
.into_iter()
.filter(|(_, rows)| !rows.is_empty())
.collect();
if ticks.is_empty() {
return self.check_alive();
}
let Some(tx) = self.tx.as_ref() else {
return Err("the archive writer thread has already been joined".into());
};
if tx.send(Msg::Wal { ticks }).is_ok() {
return Ok(());
}
Err(take_writer_error(&self.err))
}
#[cfg(any(test, feature = "test-support"))]
pub fn commits_for_test(&mut self) -> u64 {
let (tx, rx) = sync_channel(0);
let Some(sender) = self.tx.as_ref() else {
return 0;
};
if sender.send(Msg::Commits(tx)).is_err() {
return 0;
}
rx.recv().unwrap_or(0)
}
fn check_alive(&mut self) -> Result<()> {
match self.err.lock() {
Ok(guard) if guard.is_some() => Err(Error::Writer(guard.clone().expect("is_some"))),
_ => Ok(()),
}
}
#[cfg(any(test, feature = "test-support"))]
pub fn finalize_single(mut self, writer: SourceWriter, clock_offset: (i64, i64)) -> Result<()> {
let queued = writer.finalize(clock_offset);
let joined = self.join();
queued.and(joined)
}
#[cfg(any(test, feature = "test-support"))]
pub fn single(
path: &Path,
encoder: Box<dyn SegmentEncoder + Send>,
seed: SourceMeta,
) -> Result<(Self, SourceWriter)> {
let mut archive = Self::create(path, encoder)?;
let writer = archive.add_source(seed)?;
Ok((archive, writer))
}
}
impl std::fmt::Debug for Writer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Writer")
.field("path", &self.path)
.field("joined", &self.thread.is_none())
.finish_non_exhaustive()
}
}
impl Drop for Writer {
fn drop(&mut self) {
if let Err(e) = self.join() {
warn!("the archive writer failed: {e}");
}
}
}
struct Resumed {
labels: BTreeMap<String, String>,
last_ts: Option<i64>,
}
pub struct SourceWriter {
tx: SyncSender<Msg>,
source_id: i64,
floor_ts: Option<i64>,
stagger_key: String,
err: ErrorSlot,
path: PathBuf,
shadowed: ShadowCounts,
}
impl std::fmt::Debug for SourceWriter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SourceWriter")
.field("source_id", &self.source_id)
.field("path", &self.path)
.field("floor_ts", &self.floor_ts)
.finish_non_exhaustive()
}
}
impl SourceWriter {
#[cfg_attr(not(test), allow(dead_code))]
pub fn path(&self) -> &Path {
&self.path
}
pub fn source_id(&self) -> i64 {
self.source_id
}
#[cfg_attr(not(test), allow(dead_code))]
pub fn stagger_key(&self) -> &str {
&self.stagger_key
}
pub fn wal(&mut self, rows: Vec<WalRow>) -> Result<()> {
if rows.is_empty() {
return self.check_alive();
}
if let Some(floor) = self.floor_ts {
if let Some(r) = rows.iter().find(|r| r.ts <= floor) {
return Err(Error::TimelineBackwards {
source_id: self.source_id,
ts: r.ts,
floor,
});
}
}
self.send(Msg::Wal {
ticks: vec![(self.source_id, rows)],
})
}
pub fn floor_ts(&self) -> Option<i64> {
self.floor_ts
}
pub fn dropped_out_of_order(&self) -> u64 {
self.shadowed
.lock()
.ok()
.and_then(|c| c.get(&self.source_id).copied())
.unwrap_or(0)
}
pub fn seal(&mut self, batch: Vec<String>) -> Result<()> {
if batch.is_empty() {
return self.check_alive();
}
self.send(Msg::Seal {
source_id: self.source_id,
batch,
})
}
pub fn adopt_segment(
&mut self,
stream: &str,
meta: &SegmentMeta,
bytes: &[u8],
caller_index: Option<&[u8]>,
) -> Result<bool> {
let (reply_tx, reply_rx) = sync_channel(0);
if self
.tx
.send(Msg::AdoptSegment {
source_id: self.source_id,
segment: Box::new(Adopted {
stream: stream.to_string(),
meta: *meta,
bytes: bytes.to_vec(),
caller_index: caller_index.map(<[u8]>::to_vec),
}),
reply: reply_tx,
})
.is_err()
{
return Err(take_writer_error(&self.err));
}
match reply_rx.recv() {
Ok(adopted) => adopted,
Err(_) => Err(take_writer_error(&self.err)),
}
}
pub fn update_metadata(&mut self, patch: BTreeMap<String, String>) -> Result<()> {
if patch.is_empty() {
return self.check_alive();
}
self.send(Msg::UpdateMetadata {
source_id: self.source_id,
patch,
})
}
pub fn clock_offset(&mut self, ts: i64, offset_ns: i64) -> Result<()> {
self.send(Msg::ClockOffset {
source_id: self.source_id,
ts,
offset_ns,
})
}
pub fn caller_rows(&mut self, stream: impl Into<String>, rows: Vec<CallerRow>) -> Result<()> {
if rows.is_empty() {
return self.check_alive();
}
self.send(Msg::CallerRows {
source_id: self.source_id,
stream: stream.into(),
rows,
})
}
pub fn evict_before(&mut self, cutoff_ts: i64) -> Result<Evicted> {
self.evict(cutoff_ts, None)
}
pub fn evict_streams_before(&mut self, cutoff_ts: i64, keep: StreamFilter) -> Result<Evicted> {
self.evict(cutoff_ts, Some(keep))
}
fn evict(&mut self, cutoff_ts: i64, streams: Option<StreamFilter>) -> Result<Evicted> {
let (tx, rx) = sync_channel(0);
self.send(Msg::Evict {
source_id: self.source_id,
cutoff_ts,
streams,
reply: tx,
})?;
rx.recv().map_err(|_| take_writer_error(&self.err))?
}
pub fn sync(&mut self) -> Result<()> {
let (tx, rx) = sync_channel(0);
self.send(Msg::Sync(tx))?;
rx.recv().map_err(|_| take_writer_error(&self.err))
}
pub fn finalize(mut self, clock_offset: (i64, i64)) -> Result<()> {
self.send(Msg::Finalize {
source_id: self.source_id,
clock_offset,
})?;
self.sync()
}
fn check_alive(&mut self) -> Result<()> {
match self.err.lock() {
Ok(guard) if guard.is_some() => Err(Error::Writer(guard.clone().expect("is_some"))),
_ => Ok(()),
}
}
fn send(&mut self, msg: Msg) -> Result<()> {
if self.tx.send(msg).is_ok() {
return Ok(());
}
Err(take_writer_error(&self.err))
}
}
fn take_writer_error(slot: &ErrorSlot) -> Error {
match slot.lock().ok().and_then(|guard| guard.clone()) {
Some(e) => Error::Writer(e),
None => Error::WriterGone,
}
}
const RETRY_BACKOFF: [Duration; 3] = [
Duration::from_millis(10),
Duration::from_millis(50),
Duration::from_millis(250),
];
const MAX_CONSECUTIVE_DROPPED_TICKS: u32 = 30;
fn with_retries<T>(what: &str, mut op: impl FnMut() -> Result<T>) -> Result<T> {
let mut attempt = 0usize;
loop {
match op() {
Ok(v) => return Ok(v),
Err(e) if e.is_retryable() && attempt < RETRY_BACKOFF.len() => {
warn!(
"{what} failed ({e}); retrying in {:?}",
RETRY_BACKOFF[attempt]
);
std::thread::sleep(RETRY_BACKOFF[attempt]);
attempt += 1;
}
Err(e) => return Err(e),
}
}
}
#[derive(Default)]
struct WriterHealth {
consecutive_dropped: u32,
warned_collisions: BTreeSet<i64>,
warned_floors: BTreeSet<i64>,
warned_shadowed: BTreeSet<(i64, String)>,
}
impl WriterHealth {
fn committed(&mut self) {
self.consecutive_dropped = 0;
}
fn count_shadowed(
&mut self,
shadowed: &ShadowCounts,
source_id: i64,
stream: &str,
ts: i64,
watermark: i64,
) {
if let Ok(mut counts) = shadowed.lock() {
*counts.entry(source_id).or_insert(0) += 1;
}
if self.warned_shadowed.insert((source_id, stream.to_string())) {
warn!(
"source {source_id}, stream {stream}: dropped an append at {ts}, which is at \
or below the newest row already sealed there ({watermark}). Such a row is \
invisible to every read path, so it is not stored. Append in order per \
stream; later drops on this stream are counted, not logged"
);
}
}
fn dropped(&mut self, e: Error) -> Result<()> {
self.consecutive_dropped += 1;
warn!(
"a tick was dropped after retries ({e}); {} consecutive",
self.consecutive_dropped
);
if self.consecutive_dropped >= MAX_CONSECUTIVE_DROPPED_TICKS {
return Err(Error::Message(format!(
"the archive writer dropped {} consecutive ticks; last error: {e}",
self.consecutive_dropped
)));
}
Ok(())
}
}
fn commit_tick(
db: &mut ArchiveMut,
ticks: &[(i64, Vec<WalRow>)],
floors: &BTreeMap<i64, i64>,
watermarks: &BTreeMap<i64, BTreeMap<String, i64>>,
shadowed: &ShadowCounts,
health: &mut WriterHealth,
) -> Result<()> {
let mut kept: Vec<(i64, Vec<WalRow>)> = Vec::with_capacity(ticks.len());
for (source_id, rows) in ticks {
match floors.get(source_id) {
Some(floor) if rows.iter().any(|r| r.ts <= *floor) => {
if health.warned_floors.insert(*source_id) {
let ts = rows.iter().map(|r| r.ts).min().unwrap_or(*floor);
warn!(
"{}; the tick was dropped, and later ones for this source are \
not logged",
Error::TimelineBackwards {
source_id: *source_id,
ts,
floor: *floor
}
);
}
}
_ => kept.push((*source_id, rows.clone())),
}
}
let ticks: &[(i64, Vec<WalRow>)] = if kept.len() == ticks.len() {
ticks
} else {
&kept
};
let shadows = |source_id: i64, r: &WalRow| {
watermarks
.get(&source_id)
.and_then(|streams| streams.get(r.stream.as_str()))
.is_some_and(|w| r.ts <= *w)
};
let rebuilt: Vec<(i64, Vec<WalRow>)>;
let ticks: &[(i64, Vec<WalRow>)] = if ticks
.iter()
.any(|(source_id, rows)| rows.iter().any(|r| shadows(*source_id, r)))
{
let mut out: Vec<(i64, Vec<WalRow>)> = Vec::with_capacity(ticks.len());
for (source_id, rows) in ticks {
let mut keep: Vec<WalRow> = Vec::with_capacity(rows.len());
for r in rows {
if shadows(*source_id, r) {
let w = watermarks[source_id][r.stream.as_str()];
health.count_shadowed(shadowed, *source_id, &r.stream, r.ts, w);
} else {
keep.push(r.clone());
}
}
if !keep.is_empty() {
out.push((*source_id, keep));
}
}
rebuilt = out;
&rebuilt
} else {
ticks
};
if ticks.is_empty() {
return Ok(());
}
match with_retries("committing a tick", || db.insert_wal_rows_batch(ticks)) {
Ok(()) => {
health.committed();
Ok(())
}
Err(e) if e.is_constraint() => {
let mut any = false;
for (source_id, rows) in ticks {
match with_retries("committing a source's tick", || {
db.insert_wal_rows(*source_id, rows)
}) {
Ok(()) => any = true,
Err(e) if e.is_constraint() => {
if health.warned_collisions.insert(*source_id) {
warn!(
"source {source_id}: a tick was dropped because its rows \
collide with rows already committed ({e}); later collisions \
for this source are not logged"
);
}
}
Err(e) if e.is_retryable() => health.dropped(e)?,
Err(e) => return Err(e),
}
}
if any {
health.committed();
}
Ok(())
}
Err(e) if e.is_retryable() => health.dropped(e),
Err(e) => Err(e),
}
}
fn record_session(
db: &mut ArchiveMut,
source_id: i64,
clock_anchor_wall_ns: i64,
resumed_after_ts: Option<Option<i64>>,
) -> Result<()> {
use crate::keys;
let session = db.mint_uuid()?;
let metadata = db.source_metadata(source_id)?;
let mut sessions: Vec<serde_json::Value> = metadata
.get(keys::WRITER_SESSIONS)
.and_then(|v| serde_json::from_str(v).ok())
.unwrap_or_default();
let mut entry = serde_json::json!({
"session": session,
"clock_anchor_wall_ns": clock_anchor_wall_ns,
"dendro": env!("CARGO_PKG_VERSION"),
});
let mut patch = BTreeMap::new();
if let Some(after) = resumed_after_ts {
entry["resumed_after_ts"] = serde_json::json!(after);
let mut events: serde_json::Value = metadata
.get(keys::EVENTS)
.and_then(|v| serde_json::from_str(v).ok())
.unwrap_or_else(|| serde_json::json!({ "events": [] }));
if !events["events"].is_array() {
events["events"] = serde_json::json!([]);
}
events["events"]
.as_array_mut()
.expect("just ensured")
.push(serde_json::json!({
"timestamp": clock_anchor_wall_ns,
"description": "source resumed by a new writer session",
"kind": "writer_session",
"details": match after {
Some(ts) => format!("previous session's last row at {ts}"),
None => "previous session left no rows".to_string(),
},
"id": format!("writer_session:{session}"),
}));
patch.insert(keys::EVENTS.to_string(), events.to_string());
}
sessions.push(entry);
patch.insert(
keys::WRITER_SESSIONS.to_string(),
serde_json::Value::Array(sessions).to_string(),
);
db.patch_source_metadata(source_id, &patch)
}
fn resume_source(
db: &mut ArchiveMut,
source_id: i64,
clock_anchor_wall_ns: i64,
encoder: &(dyn SegmentEncoder + Send),
) -> Result<Resumed> {
let Some(src) = db.read_sources()?.into_iter().find(|s| s.id == source_id) else {
return Err(Error::Message(format!(
"no source with id {source_id} to resume"
)));
};
crate::segment::check_encoder(source_id, &src.meta.metadata, encoder)?;
let (_, last_ts) = db.source_time_span(source_id)?;
if let Some(floor) = last_ts {
if clock_anchor_wall_ns <= floor {
return Err(Error::TimelineBackwards {
source_id,
ts: clock_anchor_wall_ns,
floor,
});
}
}
db.transaction(|tx| tx.mark_incomplete(source_id))?;
record_session(db, source_id, clock_anchor_wall_ns, Some(last_ts))?;
Ok(Resumed {
labels: src.meta.labels,
last_ts,
})
}
struct Encoded {
stream: String,
seq: u64,
meta: SegmentMeta,
bytes: Vec<u8>,
caller_index: Option<Vec<u8>>,
}
pub(crate) struct Adopted {
stream: String,
meta: SegmentMeta,
bytes: Vec<u8>,
caller_index: Option<Vec<u8>>,
}
fn writer_thread(
rx: Receiver<Msg>,
mut db: ArchiveMut,
err_slot: ErrorSlot,
shadowed: ShadowCounts,
checkpoint_every: Duration,
encoder: Box<dyn SegmentEncoder + Send>,
) -> Result<()> {
match writer_loop(&rx, &mut db, &shadowed, checkpoint_every, encoder.as_ref()) {
Ok(()) => Ok(()),
Err(e) => {
let shared = Arc::new(e);
*err_slot.lock().unwrap_or_else(|e| e.into_inner()) = Some(Arc::clone(&shared));
Err(Error::Writer(shared))
}
}
}
pub const CHECKPOINT_INTERVAL: Duration = Duration::from_secs(10);
fn writer_loop(
rx: &Receiver<Msg>,
db: &mut ArchiveMut,
shadowed: &ShadowCounts,
checkpoint_every: Duration,
encoder: &(dyn SegmentEncoder + Send),
) -> Result<()> {
let mut next_seq: BTreeMap<(i64, String), u64> = db.next_seqs()?;
let mut floors: BTreeMap<i64, i64> = BTreeMap::new();
let mut watermarks: BTreeMap<i64, BTreeMap<String, i64>> = db.sealed_watermarks()?;
let mut added: usize = 0;
let mut finalized: usize = 0;
let mut last_checkpoint = Instant::now();
let mut health = WriterHealth::default();
loop {
let waited = rx.recv_timeout(checkpoint_every.saturating_sub(last_checkpoint.elapsed()));
if last_checkpoint.elapsed() >= checkpoint_every {
if let Err(e) = db.checkpoint_passive() {
warn!("failed to checkpoint the WAL: {e}");
}
last_checkpoint = Instant::now();
}
let received = match waited {
Ok(msg) => Ok(msg),
Err(RecvTimeoutError::Timeout) => continue,
Err(RecvTimeoutError::Disconnected) => Err(()),
};
match received {
Ok(Msg::Sync(reply)) => {
let _ = reply.send(());
}
Ok(Msg::AddSource { seed, uuid, reply }) => {
let inserted = db
.insert_source_with_uuid(&seed, uuid.as_deref())
.and_then(|id| {
record_session(db, id, seed.clock_anchor_wall_ns, None)?;
if let Some(version) = encoder.version() {
let mut patch = BTreeMap::new();
patch.insert(crate::keys::ENCODER.to_string(), version.to_string());
db.patch_source_metadata(id, &patch)?;
}
Ok(id)
});
if inserted.is_ok() {
added += 1;
}
let _ = reply.send(inserted);
}
Ok(Msg::ResumeSource {
source_id,
clock_anchor_wall_ns,
reply,
}) => {
let resumed = resume_source(db, source_id, clock_anchor_wall_ns, encoder);
if let Ok(Resumed {
last_ts: Some(floor),
..
}) = &resumed
{
floors.insert(source_id, *floor);
}
if resumed.is_ok() {
added += 1;
}
let _ = reply.send(resumed);
}
Ok(Msg::Wal { ticks }) => {
commit_tick(db, &ticks, &floors, &watermarks, shadowed, &mut health)?
}
#[cfg(any(test, feature = "test-support"))]
Ok(Msg::Commits(reply)) => {
let _ = reply.send(db.commits());
}
Ok(Msg::Seal { source_id, batch }) => {
match with_retries("sealing a segment batch", || {
seal_batch(
db,
source_id,
&mut next_seq,
&mut watermarks,
batch.clone(),
encoder,
)
}) {
Ok(()) => {}
Err(e) if e.is_retryable() => warn!(
"seal of {} stream(s) for source {source_id} deferred ({e}); their \
rows stay in the WAL and seal with the next batch",
batch.len()
),
Err(e) => return Err(e),
}
}
Ok(Msg::AdoptSegment {
source_id,
segment,
reply,
}) => {
let adopted = with_retries("adopting a segment", || {
adopt_segment(db, source_id, &mut next_seq, &mut watermarks, &segment)
});
let _ = reply.send(adopted);
}
Ok(Msg::Evict {
source_id,
cutoff_ts,
streams,
reply,
}) => {
let evicted = match streams {
Some(keep) => db.evict_streams_before(source_id, cutoff_ts, &*keep),
None => db.evict_before(source_id, cutoff_ts),
};
if let Ok(e) = &evicted {
if e.live_rows > 0 {
warn!(
"retention on source {source_id} deleted {} row(s) no segment \
held: the stream's seal cadence is slower than the lookback. \
Seal at least as often as you evict",
e.live_rows
);
}
}
let failed = evicted.is_err();
let _ = reply.send(evicted);
if failed {
continue;
}
if let Err(e) = reclaim_if_fragmented(db) {
warn!("reclaiming freed pages after retention failed ({e}); skipped");
}
}
Ok(Msg::UpdateMetadata { source_id, patch }) => {
if let Err(e) = with_retries("updating source metadata", || {
db.patch_source_metadata(source_id, &patch)
}) {
warn!(
"metadata update for source {source_id} dropped ({e}); keys: {:?}",
patch.keys().collect::<Vec<_>>()
);
}
}
Ok(Msg::ClockOffset {
source_id,
ts,
offset_ns,
}) => {
match with_retries("committing a clock offset", || {
db.transaction(|tx| tx.insert_clock_offset(source_id, ts, offset_ns))
}) {
Ok(()) => health.committed(),
Err(e) if e.is_retryable() => health.dropped(e)?,
Err(e) => return Err(e),
}
}
Ok(Msg::CallerRows {
source_id,
stream,
rows,
}) => {
match with_retries("committing caller rows", || {
db.insert_caller_rows(source_id, &stream, &rows)
}) {
Ok(()) => health.committed(),
Err(e) if e.is_retryable() => health.dropped(e)?,
Err(e) => return Err(e),
}
}
Ok(Msg::Finalize {
source_id,
clock_offset,
}) => {
with_retries("finalizing a source", || {
db.transaction(|tx| {
tx.insert_clock_offset(source_id, clock_offset.0, clock_offset.1)?;
tx.mark_complete(source_id)
})
})?;
finalized += 1;
}
Ok(Msg::Shutdown) => {
if added > 0 && finalized == added {
reclaim_all(db)?;
}
return Ok(());
}
Err(_) => {
if added > 0 && finalized == added {
reclaim_all(db)?;
}
return Ok(());
}
}
}
}
#[cfg_attr(not(any(test, feature = "test-support")), doc(hidden))]
#[doc(hidden)]
pub fn reclaim_if_fragmented(db: &mut ArchiveMut) -> Result<()> {
if should_reclaim(
db.pragma_u32("freelist_count")?,
db.pragma_u32("page_count")?,
) {
db.incremental_vacuum(RECLAIM_PAGES_PER_PASS)?;
}
Ok(())
}
#[cfg_attr(not(any(test, feature = "test-support")), doc(hidden))]
#[doc(hidden)]
pub fn should_reclaim(free_pages: u32, pages: u32) -> bool {
free_pages.saturating_mul(RECLAIM_FREELIST_DIVISOR) > pages
}
fn reclaim_all(db: &mut ArchiveMut) -> Result<()> {
let started = Instant::now();
while db.pragma_u32("freelist_count")? > 0 {
db.incremental_vacuum(RECLAIM_PAGES_PER_PASS)?;
if started.elapsed() >= RECLAIM_AT_CLOSE_BUDGET {
warn!(
"stopped reclaiming freed pages after {:?}; {} page(s) remain on the free list \
and will be reclaimed by a later retention pass",
RECLAIM_AT_CLOSE_BUDGET,
db.pragma_u32("freelist_count")?
);
break;
}
}
Ok(())
}
fn adopt_segment(
db: &mut ArchiveMut,
source_id: i64,
next_seq: &mut BTreeMap<(i64, String), u64>,
watermarks: &mut BTreeMap<i64, BTreeMap<String, i64>>,
segment: &Adopted,
) -> Result<bool> {
let Adopted {
stream,
meta,
bytes,
caller_index,
} = segment;
let stream = stream.as_str();
if meta.rows == 0 || meta.first_ts > meta.last_ts {
return Err(Error::Message(format!(
"source {source_id}: the segment offered for `{stream}` claims {} row(s) over [{}, {}], which is not a segment that can exist",
meta.rows, meta.first_ts, meta.last_ts
)));
}
let watermark = watermarks
.get(&source_id)
.and_then(|m| m.get(stream))
.copied();
if let Some(w) = watermark {
if meta.last_ts <= w {
return Ok(false);
}
if meta.first_ts <= w {
return Err(Error::Message(format!(
"source {source_id}: the segment offered for `{stream}` spans [{}, {}], which straddles the stream's newest sealed row ({w}); its rows at or below {w} could never be read, and splitting it would mean decoding it",
meta.first_ts, meta.last_ts
)));
}
}
let live = db.live_wal_span(source_id, stream)?;
if let Some(first) = live.first_ts {
if first <= meta.last_ts {
return Err(Error::Message(format!(
"source {source_id}: the segment offered for `{stream}` ends at {}, at or after an unsealed row already in the WAL ({first}); adopting it would shadow rows no segment holds",
meta.last_ts
)));
}
}
let seq = next_seq
.get(&(source_id, stream.to_string()))
.copied()
.unwrap_or(0);
db.insert_segment_with_index(source_id, stream, seq, meta, bytes, caller_index.as_deref())?;
next_seq.insert((source_id, stream.to_string()), seq + 1);
watermarks
.entry(source_id)
.or_default()
.insert(stream.to_string(), meta.last_ts);
Ok(true)
}
fn seal_batch(
db: &mut ArchiveMut,
source_id: i64,
next_seq: &mut BTreeMap<(i64, String), u64>,
watermarks: &mut BTreeMap<i64, BTreeMap<String, i64>>,
batch: Vec<String>,
encoder: &(dyn SegmentEncoder + Send),
) -> Result<()> {
let mut encoded = Vec::with_capacity(batch.len());
let mut observation: Option<(i64, i64)> = None;
for stream in batch {
let rows = db.live_wal(source_id, &stream)?;
if rows.is_empty() {
if db.read_segments(source_id, &stream)?.is_empty() {
warn!(
"asked to seal `{stream}`, which has no live rows and has \
never sealed a segment - is the name right?"
);
}
continue;
}
let Some(tail) = crate::segment::materialize(encoder, &stream, &rows)? else {
continue;
};
let segment_last = rows
.iter()
.rev()
.find(|r| r.ts == tail.last_ts)
.expect("the segment's last_ts is one of the rows, by the check above");
if observation.is_none_or(|(seen, _)| segment_last.ts >= seen) {
observation = Some((segment_last.ts, segment_last.wall_offset));
}
let seq = next_seq
.get(&(source_id, stream.clone()))
.copied()
.unwrap_or(0);
encoded.push(Encoded {
stream,
seq,
meta: SegmentMeta {
rows: tail.rows,
first_ts: tail.first_ts,
last_ts: tail.last_ts,
},
bytes: tail.bytes,
caller_index: tail.index,
});
}
db.transaction(|tx| {
for e in &encoded {
tx.insert_segment_with_index(
source_id,
&e.stream,
e.seq,
&e.meta,
&e.bytes,
e.caller_index.as_deref(),
)?;
}
if let Some((ts, offset)) = observation {
tx.insert_clock_offset(source_id, ts, offset)?;
}
Ok(())
})?;
for e in &encoded {
next_seq.insert((source_id, e.stream.clone()), e.seq + 1);
watermarks
.entry(source_id)
.or_default()
.insert(e.stream.clone(), e.meta.last_ts);
}
for e in &encoded {
db.prune_wal(source_id, &e.stream, e.meta.last_ts)?;
}
Ok(())
}