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, commit_manifest, 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        let durable_end = wal.as_ref().ok_or(StoreError::ReadOnly)?.durable_end();
70        let mut active = self
71            .active
72            .lock()
73            .map_err(|_| StoreError::Synchronization {
74                component: "active segment",
75            })?;
76        let current = active.as_mut().ok_or(StoreError::Closed)?;
77        let path = self.directory.join(MANIFEST_FILE);
78        check_cancel(cancel)?;
79        match vfs.open(&path) {
80            Ok(_) => {}
81            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
82                return Ok(current.generation);
83            }
84            Err(source) => return Err(StoreError::Io { path, source }),
85        }
86        let mut manifest = load_manifest(vfs, &path, durable_end).map_err(StoreError::Manifest)?;
87        loop {
88            check_cancel(cancel)?;
89            let inputs = select_inputs(vfs, &self.directory, &manifest.segments)?;
90            if inputs.len() < 2 {
91                return Ok(current.generation.max(manifest.generation));
92            }
93            let generation = current
94                .generation
95                .max(manifest.generation)
96                .checked_add(1)
97                .ok_or(StoreError::GenerationOverflow)?;
98            let id = merge_id(generation, &inputs);
99            let output_path = self.directory.join(id.file_name());
100            // Never let the segment writer's rename overwrite a live artifact
101            // or a leftover output from an interrupted invocation.
102            match vfs.open(&output_path) {
103                Ok(_) => {
104                    return Err(StoreError::Segment(SegmentError::Geometry(format!(
105                        "merge output {id} already exists; reopen before retrying"
106                    ))));
107                }
108                Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
109                Err(source) => {
110                    return Err(StoreError::Io {
111                        path: output_path,
112                        source,
113                    });
114                }
115            }
116            let readers = inputs.iter().collect::<Vec<_>>();
117            let merged = merge_segments(
118                vfs,
119                &self.directory,
120                &readers,
121                id,
122                &self.tokenizer,
123                self.durability_policy,
124            )
125            .map_err(merge_error)?;
126            let replacement = if merged.is_some() {
127                let output =
128                    SegmentReader::open(vfs, &output_path, id).map_err(StoreError::Segment)?;
129                let mut meta = output.meta().clone();
130                meta.clustering_key_range = merged_clustering(
131                    &output.columns().map_err(StoreError::Segment)?,
132                    &output.alive().map_err(StoreError::Segment)?,
133                )
134                .map_err(merge_error)?;
135                meta.epoch_id = inputs
136                    .first()
137                    .and_then(|input| {
138                        manifest
139                            .segments
140                            .iter()
141                            .find(|meta| meta.id == input.meta().id)
142                    })
143                    .and_then(|meta| meta.epoch_id);
144                Some(meta)
145            } else {
146                None // A batch containing only tombstones can be removed.
147            };
148            if cancel.is_cancelled() {
149                if replacement.is_some() {
150                    vfs.delete(&output_path).map_err(|source| StoreError::Io {
151                        path: output_path,
152                        source,
153                    })?;
154                }
155                return Err(StoreError::SealCancelled);
156            }
157            manifest
158                .segments
159                .retain(|meta| !inputs.iter().any(|input| input.meta().id == meta.id));
160            manifest.segments.extend(replacement);
161            manifest.generation = generation;
162            // Preserve log_seq: a merge absorbs no active or WAL records.
163            commit_manifest(vfs, &self.directory, &manifest, self.durability_policy)
164                .map_err(StoreError::Manifest)?;
165            let remapped = PublishedSnapshot::load_on_vfs(&self.directory, &self.accounting, vfs)?;
166            let mut published = self
167                .snapshot
168                .write()
169                .map_err(|_| StoreError::Synchronization {
170                    component: "published snapshot",
171                })?;
172            let previous = published.replace(Arc::new(remapped));
173            current.generation = generation;
174            drop(published);
175            drop(previous);
176            let old_paths = inputs
177                .iter()
178                .map(|input| self.directory.join(input.meta().id.file_name()))
179                .collect::<Vec<_>>();
180            drop(inputs);
181            // An open view retains the input paths as well as their mappings.
182            // After views close, writable-open orphan cleanup reclaims them.
183            if !self.has_snapshot_views() {
184                for path in old_paths {
185                    vfs.delete(&path)
186                        .map_err(|source| StoreError::Io { path, source })?;
187                }
188            }
189            if let SyncRequirement::Sync(kind) = self.durability_policy.directory_sync() {
190                vfs.sync(&self.directory, kind)
191                    .map_err(|source| StoreError::Io {
192                        path: self.directory.clone(),
193                        source,
194                    })?;
195            }
196        }
197    }
198}
199
200fn select_inputs(
201    vfs: &dyn Vfs,
202    directory: &std::path::Path,
203    segments: &[crate::segment::SegmentMeta],
204) -> Result<Vec<SegmentReader>, StoreError> {
205    let mut candidates = segments
206        .iter()
207        .filter(|meta| meta.scheme == 4 && meta.file_size <= MAX_INPUT_BYTES)
208        .collect::<Vec<_>>();
209    candidates.sort_by_key(|meta| (meta.file_size, meta.id));
210    // Group by immutable geometry/epoch, so one incompatible small segment
211    // does not prevent another compatible group from being merged.
212    let mut visited = Vec::new();
213    for first in &candidates {
214        let group = (first.dims, first.epoch_id);
215        if visited.contains(&group) {
216            continue;
217        }
218        visited.push(group);
219        let mut inputs = Vec::new();
220        let mut bytes = 0_u64;
221        for meta in candidates
222            .iter()
223            .filter(|meta| meta.dims == first.dims && meta.epoch_id == first.epoch_id)
224        {
225            if inputs.len() == MAX_INPUTS || meta.file_size > MAX_INPUT_BYTES - bytes {
226                break;
227            }
228            let reader = SegmentReader::open(vfs, &directory.join(meta.id.file_name()), meta.id)
229                .map_err(StoreError::Segment)?;
230            if reader
231                .directory()
232                .iter()
233                .any(|entry| entry.kind == RegionKind::GraphNodeBlocks.id())
234            {
235                continue;
236            }
237            bytes += meta.file_size;
238            inputs.push(reader);
239        }
240        if inputs.len() >= 2 {
241            inputs.sort_by_key(|input| input.meta().id);
242            return Ok(inputs);
243        }
244    }
245    Ok(Vec::new())
246}
247
248fn merge_id(generation: u64, inputs: &[SegmentReader]) -> SegmentId {
249    let mut source = Vec::with_capacity(inputs.len() * 16);
250    for input in inputs {
251        source.extend_from_slice(input.meta().id.as_bytes());
252    }
253    let hash = xxhash_rust::xxh3::xxh3_64_with_seed(&source, MERGE_ID_SEED);
254    let mut bytes = [0_u8; 16];
255    for (target, value) in bytes.iter_mut().take(8).zip(generation.to_be_bytes()) {
256        *target = value;
257    }
258    for (target, value) in bytes.iter_mut().skip(8).zip(hash.to_be_bytes()) {
259        *target = value;
260    }
261    SegmentId::from_bytes(bytes)
262}
263
264fn check_cancel(cancel: &CancelToken) -> Result<(), StoreError> {
265    if cancel.is_cancelled() {
266        Err(StoreError::SealCancelled)
267    } else {
268        Ok(())
269    }
270}
271
272fn merge_error(error: ConsolidateError) -> StoreError {
273    match error {
274        ConsolidateError::Store(error) => error,
275        ConsolidateError::Segment(error) => StoreError::Segment(error),
276        error => StoreError::Segment(SegmentError::Geometry(error.to_string())),
277    }
278}
zeppelin_embed/ingest/merge.rs

zeppelin_embed/ingest/
merge.rs
