1use super::*;
2
3impl UnifiedStore {
4 pub fn new() -> Self {
5 Self::with_config(UnifiedStoreConfig::default())
6 }
7
8 pub fn format_version(&self) -> u32 {
10 self.format_version.load(Ordering::SeqCst)
11 }
12
13 pub(crate) fn set_format_version(&self, version: u32) {
14 self.format_version.store(version, Ordering::SeqCst);
15 }
16
17 pub fn next_entity_id(&self) -> EntityId {
19 EntityId::new(self.next_entity_id.fetch_add(1, Ordering::SeqCst))
20 }
21
22 pub fn reserve_entity_ids(&self, n: u64) -> std::ops::Range<u64> {
25 let start = self.next_entity_id.fetch_add(n, Ordering::SeqCst);
26 start..start + n
27 }
28
29 pub(crate) fn register_entity_id(&self, id: EntityId) {
30 let candidate = id.raw().saturating_add(1);
31 let mut current = self.next_entity_id.load(Ordering::SeqCst);
32 while candidate > current {
33 match self.next_entity_id.compare_exchange(
34 current,
35 candidate,
36 Ordering::SeqCst,
37 Ordering::SeqCst,
38 ) {
39 Ok(_) => break,
40 Err(updated) => current = updated,
41 }
42 }
43 }
44
45 fn restore_snapshot_entity(
48 &self,
49 collection: &str,
50 mut entity: UnifiedEntity,
51 ) -> Result<EntityId, StoreError> {
52 let manager = self.get_or_create_collection(collection);
53 if entity.id.raw() == 0 {
54 entity.id = self.next_entity_id();
55 } else {
56 self.register_entity_id(entity.id);
57 }
58 if let EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
59 if *row_id == 0 {
60 *row_id = manager.next_row_id();
61 } else {
62 manager.register_row_id(*row_id);
63 }
64 }
65 entity.ensure_table_logical_id();
66 self.context_index.index_entity(collection, &entity);
67 if self.config.auto_index_refs {
68 self.index_cross_refs(&entity, collection)?;
69 }
70 let graph_label = match &entity.kind {
71 EntityKind::GraphNode(node) => Some(node.label.clone()),
72 _ => None,
73 };
74 let id = manager.restore_snapshot_entity(entity)?;
75 if let Some(label) = graph_label {
76 self.update_graph_label_index(collection, &label, id);
77 }
78 Ok(id)
79 }
80
81 pub fn load_from_file(path: &Path) -> Result<Self, Box<dyn std::error::Error>> {
93 let file = File::open(path)?;
94 let mut reader = BufReader::new(file);
95 let mut buf = Vec::new();
96 reader.read_to_end(&mut buf)?;
97 Self::load_from_bytes(&buf)
98 }
99
100 pub(crate) fn load_from_bytes(buf: &[u8]) -> Result<Self, Box<dyn std::error::Error>> {
101 Self::load_from_bytes_with_config(buf, UnifiedStoreConfig::default())
102 }
103
104 pub(crate) fn load_from_bytes_with_config(
105 buf: &[u8],
106 config: UnifiedStoreConfig,
107 ) -> Result<Self, Box<dyn std::error::Error>> {
108 use reddb_engine::storage::snapshot::{visit_snapshot, SnapshotVisitor};
109
110 struct Restore<'a>(&'a UnifiedStore);
111 impl SnapshotVisitor for Restore<'_> {
112 fn collection(&mut self, name: &str) -> Result<(), String> {
113 self.0.get_or_create_collection(name);
114 Ok(())
115 }
116
117 fn entity(
118 &mut self,
119 collection: &str,
120 entity: UnifiedEntity,
121 metadata: Option<Metadata>,
122 ) -> Result<(), String> {
123 let id = entity.id;
124 self.0
125 .restore_snapshot_entity(collection, entity)
126 .map_err(|err| err.to_string())?;
127 if let Some(metadata) = metadata {
128 if let Some(manager) = self.0.get_collection(collection) {
129 let _ = manager.set_metadata(id, metadata);
130 }
131 }
132 Ok(())
133 }
134
135 fn cross_ref(
136 &mut self,
137 cross_ref: reddb_file::NativeDumpCrossRef,
138 ) -> Result<(), String> {
139 let source_collection = self
140 .0
141 .get_any(EntityId::new(cross_ref.source_id))
142 .map(|(name, _)| name)
143 .unwrap_or_else(|| cross_ref.target_collection.clone());
144 let _ = self.0.add_cross_ref(
145 &source_collection,
146 EntityId::new(cross_ref.source_id),
147 &cross_ref.target_collection,
148 EntityId::new(cross_ref.target_id),
149 RefType::from_byte(cross_ref.ref_type),
150 1.0,
151 );
152 Ok(())
153 }
154
155 fn auxiliary(&mut self, bytes: &[u8]) -> Result<(), String> {
156 self.0.set_aux_metadata(bytes.to_vec());
157 Ok(())
158 }
159 }
160
161 let version = reddb_file::decode_native_store_header(buf).map_err(|err| err.to_string())?;
162 let mut store = Self::with_config(config);
163 let embedded_wal_path = store.config.embedded_wal_path.take();
165 store.set_format_version(version);
166 visit_snapshot(buf, &mut Restore(&store))?;
167
168 if store.format_version() < STORE_VERSION_CURRENT {
169 store.set_format_version(STORE_VERSION_CURRENT);
170 }
171 store.config.embedded_wal_path = embedded_wal_path;
172
173 Ok(store)
174 }
175
176 pub fn save_to_file(&self, path: &Path) -> Result<(), Box<dyn std::error::Error>> {
181 let buf = self.to_binary_dump_bytes();
182 reddb_file::write_native_store_bytes_atomically(path, &buf)?;
183 Ok(())
184 }
185
186 pub(crate) fn to_binary_dump_bytes(&self) -> Vec<u8> {
187 let _vault_guard = self.vault_publication_lock.lock();
188 let pending_vault_versions = self.pending_vault_versions.read();
189 let (_, excluded) = self.snapshot_manager.get().map_or_else(
190 || (0, Default::default()),
191 |manager| manager.checkpoint_excluded_state(),
192 );
193 let mut buf = Vec::new();
194
195 buf.extend_from_slice(&reddb_file::encode_native_store_header(
198 STORE_VERSION_CURRENT,
199 ));
200
201 let collections = self.collections.read();
203 reddb_file::encode_native_dump_count(&mut buf, collections.len() as u32);
204
205 let fv = STORE_VERSION_CURRENT;
206 for (name, manager) in collections.iter() {
207 let entities = manager.query_all(|entity| {
209 if pending_vault_versions.contains(&entity.id) {
210 return false;
211 }
212 !matches!(entity.data, EntityData::Row(_) | EntityData::Vector(_))
213 || !excluded.contains(&entity.xmin)
214 });
215 reddb_file::encode_native_dump_collection_header(&mut buf, name, entities.len() as u32);
216
217 for mut entity in entities {
220 if matches!(entity.data, EntityData::Row(_)) && excluded.contains(&entity.xmax) {
223 entity.xmax = 0;
224 }
225 let metadata = manager.get_metadata(entity.id);
226 let record = Self::serialize_entity_record(&entity, metadata.as_ref(), fv);
227 reddb_file::encode_native_dump_entity_record(&mut buf, &record);
228 }
229 }
230
231 let cross_refs = self.cross_refs.read();
233 let total_refs: usize = cross_refs.values().map(|v| v.len()).sum();
234 reddb_file::encode_native_dump_count(&mut buf, total_refs as u32);
235
236 for (source_id, refs) in cross_refs.iter() {
237 for (target_id, ref_type, collection) in refs {
238 reddb_file::encode_native_dump_cross_ref(
239 &mut buf,
240 source_id.raw(),
241 target_id.raw(),
242 ref_type.to_byte(),
243 collection,
244 );
245 }
246 }
247
248 {
252 let aux = self.aux_metadata.read();
253 reddb_file::encode_native_dump_entity_record(&mut buf, &aux);
254 }
255
256 self.set_format_version(STORE_VERSION_CURRENT);
257
258 reddb_file::append_native_store_crc32_footer(&mut buf);
259 buf
260 }
261
262 pub(crate) fn read_entity_binary(
264 buf: &[u8],
265 pos: &mut usize,
266 format_version: u32,
267 ) -> Result<UnifiedEntity, Box<dyn std::error::Error>> {
268 Ok(reddb_engine::storage::entity_codec::read_entity_binary(
269 buf,
270 pos,
271 format_version,
272 )?)
273 }
274
275 pub(crate) fn write_entity_binary(
277 buf: &mut Vec<u8>,
278 entity: &UnifiedEntity,
279 format_version: u32,
280 ) {
281 reddb_engine::storage::entity_codec::write_entity_binary(buf, entity, format_version);
282 }
283}
284
285#[cfg(test)]
286mod aux_metadata_dump_tests {
287 use super::*;
288 use std::collections::HashMap;
289
290 #[test]
291 fn snapshot_load_preserves_existing_wal_and_restores_future_appends() {
292 let directory = tempfile::tempdir().expect("directory");
293 let path = directory.path().join("snapshot.rdb");
294 let store = UnifiedStore::with_config(UnifiedStoreConfig::default());
295 let old =
296 UnifiedEntity::table_row(EntityId::new(1), "rows", 1, vec![Value::text("snapshot")]);
297 store.insert_auto("rows", old).expect("snapshot row");
298 let snapshot = store.to_binary_dump_bytes();
299 crate::storage::EmbeddedRdbArtifact::create_with_snapshot(&path, &snapshot)
300 .expect("artifact");
301 let newer = UnifiedEntity::table_row(
302 EntityId::new(1),
303 "rows",
304 1,
305 vec![Value::text("durable WAL")],
306 );
307 let action = StoreWalAction::upsert_entity("rows", &newer, None, STORE_VERSION_CURRENT);
308 UnifiedStore::with_config(UnifiedStoreConfig::default().with_embedded_wal_path(&path))
309 .finish_paged_write([action])
310 .expect("durable update after snapshot");
311 let before = std::fs::read(&path).expect("original artifact bytes");
312 let loaded = UnifiedStore::load_from_bytes_with_config(
313 &snapshot,
314 UnifiedStoreConfig::default().with_embedded_wal_path(&path),
315 )
316 .expect("reconstruct snapshot");
317 assert_eq!(std::fs::read(&path).expect("artifact after load"), before);
318 assert_eq!(
319 loaded.config().embedded_wal_path.as_deref(),
320 Some(path.as_path())
321 );
322
323 let generation = crate::storage::EmbeddedRdbArtifact::open(&path)
324 .expect("boundary before future write")
325 .selected_superblock
326 .generation;
327 loaded
328 .insert_auto(
329 "rows",
330 UnifiedEntity::table_row(
331 EntityId::new(2),
332 "rows",
333 2,
334 vec![Value::text("future write")],
335 ),
336 )
337 .expect("future writes still append WAL");
338 assert_eq!(
339 crate::storage::EmbeddedRdbArtifact::open(&path)
340 .expect("boundary after future write")
341 .selected_superblock
342 .generation,
343 generation + 1
344 );
345 }
346
347 #[test]
348 fn aux_metadata_round_trips_through_binary_dump() {
349 let store = UnifiedStore::with_config(UnifiedStoreConfig::default());
350 store.create_collection("k").expect("create collection");
351 store.set_aux_metadata(b"contract-blob".to_vec());
352
353 let bytes = store.to_binary_dump_bytes();
354 let reloaded = UnifiedStore::load_from_bytes(&bytes).expect("reload dump");
355
356 assert_eq!(reloaded.aux_metadata(), b"contract-blob".to_vec());
357 }
358
359 #[test]
360 fn empty_aux_metadata_round_trips_as_empty() {
361 let store = UnifiedStore::with_config(UnifiedStoreConfig::default());
362 let bytes = store.to_binary_dump_bytes();
363 let reloaded = UnifiedStore::load_from_bytes(&bytes).expect("reload dump");
364
365 assert!(reloaded.aux_metadata().is_empty());
366 }
367
368 #[test]
369 fn interned_timeseries_series_reduces_repeated_tag_storage() {
370 let tags = HashMap::from([
371 (
372 "host".to_string(),
373 "srv-very-long-hostname-0001".to_string(),
374 ),
375 ("region".to_string(), "us-east-1".to_string()),
376 ("service".to_string(), "checkout-api".to_string()),
377 ]);
378 let inline_bytes: usize = (0..100)
379 .map(|timestamp_ns| {
380 let entity = UnifiedEntity::new(
381 EntityId::new(0),
382 EntityKind::TimeSeriesPoint(Box::new(TimeSeriesPointKind {
383 series: "cpu_metrics".to_string(),
384 metric: "cpu.idle".to_string(),
385 })),
386 EntityData::TimeSeries(crate::storage::TimeSeriesData {
387 metric: "cpu.idle".to_string(),
388 series_id: None,
389 timestamp_ns,
390 value: 94.8,
391 tags: tags.clone(),
392 fields: HashMap::new(),
393 }),
394 );
395 UnifiedStore::serialize_entity_record(&entity, None, reddb_file::STORE_VERSION_V10)
396 .len()
397 })
398 .sum();
399
400 let mut dictionary_fields = HashMap::new();
401 dictionary_fields.insert("collection".to_string(), Value::text("cpu_metrics"));
402 dictionary_fields.insert("series_id".to_string(), Value::UnsignedInteger(0));
403 dictionary_fields.insert("metric".to_string(), Value::text("cpu.idle"));
404 dictionary_fields.insert(
405 "canonical_tags".to_string(),
406 Value::text(
407 r#"{"host":"srv-very-long-hostname-0001","region":"us-east-1","service":"checkout-api"}"#,
408 ),
409 );
410 dictionary_fields.insert(
411 "tags".to_string(),
412 Value::Json(
413 br#"{"host":"srv-very-long-hostname-0001","region":"us-east-1","service":"checkout-api"}"#
414 .to_vec(),
415 ),
416 );
417 let dictionary = UnifiedEntity::new(
418 EntityId::new(0),
419 EntityKind::TableRow {
420 table: Arc::from("red_timeseries_series"),
421 row_id: 0,
422 },
423 EntityData::Row(RowData {
424 columns: Vec::new(),
425 named: Some(dictionary_fields),
426 schema: None,
427 }),
428 );
429 let dictionary_bytes =
430 UnifiedStore::serialize_entity_record(&dictionary, None, STORE_VERSION_V11).len();
431 let interned_bytes = dictionary_bytes
432 + (0..100)
433 .map(|timestamp_ns| {
434 let entity = UnifiedEntity::new(
435 EntityId::new(0),
436 EntityKind::TimeSeriesPoint(Box::new(TimeSeriesPointKind {
437 series: "cpu_metrics".to_string(),
438 metric: "cpu.idle".to_string(),
439 })),
440 EntityData::TimeSeries(crate::storage::TimeSeriesData {
441 metric: "cpu.idle".to_string(),
442 series_id: Some(0),
443 timestamp_ns,
444 value: 94.8,
445 tags: HashMap::new(),
446 fields: HashMap::new(),
447 }),
448 );
449 UnifiedStore::serialize_entity_record(&entity, None, STORE_VERSION_V11).len()
450 })
451 .sum::<usize>();
452
453 assert!(
454 interned_bytes < inline_bytes,
455 "interned series storage should shrink repeated tags: inline={inline_bytes}, interned={interned_bytes}"
456 );
457 }
458}