use std::{ path::PathBuf, time::{Duration, Instant}, }; use crossbeam::channel::Receiver; use super::{ dedup::DedupWindow, error::WalError, format::{self, EventRecord, MAX_EVENTS_PER_BATCH, SessionSeqNo, SessionWalEvent}, segment::{self, SegmentWriter}, session_journal::SessionJournal, }; use crate::replication::{RegionId, ShardId}; /// A single queued append: the event plus the caller's reply channel. /// /// The reply channel receives the assigned sequence number once the batch /// containing this event is durably fsynced, the dedup sentinel `0` if the /// event was suppressed as a duplicate, or the write error if the batch failed /// to persist. It is critical that *every* queued append eventually resolves /// its reply — dropping the sender silently surfaces as `Closed` to a caller /// that is in fact still durable-or-not, so both the steady-state loop and the /// shutdown drain funnel through the same [`flush_batch`] routine. type QueuedAppend = ( EventRecord, crossbeam::channel::Sender>, ); // Narrow test-only fault hook: when set, the next `flush_batch` write fails. // // Lets a unit test drive the *error* arm of `flush_batch` deterministically // (without a real I/O fault) to prove that both the steady-state loop and the // shutdown-drain tail notify every waiting caller with an error instead of // dropping their reply channels. Thread-local so concurrent tests don't // interfere. Always compiled out of production builds. #[cfg(test)] thread_local! { static FAIL_NEXT_FLUSH: std::cell::Cell = const { std::cell::Cell::new(false) }; } /// Arm the test-only flush-failure hook for the current thread. #[cfg(test)] pub(crate) fn arm_flush_failure() { FAIL_NEXT_FLUSH.with(|c| c.set(true)); } /// Returns and clears the test-only flush-failure flag for the current thread. #[cfg(test)] fn take_flush_failure() -> bool { FAIL_NEXT_FLUSH.with(|c| c.replace(false)) } /// Commands sent from `WalHandle` to the writer thread. pub enum WalCommand { /// Append a signal event. The reply channel receives the assigned /// sequence number (or an error) once the batch containing this /// event has been durably fsynced. Append { event: EventRecord, reply: crossbeam::channel::Sender>, }, /// Delete segments whose first sequence number is less than `before_seq`. /// Runs inside the writer thread to avoid racing with concurrent writes. TruncateBefore { before_seq: u64, reply: crossbeam::channel::Sender>, }, /// Graceful shutdown: flush remaining events and exit. Shutdown, // ── Session lifecycle commands ──────────────────────────────────────── // These are fire-and-forget (no reply channel). They bypass the signal // batch system and write directly to the session journal with fsync. /// Record that a session was started. SessionStart { session_id: u64, user_id: u64, started_at_ns: u64, agent_id: String, policy_name: String, }, /// Record that a signal was written within a session. SessionSignal { session_id: u64, entity_id: u64, weight: f64, ts_ns: u64, signal_name: String, annotation: Option, /// Monotonic sequence number for this write (used for idempotent replay). /// `None` for legacy writes that predate the seqno mechanism. session_seqno: Option, /// BLAKE3-derived idempotency key for duplicate suppression. /// `None` for legacy writes. idempotency_key: Option, }, /// Record that a session was closed. SessionClose { session_id: u64 }, } /// Configuration for the group commit writer. pub struct WriterConfig { pub dir: PathBuf, pub segment_size: u64, pub batch_size: usize, pub batch_timeout: Duration, pub dedup_window: Duration, /// Path for the session journal file (optional; `None` in ephemeral mode). pub session_journal_path: Option, /// Shard identity for this writer. Written into every batch header so /// `WalShipper` and receivers can identify the source shard. /// Defaults to `ShardId::SINGLE` for single-node deployments. pub shard_id: ShardId, /// Region identity for this writer. Written into every batch header. /// Defaults to `RegionId::SINGLE` for single-node deployments. pub region_id: RegionId, } /// Validate a writer configuration before the writer thread is spawned. /// /// The wire format caps a single batch at [`MAX_EVENTS_PER_BATCH`] events /// (`format::encode_batch_with_shard` rejects anything larger). The writer /// drains up to `batch_size` events into one batch, so a `batch_size` above the /// cap would make the *first* full batch fail to encode — and because that /// failure propagates out of `run_writer`, it would terminate the writer thread /// and wedge every subsequent append forever. We reject it here so the caller of /// `WalHandle::open` fails loudly at startup instead of silently after the first /// full batch. /// /// # Errors /// /// Returns [`WalError::InvalidConfig`] if `batch_size` is `0` or exceeds /// [`MAX_EVENTS_PER_BATCH`]. pub(crate) fn validate_writer_config(batch_size: usize) -> Result<(), WalError> { let max = usize::from(MAX_EVENTS_PER_BATCH); if batch_size == 0 { return Err(WalError::InvalidConfig { message: "batch_size must be at least 1".to_string(), }); } if batch_size > max { return Err(WalError::InvalidConfig { message: format!("batch_size {batch_size} exceeds MAX_EVENTS_PER_BATCH ({max})"), }); } Ok(()) } /// The largest batch the wire format can encode in a single fsync. /// /// `run_writer` clamps its drain limit to this value as a defence-in-depth /// belt-and-braces guard: `validate_writer_config` already rejects an oversized /// `batch_size` at open time, but if a future caller constructs a `WriterConfig` /// directly (bypassing validation) the clamp still prevents the thread from ever /// building an unencodable batch and crashing. Correctness is preserved either /// way — clamping only means more, smaller batches, never a dropped event. fn effective_batch_size(batch_size: usize) -> usize { batch_size.clamp(1, usize::from(MAX_EVENTS_PER_BATCH)) } /// Encode, write, and fsync one batch of non-duplicate events, notifying every /// caller of the outcome. /// /// This is the single shared commit routine used by both the steady-state loop /// in [`run_writer`] and its shutdown-drain tail, so the two paths cannot /// diverge in how they handle a write failure (the divergence that previously /// let the drain path drop reply channels on error). Responsibilities: /// /// 1. Encode the batch at `batch_seq` with the writer's shard/region identity. /// 2. Rotate the segment first if it has reached its size threshold. /// 3. Write the encoded bytes and fsync. /// 4. On success: notify each caller of its assigned sequence number and return /// the next free sequence number. /// 5. On failure: notify *every* waiting caller with the underlying error before /// propagating it, so no caller is ever left blocked on a dropped channel. /// /// `kept_events` and `kept_replies` are 1:1 and must be non-empty (the callers /// only invoke this when there is at least one event to persist). /// /// # Errors /// /// Returns the underlying [`WalError`] from encode/rotate/write/sync. On error, /// all reply channels have already been notified with an equivalent error. /// /// Never panics on a clock anomaly: the batch timestamp is sourced from /// [`crate::schema::Timestamp::now`], which saturates a pre-Unix-epoch clock to /// the epoch and logs a warning rather than panicking. WAL ordering is by /// sequence number, not timestamp, so a clamped timestamp is informational only. fn flush_batch( segment: &mut SegmentWriter, config: &WriterConfig, batch_seq: u64, kept_events: &[EventRecord], kept_replies: Vec>>, ) -> Result { debug_assert_eq!(kept_events.len(), kept_replies.len()); // Source the batch timestamp from the canonical clock-anomaly-safe helper: // a pre-Unix-epoch wall clock (NTP step-back, dead/uninitialized RTC at boot) // saturates to the epoch and logs a warning instead of panicking the writer // thread. The timestamp is informational metadata only — the WAL is ordered // by sequence number — so a clamped value is safe. let batch_ts = crate::schema::Timestamp::now().as_nanos(); let write_result = (|| -> Result<(), WalError> { // Test-only deterministic fault: exercise the caller-notification error // arm without a real I/O failure. Compiled out of production builds. #[cfg(test)] if take_flush_failure() { return Err(WalError::Io(std::io::Error::other( "injected flush failure", ))); } let encoded = format::encode_batch_with_shard( kept_events, batch_seq, batch_ts, config.shard_id, config.region_id, )?; if segment.needs_rotation() { segment.rotate(batch_seq)?; } segment.write_batch_bytes(&encoded)?; segment.sync()?; Ok(()) })(); match write_result { Ok(()) => { let event_count = kept_events.len() as u64; tracing::debug!(seq = batch_seq, events = event_count, "wal: batch appended"); for (i, reply) in kept_replies.into_iter().enumerate() { let _ = reply.send(Ok(batch_seq + i as u64)); } Ok(batch_seq + event_count) } Err(err) => { // Notify all waiting callers with the actual error before // propagating. We cannot clone WalError, so we send a synthetic // I/O error carrying the same description; dropping the channels // would surface as a generic `Closed` and hide the real cause. let err_msg = err.to_string(); for reply in kept_replies { let _ = reply.send(Err(WalError::Io(std::io::Error::other(err_msg.clone())))); } Err(err) } } } /// Outcome of dispatching one [`WalCommand`] through [`handle_aux_command`]. /// /// The three command loops in [`run_writer`] (blocking recv, deadline drain, /// shutdown drain) share their handling of the side-effecting *auxiliary* /// commands (`TruncateBefore`, `Session*`) but differ only in their /// continue/break control flow. Routing every command through one helper that /// returns this enum means each loop matches on the enum — never on the raw /// command bodies — so the auxiliary-command logic (including the /// active-segment clamp in [`handle_aux_command`]) lives in exactly one place /// and cannot drift between the three sites. enum CommandOutcome { /// An `Append` was received: the loop must push `(event, reply)` onto its /// pending batch. Pushed(QueuedAppend), /// A side-effecting command (`TruncateBefore`, `Session*`) was fully /// handled inside the helper; the loop should keep going. Handled, /// `Shutdown` was received (or the channel disconnected): the loop must /// stop draining. Shutdown, } /// Dispatch a single received [`WalCommand`], executing any side effect. /// /// `Append` is returned as [`CommandOutcome::Pushed`] for the caller to batch; /// every other variant is fully handled here and reported as /// [`CommandOutcome::Handled`] or [`CommandOutcome::Shutdown`]. Sharing this /// dispatch across all three loops in [`run_writer`] keeps the auxiliary-command /// bodies byte-identical by construction. /// /// # Active-segment protection (`TruncateBefore`) /// /// `TruncateBefore` must NOT unlink the segment the live writer is appending to. /// After any write burst the active segment's `first_seq` sits below the /// materialized checkpoint, so a naive `delete_segments_before(checkpoint_seq)` /// would `remove_file` the very inode this writer still holds open — and on /// Linux the writer would keep appending to the now-unlinked inode, silently /// losing every post-checkpoint, already-fsync'd, acknowledged write on the next /// open. We therefore clamp the deletion floor to the live segment's /// `first_seq`, exactly as [`crate::wal::compaction::compact_wal_online`] does, /// guaranteeing the active segment always survives. The writer thread already /// owns `active_first_seq`, so the clamp is a single `min` with no extra I/O. fn handle_aux_command( cmd: WalCommand, config: &WriterConfig, active_first_seq: u64, session_journal: &mut Option, ) -> CommandOutcome { match cmd { WalCommand::Append { event, reply } => CommandOutcome::Pushed((event, reply)), WalCommand::TruncateBefore { before_seq, reply } => { // Clamp the deletion floor so the live segment (the maximum-first_seq // segment, which this writer holds open) is never unlinked out from // under our FD. See the function-level rustdoc and // `compaction::compact_wal_online` for the full hazard analysis. let floor = before_seq.min(active_first_seq); let result = segment::delete_segments_before(&config.dir, floor); let _ = reply.send(result.map(|_| ())); CommandOutcome::Handled } cmd @ (WalCommand::SessionStart { .. } | WalCommand::SessionSignal { .. } | WalCommand::SessionClose { .. }) => { handle_session_command(cmd, session_journal); CommandOutcome::Handled } WalCommand::Shutdown => CommandOutcome::Shutdown, } } /// Split a drained batch into kept events (1:1 with their replies) and duplicate /// replies, notifying duplicates immediately with the dedup sentinel `0`. /// /// Shared by the steady-state loop and the shutdown drain so dedup accounting /// cannot diverge between them. fn partition_dedup( dedup: &mut DedupWindow, batch: impl IntoIterator, ) -> ( Vec, Vec>>, ) { let mut kept_events: Vec = Vec::new(); let mut kept_replies: Vec>> = Vec::new(); // Hashes kept *within this batch*, so two identical events in the same drained // batch still dedup against each other even though the durable dedup window is // only updated AFTER a successful flush (see `run_writer`). This preserves // intra-batch suppression without violating the "mark as seen only once // durable" invariant. let mut batch_seen: std::collections::HashSet = std::collections::HashSet::new(); for (event, reply) in batch { let hash = format::event_content_hash(&event); // `dedup.contains` CHECKS membership without recording; `batch_seen.insert` // returns false when this exact event already appeared earlier in this // batch. Recording into `dedup` happens in `run_writer`, only after the // batch is durably persisted. if dedup.contains(&event) || !batch_seen.insert(hash) { // Duplicate: notify with the dedup sentinel (seq=0) immediately. let _ = reply.send(Ok(0)); } else { kept_events.push(event); kept_replies.push(reply); } } (kept_events, kept_replies) } /// The group commit writer loop. /// /// Runs on a dedicated thread. Receives events via crossbeam channel, /// accumulates them into batches, writes batches to the WAL segment, /// and fsyncs once per batch. Callers are notified of their sequence /// numbers via per-event reply channels. /// /// # Batch formation /// /// 1. Block until the first event arrives. /// 2. Drain additional events from the channel up to `batch_size` or /// until `batch_timeout` elapses (whichever comes first). /// 3. Deduplicate events, encode the batch, write to segment, fsync. /// 4. Send sequence numbers back to all waiting callers. /// /// # Resilience /// /// A flush failure in the **steady-state loop** does NOT terminate the writer. /// `flush_batch` notifies every waiting caller with the error (so they can /// retry), the batch's events are left un-recorded in the dedup window (so a /// retry is accepted, not suppressed as a phantom duplicate), and the loop keeps /// serving — a transient I/O fault (ENOSPC, EINTR, NFS blip) must never convert /// into a permanent write outage by dropping the command channel. /// /// # Errors /// /// Only the shutdown drain / final fsync propagate an error out of this function /// (the WAL is closing anyway, and callers were already notified). The encoding /// path cannot fail under normal operation — `effective_batch_size` clamps the /// drain limit to `MAX_EVENTS_PER_BATCH` so an oversized `batch_size` can never /// produce an unencodable batch. /// /// Never panics on a clock anomaly: [`flush_batch`] sources its batch timestamp /// from the clock-anomaly-safe [`crate::schema::Timestamp::now`], so a /// pre-Unix-epoch wall clock cannot kill the writer thread. // The encode/dedup/write/sync logic is extracted into `flush_batch` and // `partition_dedup`, and the per-command dispatch into `handle_aux_command`; // what remains is the steady-state loop plus the shutdown drain, each matching // only on the `CommandOutcome` enum. Splitting the control flow further would // obscure the single read-recv/drain structure. #[allow(clippy::too_many_lines)] pub fn run_writer( rx: &Receiver, config: &WriterConfig, mut segment: SegmentWriter, start_seq: u64, mut dedup: DedupWindow, ) -> Result<(), WalError> { let mut next_seq = start_seq; // Clamp the drain limit to what the wire format can encode in one batch. // `validate_writer_config` rejects an oversized `batch_size` at open time; // this is the defence-in-depth guard for any caller that builds a // `WriterConfig` directly. See `effective_batch_size`. let max_batch = effective_batch_size(config.batch_size); let mut batch: Vec = Vec::with_capacity(max_batch); let mut shutdown_requested = false; // Open the session journal if a path was provided (persistent mode). let mut session_journal: Option = config .session_journal_path .as_ref() .and_then(|p| match SessionJournal::open(p) { Ok(j) => Some(j), Err(e) => { tracing::error!(error = %e, "failed to open session journal; session WAL writes will be skipped"); None } }); loop { // Block until the first event arrives (or shutdown/disconnect). All // side-effecting commands route through `handle_aux_command` so the // `TruncateBefore` active-segment clamp and the `Session*` delegation // live in exactly one place; this loop only steers control flow. match rx.recv() { Ok(cmd) => { match handle_aux_command(cmd, config, segment.first_seq(), &mut session_journal) { CommandOutcome::Pushed(queued) => batch.push(queued), CommandOutcome::Handled => continue, CommandOutcome::Shutdown => break, } } Err(_) => break, } // Drain up to the (clamped) batch limit with a deadline. let deadline = Instant::now() + config.batch_timeout; while batch.len() < max_batch { match rx.recv_deadline(deadline) { Ok(cmd) => { match handle_aux_command(cmd, config, segment.first_seq(), &mut session_journal) { CommandOutcome::Pushed(queued) => batch.push(queued), // Side-effecting commands bypass the batch; keep draining. CommandOutcome::Handled => {} CommandOutcome::Shutdown => { shutdown_requested = true; break; } } } Err(crossbeam::channel::RecvTimeoutError::Disconnected) => { shutdown_requested = true; break; } Err(crossbeam::channel::RecvTimeoutError::Timeout) => break, } } // Deduplicate, then commit through the shared flush routine so this // path cannot diverge from the shutdown drain below. `partition_dedup` // notifies duplicate senders with the dedup sentinel; `flush_batch` // notifies every kept caller (success *or* error) before returning. // drain(..) reuses batch's heap allocation across loop iterations. #[allow(clippy::iter_with_drain)] let (kept_events, kept_replies) = partition_dedup(&mut dedup, batch.drain(..)); if !kept_events.is_empty() { match flush_batch(&mut segment, config, next_seq, &kept_events, kept_replies) { Ok(seq) => { next_seq = seq; // Record events as seen ONLY now that the batch is durably // persisted. Recording earlier (in `partition_dedup`) would // suppress a legitimate retry of a batch that failed to flush. for event in &kept_events { dedup.record(event); } } Err(e) => { // A transient I/O fault (a brief ENOSPC that an operator then // clears, an EINTR, an NFS/network-storage blip) must NOT tear // down the writer thread — that would drop the command channel // and wedge every future write forever, since the WAL is the // source of truth for all entity/signal/relationship writes. // `flush_batch` has already notified each waiting caller with // the error, so they can retry; we keep `next_seq` unchanged // (the failed batch's sequence range is free for the retry), // leave the events UNRECORDED in the dedup window so the retry // is accepted, and keep serving. tracing::error!( error = %e, seq = next_seq, events = kept_events.len(), "wal: batch flush failed; callers notified to retry, writer continuing" ); } } } if shutdown_requested { break; } } // Drain any remaining commands that arrived before senders observed // the shutdown. This ensures in-flight append() calls are not silently // dropped, which would cause callers to block forever or receive // WalError::Closed instead of a real sequence number. let mut final_batch: Vec = Vec::new(); // Same shared dispatch as the steady loop: `Append` queues into the final // batch, `TruncateBefore`/`Session*` are handled (with the active-segment // clamp), and a duplicate `Shutdown` is a no-op. The loop ends when the // channel is empty or disconnected. while let Ok(cmd) = rx.try_recv() { match handle_aux_command(cmd, config, segment.first_seq(), &mut session_journal) { CommandOutcome::Pushed(queued) => final_batch.push(queued), CommandOutcome::Handled | CommandOutcome::Shutdown => {} } } // Flush the final drain batch through the SAME shared routine the steady // loop uses. Routing both paths through `flush_batch` is what guarantees the // drain path notifies waiting callers on a write error instead of dropping // their reply channels (the divergence this consolidation fixes). A drained // batch can exceed `max_batch` if many appends queued during shutdown, so we // commit it in `max_batch`-sized chunks the wire format can encode. if !final_batch.is_empty() { let (kept_events, kept_replies) = partition_dedup(&mut dedup, final_batch); let mut events = kept_events.into_iter(); let mut replies = kept_replies.into_iter(); // `next_seq` is advanced per chunk so sequence numbers stay monotonic // across the (rare) multi-chunk drain. It is not propagated past the // writer's exit, but is read on each subsequent chunk iteration. loop { let chunk_events: Vec = events.by_ref().take(max_batch).collect(); if chunk_events.is_empty() { break; } let chunk_replies: Vec<_> = replies.by_ref().take(chunk_events.len()).collect(); next_seq = flush_batch(&mut segment, config, next_seq, &chunk_events, chunk_replies)?; } } // Final sync before exit segment.sync()?; Ok(()) } /// Write a session lifecycle command to the session journal. /// /// This function is called from the writer thread. Session commands bypass the /// signal batch system entirely. Errors are logged and swallowed -- session WAL /// writes are best-effort; the in-memory session state is the source of truth. fn handle_session_command(cmd: WalCommand, journal: &mut Option) { let Some(journal) = journal.as_mut() else { // No session journal open (should not happen in persistent mode, but // log defensively). return; }; let event = match cmd { WalCommand::SessionStart { session_id, user_id, started_at_ns, agent_id, policy_name, } => SessionWalEvent::Start { session_id, user_id, started_at_ns, agent_id, policy_name, }, WalCommand::SessionSignal { session_id, entity_id, weight, ts_ns, signal_name, annotation, session_seqno, idempotency_key, } => SessionWalEvent::Signal { session_id, entity_id, weight, ts_ns, signal_name, annotation, session_seqno: session_seqno.map(SessionSeqNo), idempotency_key, }, WalCommand::SessionClose { session_id } => SessionWalEvent::Close { session_id }, // Other commands are not handled here. _ => return, }; if let Err(e) = journal.append(&event) { tracing::warn!(error = %e, "session journal write failed"); } } #[cfg(test)] #[allow(clippy::unwrap_used, clippy::similar_names)] #[path = "writer_tests.rs"] mod tests;