tidaldb/tidal/src/text/writer.rs

use std::{collections::HashMap, sync::MutexGuard};

use tantivy::{IndexWriter, TantivyDocument, Term};

use crate::{TidalError, schema::EntityId, text::index::TantivyFields};

// Test-only fault hook: number of subsequent `commit` calls (see
// `TextIndexWriter::commit`) that should fail with a transient error before the
// real commit runs.
//
// Thread-local so each test arms its own syncer thread independently and the
// hook compiles away entirely in non-test builds. Used to prove the
// `TextIndexSyncer` (see `super::syncer::TextIndexSyncer`) survives a transient
// commit failure (disk-full / fsync IO error) and recovers on retry, the
// 3am-incident failure mode where a single commit error would otherwise kill
// the syncer thread and silently lose every future write from the text index.
#[cfg(test)]
thread_local! {
    static FORCE_COMMIT_FAILURES: std::cell::Cell<u32> = const { std::cell::Cell::new(0) };
}

/// Arm the commit-failure fault hook for the current thread.
///
/// The next `n` calls to [`commit`](TextIndexWriter::commit) on this thread
/// return `Err(TidalError::Internal)` *without* touching Tantivy, then commits
/// resume normally. Because the syncer runs on a dedicated thread, the test must
/// arm this from inside that thread (e.g. via the closure passed to `spawn`).
#[cfg(test)]
pub(crate) fn force_commit_failures(n: u32) {
    FORCE_COMMIT_FAILURES.with(|c| c.set(n));
}

/// Write operations on the Tantivy text index.
///
/// Created by [`TextIndex::writer_guard()`](super::TextIndex::writer_guard).
/// Holds the `MutexGuard` on the `IndexWriter`, preventing concurrent writes
/// (Tantivy enforces single-writer).
///
/// All write operations are batched in memory. Call [`commit()`](Self::commit)
/// to make changes durable and visible to new searchers.
pub struct TextIndexWriter<'a> {
    pub(crate) writer: MutexGuard<'a, IndexWriter>,
    pub(crate) fields: &'a TantivyFields,
}

