haematite-cli 0.21.7

use haematite::branch::history::Key as HistoryKey;
use haematite::branch::{
    CommitDurability, CommitRequest, ForkWithChanges, ForkWrite, RefRecord, current_timestamp,
    open as open_branch,
};
use haematite::db::branches::ForkKeptError;
use haematite::tree::Hash;

use super::{
    Captured, Collection, DocumentRevision, Failure, Receipt, Storage, branch_name, collection_key,
    counts, document_key, identifier, reads, receipt_key, revision_key, version,
};
use crate::serve::protocol::{Change, DocumentImport, Mutation, ReceiptQuery, Selection, Version};
use crate::serve::schema::Collection as CollectionSchema;

pub(super) type PreparedWrite = (Vec<u8>, Option<Vec<u8>>);

impl Storage {
    /// Reconcile a receipt and retry only its reference fence, not the write.
    ///
    /// # Errors
    ///
    /// Propagates any error from `identifier`, `get_bytes`, `from_slice` or `publish_record`.
    pub fn receipt(&mut self, query: &ReceiptQuery) -> Result<Option<Receipt>, Failure> {
        identifier(&query.request_id)?;
        identifier(&query.branch)?;
        if let Some(receipt) = self.discard_receipt(query)? {
            return Ok(Some(receipt));
        }
        if self.pending_discards.contains(&branch_name(&query.branch)) {
            return Err(Failure::new(
                "discard_pending",
                "branch removal requires receipt reconciliation",
            ));
        }
        // Reconciliation deliberately sees an installed but not yet published
        // head. Ordinary readers use only the separately confirmed catalogue.
        let record = self
            .refs
            .get(&branch_name(&query.branch))
            .ok_or_else(|| Failure::new("branch_missing", &query.branch))?;
        let at = Captured {
            branch: query.branch.clone(),
            version: version(record),
            shards: record.shards.clone(),
        };
        let Some(bytes) = reads::get_bytes(
            &self.reader,
            &at,
            &receipt_key(&query.branch, at.version.generation, &query.request_id),
        )?
        else {
            return Ok(None);
        };
        let receipt: Receipt = serde_json::from_slice(&bytes).map_err(Failure::storage)?;
        if receipt.request_id != query.request_id
            || receipt.branch != query.branch
            || receipt.version.generation != at.version.generation
        {
            return Err(Failure::new(
                "receipt_corrupt",
                "receipt identity does not match its stored key",
            ));
        }
        self.refs
            .confirm_version(
                &branch_name(&query.branch),
                at.version.generation,
                at.version.sequence,
            )
            .map_err(|error| Failure::new("durability_unknown", error.to_string()))?;
        self.publish_record(&branch_name(&query.branch))?;
        Ok(Some(receipt))
    }

    /// One durable request identity, validation and effect within admission.
    ///
    /// # Errors
    ///
    /// Propagates any error from `identifier`, `to_vec`, `replay` or `prepare_change`.
    pub fn mutate(&mut self, request: &Mutation) -> Result<Receipt, Failure> {
        self.mutate_with_import_bound(request, None, super::AppendPruneLimits::default())
    }

    /// Admit imports up to sixty-four pages, preserving receipt-first replay.
    ///
    /// # Errors
    /// Propagates mutation preparation, admission, replay and durable commit refusals without changing their outcome classification.
    pub fn mutate_bounded(
        &mut self,
        request: &Mutation,
        max_page: std::num::NonZeroU32,
    ) -> Result<Receipt, Failure> {
        self.mutate_with_import_bound(request, Some(max_page), super::AppendPruneLimits::default())
    }

    /// Admit a mutation with the configured retention and import bounds.
    ///
    /// # Errors
    /// Preserves receipt-first replay and propagates preparation or commit refusals.
    pub fn mutate_with_append_prune_limits(
        &mut self,
        request: &Mutation,
        max_page: std::num::NonZeroU32,
        limits: super::AppendPruneLimits,
    ) -> Result<Receipt, Failure> {
        self.mutate_with_import_bound(request, Some(max_page), limits)
    }

    /// The live record a mutation names, at exactly the version it expects.
    fn expected_record(&self, request: &Mutation) -> Result<RefRecord, Failure> {
        let record = self
            .refs
            .get(&branch_name(&request.branch))
            .ok_or_else(|| Failure::new("branch_missing", &request.branch))?;
        if version(record) != request.expected {
            return Err(Failure::new(
                "version_conflict",
                "expected branch generation or sequence differs",
            ));
        }
        if self.pending_discards.contains(&record.name) {
            return Err(Failure::new(
                "discard_pending",
                "branch removal requires receipt reconciliation",
            ));
        }
        Ok(record.clone())
    }

