zeppelin-embed 0.7.0

//! Explicit active-segment sealing and WAL-prefix absorption.

use std::path::Path;
use std::sync::Arc;

use crate::fts::sealed::{SealedSegment, SealedSegmentError};
use crate::lifecycle::durability::{DurabilityPolicy, SyncRequirement};
use crate::lifecycle::{CancelToken, Store, StoreError, StoreState};
use crate::manifest::Manifest;
use crate::manifest::io::{MANIFEST_FILE, load_manifest};
use crate::meta::{ColumnStore, ColumnStoreBuilder, Schema};
use crate::segment::writer::{
    SegmentBuild, SegmentDocumentVersions, SegmentFactors, SegmentPayloads, SegmentPostings,
    SegmentStoredMetadata, SegmentStoredText, write_segment_with_documents_payloads,
};
use crate::segment::{ClusteringKeyRange, SegmentId};
use crate::vfs::Vfs;
use crate::wal::LogSeq;

use super::ActiveState;

struct SealPayloadBuffers<'a> {
    postings: Option<Vec<u8>>,
    metadata: Option<SegmentStoredMetadata<'a>>,
    text: Option<SegmentStoredText<'a>>,
}

#[cfg(any(test, feature = "test-seams"))]
type TestSealStorage = (Vec<u8>, Vec<crate::segment::layout::Int8Factors>);

impl Store {
    /// Seals the current active segment into one appended immutable segment.
    ///
    /// An empty active segment is an idempotent no-op that returns the current
    /// generation without publishing a manifest or segment.
    pub fn seal(&self) -> Result<u64, StoreError> {
        self.seal_inner(None, self.vfs.as_ref())
    }

    /// Seals explicitly unless caller cancellation wins before manifest commit.
    ///
    /// Cancellation is checked before the empty-segment no-op.
    pub fn seal_with_cancel(&self, cancel: &CancelToken) -> Result<u64, StoreError> {
        self.seal_inner(Some(cancel), self.vfs.as_ref())
    }

    /// Test-support seam for deterministic cancellation at filesystem boundaries.
    #[doc(hidden)]
    pub fn seal_with_cancel_on_vfs(
        &self,
        cancel: &CancelToken,
        vfs: &dyn Vfs,
    ) -> Result<u64, StoreError> {
        self.seal_inner(Some(cancel), vfs)
    }

