hayson_engine/queue/

write_queue.rs

use std::collections::{HashMap, HashSet, VecDeque};
use std::fmt;
use std::io;

use serde_json::Value;

use crate::errors::{
    HaysonCipherError, HaysonPartitionError, HaysonPendingConflictError, HaysonUniqueViolationError,
    HaysonValidationError,
};
use crate::registry::composite_index::unique_key_of;
use crate::registry::table_registry::TableRegistry;
use crate::schema::{
    partition_of, primary_key_of, validate_record, HaysonRecord, TableSchema, ValidateOptions, UNPARTITIONED,
};
use crate::storage::{serialize_entry, LogFile, LogMeta, LogOp, LogSet, Slot};

use super::pending_map::{PendingMap, ALL_FIELDS};

pub trait LogAppender {
    fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>>;

    /// Appends into one named partition. Defaults to [`LogAppender::append_lines`], which is right
    /// for anything holding a single log file: the only partition such a target has is the
    /// [`UNPARTITIONED`] one every unpartitioned record already names.
    fn append_lines_in(&mut self, part: &str, lines: &[String]) -> io::Result<Vec<Slot>> {
        let _ = part;
        self.append_lines(lines)
    }

    fn sync(&mut self) -> io::Result<()>;
}

impl LogAppender for LogFile {
    fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
        self.append(lines)
    }

    fn sync(&mut self) -> io::Result<()> {
        LogFile::sync(self)
    }
}

impl LogAppender for LogSet {
    fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
        self.append(UNPARTITIONED, lines)
    }

    fn append_lines_in(&mut self, part: &str, lines: &[String]) -> io::Result<Vec<Slot>> {
        self.append(part, lines)
    }

    fn sync(&mut self) -> io::Result<()> {
        LogSet::sync(self)
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum StageError {
    Validation(HaysonValidationError),
    Unique(HaysonUniqueViolationError),
    Conflict(HaysonPendingConflictError),
    Partition(HaysonPartitionError),
}

impl fmt::Display for StageError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            StageError::Validation(e) => e.fmt(f),
            StageError::Unique(e) => e.fmt(f),
            StageError::Conflict(e) => e.fmt(f),
            StageError::Partition(e) => e.fmt(f),
        }
    }
}

impl std::error::Error for StageError {}

/// `WriteQueue::commit_sealed`'s failure: an I/O error appending to the log, or the `seal`
/// closure refusing a `cypher` field (no key to write it with).
#[derive(Debug)]
pub enum CommitError {
    Io(io::Error),
    Cipher(HaysonCipherError),
}

impl fmt::Display for CommitError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            CommitError::Io(e) => e.fmt(f),
            CommitError::Cipher(e) => e.fmt(f),
        }
    }
}

impl std::error::Error for CommitError {}

impl From<io::Error> for CommitError {
    fn from(value: io::Error) -> Self {
        CommitError::Io(value)
    }
}

impl From<HaysonCipherError> for CommitError {
    fn from(value: HaysonCipherError) -> Self {
        CommitError::Cipher(value)
    }
}

impl From<HaysonValidationError> for StageError {
    fn from(value: HaysonValidationError) -> Self {
        StageError::Validation(value)
    }
}

impl From<HaysonUniqueViolationError> for StageError {
    fn from(value: HaysonUniqueViolationError) -> Self {
        StageError::Unique(value)
    }
}

impl From<HaysonPartitionError> for StageError {
    fn from(value: HaysonPartitionError) -> Self {
        StageError::Partition(value)
    }
}