    fn mutate_with_import_bound(
        &mut self,
        request: &Mutation,
        import_max_page: Option<std::num::NonZeroU32>,
        retention_limits: super::AppendPruneLimits,
    ) -> Result<Receipt, Failure> {
        identifier(&request.request_id)?;
        identifier(&request.branch)?;
        let destination = match &request.change {
            Change::Fork(fork) => &fork.name,
            _ => &request.branch,
        };
        identifier(destination)?;
        let fingerprint =
            Hash::of(&serde_json::to_vec(request).map_err(Failure::storage)?).to_string();
        if let Some(receipt) = self.replay(request, destination, &fingerprint)? {
            return Ok(receipt);
        }
        validate_import_bound(&request.change, import_max_page)?;
        let record = self.expected_record(request)?;
        let at = Captured {
            branch: request.branch.clone(),
            version: request.expected,
            shards: record.shards.clone(),
        };
        if let Change::Fork(fork) = &request.change {
            return self.fork_request(request, &fingerprint, &fork.name, fork.from);
        }
        if let Change::Discard(discard) = &request.change {
            return self.discard_request(request, &fingerprint, &discard.name, &record);
        }
        let next = Version {
            generation: record.created,
            sequence: record.seq.checked_add(1).ok_or_else(|| {
                Failure::new("sequence_exhausted", "branch sequence cannot advance")
            })?,
        };
        let (mut writes, retention) = match &request.change {
            Change::AppendPrune(write) => {
                let (writes, result) =
                    self.prepare_append_prune(&at, write, next, retention_limits)?;
                (writes, Some(result))
            }
            change => (self.prepare_change(&at, change, next)?, None),
        };
        counts::append(&self.reader, &at, &request.change, &mut writes)?;
        // A version the advance drops from the kept window has its row first
        // (STOR-003 MD5a).
        self.rows_before_advance(&request.branch, &record)?;
        #[cfg(test)]
        super::history::after_rows::reached()?;
        let branch = open_branch(&record.name, &self.refs, &self.registry)
            .map_err(Failure::branch_commit)?;
        for (key, value) in writes {
            let shard_id = self.writer.shard_for(&key);
            match value {
                Some(value) => branch.put(shard_id, key, value),
                None => branch.delete(shard_id, key),
            }
            .map_err(Failure::branch)?;
        }
        let receipt = Receipt {
            request_id: request.request_id.clone(),
            fingerprint,
            branch: request.branch.clone(),
            version: next,
            retention,
        };
        let key = receipt_key(&request.branch, next.generation, &request.request_id);
        branch
            .put(
                self.writer.shard_for(&key),
                key,
                serde_json::to_vec(&receipt).map_err(Failure::storage)?,
            )
            .map_err(Failure::branch)?;
        let result = self.writer.commit(
            &branch,
            &self.registry,
            CommitRequest {
                durability: CommitDurability::Durable {
                    refs: &mut self.refs,
                },
                extra_parents: &[],
                timestamp: current_timestamp(),
            },
        );
        if let Err(error) = result {
            let installed = self
                .refs
                .get(&record.name)
                .is_some_and(|now| version(now) != request.expected);
            let mut failure = Failure::branch_commit(error);
            let code = if installed {
                "durability_unknown"
            } else {
                "write_refused"
            };
            code.clone_into(&mut failure.code);
            return Err(failure);
        }
        #[cfg(test)]
        lose_mutation_answer()?;
        #[cfg(test)]
        publish_pause_point();
        self.publish_record(&record.name)?;
        Ok(receipt)
    }

    fn replay(
        &mut self,
        request: &Mutation,
        destination: &str,
        fingerprint: &str,
    ) -> Result<Option<Receipt>, Failure> {
        // A fork receipt lives in its destination and is absent before creation.
        if matches!(&request.change, Change::Discard(_))
            || self.refs.get(&branch_name(destination)).is_some()
        {
            if let Some(receipt) = self.receipt(&ReceiptQuery {
                branch: destination.to_owned(),
                request_id: request.request_id.clone(),
            })? {
                return if receipt.fingerprint == fingerprint {
                    Ok(Some(receipt))
                } else {
                    Err(Failure::new(
                        "request_conflict",
                        "request identity was used with different content",
                    ))
                };
            }
            if matches!(&request.change, Change::Fork(_)) {
                return Err(Failure::new(
                    "branch_exists",
                    "destination exists without this fork receipt",
                ));
            }
        }
        Ok(None)
    }