    #[cfg(any(test, feature = "test-seams"))]
    fn test_seal_storage(
        &self,
        current: &ActiveState,
        dims: u32,
    ) -> Result<Option<TestSealStorage>, StoreError> {
        match self.vector_seal_scheme {
            Some(crate::quant::QuantScheme::Int8) => {
                let dimension = usize::try_from(dims).map_err(|_| StoreError::ActiveRowOverflow)?;
                let vectors = current.segment.vectors();
                if dimension == 0 || !vectors.len().is_multiple_of(dimension) {
                    return Err(StoreError::Segment(crate::segment::SegmentError::Geometry(
                        format!(
                            "Int8 seal vector length {} is not divisible by dimension {dimension}",
                            vectors.len()
                        ),
                    )));
                }
                let mut signed_codes = vec![0_i8; vectors.len()];
                let mut factors = Vec::with_capacity(vectors.len() / dimension);
                for (row, codes) in vectors
                    .chunks_exact(dimension)
                    .zip(signed_codes.chunks_exact_mut(dimension))
                {
                    let (scale, offset) =
                        crate::quant::quantize_int8(row, codes).map_err(|error| {
                            StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
                                "Int8 seal quantization failed: {error}"
                            )))
                        })?;
                    factors.push(crate::segment::layout::Int8Factors { scale, offset });
                }
                let codes = signed_codes
                    .into_iter()
                    .map(|code| code as u8)
                    .collect::<Vec<_>>();
                Ok(Some((codes, factors)))
            }
            Some(crate::quant::QuantScheme::Bit4) | None => Ok(None),
            Some(scheme) => Err(StoreError::Segment(crate::segment::SegmentError::Geometry(
                format!("unsupported test vector seal scheme {scheme:?}"),
            ))),
        }
    }

    fn abort_seal(
        &self,
        vfs: &dyn Vfs,
        id: SegmentId,
        _invocation: Option<u64>,
        _current: &ActiveState,
        _absorbed_through: u64,
    ) -> Result<(), StoreError> {
        cleanup_uncommitted_segment(vfs, &self.directory, id, self.durability_policy)?;
        #[cfg(any(test, feature = "test-seams"))]
        if let Some(invocation_id) = _invocation {
            let active_rows = u64::try_from(_current.segment.row_count())
                .map_err(|_| StoreError::ActiveRowOverflow)?;
            if let Some(controller) = self.ingest_retention_fault_controller.as_ref() {
                controller
                    .push_receipt(super::IngestRetentionFaultReceiptV1::seal_cancellation(
                        invocation_id,
                        active_rows,
                        _absorbed_through,
                        *id.as_bytes(),
                    ))
                    .map_err(|_| StoreError::Synchronization {
                        component: "ingest-retention controller",
                    })?;
            }
        }
        Ok(())
    }

    fn seal_inner(&self, cancel: Option<&CancelToken>, vfs: &dyn Vfs) -> Result<u64, StoreError> {
        let state = self
            .state
            .lock()
            .map_err(|_| StoreError::Synchronization { component: "state" })?;
        match *state {
            StoreState::Open => {}
            StoreState::Closing => return Err(StoreError::Closing),
            StoreState::Closed => return Err(StoreError::Closed),
        }
        let writer_lock = self
            .writer_lock
            .lock()
            .map_err(|_| StoreError::Synchronization {
                component: "writer lock",
            })?;
        if writer_lock.is_none() {
            return Err(StoreError::ReadOnly);
        }
        check_cancelled(cancel)?;
        let mut wal = self
            .wal_writer
            .lock()
            .map_err(|_| StoreError::Synchronization {
                component: "WAL writer",
            })?;
        let writer = wal.as_mut().ok_or(StoreError::ReadOnly)?;
        let mut publication = writer.manifest_publication()?;
        let absorbed_through = writer.durable_end();
        let mut active = self
            .active
            .lock()
            .map_err(|_| StoreError::Synchronization {
                component: "active segment",
            })?;
        let current = active.as_ref().ok_or(StoreError::Closed)?;
        if current.segment.is_empty() {
            // A crash-restart after a durably committed seal leaves a correctly
            // empty active segment; retrying seal is a replay of that commit.
            return Ok(current.generation);
        }
        let generation = current
            .generation
            .checked_add(1)
            .ok_or(StoreError::GenerationOverflow)?;
        #[cfg_attr(not(feature = "graph-cypher"), allow(unused_mut))]
        let mut manifest = load_current_manifest(
            vfs,
            &self.directory,
            absorbed_through,
            current.generation,
            &self.schema,
        )?;
        #[cfg(feature = "graph-cypher")]
        let graph_fold = if manifest.graph.is_some() {
            Some(
                self.prepare_native_graph_seal(&mut manifest, generation, absorbed_through)
                    .map_err(seal_graph_error)?,
            )
        } else {
            None
        };
        let columns = active_columns(&manifest.schema, &current.segment, cancel)?;
        let alive = current.segment.alive()?;
        let clustering_key_range = clustering_key_range(current.segment.timestamps(), &alive)?;
        let dims = current
            .segment
            .dims()
            .ok_or(StoreError::ActiveRowOverflow)?;
        let dims = u32::try_from(dims).map_err(|_| StoreError::ActiveRowOverflow)?;
        let id = seal_id(generation, absorbed_through);
        #[cfg(any(test, feature = "test-seams"))]
        let int8_storage = self.test_seal_storage(current, dims)?;
        #[cfg(any(test, feature = "test-seams"))]
        let (scheme, codes, factors) = match int8_storage.as_ref() {
            Some((codes, factors)) => (
                u16::from(crate::quant::QuantScheme::Int8.id()),
                codes.as_slice(),
                SegmentFactors::Int8(factors.as_slice()),
            ),
            None => (
                u16::from(crate::quant::QuantScheme::Bit4.id()),
                current.segment.codes(),
                SegmentFactors::Bit4(current.segment.factors()),
            ),
        };
        #[cfg(not(any(test, feature = "test-seams")))]
        let (scheme, codes, factors) = (
            u16::from(crate::quant::QuantScheme::Bit4.id()),
            current.segment.codes(),
            SegmentFactors::Bit4(current.segment.factors()),
        );
        let build = SegmentBuild {
            id,
            scheme,
            dims,
            codes,
            factors,
            rescore: current.segment.vectors(),
            columns: &columns,
            alive: &alive,
        };
        let documents = SegmentDocumentVersions {
            doc_ids: current.segment.doc_ids(),
            revisions: current.segment.revisions(),
        };
        let payloads = build_seal_payloads(&current.segment)?;
        publication = publication.arm();
        let written = write_segment_with_documents_payloads(
            vfs,
            &self.directory,
            build,
            SegmentPayloads {
                documents,
                metadata: payloads.metadata,
                text: payloads.text,
                postings: payloads
                    .postings
                    .as_deref()
                    .map(|bytes| SegmentPostings { bytes }),
            },
            self.durability_policy,
        );
        let mut meta = match written {
            Ok(meta) => meta,
            Err(error) => {
                cleanup_uncommitted_segment(vfs, &self.directory, id, self.durability_policy)?;
                return Err(StoreError::Segment(error));
            }
        };
        meta.clustering_key_range = clustering_key_range;
        meta.epoch_id = manifest.epoch_alias.map(|identity| identity.embedding);
        #[cfg(any(test, feature = "test-seams"))]
        let forced_late_cancellation =
            match self.ingest_retention_fault_controller.as_ref() {
                Some(controller) => controller.seal_cancellation_plan().map_err(|_| {
                    StoreError::Synchronization {
                        component: "ingest-retention controller",
                    }
                })?,
                None => None,
            };
        #[cfg(not(any(test, feature = "test-seams")))]
        let forced_late_cancellation: Option<u64> = None;
        if forced_late_cancellation.is_some() || check_cancelled(cancel).is_err() {
            self.abort_seal(vfs, id, forced_late_cancellation, current, absorbed_through)?;
            // Successful cleanup durably restores the pre-seal state; no
            // manifest or in-memory publication changed. Cleanup I/O errors
            // above still drop the armed guard and fence the shared writer.
            publication.complete();
            return Err(StoreError::SealCancelled);
        }
        let mut segments = manifest.segments;
        segments.push(meta);
        let epochs = self.epoch_registry(&manifest.epochs);
        publication
            .commit_manifest(
                vfs,
                &self.directory,
                &Manifest {
                    #[cfg(feature = "graph-cypher")]
                    graph: manifest.graph,
                    generation,
                    log_seq: absorbed_through,
                    segments,
                    epochs,
                    epoch_alias: manifest.epoch_alias,
                    schema: manifest.schema,
                },
                self.durability_policy,
            )
            .map_err(StoreError::Manifest)?;
        let remapped = crate::lifecycle::PublishedSnapshot::load_on_vfs(
            &self.directory,
            &self.accounting,
            vfs,
        )?;
        let mut published = self
            .snapshot
            .write()
            .map_err(|_| StoreError::Synchronization {
                component: "published snapshot",
            })?;
        let previous = published.replace(Arc::new(remapped));
        drop(published);
        drop(previous);
        *active = Some(ActiveState::empty(generation));
        #[cfg(feature = "graph-cypher")]
        if let Some(fold) = graph_fold {
            fold.publish(self).map_err(seal_graph_error)?;
        }
        #[cfg(feature = "graph-cypher")]
        self.publish_native_documents(active.as_ref().ok_or(StoreError::Closed)?)?;
        publication.complete();
        writer.retire_visible_through(LogSeq::new(absorbed_through))?;
        // Both watermarks above are at the durable end, so the log can
        // shrink to a header; without this wal.ze grows for the store's life.
        writer.truncate_absorbed(
            vfs,
            &self.directory,
            self.durability_policy,
            LogSeq::new(absorbed_through),
        )?;
        drop(active);
        drop(wal);
        drop(writer_lock);
        drop(state);
        Ok(generation)
    }
}

