zeppelin-embed 0.7.0

//! Whole-segment partition retention.

use std::ops::Range;
use std::sync::Arc;

use crate::lifecycle::durability::SyncRequirement;
use crate::lifecycle::{PublishedSnapshot, Store, StoreError, StoreState};
use crate::manifest::io::{MANIFEST_FILE, load_manifest};
use crate::segment::{ClusteringKeyRange, SegmentId, SegmentMeta};
use crate::vfs::Vfs;

/// Pure timestamp-window policy evaluated only when a host calls it.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct RetentionPolicy {
    window: i64,
}

/// Typed rejection of a non-positive retention window.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct RetentionPolicyError {
    window: i64,
}

impl std::fmt::Display for RetentionPolicyError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            formatter,
            "retention window {} must be positive",
            self.window
        )
    }
}

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

impl RetentionPolicy {
    /// Creates a positive retention window in the caller's timestamp unit.
    pub const fn new(window: i64) -> Result<Self, RetentionPolicyError> {
        if window <= 0 {
            Err(RetentionPolicyError { window })
        } else {
            Ok(Self { window })
        }
    }

    /// Purely decides the half-open old-key range outside the live window.
    #[must_use]
    pub const fn partition_to_drop(self, now_ts: i64) -> Range<i64> {
        i64::MIN..now_ts.saturating_sub(self.window)
    }
}

/// Result of one explicit whole-segment partition drop.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct DropPartitionReport {
    generation: u64,
    manifest_committed: bool,
    segments_dropped: Vec<SegmentId>,
    bytes_reclaimed: u64,
    orphaned_segments: Vec<SegmentId>,
    straddlers_skipped: Vec<SegmentId>,
}

impl DropPartitionReport {
    /// Returns the committed or unchanged generation observed by this call.
    #[must_use]
    pub const fn generation(&self) -> u64 {
        self.generation
    }

    /// Returns whether no manifest commit and no unlink were needed.
    #[must_use]
    pub const fn is_no_op(&self) -> bool {
        !self.manifest_committed
    }

    /// Returns immutable segments omitted by the committed manifest.
    #[must_use]
    pub fn segments_dropped(&self) -> &[SegmentId] {
        &self.segments_dropped
    }

    /// Returns the exact immutable file bytes unlinked after commit.
    #[must_use]
    pub const fn bytes_reclaimed(&self) -> u64 {
        self.bytes_reclaimed
    }

    /// Returns dropped segments whose post-commit unlink failed.
    #[must_use]
    pub fn orphaned_segments(&self) -> &[SegmentId] {
        &self.orphaned_segments
    }

    /// Returns overlapping segments retained because they crossed a boundary.
    #[must_use]
    pub fn straddlers_skipped(&self) -> &[SegmentId] {
        &self.straddlers_skipped
    }
}

impl Store {
    /// Drops immutable segments wholly contained by the half-open timestamp range.
    pub fn drop_partition(&self, key_range: Range<i64>) -> Result<DropPartitionReport, StoreError> {
        let retained_through = key_range.end;
        let report = self.drop_partition_on_vfs(key_range, self.vfs.as_ref())?;
        self.record_retention(retained_through)?;
        Ok(report)
    }

    /// Evaluates a retention policy now and explicitly invokes partition drop.
    pub fn apply_retention(
        &self,
        policy: RetentionPolicy,
        now_ts: i64,
    ) -> Result<DropPartitionReport, StoreError> {
        let range = policy.partition_to_drop(now_ts);
        #[cfg(any(test, feature = "test-seams"))]
        let receipt_plan = match self.ingest_retention_fault_controller.as_ref() {
            Some(controller) => controller
                .retention_clock_boundary_plan(now_ts, policy.window, range.end)
                .map_err(|_| StoreError::Synchronization {
                    component: "ingest-retention controller",
                })?,
            None => None,
        };
        let report = self.drop_partition(range)?;
        #[cfg(any(test, feature = "test-seams"))]
        if let Some(invocation_id) = receipt_plan
            && let Some(controller) = self.ingest_retention_fault_controller.as_ref()
        {
            let dropped_count = u64::try_from(report.segments_dropped.len())
                .map_err(|_| StoreError::PartitionBytesOverflow)?;
            let straddler_count = u64::try_from(report.straddlers_skipped.len())
                .map_err(|_| StoreError::PartitionBytesOverflow)?;
            controller
                .push_receipt(
                    super::IngestRetentionFaultReceiptV1::retention_clock_boundary(
                        invocation_id,
                        now_ts,
                        policy.window,
                        policy.partition_to_drop(now_ts).end,
                        report.generation,
                        dropped_count,
                        straddler_count,
                        report.manifest_committed,
                    ),
                )
                .map_err(|_| StoreError::Synchronization {
                    component: "ingest-retention controller",
                })?;
        }
        Ok(report)
    }