impl From<HaysonPendingConflictError> for StageError {
    fn from(value: HaysonPendingConflictError) -> Self {
        StageError::Conflict(value)
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum WriteOutcome {
    Put(HaysonRecord),
    Deleted(bool),
}

/// What [`WriteQueue::commit_sealed`] puts every queued `put` through on its way to a log line: the
/// entry's key, its plaintext record, and the field names that entry was staged as already sealed
/// with.
pub type Seal<'a> = dyn Fn(&str, &HaysonRecord, &[String]) -> Result<HaysonRecord, HaysonCipherError> + 'a;

enum EntryKind {
    Put(HaysonRecord),
    Delete,
}

struct QueuedEntry {
    key: String,
    fields: Vec<String>,
    kind: EntryKind,
    /// Field names this entry's caller declared to already hold what the `seal` closure of
    /// [`WriteQueue::commit_sealed`] would otherwise produce. Opaque here: the queue neither reads
    /// them nor knows what sealing is, it only keeps them tied to the entry they were staged with,
    /// so a second write in the same window cannot inherit them.
    sealed: Vec<String>,
    /// The log file this entry lands in — [`UNPARTITIONED`] on a table with no `partitionBy`.
    part: String,
}

/// Batches put/patch/delete operations across one commit cycle into a single `LogAppender` call,
/// same shape as the TS `WriteQueue`: a key claim (`PendingMap`) guards every queued entry until the
/// batch that carries it is committed and the registry reflects it.
pub struct WriteQueue {
    table: String,
    max_batch: usize,
    pending: PendingMap,
    queued: VecDeque<QueuedEntry>,
    /// The keys of every entry a refused [`WriteQueue::commit_sealed`] threw away, until the
    /// caller collects them through [`WriteQueue::take_abandoned_keys`].
    abandoned: Vec<String>,
    /// Keys the image file names of this table's writes; set for a table of an encrypted database.
    image_key: Option<[u8; 32]>,
}

impl WriteQueue {
    pub fn new(table: impl Into<String>) -> Self {
        Self::with_max_batch(table, 256)
    }

    /// Names the image files of every entry staged from now on with `key` — see
    /// [`crate::schema::ValidateOptions::image_key`].
    pub fn set_image_key(&mut self, key: Option<[u8; 32]>) {
        self.image_key = key;
    }

    pub fn with_max_batch(table: impl Into<String>, max_batch: usize) -> Self {
        WriteQueue {
            table: table.into(),
            max_batch: max_batch.max(1),
            pending: PendingMap::new(),
            queued: VecDeque::new(),
            abandoned: Vec::new(),
            image_key: None,
        }
    }

    pub fn depth(&self) -> usize {
        self.queued.len()
    }

    pub fn pending(&self, key: &str) -> HashSet<String> {
        self.pending.pending(key)
    }

    /// The keys the last refused commit abandoned — see [`WriteQueue::commit_sealed`] — handed over
    /// once, in staging order, so the caller can drop whatever it staged beside the queue for them.
    pub fn take_abandoned_keys(&mut self) -> Vec<String> {
        std::mem::take(&mut self.abandoned)
    }

    /// Every key with a write still queued, in staging order — what a caller holding a cache of
    /// committed records has to drop before [`WriteQueue::commit_sealed`] moves the registry, and
    /// the only place that list exists before the commit drains it.
    pub fn queued_keys(&self) -> impl Iterator<Item = &str> {
        self.queued.iter().map(|entry| entry.key.as_str())
    }

    /// Validates `document` against `schema`, checks its unique fields against `registry`, resolves
    /// the partition it belongs to, then reserves the record's fields in the pending map. The key
    /// returned is the one the caller should use for `pending()`/patch chaining.
    pub fn stage_put(
        &mut self,
        schema: &TableSchema,
        registry: &TableRegistry,
        document: &Value,
        force: bool,
    ) -> Result<String, StageError> {
        self.stage_entry(schema, registry, document, None, Vec::new(), force)
    }

    /// [`WriteQueue::stage_put`] for a record whose fields named in `sealed` must reach the log as
    /// they stand, without the `seal` closure of [`WriteQueue::commit_sealed`] touching them. The
    /// names ride with this entry alone — see [`QueuedEntry::sealed`] — and the queue makes no
    /// judgement about them; whoever hands them over has already vouched for them.
    pub fn stage_put_carrying(
        &mut self,
        schema: &TableSchema,
        registry: &TableRegistry,
        document: &Value,
        sealed: Vec<String>,
        force: bool,
    ) -> Result<String, StageError> {
        self.stage_entry(schema, registry, document, None, sealed, force)
    }

    /// [`WriteQueue::stage_put`] with the claim named by the caller instead of taken from the whole
    /// record — what a patch needs, since it only changes the fields it was given and claiming the
    /// rest would refuse a second patch touching a field this one never mentioned.
    ///
    /// `claim` is only what the pending map guards; the entry queued is still the whole record.
    pub fn stage_put_claiming(
        &mut self,
        schema: &TableSchema,
        registry: &TableRegistry,
        document: &Value,
        claim: &[String],
        force: bool,
    ) -> Result<String, StageError> {
        self.stage_entry(schema, registry, document, Some(claim), Vec::new(), force)
    }

    fn stage_entry(
        &mut self,
        schema: &TableSchema,
        registry: &TableRegistry,
        document: &Value,
        claim: Option<&[String]>,
        sealed: Vec<String>,
        force: bool,
    ) -> Result<String, StageError> {
        let record = validate_record(
            schema,
            document,
            ValidateOptions {
                image_key: self.image_key,
                ..ValidateOptions::default()
            },
        )?;
        let key = primary_key_of(schema, &record)?;
        assert_unique(schema, registry, &key, &record)?;
        let part = partition_of(schema, &record)?;
        assert_partition_stable(schema, registry, &key, &part)?;
        let fields: Vec<String> = match claim {
            Some(claim) => claim.to_vec(),
            None => record.keys().cloned().collect(),
        };
        self.reserve(&key, &fields, force)?;
        self.queued.push_back(QueuedEntry {
            key: key.clone(),
            fields,
            kind: EntryKind::Put(record),
            sealed,
            part,
        });
        Ok(key)
    }

    /// The record of the entry queued last — how the caller of [`WriteQueue::stage_put`] reads back
    /// what was validated and defaulted, without paying for a second `validate_record`. `None` when
    /// nothing is queued or the last entry queued is a delete.
    pub fn last_put_record(&self) -> Option<&HaysonRecord> {
        match self.queued.back() {
            Some(QueuedEntry {
                kind: EntryKind::Put(record),
                ..
            }) => Some(record),
            _ => None,
        }
    }

    /// Queues a tombstone into `part` — the partition the key currently lives in, which the caller
    /// reads off the registry, since a tombstone has to land in the same file as the record it
    /// retires.
    pub fn stage_delete(
        &mut self,
        key: impl Into<String>,
        force: bool,
        part: &str,
    ) -> Result<(), HaysonPendingConflictError> {
        let key = key.into();
        let fields = vec![ALL_FIELDS.to_string()];
        self.reserve(&key, &fields, force)?;
        self.queued.push_back(QueuedEntry {
            key,
            fields,
            kind: EntryKind::Delete,
            sealed: Vec::new(),
            part: part.to_string(),
        });
        Ok(())
    }

    fn reserve(&mut self, key: &str, fields: &[String], force: bool) -> Result<(), HaysonPendingConflictError> {
        let conflicts = self.pending.conflicts_of(key, fields);
        if !conflicts.is_empty() && !force {
            return Err(HaysonPendingConflictError {
                table: self.table.clone(),
                key: key.to_string(),
                fields: conflicts,
            });
        }
        self.pending.claim(key, fields);
        Ok(())
    }

    /// Commits every entry queued since the last call: one `append_lines` (plus one `sync` when
    /// requested) per batch of at most `max_batch` entries, then updates `registry` and releases
    /// the pending claims for each committed key, in the order the entries were staged. Writes no
    /// `cypher` field, so it fails if any queued record carries one — callers with a schema that
    /// declares `cypher` fields must go through [`WriteQueue::commit_sealed`] instead.
    #[allow(dead_code)]
    pub fn commit<L: LogAppender>(
        &mut self,
        log: &mut L,
        registry: &mut TableRegistry,
        seq: &mut u64,
        sync: bool,
    ) -> io::Result<Vec<WriteOutcome>> {
        self.commit_sealed(log, registry, seq, sync, &|_key, record, _sealed| Ok(record.clone()))
            .map_err(|error| match error {
                CommitError::Io(error) => error,
                CommitError::Cipher(error) => io::Error::other(error.to_string()),
            })
    }

    /// Same as [`WriteQueue::commit`], but every queued `put` is passed through `seal` before it is
    /// serialized into a log line — the plaintext record still goes to `registry` and the returned
    /// [`WriteOutcome`], only the line written to `log` carries whatever `seal` returns. `seal` is
    /// invoked once per `put` entry with its key, its plaintext record and the field names that
    /// entry was staged as already sealed with (empty for anything staged through
    /// [`WriteQueue::stage_put`]).
    ///
    /// A batch leaves the queue only once its append has landed. A refusal — the log turning the
    /// append down, `seal` finding no key, the `sync` failing — abandons that batch and every one
    /// still behind it: their claims are released and their keys kept for
    /// [`WriteQueue::take_abandoned_keys`], so nothing of a write that never reached the log
    /// survives as a pending claim a retry would conflict with. The batches that landed before it
    /// stay committed; the `_seq` values the refused batch was stamped with stay burned.
    pub fn commit_sealed<L: LogAppender>(
        &mut self,
        log: &mut L,
        registry: &mut TableRegistry,
        seq: &mut u64,
        sync: bool,
        seal: &Seal<'_>,
    ) -> Result<Vec<WriteOutcome>, CommitError> {
        let mut results = Vec::with_capacity(self.queued.len());
        while !self.queued.is_empty() {
            let take = self.max_batch.min(self.queued.len());
            let now = chrono::Utc::now().timestamp_millis();

            let landed = self.stamp_batch(take, seq, now, seal).and_then(|(lines, stamps)| {
                let batch = &self.queued.make_contiguous()[..take];
                let slots = append_by_part(log, batch, &lines)?;
                if sync {
                    log.sync()?;
                }
                Ok((slots, stamps))
            });
            let (slots, stamps) = match landed {
                Ok(landed) => landed,
                Err(error) => {
                    self.abandon_queued();
                    return Err(error);
                }
            };

            for ((entry, slot), commit_seq) in self.queued.drain(..take).zip(slots).zip(stamps) {
                match entry.kind {
                    EntryKind::Put(record) => {
                        registry.apply_put_in(&entry.key, &record, slot, commit_seq, &entry.part);
                        results.push(WriteOutcome::Put(record));
                    }
                    EntryKind::Delete => {
                        let existed = registry.slot_of(&entry.key).is_some();
                        registry.apply_delete_in(&entry.key, slot.length, commit_seq, &entry.part);
                        results.push(WriteOutcome::Deleted(existed));
                    }
                }
                self.pending.release(&entry.key, &entry.fields);
            }
        }
        Ok(results)
    }

    /// Serializes the first `take` queued entries into log lines, stamping each with the next
    /// `_seq` — the half of a commit that touches nothing on disk.
    fn stamp_batch(
        &self,
        take: usize,
        seq: &mut u64,
        now: i64,
        seal: &Seal<'_>,
    ) -> Result<(Vec<String>, Vec<u64>), CommitError> {
        let mut lines = Vec::with_capacity(take);
        let mut stamps = Vec::with_capacity(take);
        for entry in self.queued.iter().take(take) {
            *seq += 1;
            stamps.push(*seq);
            let line = match &entry.kind {
                EntryKind::Put(record) => {
                    let sealed = seal(&entry.key, record, &entry.sealed)?;
                    serialize_entry(
                        &LogMeta {
                            key: entry.key.clone(),
                            op: LogOp::Put,
                            seq: *seq,
                            ts: now,
                        },
                        Some(&sealed),
                    )
                }
                EntryKind::Delete => serialize_entry(
                    &LogMeta {
                        key: entry.key.clone(),
                        op: LogOp::Del,
                        seq: *seq,
                        ts: now,
                    },
                    None,
                ),
            };
            lines.push(line);
        }
        Ok((lines, stamps))
    }

    /// Throws away everything still queued after a batch was refused, releasing each claim and
    /// keeping each key for [`WriteQueue::take_abandoned_keys`]. The batches behind the refused one
    /// go too: they would be refused the same way, or land after the caller has already heard that
    /// its flush failed.
    fn abandon_queued(&mut self) {
        for entry in self.queued.drain(..) {
            self.pending.release(&entry.key, &entry.fields);
            self.abandoned.push(entry.key);
        }
    }
}

/// Appends one batch, one `append_lines_in` call per distinct partition it touches, and hands the
/// slots back in the batch's own order. A batch that names a single partition — every batch of an
/// unpartitioned table, and most batches of a partitioned one — makes exactly one append call, the
/// same shape this queue had before partitions existed.
fn append_by_part<L: LogAppender + ?Sized>(
    log: &mut L,
    batch: &[QueuedEntry],
    lines: &[String],
) -> io::Result<Vec<Slot>> {
    let mut grouped: Vec<(&str, Vec<usize>)> = Vec::new();
    for (position, entry) in batch.iter().enumerate() {
        match grouped.iter_mut().find(|(part, _)| *part == entry.part) {
            Some((_, positions)) => positions.push(position),
            None => grouped.push((&entry.part, vec![position])),
        }
    }

    let mut slots: Vec<Option<Slot>> = vec![None; batch.len()];
    for (part, positions) in grouped {
        let payload: Vec<String> = positions.iter().map(|position| lines[*position].clone()).collect();
        let appended = log.append_lines_in(part, &payload)?;
        for (position, slot) in positions.into_iter().zip(appended) {
            slots[position] = Some(slot);
        }
    }
    slots
        .into_iter()
        .map(|slot| slot.ok_or_else(|| io::Error::other("the log returned fewer slots than lines appended")))
        .collect()
}

/// Refuses a write that would move an existing key out of the partition it already lives in — the
/// record's old line stays where it is, so the two files would both claim the key. Mirrors the TS
/// `Table.#assertPartitionStable`: delete the record first, then write it into the new partition.
pub(crate) fn assert_partition_stable(
    schema: &TableSchema,
    registry: &TableRegistry,
    key: &str,
    part: &str,
) -> Result<(), HaysonPartitionError> {
    if part == UNPARTITIONED {
        return Ok(());
    }
    let Some(held) = registry.part_of(key) else {
        return Ok(());
    };
    if held == part {
        return Ok(());
    }
    Err(HaysonPartitionError::new(
        schema.table.clone(),
        schema.partition_by.clone().unwrap_or_default(),
        format!("key \"{key}\" already lives in partition \"{held}\" — delete it before writing it into \"{part}\""),
    ))
}

/// What a batch of operations has already made of each key it named, as seen by an operation still
/// being validated against the same batch: `Some(record)` for a key a put or patch will leave
/// behind, `None` for one a delete will take away. Empty for a write that stands alone.
pub(crate) type UnitShadow = HashMap<String, Option<HaysonRecord>>;

pub(crate) fn assert_unique(
    schema: &TableSchema,
    registry: &TableRegistry,
    key: &str,
    record: &HaysonRecord,
) -> Result<(), HaysonUniqueViolationError> {
    assert_unique_in_unit(schema, registry, key, record, &UnitShadow::new())
}

fn virtual_path_of<'a>(record: &'a HaysonRecord, field: &str) -> Option<&'a str> {
    record.get(field)?.as_object()?.get("virtualPath")?.as_str()
}