impl TextIndexWriter<'_> {
    /// Index or re-index an item.
    ///
    /// Performs a delete-then-add in the same batch so updates are atomic
    /// from the reader's perspective (both become visible after `commit()`).
    /// Only metadata keys declared as text fields in the schema are indexed;
    /// other keys are silently ignored.
    ///
    /// # Errors
    ///
    /// Returns `TidalError::Internal` if Tantivy fails to add the document.
    #[tracing::instrument(skip(self, metadata), fields(%entity_id))]
    pub fn index_item(
        &mut self,
        entity_id: EntityId,
        metadata: &HashMap<String, String>,
    ) -> crate::Result<()> {
        // Delete any previous version of this document.
        let id_term = Term::from_field_u64(self.fields.entity_id, entity_id.as_u64());
        self.writer.delete_term(id_term);

        // Build new document.
        let mut doc = TantivyDocument::default();
        doc.add_u64(self.fields.entity_id, entity_id.as_u64());

        for (key, tv_field, _field_type) in &self.fields.text_fields {
            if let Some(value) = metadata.get(key) {
                doc.add_text(*tv_field, value);
            }
        }

        self.writer.add_document(doc).map_err(|e| {
            TidalError::internal("text_index_item", format!("tantivy add_document: {e}"))
        })?;

        Ok(())
    }

    /// Schedule deletion of a document by entity ID.
    ///
    /// The delete is applied on the next `commit()`. If the entity is not in
    /// the index, this is a no-op.
    pub fn delete_item(&mut self, entity_id: EntityId) {
        let id_term = Term::from_field_u64(self.fields.entity_id, entity_id.as_u64());
        self.writer.delete_term(id_term);
    }

    /// Commit all pending writes, making them durable and visible to searchers.
    ///
    /// After this returns, new `searcher()` instances will see the committed docs.
    ///
    /// Crash recovery for the text index is NOT driven from a commit payload.
    /// The durable entity store is the source of truth, and the text index is
    /// derived state that is unconditionally rebuilt from that store at open
    /// (see [`crate::db::state_rebuild::rebuild_text_indexes_at_open`]). This
    /// mirrors the vector index, which is likewise rebuilt from the durable
    /// embeddings on every reopen rather than reconciled against a stored
    /// sequence number.
    ///
    /// # Errors
    ///
    /// Returns `TidalError::Internal` if Tantivy fails to finalize the commit.
    #[tracing::instrument(skip(self))]
    pub fn commit(&mut self) -> crate::Result<()> {
        // Test-only: simulate a transient commit failure (disk-full / fsync IO)
        // before touching Tantivy, so the pending batch is left intact and the
        // syncer can retry on the next loop iteration.
        #[cfg(test)]
        {
            let should_fail = FORCE_COMMIT_FAILURES.with(|c| {
                let remaining = c.get();
                if remaining > 0 {
                    c.set(remaining - 1);
                    true
                } else {
                    false
                }
            });
            if should_fail {
                return Err(TidalError::internal(
                    "text_commit",
                    "injected transient commit failure".to_owned(),
                ));
            }
        }

        self.writer
            .commit()
            .map_err(|e| TidalError::internal("text_commit", format!("tantivy commit: {e}")))?;
        Ok(())
    }

    /// Delete all documents from the index.
    ///
    /// Used by [`TextIndex::rebuild_from`](super::TextIndex::rebuild_from) to
    /// start fresh before re-indexing from the entity store.
    ///
    /// # Errors
    ///
    /// Returns `TidalError::Internal` if Tantivy fails to delete documents.
    pub fn delete_all(&mut self) -> crate::Result<()> {
        self.writer.delete_all_documents().map_err(|e| {
            TidalError::internal("text_delete_all", format!("tantivy delete_all: {e}"))
        })?;
        Ok(())
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use tantivy::{collector::TopDocs, query::QueryParser, schema::Value};

    use super::*;
    use crate::{
        schema::{TextFieldDef, TextFieldType},
        text::TextIndex,
    };

    fn sample_text_fields() -> Vec<TextFieldDef> {
        vec![
            TextFieldDef {
                key: "title".to_owned(),
                field_type: TextFieldType::Text,
            },
            TextFieldDef {
                key: "description".to_owned(),
                field_type: TextFieldType::Text,
            },
        ]
    }

    fn make_metadata(pairs: &[(&str, &str)]) -> HashMap<String, String> {
        pairs
            .iter()
            .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
            .collect()
    }

    #[test]
    fn index_and_search() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(
                EntityId::new(1),
                &make_metadata(&[("title", "smooth jazz classics")]),
            )
            .unwrap();
            w.index_item(
                EntityId::new(2),
                &make_metadata(&[("title", "rock anthems")]),
            )
            .unwrap();
            w.index_item(
                EntityId::new(3),
                &make_metadata(&[("title", "jazz fusion favorites")]),
            )
            .unwrap();
            w.commit().unwrap();
        }

        idx.reader.reload().unwrap();
        let searcher = idx.reader.searcher();
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("jazz").unwrap();
        let top_docs = searcher.search(&query, &TopDocs::with_limit(10)).unwrap();

        assert_eq!(top_docs.len(), 2, "two docs contain 'jazz'");

        let found_ids: Vec<u64> = top_docs
            .iter()
            .map(|(_score, doc_address)| {
                let doc: TantivyDocument = searcher.doc(*doc_address).unwrap();
                doc.get_first(idx.fields().entity_id)
                    .and_then(|v| v.as_u64())
                    .unwrap()
            })
            .collect();

        assert!(found_ids.contains(&1));
        assert!(found_ids.contains(&3));
    }

    #[test]
    fn delete_removes_document() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(EntityId::new(1), &make_metadata(&[("title", "alpha")]))
                .unwrap();
            w.index_item(EntityId::new(2), &make_metadata(&[("title", "beta")]))
                .unwrap();
            w.commit().unwrap();
        }

        idx.reader.reload().unwrap();
        assert_eq!(idx.reader.searcher().num_docs(), 2);

        {
            let mut w = idx.writer_guard().unwrap();
            w.delete_item(EntityId::new(1));
            w.commit().unwrap();
        }

        idx.reader.reload().unwrap();
        assert_eq!(idx.reader.searcher().num_docs(), 1);
    }

    #[test]
    fn update_replaces_document() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        // Index with title A.
        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(
                EntityId::new(1),
                &make_metadata(&[("title", "original alpha")]),
            )
            .unwrap();
            w.commit().unwrap();
        }

        // Re-index same entity with title B.
        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(
                EntityId::new(1),
                &make_metadata(&[("title", "replacement beta")]),
            )
            .unwrap();
            w.commit().unwrap();
        }

        idx.reader.reload().unwrap();
        let searcher = idx.reader.searcher();

        // Only one document should exist.
        assert_eq!(searcher.num_docs(), 1);

        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);

        // Search for old title: not found.
        let query_old = qp.parse_query("original").unwrap();
        let results_old = searcher
            .search(&query_old, &TopDocs::with_limit(10))
            .unwrap();
        assert!(results_old.is_empty(), "old title should not be found");

        // Search for new title: found.
        let query_new = qp.parse_query("replacement").unwrap();
        let results_new = searcher
            .search(&query_new, &TopDocs::with_limit(10))
            .unwrap();
        assert_eq!(results_new.len(), 1, "new title should be found");

        let doc: TantivyDocument = searcher.doc(results_new[0].1).unwrap();
        let eid = doc
            .get_first(idx.fields().entity_id)
            .and_then(|v| v.as_u64())
            .unwrap();
        assert_eq!(eid, 1);
    }

    #[test]
    fn unknown_metadata_keys_ignored() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(
                EntityId::new(1),
                &make_metadata(&[
                    ("title", "hello world"),
                    ("nonexistent_key", "should be ignored"),
                    ("another_bogus", "also ignored"),
                ]),
            )
            .unwrap();
            w.commit().unwrap();
        }

        idx.reader.reload().unwrap();
        let searcher = idx.reader.searcher();
        assert_eq!(searcher.num_docs(), 1, "document should be indexed");

        // The entity_id should still be readable.
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("hello").unwrap();
        let results = searcher.search(&query, &TopDocs::with_limit(10)).unwrap();
        assert_eq!(results.len(), 1);
    }

    #[test]
    fn delete_nonexistent_is_noop() {
        let idx = TextIndex::ephemeral(&sample_text_fields()).unwrap();

        {
            let mut w = idx.writer_guard().unwrap();
            // Delete entity that was never indexed — should not panic or error.
            w.delete_item(EntityId::new(999));
            w.commit().unwrap();
        }

        idx.reader.reload().unwrap();
        assert_eq!(idx.reader.searcher().num_docs(), 0);
    }
}