    pub(super) fn prepare_change(
        &self,
        at: &Captured,
        change: &Change,
        next: Version,
    ) -> Result<Vec<PreparedWrite>, Failure> {
        match change {
            Change::Collection(write) => {
                identifier(&write.collection)?;
                if let Some(schema) = &write.schema {
                    CollectionSchema::compile(schema).map_err(schema_failure)?;
                }
                let collection = Collection {
                    name: write.collection.clone(),
                    schema: write.schema.clone(),
                    changed_in: at.branch.clone(),
                    changed_at: next,
                };
                Ok(vec![(
                    collection_key(&write.collection),
                    Some(serde_json::to_vec(&collection).map_err(Failure::storage)?),
                )])
            }
            Change::Put(write) => {
                identifier(&write.document_id)?;
                let collection = reads::collection(&self.reader, at, &write.collection)?
                    .ok_or_else(|| Failure::new("collection_missing", &write.collection))?;
                if let Some(schema) = &collection.schema {
                    CollectionSchema::compile(schema)
                        .map_err(schema_failure)?
                        .validate(
                            &serde_json::from_str(write.document.get())
                                .map_err(Failure::storage)?,
                        )
                        .map_err(schema_failure)?;
                }
                let revision = DocumentRevision {
                    collection_branch: collection.changed_in,
                    collection_version: collection.changed_at,
                    written_in: at.branch.clone(),
                    written_at: next,
                };
                Ok(vec![
                    (
                        document_key(&write.collection, &write.document_id),
                        Some(write.document.get().as_bytes().to_vec()),
                    ),
                    (
                        revision_key(&write.collection, &write.document_id),
                        Some(serde_json::to_vec(&revision).map_err(Failure::storage)?),
                    ),
                ])
            }
            Change::AppendPrune(_) => Err(Failure::new(
                "append_prune_preparation_required",
                "retention uses bounded mutation preparation",
            )),
            Change::Delete(write) => {
                identifier(&write.collection)?;
                identifier(&write.document_id)?;
                // A deleted document takes its edges, in both directions, and its
                // vectors with it in the same version (D7).
                let mut writes = vec![
                    (document_key(&write.collection, &write.document_id), None),
                    (revision_key(&write.collection, &write.document_id), None),
                ];
                writes.extend(super::edges::removals_for_document(
                    &self.reader,
                    at,
                    &write.collection,
                    &write.document_id,
                )?);
                writes.extend(super::vectors::removals_for_document(
                    &self.reader,
                    at,
                    &write.collection,
                    &write.document_id,
                )?);
                Ok(writes)
            }
            Change::DeleteMany(delete) => self.prepare_delete_many(at, delete, next),
            Change::Merge(merge) => self.merge_writes(at, merge),
            Change::Fork(_) | Change::Discard(_) => Err(Failure::new(
                "invalid_operation",
                "branch control uses the publication path",
            )),
            Change::EdgePut(edge) => super::edges::put(&self.reader, at, edge),
            Change::EdgeDelete(edge) => super::edges::delete(edge),
            Change::EdgePutMany(import) => super::edges::put_many(&self.reader, at, import),
            Change::IndexDeclare(field) => {
                super::field_index::declare(&self.reader, at, field, next)
            }
            Change::VectorDeclare(field) => super::vectors::declare(&self.reader, at, field),
            Change::VectorPut(vector) => super::vectors::put(&self.reader, at, vector),
            Change::VectorPutMany(import) => {
                super::vector_import::prepare(&self.reader, at, import)
            }
            Change::VectorDelete(vector) => super::vectors::delete(vector),
            Change::PutMany(import) => self.prepare_import(at, import, next),
            Change::RelationPut(relation) => super::relations::put(&self.reader, at, relation),
            Change::RelationDelete(relation) => super::relations::delete(relation),
        }
    }

