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

zeppelin_embed/ingest/
seal.rs
