Skip to main content

nodedb_lite/engine/array/
segments.rs

1//! Segment read/write helpers for the Lite array engine.
2//!
3//! On pagedb-backed storage (`as_array_segment_ext()` returns `Some`), tile
4//! data is stored in pagedb encrypted segments under `arr/tile/{name}/{id}`.
5//! On WASM (where `as_array_segment_ext()` returns `None`), the legacy KV blob path is used:
6//! bytes stored in the `Array` namespace under `segment:{name}:{id}`.
7//!
8//! The on-disk bytes are identical in both paths — the exact output of
9//! `nodedb_array::SegmentWriter::finish()`, which includes the header, tile
10//! frames, and footer.  `SegmentReader::open` can parse them directly.
11
12use std::sync::Arc;
13
14use nodedb_array::segment::extract_cell_bytes;
15use nodedb_array::{SegmentReader, SegmentWriter, SparseTile, TilePayload};
16use nodedb_types::Namespace;
17
18use crate::engine::array::manifest::segment_key;
19use crate::error::LiteError;
20use crate::storage::engine::StorageEngine;
21
22/// Flush a batch of `(TileId, SparseTile)` pairs into a new segment and
23/// persist the bytes.
24///
25/// When `as_array_segment_ext()` is available (pagedb backend), the bytes are
26/// stored as an encrypted pagedb segment.  Otherwise they are stored as a KV
27/// blob in the `Array` namespace.
28///
29/// Returns the serialized segment bytes so the caller can record the
30/// `byte_len` in the manifest without a second storage read.
31pub async fn write_segment<S: StorageEngine>(
32    storage: &Arc<S>,
33    name: &str,
34    seg_id: u64,
35    schema_hash: u64,
36    tiles: &[(nodedb_array::types::TileId, &SparseTile)],
37) -> Result<Vec<u8>, LiteError> {
38    let mut writer = SegmentWriter::new(schema_hash);
39    for (tile_id, tile) in tiles {
40        writer
41            .append_sparse(*tile_id, tile)
42            .map_err(|e| LiteError::Storage {
43                detail: format!("append_sparse: {e}"),
44            })?;
45    }
46    let bytes = writer.finish(None).map_err(|e| LiteError::Storage {
47        detail: format!("segment finish: {e}"),
48    })?;
49
50    #[cfg(not(target_arch = "wasm32"))]
51    if let Some(ext) = storage.as_array_segment_ext() {
52        ext.write_array_segment(name, seg_id, &bytes).await?;
53        return Ok(bytes);
54    }
55
56    let key = segment_key(name, seg_id);
57    storage.put(Namespace::Array, &key, &bytes).await?;
58    Ok(bytes)
59}
60
61/// Load segment bytes for `seg_id`.
62///
63/// When `as_array_segment_ext()` is available (pagedb backend), the bytes are
64/// read from the encrypted pagedb segment.  Otherwise they are read from the
65/// KV blob in the `Array` namespace.
66pub async fn load_segment<S: StorageEngine>(
67    storage: &Arc<S>,
68    name: &str,
69    seg_id: u64,
70) -> Result<Vec<u8>, LiteError> {
71    // On pagedb-backed storage, attempt to read from the encrypted segment.
72    // Fall through to the KV path if the segment is not found there — this
73    // handles data written via the legacy KV path (e.g. pre-migration data
74    // or tests that write directly to KV).
75    #[cfg(not(target_arch = "wasm32"))]
76    if let Some(ext) = storage.as_array_segment_ext()
77        && let Some(bytes) = ext.open_array_segment(name, seg_id).await?
78    {
79        return Ok(bytes.into_vec());
80    }
81    // Not in pagedb — fall through to KV lookup below.
82
83    let key = segment_key(name, seg_id);
84    storage
85        .get(Namespace::Array, &key)
86        .await?
87        .ok_or_else(|| LiteError::Storage {
88            detail: format!("segment {name}/{seg_id} not found"),
89        })
90}
91
92/// Open a `SegmentReader` over the given byte slice.
93pub fn open_reader(bytes: &[u8]) -> Result<SegmentReader<'_>, LiteError> {
94    SegmentReader::open(bytes).map_err(|e| LiteError::Storage {
95        detail: format!("SegmentReader::open: {e}"),
96    })
97}
98
99/// Delete the segment for `seg_id` from storage.
100///
101/// On pagedb backend, the segment is tombstoned and reaped by the next GC
102/// cycle.  On KV backends, the blob is deleted immediately.
103pub async fn delete_segment<S: StorageEngine>(
104    storage: &Arc<S>,
105    name: &str,
106    seg_id: u64,
107) -> Result<(), LiteError> {
108    // On pagedb-backed storage, tombstone the encrypted segment and also
109    // remove any legacy KV entry for the same segment (dual cleanup ensures
110    // no stale KV blobs after migration).
111    #[cfg(not(target_arch = "wasm32"))]
112    if let Some(ext) = storage.as_array_segment_ext() {
113        ext.delete_array_segment(name, seg_id).await?;
114        // Best-effort cleanup of any legacy KV entry.
115        let _ = storage
116            .delete(Namespace::Array, &segment_key(name, seg_id))
117            .await;
118        return Ok(());
119    }
120
121    storage
122        .delete(Namespace::Array, &segment_key(name, seg_id))
123        .await
124}
125
126/// Collect all tile versions for `hilbert_prefix` at or before `system_as_of`
127/// across all segments, filtering to cells at `coord`. Returns raw cell bytes
128/// for each `(TileId, cell_bytes)` pair compatible with the ceiling resolver.
129///
130/// If a coord is not in the tile, `extract_cell_bytes` returns `None` and
131/// that tile is silently skipped. The caller must aggregate across segments.
132pub async fn iter_cell_versions_across_segments<S: StorageEngine>(
133    storage: &Arc<S>,
134    name: &str,
135    seg_ids: impl Iterator<Item = u64>,
136    hilbert_prefix: u64,
137    system_as_of: i64,
138    coord: &[nodedb_array::types::coord::value::CoordValue],
139) -> Result<Vec<(nodedb_array::types::TileId, Vec<u8>)>, LiteError> {
140    let mut out = Vec::new();
141    for seg_id in seg_ids {
142        let bytes = load_segment(storage, name, seg_id).await?;
143        let reader = SegmentReader::open(&bytes).map_err(|e| LiteError::Storage {
144            detail: format!("open reader seg {seg_id}: {e}"),
145        })?;
146        let versions = reader
147            .iter_tile_versions(hilbert_prefix, system_as_of)
148            .map_err(|e| LiteError::Storage {
149                detail: format!("iter_tile_versions seg {seg_id}: {e}"),
150            })?;
151        for res in versions {
152            let (tile_id, payload) = res.map_err(|e| LiteError::Storage {
153                detail: format!("iter_tile_versions: {e}"),
154            })?;
155            if let TilePayload::Sparse(tile) = &payload {
156                match extract_cell_bytes(tile, coord) {
157                    Ok(Some(cell_bytes)) => out.push((tile_id, cell_bytes)),
158                    Ok(None) => {}
159                    Err(e) => {
160                        return Err(LiteError::Storage {
161                            detail: format!("extract_cell_bytes: {e}"),
162                        });
163                    }
164                }
165            }
166        }
167    }
168    Ok(out)
169}
170
171/// Collect all live cells from all segments for the given `hilbert_prefix`
172/// at or before `system_as_of` — used by `array_slice`.
173pub async fn collect_tile_versions_across_segments<S: StorageEngine>(
174    storage: &Arc<S>,
175    name: &str,
176    seg_ids: &[u64],
177    hilbert_prefix: u64,
178    system_as_of: i64,
179) -> Result<Vec<(nodedb_array::types::TileId, TilePayload)>, LiteError> {
180    let mut out = Vec::new();
181    for &seg_id in seg_ids {
182        let bytes = load_segment(storage, name, seg_id).await?;
183        let reader = SegmentReader::open(&bytes).map_err(|e| LiteError::Storage {
184            detail: format!("open reader seg {seg_id}: {e}"),
185        })?;
186        let versions = reader
187            .iter_tile_versions(hilbert_prefix, system_as_of)
188            .map_err(|e| LiteError::Storage {
189                detail: format!("iter_tile_versions seg {seg_id}: {e}"),
190            })?;
191        for res in versions {
192            let (tile_id, payload) = res.map_err(|e| LiteError::Storage {
193                detail: format!("decode tile seg {seg_id}: {e}"),
194            })?;
195            out.push((tile_id, payload));
196        }
197    }
198    Ok(out)
199}
200
201/// Compile a new merged segment from the provided `(TileId, SparseTile)` list
202/// and persist it, then delete the old segments.
203///
204/// Used by retention compaction. Returns the new segment ID and byte length.
205pub async fn rewrite_segment<S: StorageEngine>(
206    storage: &Arc<S>,
207    name: &str,
208    new_seg_id: u64,
209    schema_hash: u64,
210    tiles: &[(nodedb_array::types::TileId, SparseTile)],
211    old_seg_ids: &[u64],
212) -> Result<Vec<u8>, LiteError> {
213    let refs: Vec<_> = tiles.iter().map(|(id, tile)| (*id, tile)).collect();
214    let bytes = write_segment(storage, name, new_seg_id, schema_hash, &refs).await?;
215    for &old_id in old_seg_ids {
216        delete_segment(storage, name, old_id).await?;
217    }
218    Ok(bytes)
219}
220
221/// Read all decoded cells matching `hilbert_prefix` at or before `system_as_of`
222/// from a single segment's bytes (already loaded). Returns `(TileId, TilePayload)`.
223pub fn tile_versions_from_bytes(
224    bytes: &[u8],
225    hilbert_prefix: u64,
226    system_as_of: i64,
227) -> Result<Vec<(nodedb_array::types::TileId, TilePayload)>, LiteError> {
228    let reader = SegmentReader::open(bytes).map_err(|e| LiteError::Storage {
229        detail: format!("SegmentReader::open: {e}"),
230    })?;
231    let versions = reader
232        .iter_tile_versions(hilbert_prefix, system_as_of)
233        .map_err(|e| LiteError::Storage {
234            detail: format!("iter_tile_versions: {e}"),
235        })?;
236    versions
237        .map(|r| {
238            r.map_err(|e| LiteError::Storage {
239                detail: format!("tile decode: {e}"),
240            })
241        })
242        .collect()
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248    use crate::storage::pagedb_storage::PagedbStorageMem;
249    use nodedb_array::schema::ArraySchemaBuilder;
250    use nodedb_array::schema::attr_spec::{AttrSpec, AttrType};
251    use nodedb_array::schema::dim_spec::{DimSpec, DimType};
252    use nodedb_array::tile::sparse_tile::{SparseRow, SparseTileBuilder};
253    use nodedb_array::types::TileId;
254    use nodedb_array::types::cell_value::value::CellValue;
255    use nodedb_array::types::coord::value::CoordValue;
256    use nodedb_array::types::domain::{Domain, DomainBound};
257    use nodedb_types::{OPEN_UPPER, Surrogate};
258    use std::sync::Arc;
259
260    fn schema() -> nodedb_array::schema::ArraySchema {
261        ArraySchemaBuilder::new("g")
262            .dim(DimSpec::new(
263                "x",
264                DimType::Int64,
265                Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)),
266            ))
267            .attr(AttrSpec::new("v", AttrType::Int64, true))
268            .tile_extents(vec![4])
269            .build()
270            .unwrap()
271    }
272
273    const STORED_SURROGATE: Surrogate = Surrogate::new(2);
274
275    fn make_tile(s: &nodedb_array::schema::ArraySchema) -> SparseTile {
276        let mut b = SparseTileBuilder::new(s);
277        b.push_row(SparseRow::live(
278            &[CoordValue::Int64(1)],
279            &[CellValue::Int64(42)],
280            STORED_SURROGATE,
281            0,
282            OPEN_UPPER,
283        ))
284        .unwrap();
285        b.build()
286    }
287
288    #[tokio::test]
289    async fn write_and_load_segment() {
290        let s = schema();
291        let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
292        let tile = make_tile(&s);
293        let tile_id = TileId::snapshot(0);
294        write_segment(&storage, "g", 0, 0xABCD, &[(tile_id, &tile)])
295            .await
296            .unwrap();
297
298        let bytes = load_segment(&storage, "g", 0).await.unwrap();
299        let reader = open_reader(&bytes).unwrap();
300        assert_eq!(reader.tile_count(), 1);
301    }
302
303    #[tokio::test]
304    async fn delete_segment_removes_key() {
305        let s = schema();
306        let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
307        let tile = make_tile(&s);
308        write_segment(&storage, "g", 0, 0, &[(TileId::snapshot(0), &tile)])
309            .await
310            .unwrap();
311        delete_segment(&storage, "g", 0).await.unwrap();
312        assert!(load_segment(&storage, "g", 0).await.is_err());
313    }
314
315    /// Verify that `write_segment` dispatches through the pagedb segment path
316    /// when the storage backend supports `as_array_segment_ext()`.
317    ///
318    /// The segment bytes must be absent from the KV namespace (proving the
319    /// pagedb path was taken), and must be readable back via `load_segment`.
320    #[cfg(not(target_arch = "wasm32"))]
321    #[tokio::test]
322    async fn pagedb_path_writes_to_segment_not_kv() {
323        use nodedb_types::Namespace;
324
325        let s = schema();
326        let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
327        let tile = make_tile(&s);
328        let tile_id = TileId::snapshot(0);
329
330        write_segment(&storage, "arr_test", 7, 0xCAFE, &[(tile_id, &tile)])
331            .await
332            .unwrap();
333
334        // The KV namespace must be empty — no blob was written there.
335        let kv_key = crate::engine::array::manifest::segment_key("arr_test", 7);
336        let kv_val = storage.get(Namespace::Array, &kv_key).await.unwrap();
337        assert!(
338            kv_val.is_none(),
339            "expected no KV entry on pagedb-backed storage"
340        );
341
342        // But load_segment must succeed via the pagedb path.
343        let bytes = load_segment(&storage, "arr_test", 7).await.unwrap();
344        let reader = open_reader(&bytes).unwrap();
345        assert_eq!(reader.tile_count(), 1);
346    }
347
348    /// Format bit-identity: bytes written by `write_segment` must parse via
349    /// `nodedb_array::SegmentReader::open` — same as the Origin-side reader.
350    #[cfg(not(target_arch = "wasm32"))]
351    #[tokio::test]
352    async fn format_bit_identity_with_origin_reader() {
353        let s = schema();
354        let storage = Arc::new(PagedbStorageMem::open_in_memory().await.unwrap());
355        let tile = make_tile(&s);
356        let tile_id = TileId::new(1, 100);
357
358        let written_bytes =
359            write_segment(&storage, "identity_test", 0, 0xDEAD, &[(tile_id, &tile)])
360                .await
361                .unwrap();
362
363        // The bytes returned by write_segment must be parseable by SegmentReader
364        // (the same path used by Origin's SegmentHandle).
365        let reader = nodedb_array::SegmentReader::open(&written_bytes).unwrap();
366        assert_eq!(reader.tile_count(), 1);
367        assert_eq!(reader.tiles()[0].tile_id, tile_id);
368        assert_eq!(reader.schema_hash(), 0xDEAD);
369
370        // The stored live row keeps its bound surrogate through the round trip.
371        let TilePayload::Sparse(read_back) = reader.read_tile(0).unwrap() else {
372            panic!("expected a sparse tile");
373        };
374        read_back.check_stored_identities().unwrap();
375        assert_eq!(read_back.live_surrogate(0).unwrap(), STORED_SURROGATE);
376
377        // The bytes retrieved via load_segment must also be parseable.
378        let loaded = load_segment(&storage, "identity_test", 0).await.unwrap();
379        assert_eq!(loaded, written_bytes, "loaded bytes must be bit-identical");
380    }
381}