    /// Test-support seam for counting filesystem work during partition drop.
    #[doc(hidden)]
    pub fn drop_partition_on_vfs(
        &self,
        key_range: Range<i64>,
        vfs: &dyn Vfs,
    ) -> Result<DropPartitionReport, StoreError> {
        let _data_read_audit =
            crate::segment::reader::install_data_read_audit(vfs.segment_data_read_counter());
        let _maintenance = self
            .maintenance
            .lock()
            .map_err(|_| StoreError::Synchronization {
                component: "partition retention",
            })?;
        #[cfg(feature = "graph-cypher")]
        let mut native_writer = self.native_document_writer()?;
        #[cfg(feature = "graph-cypher")]
        if let Some(writer) = native_writer.as_mut().and_then(|slot| slot.as_mut()) {
            let ids = {
                let snapshot = self
                    .snapshot
                    .read()
                    .map_err(|_| StoreError::Synchronization {
                        component: "published snapshot",
                    })?;
                let snapshot = snapshot.as_ref().ok_or(StoreError::Closed)?;
                let mut ids = Vec::new();
                if key_range.start < key_range.end {
                    for segment in snapshot.segments() {
                        if matches!(segment.meta().clustering_key_range,
                            ClusteringKeyRange::Bounded { min_ts, max_ts }
                                if min_ts >= key_range.start && max_ts < key_range.end)
                        {
                            for row in segment.query_alive()?.alive_bitmap().iter() {
                                if let Some(version) = segment
                                    .document_version(row as usize)
                                    .map_err(StoreError::Segment)?
                                {
                                    ids.push(version.doc_id());
                                }
                            }
                        }
                    }
                }
                ids
            };
            if !ids.is_empty()
                && self
                    .documents_have_native_nodes(&ids)
                    .map_err(|error| StoreError::DocumentMutation(Box::new(error)))?
            {
                self.delete_native_documents_locked(&super::DeleteBatch::new(ids), writer)
                    .map_err(|error| StoreError::DocumentMutation(Box::new(error)))?;
            }
        }
        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);
        }
        let wal = self
            .wal_writer
            .lock()
            .map_err(|_| StoreError::Synchronization {
                component: "WAL writer",
            })?;
        let mut publication = wal
            .as_ref()
            .ok_or(StoreError::ReadOnly)?
            .manifest_publication()?;
        let durable_end = wal.as_ref().ok_or(StoreError::ReadOnly)?.durable_end();
        let mut active = self
            .active
            .lock()
            .map_err(|_| StoreError::Synchronization {
                component: "active segment",
            })?;
        let active_state = active.as_mut().ok_or(StoreError::Closed)?;
        if key_range.start >= key_range.end {
            return Ok(no_op_report(active_state.generation, Vec::new()));
        }
        let manifest_path = self.directory.join(MANIFEST_FILE);
        let mut manifest = match vfs.open(&manifest_path) {
            Ok(_) => {
                load_manifest(vfs, &manifest_path, durable_end).map_err(StoreError::Manifest)?
            }
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
                return Ok(no_op_report(active_state.generation, Vec::new()));
            }
            Err(source) => {
                return Err(StoreError::Io {
                    path: manifest_path,
                    source,
                });
            }
        };
        let selection = select_segments(manifest.segments, &key_range)?;
        if selection.dropped.is_empty() {
            return Ok(no_op_report(
                active_state.generation.max(manifest.generation),
                selection.straddlers,
            ));
        }
        let generation = active_state
            .generation
            .max(manifest.generation)
            .checked_add(1)
            .ok_or(StoreError::GenerationOverflow)?;
        manifest.generation = generation;
        manifest.segments = selection.retained;
        #[cfg(feature = "graph-cypher")]
        manifest
            .record_generation_bump(durable_end)
            .map_err(StoreError::Manifest)?;
        manifest.epochs = self.epoch_registry(&manifest.epochs);
        // Construct and validate the replacement snapshot before crossing the
        // manifest commit point. After commit, publication is an in-memory
        // pointer swap and cannot fail on segment I/O or accounting budget.
        let remapped =
            PublishedSnapshot::from_manifest(vfs, &self.directory, &manifest, &self.accounting)?;
        publication
            .commit_manifest(vfs, &self.directory, &manifest, self.durability_policy)
            .map_err(StoreError::Manifest)?;

        // Store admission holds `state` across the manifest commit and this
        // publication, so a query admitted after the commit can only pin the
        // replacement snapshot that omits every dropped segment.
        let mut published = self
            .snapshot
            .write()
            .map_err(|_| StoreError::Synchronization {
                component: "published snapshot",
            })?;
        let previous = published.replace(Arc::new(remapped));
        active_state.generation = generation;
        #[cfg(feature = "graph-cypher")]
        {
            drop(published);
            self.publish_native_documents(active_state)?;
        }
        publication.complete();
        #[cfg(not(feature = "graph-cypher"))]
        drop(published);
        drop(previous);

        let mut bytes_reclaimed = selection.bytes_reclaimed;
        let mut orphaned_segments = Vec::new();
        for segment in &selection.dropped {
            let path = self.directory.join(segment.id.file_name());
            // The manifest omission is already committed, so an unlink failure
            // leaves an unreachable orphan rather than rolling back the drop.
            // Queries admitted before the commit own their snapshot's open
            // segment descriptor and mmap. POSIX unlink removes only the path;
            // the inode remains alive until the final pinned descriptor/mapping
            // is released, so their in-flight reads remain valid.
            if self.has_snapshot_views() || vfs.delete(&path).is_err() {
                bytes_reclaimed = bytes_reclaimed.saturating_sub(segment.file_size);
                orphaned_segments.push(segment.id);
            }
        }
        if let SyncRequirement::Sync(kind) = self.durability_policy.directory_sync() {
            // The manifest commit already crossed its durable directory-sync
            // boundary. This later sync only persists orphan cleanup and must
            // not deny the generation already published.
            let _ = vfs.sync(&self.directory, kind);
        }
        let dropped_ids = selection.dropped.iter().map(|segment| segment.id).collect();
        drop(active);
        drop(wal);
        drop(writer_lock);
        drop(state);
        Ok(DropPartitionReport {
            generation,
            manifest_committed: true,
            segments_dropped: dropped_ids,
            bytes_reclaimed,
            orphaned_segments,
            straddlers_skipped: selection.straddlers,
        })
    }
}