    fn prepare_import(
        &self,
        at: &Captured,
        import: &DocumentImport,
        next: Version,
    ) -> Result<Vec<PreparedWrite>, Failure> {
        if import.documents.is_empty() {
            return Err(Failure {
                code: "import_empty".to_owned(),
                message: "an import must contain at least one document".to_owned(),
                path: Some("/documents".to_owned()),
                cause: None,
            });
        }
        let collection = reads::collection(&self.reader, at, &import.collection)
            .map_err(|mut failure| {
                failure.path = Some("/collection".to_owned());
                failure
            })?
            .ok_or_else(|| Failure {
                code: "collection_missing".to_owned(),
                message: import.collection.clone(),
                path: Some("/collection".to_owned()),
                cause: None,
            })?;
        let schema = collection
            .schema
            .as_ref()
            .map(CollectionSchema::compile)
            .transpose()
            .map_err(schema_failure)?;
        validate_import_rows(import, schema.as_ref())?;
        let revision = serde_json::to_vec(&DocumentRevision {
            collection_branch: collection.changed_in,
            collection_version: collection.changed_at,
            written_in: at.branch.clone(),
            written_at: next,
        })
        .map_err(Failure::storage)?;
        let mut writes = Vec::new();
        for document in &import.documents {
            writes.push((
                document_key(&import.collection, &document.document_id),
                Some(document.document.get().as_bytes().to_vec()),
            ));
            writes.push((
                revision_key(&import.collection, &document.document_id),
                Some(revision.clone()),
            ));
        }
        Ok(writes)
    }

    fn fork_request(
        &mut self,
        request: &Mutation,
        fingerprint: &str,
        destination: &str,
        from: Version,
    ) -> Result<Receipt, Failure> {
        if self.pending_discards.contains(&branch_name(destination)) {
            return Err(Failure::new(
                "discard_pending",
                "destination removal requires receipt reconciliation",
            ));
        }
        self.capture(&Selection {
            branch: request.branch.clone(),
            version: Some(from),
        })?;
        let current = self
            .refs
            .get(&branch_name(&request.branch))
            .ok_or_else(|| Failure::new("branch_missing", &request.branch))?
            .clone();
        // `fork_kept` reads the row of the version forked from: the head's is
        // written only when the head is that version (STOR-003 MD5b).
        if from == version(&current) {
            self.retain_current(&request.branch, &current)?;
        }
        // A kept version forked from has its row before `fork_kept` reads it
        // (STOR-003 MD5a).
        self.row_for_kept(&request.branch, &current, from)?;
        #[cfg(test)]
        super::history::after_rows::reached()?;
        let service =
            open_branch(super::SERVICE, &self.refs, &self.registry).map_err(|source| {
                fork_kept_failure(&ForkKeptError::ServiceBindingStale { source }, false)
            })?;
        let history_key = HistoryKey {
            branch: Hash::of(request.branch.as_bytes()),
            generation: from.generation,
            sequence: from.sequence,
        };
        let created = current_timestamp();
        let receipt = Receipt {
            request_id: request.request_id.clone(),
            fingerprint: fingerprint.to_owned(),
            branch: destination.to_owned(),
            version: Version {
                generation: created,
                sequence: 0,
            },
            retention: None,
        };
        let key = receipt_key(destination, created, &request.request_id);
        let writes = [
            ForkWrite {
                shard_id: self.writer.shard_for(&key),
                key,
                value: Some(serde_json::to_vec(&receipt).map_err(Failure::storage)?),
            },
            super::branches::provenance_write(&self.writer, &request.branch, from)?,
        ];
        let name = branch_name(destination);
        self.writer
            .fork_kept(
                &service,
                history_key,
                &mut self.refs,
                &self.registry,
                ForkWithChanges {
                    name: &name,
                    timestamp: created,
                    writes: &writes,
                },
            )
            .map_err(|error| fork_kept_failure(&error, self.refs.get(&name).is_some()))?;
        #[cfg(test)]
        lose_mutation_answer()?;
        self.publish_record(&name)?;
        Ok(receipt)
    }

    /// Stop this thread's next committed mutation before it publishes.
    ///
    /// The pause point signals on `paused` once the store commit has returned,
    /// then waits on `release` while still holding the mutation admission.
    #[cfg(test)]
    pub fn arm_publish_pause(
        paused: std::sync::mpsc::Sender<()>,
        release: std::sync::mpsc::Receiver<()>,
    ) {
        PUBLISH_PAUSE.with(|slot| {
            *slot.borrow_mut() = Some(PublishPause { paused, release });
        });
    }

    #[cfg(test)]
    pub(crate) fn lose_next_mutation_answer() {
        Self::lose_next_mutation_answer_with_message(
            "injected lost answer after reference commit".to_owned(),
        );
    }

    #[cfg(test)]
    pub(crate) fn lose_next_mutation_answer_with_message(message: String) {
        LOST_MUTATION_ANSWER.with(|slot| {
            *slot.borrow_mut() = Some(message);
        });
    }
}

#[cfg(test)]
thread_local! {
    static LOST_MUTATION_ANSWER: std::cell::RefCell<Option<String>> = const { std::cell::RefCell::new(None) };
}

#[cfg(test)]
fn lose_mutation_answer() -> Result<(), Failure> {
    if let Some(message) = LOST_MUTATION_ANSWER.with(|slot| slot.borrow_mut().take()) {
        return Err(Failure::new("durability_unknown", message));
    }
    Ok(())
}