#[cfg(feature = "graph-cypher")]
fn seal_graph_error(error: crate::lifecycle::native_graph::NativeGraphError) -> StoreError {
    match error {
        crate::lifecycle::native_graph::NativeGraphError::Store(error) => error,
        error => StoreError::Manifest(crate::manifest::ManifestError::Decode(format!(
            "seal graph fold: {error}",
        ))),
    }
}

fn build_seal_payloads(
    segment: &super::ActiveSegment,
) -> Result<SealPayloadBuffers<'_>, StoreError> {
    let postings = if segment.has_text() {
        let sealed = SealedSegment::seal(segment.lexical())
            .map_err(SealedSegmentError::from)
            .map_err(crate::segment::SegmentError::from)
            .map_err(StoreError::Segment)?;
        Some(
            sealed
                .encode_region()
                .map_err(crate::segment::SegmentError::from)
                .map_err(StoreError::Segment)?,
        )
    } else {
        None
    };
    let metadata = (!segment.metadata_bytes().is_empty()).then_some(SegmentStoredMetadata {
        end_offsets: segment.metadata_end_offsets(),
        bytes: segment.metadata_bytes(),
    });
    let text = segment.has_text().then_some(SegmentStoredText {
        present: segment.text_present(),
        end_offsets: segment.text_end_offsets(),
        bytes: segment.text_bytes(),
    });
    Ok(SealPayloadBuffers {
        postings,
        metadata,
        text,
    })
}

