Skip to main content

a3s_memory/repository/
snapshot.rs

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
8/// Stable identity of the exact namespace-snapshot algorithm.
9pub const MEMORY_NAMESPACE_SNAPSHOT_PROFILE_V1: &str = "a3s.memory.namespace-snapshot.sha256.v1";
10
11/// Hard upper bound for one exact namespace snapshot.
12pub const MAX_SNAPSHOT_NODES: usize = 100_000;
13
14/// Hard upper bound for the canonical payload of one exact snapshot.
15pub const MAX_SNAPSHOT_BYTES: usize = 256 * 1024 * 1024;
16
17const SNAPSHOT_DIGEST_DOMAIN: &str = "a3s.memory.namespace-snapshot.v1";
18
19/// Caller-selected scope and hard node/byte budgets for one exact repository view.
20#[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    /// Select the complete current Active view of one exact namespace.
31    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/// Complete, deterministically ordered current view selected by one request.
79#[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    /// Construct and hash a complete caller-provided view.
92    ///
93    /// Custom repositories should use this constructor rather than assembling
94    /// response fields independently.
95    pub fn try_new(
96        request: MemorySnapshotRequest,
97        nodes: Vec<MemoryNode>,
98    ) -> Result<Self, MemoryRepositoryError> {
99        build_snapshot(request, nodes)
100    }
101
102    /// Recompute and verify every response field against the original request.
103    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}