zeppelin_embed/ingest/
merge.rs

1//! Host-scheduled, bounded copy-on-write merging of small scan segments.
2
3use std::sync::Arc;
4
5use crate::graph::consolidate::{ConsolidateError, merge_segments, merged_clustering};
6use crate::lifecycle::durability::SyncRequirement;
7use crate::lifecycle::{CancelToken, PublishedSnapshot, Store, StoreError, StoreState};
8use crate::manifest::io::{MANIFEST_FILE, load_manifest};
9use crate::segment::layout::RegionKind;
10use crate::segment::reader::SegmentReader;
11use crate::segment::{SegmentError, SegmentId};
12use crate::vfs::Vfs;
13
14// Admission bounds the existing gather/re-encode implementation independently
15// of corpus size. Decoded structures require more memory than the input bytes.
16const MAX_INPUT_BYTES: u64 = 8 * 1024 * 1024;
17const MAX_INPUTS: usize = 16;
18const MERGE_ID_SEED: u64 = 0x5a45_0234_6d65_7267;
19
20impl Store {
21    /// Merges small compatible sealed scan segments during host-selected idle time.
22    ///
23    /// Each atomic batch admits at most 16 segments and 8 MiB of input files.
24    /// Decoding and rebuilding require additional bounded working memory. Large
25    /// segments and graph segments stay separate. Repeats while a batch fits;
26    /// this synchronous call blocks writers and does not seal active writes.
27    /// Open snapshot views defer input unlinking until later orphan cleanup.
28    /// Returns the last publication generation, unchanged if no merge is due.
29    pub fn merge_sealed(&self) -> Result<u64, StoreError> {
30        self.merge_sealed_with_cancel(&CancelToken::new())
31    }
32
33    /// Cancels between batches or before publication; completed batches remain.
34    pub fn merge_sealed_with_cancel(&self, cancel: &CancelToken) -> Result<u64, StoreError> {
35        self.merge_sealed_with_cancel_on_vfs(cancel, self.vfs.as_ref())
36    }
37
38    /// Filesystem seam for deterministic interrupted-publication tests.
39    #[doc(hidden)]
40    pub fn merge_sealed_with_cancel_on_vfs(
41        &self,
42        cancel: &CancelToken,
43        vfs: &dyn Vfs,
44    ) -> Result<u64, StoreError> {
45        let state = self
46            .state
47            .lock()
48            .map_err(|_| StoreError::Synchronization { component: "state" })?;
49        match *state {
50            StoreState::Open => {}
51            StoreState::Closing => return Err(StoreError::Closing),
52            StoreState::Closed => return Err(StoreError::Closed),
53        }
54        let writer = self
55            .writer_lock
56            .lock()
57            .map_err(|_| StoreError::Synchronization {
58                component: "writer lock",
59            })?;
60        if writer.is_none() {
61            return Err(StoreError::ReadOnly);
62        }
63        let wal = self
64            .wal_writer
65            .lock()
66            .map_err(|_| StoreError::Synchronization {
67                component: "WAL writer",
68            })?;
69        wal.as_ref()
70            .ok_or(StoreError::ReadOnly)?
71            .manifest_publication()?
72            .complete();
73        let durable_end = wal.as_ref().ok_or(StoreError::ReadOnly)?.durable_end();
74        let mut active = self
75            .active
76            .lock()
77            .map_err(|_| StoreError::Synchronization {
78                component: "active segment",
79            })?;
80        let current = active.as_mut().ok_or(StoreError::Closed)?;
81        let path = self.directory.join(MANIFEST_FILE);
82        check_cancel(cancel)?;
83        match vfs.open(&path) {
84            Ok(_) => {}
85            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
86                return Ok(current.generation);
87            }
88            Err(source) => return Err(StoreError::Io { path, source }),
89        }
90        let mut manifest = load_manifest(vfs, &path, durable_end).map_err(StoreError::Manifest)?;
91        loop {
92            check_cancel(cancel)?;
93            let inputs = select_inputs(vfs, &self.directory, &manifest.segments)?;
94            if inputs.len() < 2 {
95                return Ok(current.generation.max(manifest.generation));
96            }
97            let mut publication = wal
98                .as_ref()
99                .ok_or(StoreError::ReadOnly)?
100                .manifest_publication()?;
101            let generation = current
102                .generation
103                .max(manifest.generation)
104                .checked_add(1)
105                .ok_or(StoreError::GenerationOverflow)?;
106            let id = merge_id(generation, &inputs);
107            let output_path = self.directory.join(id.file_name());
108            // Never let the segment writer's rename overwrite a live artifact
109            // or a leftover output from an interrupted invocation.
110            match vfs.open(&output_path) {
111                Ok(_) => {
112                    return Err(StoreError::Segment(SegmentError::Geometry(format!(
113                        "merge output {id} already exists; reopen before retrying"
114                    ))));
115                }
116                Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
117                Err(source) => {
118                    return Err(StoreError::Io {
119                        path: output_path,
120                        source,
121                    });
122                }
123            }
124            let readers = inputs.iter().collect::<Vec<_>>();
125            let merged = merge_segments(
126                vfs,
127                &self.directory,
128                &readers,
129                id,
130                &self.tokenizer,
131                self.durability_policy,
132            )
133            .map_err(merge_error)?;
134            let replacement = if merged.is_some() {
135                let output =
136                    SegmentReader::open(vfs, &output_path, id).map_err(StoreError::Segment)?;
137                let mut meta = output.meta().clone();
138                meta.clustering_key_range = merged_clustering(
139                    &output.columns().map_err(StoreError::Segment)?,
140                    &output.alive().map_err(StoreError::Segment)?,
141                )
142                .map_err(merge_error)?;
143                meta.epoch_id = inputs
144                    .first()
145                    .and_then(|input| {
146                        manifest
147                            .segments
148                            .iter()
149                            .find(|meta| meta.id == input.meta().id)
150                    })
151                    .and_then(|meta| meta.epoch_id);
152                Some(meta)
153            } else {
154                None // A batch containing only tombstones can be removed.
155            };
156            if cancel.is_cancelled() {
157                if replacement.is_some() {
158                    vfs.delete(&output_path).map_err(|source| StoreError::Io {
159                        path: output_path,
160                        source,
161                    })?;
162                }
163                return Err(StoreError::SealCancelled);
164            }
165            manifest
166                .segments
167                .retain(|meta| !inputs.iter().any(|input| input.meta().id == meta.id));
168            manifest.segments.extend(replacement);
169            manifest.generation = generation;
170            #[cfg(feature = "graph-cypher")]
171            manifest
172                .record_generation_bump(durable_end)
173                .map_err(StoreError::Manifest)?;
174            // Preserve log_seq: a merge absorbs no active or WAL records.
175            publication
176                .commit_manifest(vfs, &self.directory, &manifest, self.durability_policy)
177                .map_err(StoreError::Manifest)?;
178            let remapped = PublishedSnapshot::load_on_vfs(&self.directory, &self.accounting, vfs)?;
179            let mut published = self
180                .snapshot
181                .write()
182                .map_err(|_| StoreError::Synchronization {
183                    component: "published snapshot",
184                })?;
185            let previous = published.replace(Arc::new(remapped));
186            current.generation = generation;
187            publication.complete();
188            drop(published);
189            drop(previous);
190            let old_paths = inputs
191                .iter()
192                .map(|input| self.directory.join(input.meta().id.file_name()))
193                .collect::<Vec<_>>();
194            drop(inputs);
195            // An open view retains the input paths as well as their mappings.
196            // After views close, writable-open orphan cleanup reclaims them.
197            if !self.has_snapshot_views() {
198                for path in old_paths {
199                    vfs.delete(&path)
200                        .map_err(|source| StoreError::Io { path, source })?;
201                }
202            }
203            if let SyncRequirement::Sync(kind) = self.durability_policy.directory_sync() {
204                vfs.sync(&self.directory, kind)
205                    .map_err(|source| StoreError::Io {
206                        path: self.directory.clone(),
207                        source,
208                    })?;
209            }
210        }
211    }
212}
213
214fn select_inputs(
215    vfs: &dyn Vfs,
216    directory: &std::path::Path,
217    segments: &[crate::segment::SegmentMeta],
218) -> Result<Vec<SegmentReader>, StoreError> {
219    let mut candidates = segments
220        .iter()
221        .filter(|meta| meta.scheme == 4 && meta.file_size <= MAX_INPUT_BYTES)
222        .collect::<Vec<_>>();
223    candidates.sort_by_key(|meta| (meta.file_size, meta.id));
224    // Group by immutable geometry/epoch, so one incompatible small segment
225    // does not prevent another compatible group from being merged.
226    let mut visited = Vec::new();
227    for first in &candidates {
228        let group = (first.dims, first.epoch_id);
229        if visited.contains(&group) {
230            continue;
231        }
232        visited.push(group);
233        let mut inputs = Vec::new();
234        let mut bytes = 0_u64;
235        for meta in candidates
236            .iter()
237            .filter(|meta| meta.dims == first.dims && meta.epoch_id == first.epoch_id)
238        {
239            if inputs.len() == MAX_INPUTS || meta.file_size > MAX_INPUT_BYTES - bytes {
240                break;
241            }
242            let reader = SegmentReader::open(vfs, &directory.join(meta.id.file_name()), meta.id)
243                .map_err(StoreError::Segment)?;
244            if reader
245                .directory()
246                .iter()
247                .any(|entry| entry.kind == RegionKind::GraphNodeBlocks.id())
248            {
249                continue;
250            }
251            bytes += meta.file_size;
252            inputs.push(reader);
253        }
254        if inputs.len() >= 2 {
255            inputs.sort_by_key(|input| input.meta().id);
256            return Ok(inputs);
257        }
258    }
259    Ok(Vec::new())
260}
261
262fn merge_id(generation: u64, inputs: &[SegmentReader]) -> SegmentId {
263    let mut source = Vec::with_capacity(inputs.len() * 16);
264    for input in inputs {
265        source.extend_from_slice(input.meta().id.as_bytes());
266    }
267    let hash = xxhash_rust::xxh3::xxh3_64_with_seed(&source, MERGE_ID_SEED);
268    let mut bytes = [0_u8; 16];
269    for (target, value) in bytes.iter_mut().take(8).zip(generation.to_be_bytes()) {
270        *target = value;
271    }
272    for (target, value) in bytes.iter_mut().skip(8).zip(hash.to_be_bytes()) {
273        *target = value;
274    }
275    SegmentId::from_bytes(bytes)
276}
277
278fn check_cancel(cancel: &CancelToken) -> Result<(), StoreError> {
279    if cancel.is_cancelled() {
280        Err(StoreError::SealCancelled)
281    } else {
282        Ok(())
283    }
284}
285
286fn merge_error(error: ConsolidateError) -> StoreError {
287    match error {
288        ConsolidateError::Store(error) => error,
289        ConsolidateError::Segment(error) => StoreError::Segment(error),
290        error => StoreError::Segment(SegmentError::Geometry(error.to_string())),
291    }
292}
zeppelin_embed/ingest/merge.rs

zeppelin_embed/ingest/
merge.rs