pub(crate) fn load_current_manifest(
    vfs: &dyn Vfs,
    directory: &Path,
    durable_end: u64,
    generation: u64,
    schema: &Schema,
) -> Result<Manifest, StoreError> {
    let path = directory.join(MANIFEST_FILE);
    match vfs.open(&path) {
        Ok(_) => load_manifest(vfs, &path, durable_end).map_err(StoreError::Manifest),
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Manifest {
            #[cfg(feature = "graph-cypher")]
            graph: None,
            generation,
            log_seq: 0,
            segments: Vec::new(),
            epochs: Vec::new(),
            epoch_alias: None,
            schema: schema.clone(),
        }),
        Err(source) => Err(StoreError::Io { path, source }),
    }
}

fn active_columns(
    schema: &Schema,
    active: &super::ActiveSegment,
    cancel: Option<&CancelToken>,
) -> Result<ColumnStore, StoreError> {
    let mut builder = ColumnStoreBuilder::new(schema.clone());
    for (row, timestamp) in active.timestamps().iter().copied().enumerate() {
        if row.is_multiple_of(64) {
            check_cancelled(cancel)?;
        }
        let values = active.column_values(row)?;
        builder
            .push_row(timestamp, &super::column_inputs(&values))
            .map_err(|error| {
                StoreError::Segment(crate::segment::SegmentError::Columns(error.to_string()))
            })?;
    }
    builder.finish().map_err(|error| {
        StoreError::Segment(crate::segment::SegmentError::Columns(error.to_string()))
    })
}

fn clustering_key_range(
    timestamps: &[i64],
    alive: &crate::meta::AliveSet,
) -> Result<ClusteringKeyRange, StoreError> {
    let mut bounds: Option<(i64, i64)> = None;
    for row in alive.iter_alive() {
        let index = usize::try_from(row).map_err(|_| StoreError::ActiveRowOverflow)?;
        let timestamp = timestamps
            .get(index)
            .copied()
            .ok_or(StoreError::ActiveRowOverflow)?;
        bounds = Some(match bounds {
            None => (timestamp, timestamp),
            Some((minimum, maximum)) => (minimum.min(timestamp), maximum.max(timestamp)),
        });
    }
    Ok(match bounds {
        Some((min_ts, max_ts)) => ClusteringKeyRange::Bounded { min_ts, max_ts },
        // An all-tombstoned segment has the empty key set. It is represented
        // explicitly rather than overloading a numeric sentinel or Unstamped.
        None => ClusteringKeyRange::Empty,
    })
}

fn check_cancelled(cancel: Option<&CancelToken>) -> Result<(), StoreError> {
    if cancel.is_some_and(CancelToken::is_cancelled) {
        Err(StoreError::SealCancelled)
    } else {
        Ok(())
    }
}

fn seal_id(generation: u64, absorbed_through: u64) -> SegmentId {
    let mut bytes = [0_u8; 16];
    if let Some(prefix) = bytes.get_mut(..8) {
        prefix.copy_from_slice(&generation.to_be_bytes());
    }
    if let Some(suffix) = bytes.get_mut(8..) {
        suffix.copy_from_slice(&absorbed_through.to_be_bytes());
    }
    SegmentId::from_bytes(bytes)
}

fn cleanup_uncommitted_segment(
    vfs: &dyn Vfs,
    directory: &Path,
    id: SegmentId,
    policy: DurabilityPolicy,
) -> Result<(), StoreError> {
    let final_path = directory.join(id.file_name());
    let temporary_path = directory.join(format!(".{}.tmp", id.file_name()));
    let mut deleted = false;
    for path in [temporary_path, final_path] {
        match vfs.open(&path) {
            Ok(_) => {
                vfs.delete(&path)
                    .map_err(|source| StoreError::Io { path, source })?;
                deleted = true;
            }
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
            Err(source) => return Err(StoreError::Io { path, source }),
        }
    }
    if deleted && let SyncRequirement::Sync(kind) = policy.directory_sync() {
        vfs.sync(directory, kind).map_err(|source| StoreError::Io {
            path: directory.to_path_buf(),
            source,
        })?;
    }
    Ok(())
}

#[cfg(test)]
#[allow(clippy::expect_used, clippy::panic)]
mod wal_truncation_crash_tests {
    use std::path::Path;
    use std::sync::Arc;

