nodedb_lite/engine/spatial/
checkpoint.rs

1//! Checkpoint serialization and restoration for [`SpatialIndexManager`].
2//!
3//! Persists the full in-memory spatial state to `Namespace::Spatial` so that a
4//! cold open can load the index without rebuilding from CRDT documents.
5//!
6//! ## Key layout under `Namespace::Spatial`
7//!
8//! | Key                                    | Value                                               |
9//! |----------------------------------------|-----------------------------------------------------|
10//! | `spatial:_collections`                 | MessagePack `Vec<(String, String)>` — (collection, field) pairs |
11//! | `spatial:{collection}:{field}:docmap`  | MessagePack `Vec<(String, u64)>` — doc_id → entry_id |
12//! | `spatial:_next_id`                     | MessagePack `u64` — next entry ID                  |
13//!
14//! The R-tree blob (`spatial:{collection}:{field}:rtree`) is stored in a pagedb
15//! segment when `as_spatial_segment_ext()` is available, or falls back to the
16//! `Namespace::Spatial` KV path (e.g. WASM). In both cases the
17//! bytes stored are the CRC32C-wrapped R-tree checkpoint produced by
18//! `crate::storage::checksum::wrap`.
19
20use std::collections::HashMap;
21
22use nodedb_types::Namespace;
23use nodedb_types::error::{NodeDbError, NodeDbResult};
24
25use crate::storage::engine::{StorageEngine, WriteOp};
26
27/// Flush the full spatial state to storage.
28///
29/// Persists each R-tree checkpoint (CRC32C-wrapped) plus the `doc_id → entry_id`
30/// mapping so that cold opens can restore exact index state.
31pub(crate) async fn flush_spatial<S>(
32    storage: &S,
33    checkpoints: &[(String, String, Vec<u8>)],
34    doc_to_entry: &HashMap<(String, String), u64>,
35    next_id: u64,
36) -> NodeDbResult<()>
37where
38    S: StorageEngine,
39{
40    let mut ops: Vec<WriteOp> = Vec::new();
41
42    // ── Collection list ───────────────────────────────────────────────────────
43    let index_keys: Vec<(String, String)> = checkpoints
44        .iter()
45        .map(|(c, f, _)| (c.clone(), f.clone()))
46        .collect();
47    let keys_bytes = zerompk::to_msgpack_vec(&index_keys)
48        .map_err(|e| NodeDbError::serialization("msgpack", e))?;
49    ops.push(WriteOp::Put {
50        ns: Namespace::Spatial,
51        key: b"spatial:_collections".to_vec(),
52        value: keys_bytes,
53    });
54
55    // ── Next entry ID ─────────────────────────────────────────────────────────
56    let next_id_bytes =
57        zerompk::to_msgpack_vec(&next_id).map_err(|e| NodeDbError::serialization("msgpack", e))?;
58    ops.push(WriteOp::Put {
59        ns: Namespace::Spatial,
60        key: b"spatial:_next_id".to_vec(),
61        value: next_id_bytes,
62    });
63
64    // ── Per-index doc-map (always on B+ tree) ─────────────────────────────────
65    for (collection, field, _rtree_bytes) in checkpoints {
66        let docmap_key = format!("spatial:{collection}:{field}:docmap");
67        let pairs: Vec<(String, u64)> = doc_to_entry
68            .iter()
69            .filter(|((coll, _doc_id), _)| coll == collection)
70            .map(|((_coll, doc_id), &entry_id)| (doc_id.clone(), entry_id))
71            .collect();
72        let docmap_bytes = zerompk::to_msgpack_vec(&pairs)
73            .map_err(|e| NodeDbError::serialization("msgpack", e))?;
74        ops.push(WriteOp::Put {
75            ns: Namespace::Spatial,
76            key: docmap_key.into_bytes(),
77            value: docmap_bytes,
78        });
79    }
80
81    // ── Commit catalog + docmap to B+ tree ────────────────────────────────────
82    storage
83        .batch_write(&ops)
84        .await
85        .map_err(NodeDbError::storage)?;
86
87    // ── Per-index R-tree bytes: segment when available, KV fallback ───────────
88    #[cfg(not(target_arch = "wasm32"))]
89    if let Some(seg) = storage.as_spatial_segment_ext() {
90        for (collection, field, rtree_bytes) in checkpoints {
91            let wrapped = crate::storage::checksum::wrap(rtree_bytes);
92            seg.write_spatial_segment(collection, field, &wrapped)
93                .await
94                .map_err(NodeDbError::storage)?;
95        }
96        return Ok(());
97    }
98
99    // Legacy KV path (WASM fallback).
100    let mut rtree_ops: Vec<WriteOp> = Vec::with_capacity(checkpoints.len());
101    for (collection, field, rtree_bytes) in checkpoints {
102        let rtree_key = format!("spatial:{collection}:{field}:rtree");
103        rtree_ops.push(WriteOp::Put {
104            ns: Namespace::Spatial,
105            key: rtree_key.into_bytes(),
106            value: crate::storage::checksum::wrap(rtree_bytes),
107        });
108    }
109    if !rtree_ops.is_empty() {
110        storage
111            .batch_write(&rtree_ops)
112            .await
113            .map_err(NodeDbError::storage)?;
114    }
115
116    Ok(())
117}
118
119/// Restore spatial state from storage on cold open.
120///
121/// Returns `(checkpoints, doc_to_entry, next_id)`.
122/// Returns an empty state if no checkpoint is found.
123pub(crate) async fn restore_spatial<S>(
124    storage: &S,
125) -> NodeDbResult<(
126    Vec<(String, String, Vec<u8>)>,
127    HashMap<(String, String), u64>,
128    u64,
129)>
130where
131    S: StorageEngine,
132{
133    // ── Read collection list ──────────────────────────────────────────────────
134    let Some(keys_bytes) = storage
135        .get(Namespace::Spatial, b"spatial:_collections")
136        .await?
137    else {
138        return Ok((Vec::new(), HashMap::new(), 1));
139    };
140
141    let Ok(index_keys) = zerompk::from_msgpack::<Vec<(String, String)>>(&keys_bytes) else {
142        tracing::warn!("spatial checkpoint: failed to decode collection list — starting fresh");
143        return Ok((Vec::new(), HashMap::new(), 1));
144    };
145
146    if index_keys.is_empty() {
147        return Ok((Vec::new(), HashMap::new(), 1));
148    }
149
150    // ── Read next_id ──────────────────────────────────────────────────────────
151    let next_id = if let Some(bytes) = storage.get(Namespace::Spatial, b"spatial:_next_id").await? {
152        zerompk::from_msgpack::<u64>(&bytes).unwrap_or(1)
153    } else {
154        1
155    };
156
157    // ── Per-index R-tree bytes and doc-map ────────────────────────────────────
158    let mut checkpoints: Vec<(String, String, Vec<u8>)> = Vec::new();
159    let mut doc_to_entry: HashMap<(String, String), u64> = HashMap::new();
160
161    for (collection, field) in &index_keys {
162        // Try pagedb segment first (non-WASM), then fall back to KV blob.
163        let rtree_envelope: Option<Vec<u8>> = {
164            #[cfg(not(target_arch = "wasm32"))]
165            {
166                if let Some(seg) = storage.as_spatial_segment_ext() {
167                    match seg.open_spatial_segment(collection, field).await {
168                        Ok(Some(boxed)) => Some(boxed.into_vec()),
169                        Ok(None) => {
170                            // Segment absent — fall through to KV blob.
171                            None
172                        }
173                        Err(e) => {
174                            tracing::warn!(
175                                collection = %collection,
176                                field = %field,
177                                error = %e,
178                                "spatial segment open failed — falling back to KV blob"
179                            );
180                            None
181                        }
182                    }
183                } else {
184                    None
185                }
186            }
187            #[cfg(target_arch = "wasm32")]
188            {
189                None
190            }
191        };
192
193        // If segment was not found (absent or not supported), try KV blob.
194        let envelope = if let Some(env) = rtree_envelope {
195            env
196        } else {
197            let rtree_key = format!("spatial:{collection}:{field}:rtree");
198            match storage.get(Namespace::Spatial, rtree_key.as_bytes()).await {
199                Ok(Some(env)) => env,
200                Ok(None) => continue,
201                Err(_) => continue,
202            }
203        };
204
205        match crate::storage::checksum::unwrap(&envelope) {
206            Some(bytes) => {
207                checkpoints.push((collection.clone(), field.clone(), bytes));
208            }
209            None => {
210                tracing::error!(
211                    collection = %collection,
212                    field = %field,
213                    "spatial R-tree CRC32C mismatch — discarding"
214                );
215                // Best-effort cleanup of the stale KV blob (segment path has
216                // no stale entry to clean in this branch).
217                let rtree_key = format!("spatial:{collection}:{field}:rtree");
218                let _ = storage
219                    .delete(Namespace::Spatial, rtree_key.as_bytes())
220                    .await;
221                continue;
222            }
223        }
224
225        // ── Restore doc_id → entry_id mapping ────────────────────────────────
226        let docmap_key = format!("spatial:{collection}:{field}:docmap");
227        if let Ok(Some(docmap_bytes)) = storage.get(Namespace::Spatial, docmap_key.as_bytes()).await
228            && let Ok(pairs) = zerompk::from_msgpack::<Vec<(String, u64)>>(&docmap_bytes)
229        {
230            for (doc_id, entry_id) in pairs {
231                doc_to_entry.insert((collection.clone(), doc_id), entry_id);
232            }
233        }
234    }
235
236    tracing::debug!(
237        index_count = checkpoints.len(),
238        doc_entry_count = doc_to_entry.len(),
239        next_id,
240        "spatial checkpoint restored"
241    );
242
243    Ok((checkpoints, doc_to_entry, next_id))
244}
nodedb_lite/engine/spatial/checkpoint.rs

nodedb_lite/engine/spatial/
checkpoint.rs
