Skip to main content

zeppelin_embed/ingest/
seal.rs

1//! Explicit active-segment sealing and WAL-prefix absorption.
2
3use std::path::Path;
4use std::sync::Arc;
5
6use crate::fts::sealed::{SealedSegment, SealedSegmentError};
7use crate::lifecycle::durability::{DurabilityPolicy, SyncRequirement};
8use crate::lifecycle::{CancelToken, Store, StoreError, StoreState};
9use crate::manifest::Manifest;
10use crate::manifest::io::{MANIFEST_FILE, load_manifest};
11use crate::meta::{ColumnStore, ColumnStoreBuilder, Schema};
12use crate::segment::writer::{
13    SegmentBuild, SegmentDocumentVersions, SegmentFactors, SegmentPayloads, SegmentPostings,
14    SegmentStoredMetadata, SegmentStoredText, write_segment_with_documents_payloads,
15};
16use crate::segment::{ClusteringKeyRange, SegmentId};
17use crate::vfs::Vfs;
18use crate::wal::LogSeq;
19
20use super::ActiveState;
21
22struct SealPayloadBuffers<'a> {
23    postings: Option<Vec<u8>>,
24    metadata: Option<SegmentStoredMetadata<'a>>,
25    text: Option<SegmentStoredText<'a>>,
26}
27
28#[cfg(any(test, feature = "test-seams"))]
29type TestSealStorage = (Vec<u8>, Vec<crate::segment::layout::Int8Factors>);
30
31impl Store {
32    /// Seals the current active segment into one appended immutable segment.
33    ///
34    /// An empty active segment is an idempotent no-op that returns the current
35    /// generation without publishing a manifest or segment.
36    pub fn seal(&self) -> Result<u64, StoreError> {
37        self.seal_inner(None, self.vfs.as_ref())
38    }
39
40    /// Seals explicitly unless caller cancellation wins before manifest commit.
41    ///
42    /// Cancellation is checked before the empty-segment no-op.
43    pub fn seal_with_cancel(&self, cancel: &CancelToken) -> Result<u64, StoreError> {
44        self.seal_inner(Some(cancel), self.vfs.as_ref())
45    }
46
47    /// Test-support seam for deterministic cancellation at filesystem boundaries.
48    #[doc(hidden)]
49    pub fn seal_with_cancel_on_vfs(
50        &self,
51        cancel: &CancelToken,
52        vfs: &dyn Vfs,
53    ) -> Result<u64, StoreError> {
54        self.seal_inner(Some(cancel), vfs)
55    }
56
57    #[cfg(any(test, feature = "test-seams"))]
58    fn test_seal_storage(
59        &self,
60        current: &ActiveState,
61        dims: u32,
62    ) -> Result<Option<TestSealStorage>, StoreError> {
63        match self.vector_seal_scheme {
64            Some(crate::quant::QuantScheme::Int8) => {
65                let dimension = usize::try_from(dims).map_err(|_| StoreError::ActiveRowOverflow)?;
66                let vectors = current.segment.vectors();
67                if dimension == 0 || !vectors.len().is_multiple_of(dimension) {
68                    return Err(StoreError::Segment(crate::segment::SegmentError::Geometry(
69                        format!(
70                            "Int8 seal vector length {} is not divisible by dimension {dimension}",
71                            vectors.len()
72                        ),
73                    )));
74                }
75                let mut signed_codes = vec![0_i8; vectors.len()];
76                let mut factors = Vec::with_capacity(vectors.len() / dimension);
77                for (row, codes) in vectors
78                    .chunks_exact(dimension)
79                    .zip(signed_codes.chunks_exact_mut(dimension))
80                {
81                    let (scale, offset) =
82                        crate::quant::quantize_int8(row, codes).map_err(|error| {
83                            StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
84                                "Int8 seal quantization failed: {error}"
85                            )))
86                        })?;
87                    factors.push(crate::segment::layout::Int8Factors { scale, offset });
88                }
89                let codes = signed_codes
90                    .into_iter()
91                    .map(|code| code as u8)
92                    .collect::<Vec<_>>();
93                Ok(Some((codes, factors)))
94            }
95            Some(crate::quant::QuantScheme::Bit4) | None => Ok(None),
96            Some(scheme) => Err(StoreError::Segment(crate::segment::SegmentError::Geometry(
97                format!("unsupported test vector seal scheme {scheme:?}"),
98            ))),
99        }
100    }
101
102    fn abort_seal(
103        &self,
104        vfs: &dyn Vfs,
105        id: SegmentId,
106        _invocation: Option<u64>,
107        _current: &ActiveState,
108        _absorbed_through: u64,
109    ) -> Result<(), StoreError> {
110        cleanup_uncommitted_segment(vfs, &self.directory, id, self.durability_policy)?;
111        #[cfg(any(test, feature = "test-seams"))]
112        if let Some(invocation_id) = _invocation {
113            let active_rows = u64::try_from(_current.segment.row_count())
114                .map_err(|_| StoreError::ActiveRowOverflow)?;
115            if let Some(controller) = self.ingest_retention_fault_controller.as_ref() {
116                controller
117                    .push_receipt(super::IngestRetentionFaultReceiptV1::seal_cancellation(
118                        invocation_id,
119                        active_rows,
120                        _absorbed_through,
121                        *id.as_bytes(),
122                    ))
123                    .map_err(|_| StoreError::Synchronization {
124                        component: "ingest-retention controller",
125                    })?;
126            }
127        }
128        Ok(())
129    }
130
131    fn seal_inner(&self, cancel: Option<&CancelToken>, vfs: &dyn Vfs) -> Result<u64, StoreError> {
132        let state = self
133            .state
134            .lock()
135            .map_err(|_| StoreError::Synchronization { component: "state" })?;
136        match *state {
137            StoreState::Open => {}
138            StoreState::Closing => return Err(StoreError::Closing),
139            StoreState::Closed => return Err(StoreError::Closed),
140        }
141        let writer_lock = self
142            .writer_lock
143            .lock()
144            .map_err(|_| StoreError::Synchronization {
145                component: "writer lock",
146            })?;
147        if writer_lock.is_none() {
148            return Err(StoreError::ReadOnly);
149        }
150        check_cancelled(cancel)?;
151        let mut wal = self
152            .wal_writer
153            .lock()
154            .map_err(|_| StoreError::Synchronization {
155                component: "WAL writer",
156            })?;
157        let writer = wal.as_mut().ok_or(StoreError::ReadOnly)?;
158        let mut active = self
159            .active
160            .lock()
161            .map_err(|_| StoreError::Synchronization {
162                component: "active segment",
163            })?;
164        let generation = self.seal_active_locked(writer, &mut active, cancel, vfs)?;
165        drop(active);
166        drop(wal);
167        drop(writer_lock);
168        drop(state);
169        Ok(generation)
170    }
171
172    /// The caller holds WAL and active locks, excluding every publisher.
173    pub(crate) fn seal_active_locked(
174        &self,
175        writer: &mut super::StoreWal,
176        active: &mut Option<ActiveState>,
177        cancel: Option<&CancelToken>,
178        vfs: &dyn Vfs,
179    ) -> Result<u64, StoreError> {
180        let mut publication = writer.manifest_publication()?;
181        let absorbed_through = writer.durable_end();
182        let current = active.as_ref().ok_or(StoreError::Closed)?;
183        if current.segment.is_empty() {
184            // A crash-restart after a durably committed seal leaves a correctly
185            // empty active segment; retrying seal is a replay of that commit.
186            return Ok(current.generation);
187        }
188        let generation = current
189            .generation
190            .checked_add(1)
191            .ok_or(StoreError::GenerationOverflow)?;
192        #[cfg_attr(not(feature = "graph-cypher"), allow(unused_mut))]
193        let mut manifest = load_current_manifest(
194            vfs,
195            &self.directory,
196            absorbed_through,
197            current.generation,
198            &self.schema,
199        )?;
200        #[cfg(feature = "graph-cypher")]
201        let graph_fold = if manifest.graph.is_some() {
202            Some(
203                self.prepare_native_graph_seal(&mut manifest, generation, absorbed_through)
204                    .map_err(seal_graph_error)?,
205            )
206        } else {
207            None
208        };
209        let columns = active_columns(&manifest.schema, &current.segment, cancel)?;
210        let alive = current.segment.alive()?;
211        let clustering_key_range = clustering_key_range(current.segment.timestamps(), &alive)?;
212        let dims = current
213            .segment
214            .dims()
215            .ok_or(StoreError::ActiveRowOverflow)?;
216        let dims = u32::try_from(dims).map_err(|_| StoreError::ActiveRowOverflow)?;
217        let id = seal_id(generation, absorbed_through);
218        #[cfg(any(test, feature = "test-seams"))]
219        let int8_storage = self.test_seal_storage(current, dims)?;
220        #[cfg(any(test, feature = "test-seams"))]
221        let (scheme, codes, factors) = match int8_storage.as_ref() {
222            Some((codes, factors)) => (
223                u16::from(crate::quant::QuantScheme::Int8.id()),
224                codes.as_slice(),
225                SegmentFactors::Int8(factors.as_slice()),
226            ),
227            None => (
228                u16::from(crate::quant::QuantScheme::Bit4.id()),
229                current.segment.codes(),
230                SegmentFactors::Bit4(current.segment.factors()),
231            ),
232        };
233        #[cfg(not(any(test, feature = "test-seams")))]
234        let (scheme, codes, factors) = (
235            u16::from(crate::quant::QuantScheme::Bit4.id()),
236            current.segment.codes(),
237            SegmentFactors::Bit4(current.segment.factors()),
238        );
239        let build = SegmentBuild {
240            id,
241            scheme,
242            dims,
243            codes,
244            factors,
245            rescore: current.segment.vectors(),
246            columns: &columns,
247            alive: &alive,
248        };
249        let documents = SegmentDocumentVersions {
250            doc_ids: current.segment.doc_ids(),
251            revisions: current.segment.revisions(),
252        };
253        let payloads = build_seal_payloads(&current.segment)?;
254        publication = publication.arm();
255        let written = write_segment_with_documents_payloads(
256            vfs,
257            &self.directory,
258            build,
259            SegmentPayloads {
260                documents,
261                metadata: payloads.metadata,
262                text: payloads.text,
263                postings: payloads
264                    .postings
265                    .as_deref()
266                    .map(|bytes| SegmentPostings { bytes }),
267            },
268            self.durability_policy,
269        );
270        let mut meta = match written {
271            Ok(meta) => meta,
272            Err(error) => {
273                cleanup_uncommitted_segment(vfs, &self.directory, id, self.durability_policy)?;
274                return Err(StoreError::Segment(error));
275            }
276        };
277        meta.clustering_key_range = clustering_key_range;
278        meta.epoch_id = manifest.epoch_alias.map(|identity| identity.embedding);
279        #[cfg(any(test, feature = "test-seams"))]
280        let forced_late_cancellation =
281            match self.ingest_retention_fault_controller.as_ref() {
282                Some(controller) => controller.seal_cancellation_plan().map_err(|_| {
283                    StoreError::Synchronization {
284                        component: "ingest-retention controller",
285                    }
286                })?,
287                None => None,
288            };
289        #[cfg(not(any(test, feature = "test-seams")))]
290        let forced_late_cancellation: Option<u64> = None;
291        if forced_late_cancellation.is_some() || check_cancelled(cancel).is_err() {
292            self.abort_seal(vfs, id, forced_late_cancellation, current, absorbed_through)?;
293            // Successful cleanup durably restores the pre-seal state; no
294            // manifest or in-memory publication changed. Cleanup I/O errors
295            // above still drop the armed guard and fence the shared writer.
296            publication.complete();
297            return Err(StoreError::SealCancelled);
298        }
299        let mut segments = manifest.segments;
300        segments.push(meta);
301        let epochs = self.epoch_registry(&manifest.epochs);
302        publication
303            .commit_manifest(
304                vfs,
305                &self.directory,
306                &Manifest {
307                    #[cfg(feature = "graph-cypher")]
308                    graph: manifest.graph,
309                    generation,
310                    log_seq: absorbed_through,
311                    segments,
312                    epochs,
313                    epoch_alias: manifest.epoch_alias,
314                    schema: manifest.schema,
315                },
316                self.durability_policy,
317            )
318            .map_err(StoreError::Manifest)?;
319        let remapped = crate::lifecycle::PublishedSnapshot::load_on_vfs(
320            &self.directory,
321            &self.accounting,
322            vfs,
323        )?;
324        let mut published = self
325            .snapshot
326            .write()
327            .map_err(|_| StoreError::Synchronization {
328                component: "published snapshot",
329            })?;
330        let previous = published.replace(Arc::new(remapped));
331        drop(published);
332        drop(previous);
333        *active = Some(ActiveState::empty(generation));
334        #[cfg(feature = "graph-cypher")]
335        if let Some(fold) = graph_fold {
336            fold.publish(self).map_err(seal_graph_error)?;
337        }
338        #[cfg(feature = "graph-cypher")]
339        self.publish_native_documents(active.as_ref().ok_or(StoreError::Closed)?)?;
340        publication.complete();
341        writer.retire_visible_through(LogSeq::new(absorbed_through))?;
342        // Both watermarks above are at the durable end, so the log can
343        // shrink to a header; without this wal.ze grows for the store's life.
344        writer.truncate_absorbed(
345            vfs,
346            &self.directory,
347            self.durability_policy,
348            LogSeq::new(absorbed_through),
349        )?;
350        Ok(generation)
351    }
352}
353
354#[cfg(feature = "graph-cypher")]
355fn seal_graph_error(error: crate::lifecycle::native_graph::NativeGraphError) -> StoreError {
356    match error {
357        crate::lifecycle::native_graph::NativeGraphError::Store(error) => error,
358        error => StoreError::Manifest(crate::manifest::ManifestError::Decode(format!(
359            "seal graph fold: {error}",
360        ))),
361    }
362}
363
364fn build_seal_payloads(
365    segment: &super::ActiveSegment,
366) -> Result<SealPayloadBuffers<'_>, StoreError> {
367    let postings = if segment.has_text() {
368        let sealed = SealedSegment::seal(segment.lexical())
369            .map_err(SealedSegmentError::from)
370            .map_err(crate::segment::SegmentError::from)
371            .map_err(StoreError::Segment)?;
372        Some(
373            sealed
374                .encode_region()
375                .map_err(crate::segment::SegmentError::from)
376                .map_err(StoreError::Segment)?,
377        )
378    } else {
379        None
380    };
381    let metadata = (!segment.metadata_bytes().is_empty()).then_some(SegmentStoredMetadata {
382        end_offsets: segment.metadata_end_offsets(),
383        bytes: segment.metadata_bytes(),
384    });
385    let text = segment.has_text().then_some(SegmentStoredText {
386        present: segment.text_present(),
387        end_offsets: segment.text_end_offsets(),
388        bytes: segment.text_bytes(),
389    });
390    Ok(SealPayloadBuffers {
391        postings,
392        metadata,
393        text,
394    })
395}
396
397pub(crate) fn load_current_manifest(
398    vfs: &dyn Vfs,
399    directory: &Path,
400    durable_end: u64,
401    generation: u64,
402    schema: &Schema,
403) -> Result<Manifest, StoreError> {
404    let path = directory.join(MANIFEST_FILE);
405    match vfs.open(&path) {
406        Ok(_) => load_manifest(vfs, &path, durable_end).map_err(StoreError::Manifest),
407        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Manifest {
408            #[cfg(feature = "graph-cypher")]
409            graph: None,
410            generation,
411            log_seq: 0,
412            segments: Vec::new(),
413            epochs: Vec::new(),
414            epoch_alias: None,
415            schema: schema.clone(),
416        }),
417        Err(source) => Err(StoreError::Io { path, source }),
418    }
419}
420
421fn active_columns(
422    schema: &Schema,
423    active: &super::ActiveSegment,
424    cancel: Option<&CancelToken>,
425) -> Result<ColumnStore, StoreError> {
426    let mut builder = ColumnStoreBuilder::new(schema.clone());
427    for (row, timestamp) in active.timestamps().iter().copied().enumerate() {
428        if row.is_multiple_of(64) {
429            check_cancelled(cancel)?;
430        }
431        let values = active.column_values(row)?;
432        builder
433            .push_row(timestamp, &super::column_inputs(&values))
434            .map_err(|error| {
435                StoreError::Segment(crate::segment::SegmentError::Columns(error.to_string()))
436            })?;
437    }
438    builder.finish().map_err(|error| {
439        StoreError::Segment(crate::segment::SegmentError::Columns(error.to_string()))
440    })
441}
442
443fn clustering_key_range(
444    timestamps: &[i64],
445    alive: &crate::meta::AliveSet,
446) -> Result<ClusteringKeyRange, StoreError> {
447    let mut bounds: Option<(i64, i64)> = None;
448    for row in alive.iter_alive() {
449        let index = usize::try_from(row).map_err(|_| StoreError::ActiveRowOverflow)?;
450        let timestamp = timestamps
451            .get(index)
452            .copied()
453            .ok_or(StoreError::ActiveRowOverflow)?;
454        bounds = Some(match bounds {
455            None => (timestamp, timestamp),
456            Some((minimum, maximum)) => (minimum.min(timestamp), maximum.max(timestamp)),
457        });
458    }
459    Ok(match bounds {
460        Some((min_ts, max_ts)) => ClusteringKeyRange::Bounded { min_ts, max_ts },
461        // An all-tombstoned segment has the empty key set. It is represented
462        // explicitly rather than overloading a numeric sentinel or Unstamped.
463        None => ClusteringKeyRange::Empty,
464    })
465}
466
467fn check_cancelled(cancel: Option<&CancelToken>) -> Result<(), StoreError> {
468    if cancel.is_some_and(CancelToken::is_cancelled) {
469        Err(StoreError::SealCancelled)
470    } else {
471        Ok(())
472    }
473}
474
475fn seal_id(generation: u64, absorbed_through: u64) -> SegmentId {
476    let mut bytes = [0_u8; 16];
477    if let Some(prefix) = bytes.get_mut(..8) {
478        prefix.copy_from_slice(&generation.to_be_bytes());
479    }
480    if let Some(suffix) = bytes.get_mut(8..) {
481        suffix.copy_from_slice(&absorbed_through.to_be_bytes());
482    }
483    SegmentId::from_bytes(bytes)
484}
485
486fn cleanup_uncommitted_segment(
487    vfs: &dyn Vfs,
488    directory: &Path,
489    id: SegmentId,
490    policy: DurabilityPolicy,
491) -> Result<(), StoreError> {
492    let final_path = directory.join(id.file_name());
493    let temporary_path = directory.join(format!(".{}.tmp", id.file_name()));
494    let mut deleted = false;
495    for path in [temporary_path, final_path] {
496        match vfs.open(&path) {
497            Ok(_) => {
498                vfs.delete(&path)
499                    .map_err(|source| StoreError::Io { path, source })?;
500                deleted = true;
501            }
502            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
503            Err(source) => return Err(StoreError::Io { path, source }),
504        }
505    }
506    if deleted && let SyncRequirement::Sync(kind) = policy.directory_sync() {
507        vfs.sync(directory, kind).map_err(|source| StoreError::Io {
508            path: directory.to_path_buf(),
509            source,
510        })?;
511    }
512    Ok(())
513}
514
515#[cfg(test)]
516#[allow(clippy::expect_used, clippy::panic)]
517mod wal_truncation_crash_tests {
518    use std::path::Path;
519    use std::sync::Arc;
520
521    use crate::ingest::StoreWal;
522    use crate::lifecycle::durability::{CommitTier, DurabilityMode, DurabilityPolicy};
523    use crate::lifecycle::stats::Accounting;
524    use crate::vfs::crash::{CrashVfs, MemoryVfs};
525    use crate::wal::{LogSeq, WalReader, encode_wal_image};
526
527    /// Every power-cut image of the rotation, including torn, reordered and
528    /// unsynced-rename states, must still open with the absorbed boundary
529    /// intact and hand out sequence 4 next. An empty or headerless wal.ze
530    /// would restart numbering at 1, below the manifest, and lose writes.
531    #[test]
532    fn every_crash_state_of_the_wal_truncation_resumes_after_the_boundary() {
533        let directory = Path::new("/wal-truncation");
534        let wal = directory.join("wal.ze");
535        let records = (0..3_u8)
536            .map(|index| (99, vec![index; 8]))
537            .collect::<Vec<_>>();
538        let memory = MemoryVfs::new();
539        memory
540            .insert(
541                wal.clone(),
542                encode_wal_image(LogSeq::new(1), &records).expect("old WAL"),
543            )
544            .expect("seed old WAL");
545        let crash = CrashVfs::new(memory).expect("crash recorder");
546        let policy = DurabilityPolicy::new(DurabilityMode::Durable, CommitTier::Ordered)
547            .expect("ordered policy");
548        let accounting = Arc::new(Accounting::new(u64::MAX, u64::MAX));
549        let recovered = WalReader::open(&crash, &wal)
550            .expect("open old WAL")
551            .into_clean()
552            .expect("clean old WAL");
553        let mut writer = StoreWal::resume(&crash, &wal, recovered, policy, 3, &accounting)
554            .expect("resume old WAL");
555
556        writer
557            .truncate_absorbed(&crash, directory, policy, LogSeq::new(3))
558            .expect("truncate");
559
560        let states = crash.crash_states().expect("crash states");
561        assert!(!states.was_capped());
562        assert!(states.len() > 10, "only {} crash states", states.len());
563        for state in states.iter() {
564            let recovered = WalReader::open(state.vfs(), &wal)
565                .expect("read WAL")
566                .into_clean()
567                .unwrap_or_else(|error| panic!("{:?}: WAL unreadable: {error}", state.kind()));
568            let sequences = recovered
569                .records()
570                .iter()
571                .map(|record| record.seq.get())
572                .collect::<Vec<_>>();
573            assert!(
574                sequences == [1, 2, 3] || sequences.is_empty(),
575                "{:?}: WAL holds {sequences:?}",
576                state.kind()
577            );
578            let mut resumed =
579                StoreWal::resume(state.vfs(), &wal, recovered, policy, 3, &accounting)
580                    .unwrap_or_else(|error| panic!("{:?}: resume: {error}", state.kind()));
581            assert_eq!(
582                resumed.commit(99, b"next").expect("commit after crash"),
583                LogSeq::new(4),
584                "{:?}",
585                state.kind()
586            );
587            #[cfg(feature = "graph-cypher")]
588            {
589                use crate::ingest::wal_payload::{GRAPH_COMMIT_V1, encode_graph_commit};
590                let graph = encode_graph_commit(
591                    include_bytes!("../../tests/fixtures/graph-wal/complete-v1.bin")
592                        .get(64..2105)
593                        .expect("graph envelope"),
594                )
595                .expect("graph payload");
596                assert_eq!(
597                    resumed
598                        .commit(GRAPH_COMMIT_V1, &graph)
599                        .expect("graph after fold"),
600                    LogSeq::new(5)
601                );
602                let before = crate::vfs::Vfs::read(state.vfs(), &wal).expect("WAL before refusal");
603                assert!(
604                    resumed
605                        .truncate_absorbed(state.vfs(), directory, policy, LogSeq::new(3))
606                        .is_err()
607                );
608                assert_eq!(
609                    crate::vfs::Vfs::read(state.vfs(), &wal).expect("WAL after refusal"),
610                    before
611                );
612                assert_eq!(
613                    resumed
614                        .commit(99, b"after graph")
615                        .expect("definite refusal keeps writer usable"),
616                    LogSeq::new(6)
617                );
618                let recovered = WalReader::open(state.vfs(), &wal)
619                    .expect("graph WAL")
620                    .into_clean()
621                    .expect("clean graph WAL");
622                let record = recovered
623                    .records()
624                    .iter()
625                    .find(|record| record.seq.get() == 5)
626                    .expect("unabsorbed graph retained");
627                assert_eq!(record.op, GRAPH_COMMIT_V1);
628                assert_eq!(record.payload().expect("graph payload"), graph);
629                let mut next =
630                    StoreWal::resume(state.vfs(), &wal, recovered, policy, 3, &accounting)
631                        .expect("resume graph tail");
632                assert_eq!(
633                    next.commit(99, b"next reopen").expect("write after reopen"),
634                    LogSeq::new(7)
635                );
636            }
637        }
638    }
639}
640
641#[cfg(test)]
642#[path = "seal_id_tests.rs"]
643mod seal_id_tests;