fn violation(schema: &TableSchema, field: String, holder: String) -> HaysonUniqueViolationError {
    HaysonUniqueViolationError {
        table: schema.table.clone(),
        field,
        holder,
    }
}

/// True when the committed `holder` of a value has already been rewritten by this batch into
/// something that no longer holds it — the committed index still names it, but the batch it belongs
/// to will release it before anything lands, so it may not refuse `mine`.
fn shadow_released(
    shadow: &UnitShadow,
    holder: &str,
    mine: &str,
    of: impl Fn(&HaysonRecord) -> Option<String>,
) -> bool {
    match shadow.get(holder) {
        None => false,
        Some(None) => true,
        Some(Some(record)) => of(record).as_deref() != Some(mine),
    }
}

/// Every uniqueness rule the table enforces, checked against the committed registry **and** against
/// `shadow` — the records earlier operations of the same batch will leave behind. A unit of work
/// appends nothing until it commits, so without `shadow` two operations of one unit can each pass
/// against the committed state and still duplicate each other (ADR-3h3r).
pub(crate) fn assert_unique_in_unit(
    schema: &TableSchema,
    registry: &TableRegistry,
    key: &str,
    record: &HaysonRecord,
    shadow: &UnitShadow,
) -> Result<(), HaysonUniqueViolationError> {
    for index in registry.indexes() {
        if !index.unique {
            continue;
        }
        if index.fields.len() == 1 && index.fields[0] == schema.primary_key {
            continue;
        }
        let unique_key = |candidate: &HaysonRecord| unique_key_of(candidate, &index.fields);
        let Some(mine) = unique_key(record) else { continue };
        for (other_key, other) in shadow {
            if other_key == key {
                continue;
            }
            let Some(other) = other else { continue };
            if unique_key(other).as_deref() == Some(mine.as_str()) {
                return Err(violation(schema, index.fields.join(", "), other_key.clone()));
            }
        }
        if let Some(holder) = registry.unique_holder_of(index, record) {
            if holder != key && !shadow_released(shadow, &holder, &mine, unique_key) {
                return Err(violation(schema, index.fields.join(", "), holder));
            }
        }
    }
    for field in registry.virtual_path_fields() {
        let path_of = |candidate: &HaysonRecord| virtual_path_of(candidate, field).map(str::to_string);
        let Some(mine) = virtual_path_of(record, field) else {
            continue;
        };
        for (other_key, other) in shadow {
            if other_key == key {
                continue;
            }
            let Some(other) = other else { continue };
            if virtual_path_of(other, field) == Some(mine) {
                return Err(violation(schema, field.clone(), other_key.clone()));
            }
        }
        if let Some(holder) = registry.virtual_path_holder(field, mine) {
            if holder != key && !shadow_released(shadow, &holder, mine, path_of) {
                return Err(violation(schema, field.clone(), holder));
            }
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::parse_schema;
    use serde_json::json;

    fn users_schema() -> TableSchema {
        parse_schema(&json!({
            "schemaVersion": 1,
            "table": "users",
            "primaryKey": "id",
            "fields": {
                "id": { "type": "string" },
                "name": { "type": "string", "required": true },
                "email": { "type": "string", "unique": true },
            },
        }))
        .unwrap()
    }

    struct CountingLog {
        calls: usize,
        synced: usize,
        appended: Vec<String>,
    }

    impl CountingLog {
        fn new() -> Self {
            CountingLog {
                calls: 0,
                synced: 0,
                appended: Vec::new(),
            }
        }
    }

    impl LogAppender for CountingLog {
        fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
            self.calls += 1;
            let start = self.appended.len() as u64;
            let slots = lines
                .iter()
                .enumerate()
                .map(|(i, line)| Slot {
                    offset: start + i as u64,
                    length: line.len() as u64,
                })
                .collect();
            self.appended.extend(lines.iter().cloned());
            Ok(slots)
        }

        fn sync(&mut self) -> io::Result<()> {
            self.synced += 1;
            Ok(())
        }
    }

    #[test]
    fn commits_every_entry_queued_since_the_last_call_in_one_append() {
        let schema = users_schema();
        let mut registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::new("users");
        let mut log = CountingLog::new();
        let mut seq = 0u64;

        queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Ania" }), false)
            .unwrap();
        queue
            .stage_put(&schema, &registry, &json!({ "id": "u2", "name": "Piotr" }), false)
            .unwrap();
        assert_eq!(queue.depth(), 2);

        let results = queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();

        assert_eq!(log.calls, 1);
        assert_eq!(log.synced, 1);
        assert_eq!(results.len(), 2);
        assert_eq!(queue.depth(), 0);
        assert_eq!(registry.records(), 2);
    }

    #[test]
    fn a_second_write_to_the_same_field_conflicts_unless_forced() {
        let schema = users_schema();
        let registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::new("users");

        queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Ania" }), false)
            .unwrap();
        let err = queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Zosia" }), false)
            .unwrap_err();
        assert!(matches!(err, StageError::Conflict(_)));

        queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Zosia" }), true)
            .unwrap();
        assert_eq!(queue.depth(), 2);
    }

    #[test]
    fn a_delete_claims_every_field_and_conflicts_with_a_write_on_any_of_them() {
        let schema = users_schema();
        let registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::new("users");

        queue.stage_delete("u1", false, "").unwrap();
        let err = queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Ania" }), false)
            .unwrap_err();
        assert!(matches!(err, StageError::Conflict(_)));
    }

    #[test]
    fn releases_the_claim_only_after_the_batch_commits() {
        let schema = users_schema();
        let mut registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::new("users");
        let mut log = CountingLog::new();
        let mut seq = 0u64;

        queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Ania" }), false)
            .unwrap();
        assert_eq!(queue.pending("u1"), HashSet::from(["id".into(), "name".into()]));

        queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
        assert!(queue.pending("u1").is_empty());
    }

    /// A log that refuses every append from the `failing_from`-th call on — a disk that filled up,
    /// a file swapped out under the handle.
    struct RefusingLog {
        inner: CountingLog,
        failing_from: usize,
    }

    impl LogAppender for RefusingLog {
        fn append_lines(&mut self, lines: &[String]) -> io::Result<Vec<Slot>> {
            if self.inner.calls + 1 >= self.failing_from {
                self.inner.calls += 1;
                return Err(io::Error::other("refused"));
            }
            self.inner.append_lines(lines)
        }

        fn sync(&mut self) -> io::Result<()> {
            self.inner.sync()
        }
    }

    #[test]
    fn a_refused_append_leaves_nothing_queued_and_nothing_claimed() {
        let schema = users_schema();
        let mut registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::new("users");
        let mut log = RefusingLog {
            inner: CountingLog::new(),
            failing_from: 1,
        };
        let mut seq = 0u64;
        queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Ania" }), false)
            .unwrap();

        queue.commit(&mut log, &mut registry, &mut seq, true).unwrap_err();

        assert_eq!(queue.depth(), 0, "the refused entry stayed queued");
        assert!(queue.pending("u1").is_empty(), "the refused entry kept its claim");
        assert_eq!(queue.take_abandoned_keys(), vec!["u1".to_string()]);
        assert!(queue.take_abandoned_keys().is_empty(), "the keys are handed over once");
        assert_eq!(registry.records(), 0);
        queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Ania" }), false)
            .expect("the same write can be staged again without force");
    }

    #[test]
    fn a_refusal_mid_way_keeps_the_batches_that_landed_and_abandons_the_rest() {
        let schema = users_schema();
        let mut registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::with_max_batch("users", 1);
        let mut log = RefusingLog {
            inner: CountingLog::new(),
            failing_from: 2,
        };
        let mut seq = 0u64;
        for (id, name) in [("u1", "Ania"), ("u2", "Piotr"), ("u3", "Zosia")] {
            queue
                .stage_put(&schema, &registry, &json!({ "id": id, "name": name }), false)
                .unwrap();
        }

        queue.commit(&mut log, &mut registry, &mut seq, true).unwrap_err();

        assert_eq!(
            registry.records(),
            1,
            "the batch that landed before the refusal is committed"
        );
        assert_eq!(queue.depth(), 0);
        assert_eq!(queue.take_abandoned_keys(), vec!["u2".to_string(), "u3".to_string()]);
        for key in ["u1", "u2", "u3"] {
            assert!(queue.pending(key).is_empty(), "{key} kept its claim");
        }
    }

    #[test]
    fn rejects_a_second_record_claiming_a_unique_field_already_held() {
        let schema = users_schema();
        let mut registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::new("users");
        let mut log = CountingLog::new();
        let mut seq = 0u64;

        queue
            .stage_put(
                &schema,
                &registry,
                &json!({ "id": "u1", "name": "Ania", "email": "a@x.io" }),
                false,
            )
            .unwrap();
        queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();

        let err = queue
            .stage_put(
                &schema,
                &registry,
                &json!({ "id": "u2", "name": "Zosia", "email": "a@x.io" }),
                false,
            )
            .unwrap_err();
        assert!(matches!(err, StageError::Unique(_)));
    }

    fn docs_schema() -> TableSchema {
        parse_schema(&json!({
            "schemaVersion": 1,
            "table": "docs",
            "primaryKey": "id",
            "fields": {
                "id": { "type": "string" },
                "page": { "type": "markdown" },
            },
        }))
        .unwrap()
    }

    #[test]
    fn rejects_a_second_record_claiming_an_already_taken_markdown_virtual_path() {
        let schema = docs_schema();
        let mut registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::new("docs");
        let mut log = CountingLog::new();
        let mut seq = 0u64;

        queue
            .stage_put(
                &schema,
                &registry,
                &json!({ "id": "d1", "page": { "content": "a", "virtualPath": "/docs/a.md" } }),
                false,
            )
            .unwrap();
        queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();

        let err = queue
            .stage_put(
                &schema,
                &registry,
                &json!({ "id": "d2", "page": { "content": "b", "virtualPath": "/docs/a.md" } }),
                false,
            )
            .unwrap_err();
        assert!(matches!(err, StageError::Unique(_)));
    }

    #[test]
    fn dead_bytes_grow_when_an_existing_key_is_overwritten_or_removed() {
        let schema = users_schema();
        let mut registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::new("users");
        let mut log = CountingLog::new();
        let mut seq = 0u64;

        queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Ania" }), false)
            .unwrap();
        queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
        assert_eq!(registry.stats().dead_bytes, 0);

        queue
            .stage_put(&schema, &registry, &json!({ "id": "u1", "name": "Zosia" }), false)
            .unwrap();
        queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
        let after_overwrite = registry.stats().dead_bytes;
        assert!(after_overwrite > 0);

        queue.stage_delete("u1", false, "").unwrap();
        queue.commit(&mut log, &mut registry, &mut seq, true).unwrap();
        assert!(registry.stats().dead_bytes > after_overwrite);
        assert_eq!(registry.records(), 0);
    }

    #[test]
    fn splits_a_batch_larger_than_max_batch_into_several_appends() {
        let schema = users_schema();
        let mut registry = TableRegistry::new(schema.clone());
        let mut queue = WriteQueue::with_max_batch("users", 2);
        let mut log = CountingLog::new();
        let mut seq = 0u64;

        for i in 0..5 {
            queue
                .stage_put(
                    &schema,
                    &registry,
                    &json!({ "id": format!("u{i}"), "name": "x" }),
                    false,
                )
                .unwrap();
        }
        let results = queue.commit(&mut log, &mut registry, &mut seq, false).unwrap();
        assert_eq!(results.len(), 5);
        assert_eq!(log.calls, 3);
        assert_eq!(log.synced, 0);
    }
}