    use crate::ingest::StoreWal;
    use crate::lifecycle::durability::{CommitTier, DurabilityMode, DurabilityPolicy};
    use crate::lifecycle::stats::Accounting;
    use crate::vfs::crash::{CrashVfs, MemoryVfs};
    use crate::wal::{LogSeq, WalReader, encode_wal_image};

    /// Every power-cut image of the rotation, including torn, reordered and
    /// unsynced-rename states, must still open with the absorbed boundary
    /// intact and hand out sequence 4 next. An empty or headerless wal.ze
    /// would restart numbering at 1, below the manifest, and lose writes.
    #[test]
    fn every_crash_state_of_the_wal_truncation_resumes_after_the_boundary() {
        let directory = Path::new("/wal-truncation");
        let wal = directory.join("wal.ze");
        let records = (0..3_u8)
            .map(|index| (99, vec![index; 8]))
            .collect::<Vec<_>>();
        let memory = MemoryVfs::new();
        memory
            .insert(
                wal.clone(),
                encode_wal_image(LogSeq::new(1), &records).expect("old WAL"),
            )
            .expect("seed old WAL");
        let crash = CrashVfs::new(memory).expect("crash recorder");
        let policy = DurabilityPolicy::new(DurabilityMode::Durable, CommitTier::Ordered)
            .expect("ordered policy");
        let accounting = Arc::new(Accounting::new(u64::MAX, u64::MAX));
        let recovered = WalReader::open(&crash, &wal)
            .expect("open old WAL")
            .into_clean()
            .expect("clean old WAL");
        let mut writer = StoreWal::resume(&crash, &wal, recovered, policy, 3, &accounting)
            .expect("resume old WAL");

        writer
            .truncate_absorbed(&crash, directory, policy, LogSeq::new(3))
            .expect("truncate");

        let states = crash.crash_states().expect("crash states");
        assert!(!states.was_capped());
        assert!(states.len() > 10, "only {} crash states", states.len());
        for state in states.iter() {
            let recovered = WalReader::open(state.vfs(), &wal)
                .expect("read WAL")
                .into_clean()
                .unwrap_or_else(|error| panic!("{:?}: WAL unreadable: {error}", state.kind()));
            let sequences = recovered
                .records()
                .iter()
                .map(|record| record.seq.get())
                .collect::<Vec<_>>();
            assert!(
                sequences == [1, 2, 3] || sequences.is_empty(),
                "{:?}: WAL holds {sequences:?}",
                state.kind()
            );
            let mut resumed =
                StoreWal::resume(state.vfs(), &wal, recovered, policy, 3, &accounting)
                    .unwrap_or_else(|error| panic!("{:?}: resume: {error}", state.kind()));
            assert_eq!(
                resumed.commit(99, b"next").expect("commit after crash"),
                LogSeq::new(4),
                "{:?}",
                state.kind()
            );
            #[cfg(feature = "graph-cypher")]
            {
                use crate::ingest::wal_payload::{GRAPH_COMMIT_V1, encode_graph_commit};
                let graph = encode_graph_commit(
                    include_bytes!("../../tests/fixtures/graph-wal/complete-v1.bin")
                        .get(64..2105)
                        .expect("graph envelope"),
                )
                .expect("graph payload");
                assert_eq!(
                    resumed
                        .commit(GRAPH_COMMIT_V1, &graph)
                        .expect("graph after fold"),
                    LogSeq::new(5)
                );
                let before = crate::vfs::Vfs::read(state.vfs(), &wal).expect("WAL before refusal");
                assert!(
                    resumed
                        .truncate_absorbed(state.vfs(), directory, policy, LogSeq::new(3))
                        .is_err()
                );
                assert_eq!(
                    crate::vfs::Vfs::read(state.vfs(), &wal).expect("WAL after refusal"),
                    before
                );
                assert_eq!(
                    resumed
                        .commit(99, b"after graph")
                        .expect("definite refusal keeps writer usable"),
                    LogSeq::new(6)
                );
                let recovered = WalReader::open(state.vfs(), &wal)
                    .expect("graph WAL")
                    .into_clean()
                    .expect("clean graph WAL");
                let record = recovered
                    .records()
                    .iter()
                    .find(|record| record.seq.get() == 5)
                    .expect("unabsorbed graph retained");
                assert_eq!(record.op, GRAPH_COMMIT_V1);
                assert_eq!(record.payload().expect("graph payload"), graph);
                let mut next =
                    StoreWal::resume(state.vfs(), &wal, recovered, policy, 3, &accounting)
                        .expect("resume graph tail");
                assert_eq!(
                    next.commit(99, b"next reopen").expect("write after reopen"),
                    LogSeq::new(7)
                );
            }
        }
    }
}

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