1use super::validation::validate_count;
2use super::{MemoryNamespace, MemoryNode, MemoryRepositoryError, MemoryStatus};
3use serde::{Deserialize, Serialize};
4use sha2::{Digest, Sha256};
5use std::collections::{BTreeMap, BTreeSet};
6use std::io::Write;
7
8pub const MEMORY_NAMESPACE_SNAPSHOT_PROFILE_V1: &str = "a3s.memory.namespace-snapshot.sha256.v1";
10
11pub const MAX_SNAPSHOT_NODES: usize = 100_000;
13
14pub const MAX_SNAPSHOT_BYTES: usize = 256 * 1024 * 1024;
16
17const SNAPSHOT_DIGEST_DOMAIN: &str = "a3s.memory.namespace-snapshot.v1";
18
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21#[serde(rename_all = "camelCase")]
22pub struct MemorySnapshotRequest {
23 pub namespace: MemoryNamespace,
24 pub statuses: BTreeSet<MemoryStatus>,
25 pub max_nodes: usize,
26 pub max_bytes: usize,
27}
28
29impl MemorySnapshotRequest {
30 pub fn new(namespace: MemoryNamespace, max_nodes: usize, max_bytes: usize) -> Self {
32 Self {
33 namespace,
34 statuses: BTreeSet::from([MemoryStatus::Active]),
35 max_nodes,
36 max_bytes,
37 }
38 }
39
40 pub fn with_statuses(mut self, statuses: impl IntoIterator<Item = MemoryStatus>) -> Self {
41 self.statuses = statuses.into_iter().collect();
42 self
43 }
44
45 pub(crate) fn validate(&self) -> Result<(), MemoryRepositoryError> {
46 self.namespace.validate()?;
47 if self.statuses.is_empty() {
48 return Err(MemoryRepositoryError::invalid(
49 "snapshot.statuses",
50 "must contain at least one status",
51 ));
52 }
53 if self.max_nodes == 0 {
54 return Err(MemoryRepositoryError::invalid(
55 "snapshot.maxNodes",
56 "must be greater than zero",
57 ));
58 }
59 validate_count(
60 "namespace snapshot nodes",
61 self.max_nodes,
62 MAX_SNAPSHOT_NODES,
63 )?;
64 if self.max_bytes == 0 {
65 return Err(MemoryRepositoryError::invalid(
66 "snapshot.maxBytes",
67 "must be greater than zero",
68 ));
69 }
70 validate_count(
71 "namespace snapshot bytes",
72 self.max_bytes,
73 MAX_SNAPSHOT_BYTES,
74 )
75 }
76}
77
78#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
80#[serde(rename_all = "camelCase")]
81pub struct MemoryNamespaceSnapshot {
82 profile: String,
83 namespace: MemoryNamespace,
84 statuses: BTreeSet<MemoryStatus>,
85 nodes: Vec<MemoryNode>,
86 byte_count: usize,
87 digest: String,
88}
89
90impl MemoryNamespaceSnapshot {
91 pub fn try_new(
96 request: MemorySnapshotRequest,
97 nodes: Vec<MemoryNode>,
98 ) -> Result<Self, MemoryRepositoryError> {
99 build_snapshot(request, nodes)
100 }
101
102 pub fn verify(&self, request: &MemorySnapshotRequest) -> Result<(), MemoryRepositoryError> {
104 let expected = build_snapshot(request.clone(), self.nodes.clone())?;
105 if expected != *self {
106 return Err(MemoryRepositoryError::invariant(
107 "namespace snapshot identity or shape does not match its request",
108 ));
109 }
110 Ok(())
111 }
112
113 pub fn profile(&self) -> &str {
114 &self.profile
115 }
116
117 pub fn namespace(&self) -> &MemoryNamespace {
118 &self.namespace
119 }
120
121 pub fn statuses(&self) -> &BTreeSet<MemoryStatus> {
122 &self.statuses
123 }
124
125 pub fn nodes(&self) -> &[MemoryNode] {
126 &self.nodes
127 }
128
129 pub fn digest(&self) -> &str {
130 &self.digest
131 }
132
133 pub fn byte_count(&self) -> usize {
134 self.byte_count
135 }
136
137 pub fn into_nodes(self) -> Vec<MemoryNode> {
138 self.nodes
139 }
140}
141
142pub(crate) fn snapshot_from_map(
143 namespace_nodes: Option<&BTreeMap<String, MemoryNode>>,
144 request: MemorySnapshotRequest,
145) -> Result<MemoryNamespaceSnapshot, MemoryRepositoryError> {
146 request.validate()?;
147 let actual = namespace_nodes
148 .into_iter()
149 .flat_map(BTreeMap::values)
150 .filter(|node| request.statuses.contains(&node.status))
151 .count();
152 if actual > request.max_nodes {
153 return Err(MemoryRepositoryError::LimitExceeded {
154 resource: "namespace snapshot nodes".into(),
155 limit: request.max_nodes,
156 actual,
157 });
158 }
159 let nodes = namespace_nodes
160 .into_iter()
161 .flat_map(BTreeMap::values)
162 .filter(|node| request.statuses.contains(&node.status))
163 .collect::<Vec<_>>();
164 for node in &nodes {
165 if node.namespace != request.namespace {
166 return Err(MemoryRepositoryError::NamespaceMismatch {
167 context: "namespace snapshot".into(),
168 });
169 }
170 }
171 let (digest, byte_count) = snapshot_identity(&request, &nodes)?;
172 Ok(MemoryNamespaceSnapshot {
173 profile: MEMORY_NAMESPACE_SNAPSHOT_PROFILE_V1.to_string(),
174 namespace: request.namespace,
175 statuses: request.statuses,
176 nodes: nodes.into_iter().cloned().collect(),
177 byte_count,
178 digest,
179 })
180}
181
182pub(crate) fn snapshot_from_nodes(
183 request: MemorySnapshotRequest,
184 nodes: Vec<MemoryNode>,
185) -> Result<MemoryNamespaceSnapshot, MemoryRepositoryError> {
186 MemoryNamespaceSnapshot::try_new(request, nodes)
187}
188
189fn build_snapshot(
190 request: MemorySnapshotRequest,
191 mut nodes: Vec<MemoryNode>,
192) -> Result<MemoryNamespaceSnapshot, MemoryRepositoryError> {
193 request.validate()?;
194 nodes.sort_by(|left, right| left.id.cmp(&right.id));
195 if nodes.len() > request.max_nodes {
196 return Err(MemoryRepositoryError::LimitExceeded {
197 resource: "namespace snapshot nodes".into(),
198 limit: request.max_nodes,
199 actual: nodes.len(),
200 });
201 }
202 let mut previous_id: Option<&str> = None;
203 for node in &nodes {
204 if node.namespace != request.namespace || !request.statuses.contains(&node.status) {
205 return Err(MemoryRepositoryError::NamespaceMismatch {
206 context: "namespace snapshot".into(),
207 });
208 }
209 if previous_id == Some(node.id.as_str()) {
210 return Err(MemoryRepositoryError::invariant(format!(
211 "namespace snapshot contains duplicate node {}",
212 node.id
213 )));
214 }
215 previous_id = Some(&node.id);
216 }
217
218 let (digest, byte_count) = snapshot_identity(&request, &nodes)?;
219 Ok(MemoryNamespaceSnapshot {
220 profile: MEMORY_NAMESPACE_SNAPSHOT_PROFILE_V1.to_string(),
221 namespace: request.namespace,
222 statuses: request.statuses,
223 nodes,
224 byte_count,
225 digest,
226 })
227}
228
229fn snapshot_identity<T: Serialize>(
230 request: &MemorySnapshotRequest,
231 nodes: &[T],
232) -> Result<(String, usize), MemoryRepositoryError> {
233 #[derive(Serialize)]
234 #[serde(rename_all = "camelCase")]
235 struct DigestPayload<'a, T: Serialize> {
236 profile: &'static str,
237 namespace: &'a MemoryNamespace,
238 statuses: &'a BTreeSet<MemoryStatus>,
239 nodes: &'a [T],
240 }
241
242 let mut writer = BoundedDigestWriter::new(request.max_bytes);
243 let result = serde_json::to_writer(
244 &mut writer,
245 &DigestPayload {
246 profile: MEMORY_NAMESPACE_SNAPSHOT_PROFILE_V1,
247 namespace: &request.namespace,
248 statuses: &request.statuses,
249 nodes,
250 },
251 );
252 if writer.exceeded {
253 return Err(MemoryRepositoryError::LimitExceeded {
254 resource: "namespace snapshot bytes".into(),
255 limit: request.max_bytes,
256 actual: request.max_bytes.saturating_add(1),
257 });
258 }
259 result.map_err(|error| {
260 MemoryRepositoryError::invariant(format!(
261 "namespace snapshot could not be encoded: {error}"
262 ))
263 })?;
264 Ok((
265 format!("sha256:{:x}", writer.hasher.finalize()),
266 writer.byte_count,
267 ))
268}
269
270struct BoundedDigestWriter {
271 hasher: Sha256,
272 byte_count: usize,
273 limit: usize,
274 exceeded: bool,
275}
276
277impl BoundedDigestWriter {
278 fn new(limit: usize) -> Self {
279 let mut hasher = Sha256::new();
280 hasher.update(SNAPSHOT_DIGEST_DOMAIN.as_bytes());
281 hasher.update([0]);
282 Self {
283 hasher,
284 byte_count: 0,
285 limit,
286 exceeded: false,
287 }
288 }
289}
290
291impl Write for BoundedDigestWriter {
292 fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
293 let next = self
294 .byte_count
295 .checked_add(buffer.len())
296 .ok_or_else(|| std::io::Error::other("namespace snapshot byte count overflowed"))?;
297 if next > self.limit {
298 self.exceeded = true;
299 return Err(std::io::Error::other(
300 "namespace snapshot byte budget exceeded",
301 ));
302 }
303 self.hasher.update(buffer);
304 self.byte_count = next;
305 Ok(buffer.len())
306 }
307
308 fn flush(&mut self) -> std::io::Result<()> {
309 Ok(())
310 }
311}