selene_graph/
shared_snapshot.rs1use selene_persist::{SnapshotBuilder, SnapshotConfig, SnapshotFinalizeOutcome};
4
5use crate::{GraphResult, SharedGraph};
6
7impl SharedGraph {
8 pub fn write_snapshot(&self, config: SnapshotConfig) -> GraphResult<SnapshotFinalizeOutcome> {
24 let mut builder = SnapshotBuilder::new(config);
25 for provider in self.index_providers() {
26 let provider_tag = provider.provider_tag();
27 for sub_tag in provider.declared_sub_tags() {
28 let bytes = provider.write_section(*sub_tag)?;
29 builder.add_section(provider_tag.0, sub_tag.0, bytes)?;
30 }
31 }
32 builder.finalize().map_err(Into::into)
33 }
34}
35
36#[cfg(test)]
37mod tests {
38 use std::path::PathBuf;
39 use std::sync::Arc;
40 use std::time::{SystemTime, UNIX_EPOCH};
41
42 use selene_core::{Change, GraphId};
43 use selene_persist::{SectionCompression, SnapshotConfig, SnapshotReader, snapshot_path};
44
45 use crate::SharedGraph;
46 use crate::index_provider::{IndexProvider, ProviderError, ProviderTag, SubTag};
47
48 const TEST_PROVIDER: [u8; 4] = *b"TST1";
49 const TEST_SUB: [u8; 4] = *b"BODY";
50 const TEST_SUB_TAGS: &[SubTag] = &[SubTag(TEST_SUB)];
51
52 struct SnapshotOnlyProvider;
53
54 impl IndexProvider for SnapshotOnlyProvider {
55 fn provider_tag(&self) -> ProviderTag {
56 ProviderTag(TEST_PROVIDER)
57 }
58
59 fn read_section(&self, _sub_tag: SubTag, _bytes: &[u8]) -> Result<(), ProviderError> {
60 Ok(())
61 }
62
63 fn write_section(&self, sub_tag: SubTag) -> Result<Vec<u8>, ProviderError> {
64 assert_eq!(sub_tag, SubTag(TEST_SUB));
65 Ok(b"provider-body".to_vec())
66 }
67
68 fn on_change(&self, _change: &Change) -> Result<(), ProviderError> {
69 Ok(())
70 }
71
72 fn declared_sub_tags(&self) -> &[SubTag] {
73 TEST_SUB_TAGS
74 }
75 }
76
77 #[test]
78 fn write_snapshot_includes_registered_provider_sections() {
79 let dir = temp_dir("shared-snapshot");
80 let provider = Arc::new(SnapshotOnlyProvider);
81 let shared = SharedGraph::builder(GraphId::new(82_001))
82 .with_provider(provider as Arc<dyn IndexProvider>)
83 .build()
84 .unwrap();
85
86 let outcome = shared
87 .write_snapshot(SnapshotConfig {
88 dir: dir.clone(),
89 sequence: 7,
90 compression: SectionCompression::None,
91 fsync: false,
92 })
93 .unwrap();
94
95 assert_eq!(outcome.snapshot_seq, 7);
96 assert!(outcome.section_count > 1);
97 let mut reader = SnapshotReader::open(&snapshot_path(&dir, 7)).unwrap();
98 assert_eq!(
99 reader.read_section(TEST_PROVIDER, TEST_SUB).unwrap(),
100 b"provider-body"
101 );
102 }
103
104 fn temp_dir(name: &str) -> PathBuf {
105 let nanos = SystemTime::now()
106 .duration_since(UNIX_EPOCH)
107 .unwrap()
108 .as_nanos();
109 let dir = std::env::temp_dir().join(format!(
110 "selene-graph-{name}-{}-{nanos}",
111 std::process::id()
112 ));
113 let _ = std::fs::remove_dir_all(&dir);
114 std::fs::create_dir(&dir).unwrap();
115 dir
116 }
117}