fn no_op_report(generation: u64, straddlers_skipped: Vec<SegmentId>) -> DropPartitionReport {
    DropPartitionReport {
        generation,
        manifest_committed: false,
        segments_dropped: Vec::new(),
        bytes_reclaimed: 0,
        orphaned_segments: Vec::new(),
        straddlers_skipped,
    }
}

struct SegmentSelection {
    retained: Vec<SegmentMeta>,
    dropped: Vec<SegmentMeta>,
    straddlers: Vec<SegmentId>,
    bytes_reclaimed: u64,
}

fn select_segments(
    segments: Vec<SegmentMeta>,
    requested: &Range<i64>,
) -> Result<SegmentSelection, StoreError> {
    let mut retained = Vec::with_capacity(segments.len());
    let mut dropped = Vec::new();
    let mut straddlers = Vec::new();
    let mut bytes_reclaimed = 0_u64;
    for segment in segments {
        match segment.clustering_key_range {
            // The empty set is safely contained by every non-empty requested
            // range, so all-tombstoned segments can be reclaimed without a
            // numeric sentinel pretending to be a timestamp.
            ClusteringKeyRange::Empty => {
                bytes_reclaimed = bytes_reclaimed
                    .checked_add(segment.file_size)
                    .ok_or(StoreError::PartitionBytesOverflow)?;
                dropped.push(segment);
            }
            ClusteringKeyRange::Bounded { min_ts, max_ts }
                if min_ts >= requested.start && max_ts < requested.end =>
            {
                bytes_reclaimed = bytes_reclaimed
                    .checked_add(segment.file_size)
                    .ok_or(StoreError::PartitionBytesOverflow)?;
                dropped.push(segment);
            }
            ClusteringKeyRange::Bounded { min_ts, max_ts }
                if max_ts >= requested.start && min_ts < requested.end =>
            {
                straddlers.push(segment.id);
                retained.push(segment);
            }
            ClusteringKeyRange::Unstamped | ClusteringKeyRange::Bounded { .. } => {
                retained.push(segment);
            }
        }
    }
    Ok(SegmentSelection {
        retained,
        dropped,
        straddlers,
        bytes_reclaimed,
    })
}