fn fork_kept_failure(error: &ForkKeptError, installed: bool) -> Failure {
    let code = match error {
        ForkKeptError::ServiceBindingStale { .. } => "kept_service_binding_stale",
        ForkKeptError::ServiceInvalid { .. } => "kept_service_invalid",
        ForkKeptError::ShardSetIncomplete { .. } => "kept_shard_set_incomplete",
        ForkKeptError::RowMissing { .. } => "kept_row_missing",
        ForkKeptError::RowInvalid { .. } => "kept_row_invalid",
        ForkKeptError::RowRead { .. } => "kept_row_read",
        ForkKeptError::RootUnavailable { .. } => "kept_root_unavailable",
        ForkKeptError::Fork(_) if installed => "durability_unknown",
        ForkKeptError::Fork(_) => "write_refused",
    };
    Failure::new(code, error.to_string())
}

#[cfg(test)]
#[path = "mutations_kept_tests.rs"]
mod kept_tests;

fn validate_import_rows(
    import: &DocumentImport,
    schema: Option<&CollectionSchema>,
) -> Result<(), Failure> {
    let mut ids = std::collections::BTreeSet::new();
    for (index, document) in import.documents.iter().enumerate() {
        identifier(&document.document_id).map_err(|mut failure| {
            failure.path = Some(format!("/documents/{index}/document_id"));
            failure
        })?;
        if !ids.insert(document.document_id.as_str()) {
            return Err(Failure {
                code: "import_duplicate".to_owned(),
                message: format!("repeated document_id {}", document.document_id),
                path: Some(format!("/documents/{index}/document_id")),
                cause: None,
            });
        }
        if let Some(schema) = schema {
            let value = serde_json::from_str(document.document.get()).map_err(|error| Failure {
                code: "invalid_document".to_owned(),
                message: error.to_string(),
                path: Some(format!("/documents/{index}/document")),
                cause: None,
            })?;
            schema.validate(&value).map_err(|error| Failure {
                code: error.code().to_owned(),
                message: error.reason,
                path: Some(format!("/documents/{index}/document{}", error.path)),
                cause: None,
            })?;
        }
    }
    Ok(())
}

fn validate_import_bound(
    change: &Change,
    max_page: Option<std::num::NonZeroU32>,
) -> Result<(), Failure> {
    let (rows, kind) = match change {
        Change::PutMany(import) => (import.documents.len(), "documents"),
        Change::DeleteMany(delete) => (delete.document_ids.len(), "document_ids"),
        Change::EdgePutMany(import) => (import.edges.len(), "edges"),
        Change::VectorPutMany(import) => (import.vectors.len(), "vectors"),
        _ => return Ok(()),
    };
    let max_page = max_page.ok_or_else(|| {
        Failure::new(
            "import_limit_required",
            "an import requires the service page bound",
        )
    })?;
    let bound = u64::from(max_page.get())
        .checked_mul(64)
        .ok_or_else(|| Failure::new("import_limit", "import row bound exceeds u64"))?;
    if u64::try_from(rows).map_err(Failure::storage)? > bound {
        return Err(Failure {
            code: "import_limit".to_owned(),
            message: format!("maximum import size is {bound} {kind}"),
            path: Some(format!("/{kind}/{bound}")),
            cause: None,
        });
    }
    Ok(())
}

fn schema_failure(error: crate::serve::schema::Violation) -> Failure {
    Failure {
        code: error.code().to_owned(),
        message: error.reason,
        path: Some(error.path),
        cause: None,
    }
}

#[cfg(test)]
struct PublishPause {
    paused: std::sync::mpsc::Sender<()>,
    release: std::sync::mpsc::Receiver<()>,
}

#[cfg(test)]
thread_local! {
    static PUBLISH_PAUSE: std::cell::RefCell<Option<PublishPause>> =
        const { std::cell::RefCell::new(None) };
}

/// Between the durable store commit and the catalogue publication. The ten
/// second bound turns a lost signal into a named failure, not a hung test.
#[cfg(test)]
fn publish_pause_point() {
    let Some(pause) = PUBLISH_PAUSE.with(|slot| slot.borrow_mut().take()) else {
        return;
    };
    assert!(
        pause.paused.send(()).is_ok(),
        "paused signal receiver was dropped"
    );
    let wait = std::time::Duration::from_secs(10);
    let released = pause.release.recv_timeout(wait);
    assert!(released.is_ok(), "release signal not received within 10 s");
}