tidaldb/tidal/src/wal/writer.rs

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<Result<u64, WalError>>,
);

// 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<bool> = 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<Result<u64, WalError>>,
    },
    /// 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<Result<(), WalError>>,
    },
    /// 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<String>,
        /// Monotonic sequence number for this write (used for idempotent replay).
        /// `None` for legacy writes that predate the seqno mechanism.
        session_seqno: Option<u64>,
        /// BLAKE3-derived idempotency key for duplicate suppression.
        /// `None` for legacy writes.
        idempotency_key: Option<u128>,
    },
    /// 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<PathBuf>,
    /// 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<crossbeam::channel::Sender<Result<u64, WalError>>>,
) -> Result<u64, WalError> {
    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<SessionJournal>,
) -> 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<Item = QueuedAppend>,
) -> (
    Vec<EventRecord>,
    Vec<crossbeam::channel::Sender<Result<u64, WalError>>>,
) {
    let mut kept_events: Vec<EventRecord> = Vec::new();
    let mut kept_replies: Vec<crossbeam::channel::Sender<Result<u64, WalError>>> = 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<u128> = 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<WalCommand>,
    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<QueuedAppend> = 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<SessionJournal> = 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<QueuedAppend> = 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<EventRecord> = 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<SessionJournal>) {
    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;