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 publication = writer.manifest_publication()?;
159        let absorbed_through = writer.durable_end();
160        let mut active = self
161            .active
162            .lock()
163            .map_err(|_| StoreError::Synchronization {
164                component: "active segment",
165            })?;
166        let current = active.as_ref().ok_or(StoreError::Closed)?;
167        if current.segment.is_empty() {
168            // A crash-restart after a durably committed seal leaves a correctly
169            // empty active segment; retrying seal is a replay of that commit.
170            return Ok(current.generation);
171        }
172        let generation = current
173            .generation
174            .checked_add(1)
175            .ok_or(StoreError::GenerationOverflow)?;
176        #[cfg_attr(not(feature = "graph-cypher"), allow(unused_mut))]
177        let mut manifest = load_current_manifest(
178            vfs,
179            &self.directory,
180            absorbed_through,
181            current.generation,
182            &self.schema,
183        )?;
184        #[cfg(feature = "graph-cypher")]
185        let graph_fold = if manifest.graph.is_some() {
186            Some(
187                self.prepare_native_graph_seal(&mut manifest, generation, absorbed_through)
188                    .map_err(seal_graph_error)?,
189            )
190        } else {
191            None
192        };
193        let columns = active_columns(&manifest.schema, &current.segment, cancel)?;
194        let alive = current.segment.alive()?;
195        let clustering_key_range = clustering_key_range(current.segment.timestamps(), &alive)?;
196        let dims = current
197            .segment
198            .dims()
199            .ok_or(StoreError::ActiveRowOverflow)?;
200        let dims = u32::try_from(dims).map_err(|_| StoreError::ActiveRowOverflow)?;
201        let id = seal_id(generation, absorbed_through);
202        #[cfg(any(test, feature = "test-seams"))]
203        let int8_storage = self.test_seal_storage(current, dims)?;
204        #[cfg(any(test, feature = "test-seams"))]
205        let (scheme, codes, factors) = match int8_storage.as_ref() {
206            Some((codes, factors)) => (
207                u16::from(crate::quant::QuantScheme::Int8.id()),
208                codes.as_slice(),
209                SegmentFactors::Int8(factors.as_slice()),
210            ),
211            None => (
212                u16::from(crate::quant::QuantScheme::Bit4.id()),
213                current.segment.codes(),
214                SegmentFactors::Bit4(current.segment.factors()),
215            ),
216        };
217        #[cfg(not(any(test, feature = "test-seams")))]
218        let (scheme, codes, factors) = (
219            u16::from(crate::quant::QuantScheme::Bit4.id()),
220            current.segment.codes(),
221            SegmentFactors::Bit4(current.segment.factors()),
222        );
223        let build = SegmentBuild {
224            id,
225            scheme,
226            dims,
227            codes,
228            factors,
229            rescore: current.segment.vectors(),
230            columns: &columns,
231            alive: &alive,
232        };
233        let documents = SegmentDocumentVersions {
234            doc_ids: current.segment.doc_ids(),
235            revisions: current.segment.revisions(),
236        };
237        let payloads = build_seal_payloads(&current.segment)?;
238        publication = publication.arm();
239        let written = write_segment_with_documents_payloads(
240            vfs,
241            &self.directory,
242            build,
243            SegmentPayloads {
244                documents,
245                metadata: payloads.metadata,
246                text: payloads.text,
247                postings: payloads
248                    .postings
249                    .as_deref()
250                    .map(|bytes| SegmentPostings { bytes }),
251            },
252            self.durability_policy,
253        );
254        let mut meta = match written {
255            Ok(meta) => meta,
256            Err(error) => {
257                cleanup_uncommitted_segment(vfs, &self.directory, id, self.durability_policy)?;
258                return Err(StoreError::Segment(error));
259            }
260        };
261        meta.clustering_key_range = clustering_key_range;
262        meta.epoch_id = manifest.epoch_alias.map(|identity| identity.embedding);
263        #[cfg(any(test, feature = "test-seams"))]
264        let forced_late_cancellation =
265            match self.ingest_retention_fault_controller.as_ref() {
266                Some(controller) => controller.seal_cancellation_plan().map_err(|_| {
267                    StoreError::Synchronization {
268                        component: "ingest-retention controller",
269                    }
270                })?,
271                None => None,
272            };
273        #[cfg(not(any(test, feature = "test-seams")))]
274        let forced_late_cancellation: Option<u64> = None;
275        if forced_late_cancellation.is_some() || check_cancelled(cancel).is_err() {
276            self.abort_seal(vfs, id, forced_late_cancellation, current, absorbed_through)?;
277            // Successful cleanup durably restores the pre-seal state; no
278            // manifest or in-memory publication changed. Cleanup I/O errors
279            // above still drop the armed guard and fence the shared writer.
280            publication.complete();
281            return Err(StoreError::SealCancelled);
282        }
283        let mut segments = manifest.segments;
284        segments.push(meta);
285        let epochs = self.epoch_registry(&manifest.epochs);
286        publication
287            .commit_manifest(
288                vfs,
289                &self.directory,
290                &Manifest {
291                    #[cfg(feature = "graph-cypher")]
292                    graph: manifest.graph,
293                    generation,
294                    log_seq: absorbed_through,
295                    segments,
296                    epochs,
297                    epoch_alias: manifest.epoch_alias,
298                    schema: manifest.schema,
299                },
300                self.durability_policy,
301            )
302            .map_err(StoreError::Manifest)?;
303        let remapped = crate::lifecycle::PublishedSnapshot::load_on_vfs(
304            &self.directory,
305            &self.accounting,
306            vfs,
307        )?;
308        let mut published = self
309            .snapshot
310            .write()
311            .map_err(|_| StoreError::Synchronization {
312                component: "published snapshot",
313            })?;
314        let previous = published.replace(Arc::new(remapped));
315        drop(published);
316        drop(previous);
317        *active = Some(ActiveState::empty(generation));
318        #[cfg(feature = "graph-cypher")]
319        if let Some(fold) = graph_fold {
320            fold.publish(self).map_err(seal_graph_error)?;
321        }
322        #[cfg(feature = "graph-cypher")]
323        self.publish_native_documents(active.as_ref().ok_or(StoreError::Closed)?)?;
324        publication.complete();
325        writer.retire_visible_through(LogSeq::new(absorbed_through))?;
326        // Both watermarks above are at the durable end, so the log can
327        // shrink to a header; without this wal.ze grows for the store's life.
328        writer.truncate_absorbed(
329            vfs,
330            &self.directory,
331            self.durability_policy,
332            LogSeq::new(absorbed_through),
333        )?;
334        drop(active);
335        drop(wal);
336        drop(writer_lock);
337        drop(state);
338        Ok(generation)
339    }
340}
341
342#[cfg(feature = "graph-cypher")]
343fn seal_graph_error(error: crate::lifecycle::native_graph::NativeGraphError) -> StoreError {
344    match error {
345        crate::lifecycle::native_graph::NativeGraphError::Store(error) => error,
346        error => StoreError::Manifest(crate::manifest::ManifestError::Decode(format!(
347            "seal graph fold: {error}",
348        ))),
349    }
350}
351
352fn build_seal_payloads(
353    segment: &super::ActiveSegment,
354) -> Result<SealPayloadBuffers<'_>, StoreError> {
355    let postings = if segment.has_text() {
356        let sealed = SealedSegment::seal(segment.lexical())
357            .map_err(SealedSegmentError::from)
358            .map_err(crate::segment::SegmentError::from)
359            .map_err(StoreError::Segment)?;
360        Some(
361            sealed
362                .encode_region()
363                .map_err(crate::segment::SegmentError::from)
364                .map_err(StoreError::Segment)?,
365        )
366    } else {
367        None
368    };
369    let metadata = (!segment.metadata_bytes().is_empty()).then_some(SegmentStoredMetadata {
370        end_offsets: segment.metadata_end_offsets(),
371        bytes: segment.metadata_bytes(),
372    });
373    let text = segment.has_text().then_some(SegmentStoredText {
374        present: segment.text_present(),
375        end_offsets: segment.text_end_offsets(),
376        bytes: segment.text_bytes(),
377    });
378    Ok(SealPayloadBuffers {
379        postings,
380        metadata,
381        text,
382    })
383}
384
385pub(crate) fn load_current_manifest(
386    vfs: &dyn Vfs,
387    directory: &Path,
388    durable_end: u64,
389    generation: u64,
390    schema: &Schema,
391) -> Result<Manifest, StoreError> {
392    let path = directory.join(MANIFEST_FILE);
393    match vfs.open(&path) {
394        Ok(_) => load_manifest(vfs, &path, durable_end).map_err(StoreError::Manifest),
395        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Manifest {
396            #[cfg(feature = "graph-cypher")]
397            graph: None,
398            generation,
399            log_seq: 0,
400            segments: Vec::new(),
401            epochs: Vec::new(),
402            epoch_alias: None,
403            schema: schema.clone(),
404        }),
405        Err(source) => Err(StoreError::Io { path, source }),
406    }
407}
408
409fn active_columns(
410    schema: &Schema,
411    active: &super::ActiveSegment,
412    cancel: Option<&CancelToken>,
413) -> Result<ColumnStore, StoreError> {
414    let mut builder = ColumnStoreBuilder::new(schema.clone());
415    for (row, timestamp) in active.timestamps().iter().copied().enumerate() {
416        if row.is_multiple_of(64) {
417            check_cancelled(cancel)?;
418        }
419        let values = active.column_values(row)?;
420        builder
421            .push_row(timestamp, &super::column_inputs(&values))
422            .map_err(|error| {
423                StoreError::Segment(crate::segment::SegmentError::Columns(error.to_string()))
424            })?;
425    }
426    builder.finish().map_err(|error| {
427        StoreError::Segment(crate::segment::SegmentError::Columns(error.to_string()))
428    })
429}
430
431fn clustering_key_range(
432    timestamps: &[i64],
433    alive: &crate::meta::AliveSet,
434) -> Result<ClusteringKeyRange, StoreError> {
435    let mut bounds: Option<(i64, i64)> = None;
436    for row in alive.iter_alive() {
437        let index = usize::try_from(row).map_err(|_| StoreError::ActiveRowOverflow)?;
438        let timestamp = timestamps
439            .get(index)
440            .copied()
441            .ok_or(StoreError::ActiveRowOverflow)?;
442        bounds = Some(match bounds {
443            None => (timestamp, timestamp),
444            Some((minimum, maximum)) => (minimum.min(timestamp), maximum.max(timestamp)),
445        });
446    }
447    Ok(match bounds {
448        Some((min_ts, max_ts)) => ClusteringKeyRange::Bounded { min_ts, max_ts },
449        // An all-tombstoned segment has the empty key set. It is represented
450        // explicitly rather than overloading a numeric sentinel or Unstamped.
451        None => ClusteringKeyRange::Empty,
452    })
453}
454
455fn check_cancelled(cancel: Option<&CancelToken>) -> Result<(), StoreError> {
456    if cancel.is_some_and(CancelToken::is_cancelled) {
457        Err(StoreError::SealCancelled)
458    } else {
459        Ok(())
460    }
461}
462
463fn seal_id(generation: u64, absorbed_through: u64) -> SegmentId {
464    let mut bytes = [0_u8; 16];
465    if let Some(prefix) = bytes.get_mut(..8) {
466        prefix.copy_from_slice(&generation.to_be_bytes());
467    }
468    if let Some(suffix) = bytes.get_mut(8..) {
469        suffix.copy_from_slice(&absorbed_through.to_be_bytes());
470    }
471    SegmentId::from_bytes(bytes)
472}
473
474fn cleanup_uncommitted_segment(
475    vfs: &dyn Vfs,
476    directory: &Path,
477    id: SegmentId,
478    policy: DurabilityPolicy,
479) -> Result<(), StoreError> {
480    let final_path = directory.join(id.file_name());
481    let temporary_path = directory.join(format!(".{}.tmp", id.file_name()));
482    let mut deleted = false;
483    for path in [temporary_path, final_path] {
484        match vfs.open(&path) {
485            Ok(_) => {
486                vfs.delete(&path)
487                    .map_err(|source| StoreError::Io { path, source })?;
488                deleted = true;
489            }
490            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
491            Err(source) => return Err(StoreError::Io { path, source }),
492        }
493    }
494    if deleted && let SyncRequirement::Sync(kind) = policy.directory_sync() {
495        vfs.sync(directory, kind).map_err(|source| StoreError::Io {
496            path: directory.to_path_buf(),
497            source,
498        })?;
499    }
500    Ok(())
501}
502
503#[cfg(test)]
504#[allow(clippy::expect_used, clippy::panic)]
505mod wal_truncation_crash_tests {
506    use std::path::Path;
507    use std::sync::Arc;
508
509    use crate::ingest::StoreWal;
510    use crate::lifecycle::durability::{CommitTier, DurabilityMode, DurabilityPolicy};
511    use crate::lifecycle::stats::Accounting;
512    use crate::vfs::crash::{CrashVfs, MemoryVfs};
513    use crate::wal::{LogSeq, WalReader, encode_wal_image};
514
515    /// Every power-cut image of the rotation, including torn, reordered and
516    /// unsynced-rename states, must still open with the absorbed boundary
517    /// intact and hand out sequence 4 next. An empty or headerless wal.ze
518    /// would restart numbering at 1, below the manifest, and lose writes.
519    #[test]
520    fn every_crash_state_of_the_wal_truncation_resumes_after_the_boundary() {
521        let directory = Path::new("/wal-truncation");
522        let wal = directory.join("wal.ze");
523        let records = (0..3_u8)
524            .map(|index| (99, vec![index; 8]))
525            .collect::<Vec<_>>();
526        let memory = MemoryVfs::new();
527        memory
528            .insert(
529                wal.clone(),
530                encode_wal_image(LogSeq::new(1), &records).expect("old WAL"),
531            )
532            .expect("seed old WAL");
533        let crash = CrashVfs::new(memory).expect("crash recorder");
534        let policy = DurabilityPolicy::new(DurabilityMode::Durable, CommitTier::Ordered)
535            .expect("ordered policy");
536        let accounting = Arc::new(Accounting::new(u64::MAX, u64::MAX));
537        let recovered = WalReader::open(&crash, &wal)
538            .expect("open old WAL")
539            .into_clean()
540            .expect("clean old WAL");
541        let mut writer = StoreWal::resume(&crash, &wal, recovered, policy, 3, &accounting)
542            .expect("resume old WAL");
543
544        writer
545            .truncate_absorbed(&crash, directory, policy, LogSeq::new(3))
546            .expect("truncate");
547
548        let states = crash.crash_states().expect("crash states");
549        assert!(!states.was_capped());
550        assert!(states.len() > 10, "only {} crash states", states.len());
551        for state in states.iter() {
552            let recovered = WalReader::open(state.vfs(), &wal)
553                .expect("read WAL")
554                .into_clean()
555                .unwrap_or_else(|error| panic!("{:?}: WAL unreadable: {error}", state.kind()));
556            let sequences = recovered
557                .records()
558                .iter()
559                .map(|record| record.seq.get())
560                .collect::<Vec<_>>();
561            assert!(
562                sequences == [1, 2, 3] || sequences.is_empty(),
563                "{:?}: WAL holds {sequences:?}",
564                state.kind()
565            );
566            let mut resumed =
567                StoreWal::resume(state.vfs(), &wal, recovered, policy, 3, &accounting)
568                    .unwrap_or_else(|error| panic!("{:?}: resume: {error}", state.kind()));
569            assert_eq!(
570                resumed.commit(99, b"next").expect("commit after crash"),
571                LogSeq::new(4),
572                "{:?}",
573                state.kind()
574            );
575            #[cfg(feature = "graph-cypher")]
576            {
577                use crate::ingest::wal_payload::{GRAPH_COMMIT_V1, encode_graph_commit};
578                let graph = encode_graph_commit(
579                    include_bytes!("../../tests/fixtures/graph-wal/complete-v1.bin")
580                        .get(64..2105)
581                        .expect("graph envelope"),
582                )
583                .expect("graph payload");
584                assert_eq!(
585                    resumed
586                        .commit(GRAPH_COMMIT_V1, &graph)
587                        .expect("graph after fold"),
588                    LogSeq::new(5)
589                );
590                let before = crate::vfs::Vfs::read(state.vfs(), &wal).expect("WAL before refusal");
591                assert!(
592                    resumed
593                        .truncate_absorbed(state.vfs(), directory, policy, LogSeq::new(3))
594                        .is_err()
595                );
596                assert_eq!(
597                    crate::vfs::Vfs::read(state.vfs(), &wal).expect("WAL after refusal"),
598                    before
599                );
600                assert_eq!(
601                    resumed
602                        .commit(99, b"after graph")
603                        .expect("definite refusal keeps writer usable"),
604                    LogSeq::new(6)
605                );
606                let recovered = WalReader::open(state.vfs(), &wal)
607                    .expect("graph WAL")
608                    .into_clean()
609                    .expect("clean graph WAL");
610                let record = recovered
611                    .records()
612                    .iter()
613                    .find(|record| record.seq.get() == 5)
614                    .expect("unabsorbed graph retained");
615                assert_eq!(record.op, GRAPH_COMMIT_V1);
616                assert_eq!(record.payload().expect("graph payload"), graph);
617                let mut next =
618                    StoreWal::resume(state.vfs(), &wal, recovered, policy, 3, &accounting)
619                        .expect("resume graph tail");
620                assert_eq!(
621                    next.commit(99, b"next reopen").expect("write after reopen"),
622                    LogSeq::new(7)
623                );
624            }
625        }
626    }
627}
628
629#[cfg(test)]
630#[path = "seal_id_tests.rs"]
631mod seal_id_tests;