1use std::collections::HashMap;
11use std::sync::Arc;
12
13use xds_core::{BoxResource, TypeUrl};
14
15#[derive(Debug, Clone, Default)]
17pub struct SnapshotResources {
18 version: String,
20 resources: HashMap<String, BoxResource>,
22}
23
24impl SnapshotResources {
25 pub fn new(version: impl Into<String>) -> Self {
27 Self {
28 version: version.into(),
29 resources: HashMap::new(),
30 }
31 }
32
33 #[inline]
35 pub fn version(&self) -> &str {
36 &self.version
37 }
38
39 #[inline]
41 pub fn len(&self) -> usize {
42 self.resources.len()
43 }
44
45 #[inline]
47 pub fn is_empty(&self) -> bool {
48 self.resources.is_empty()
49 }
50
51 #[inline]
53 pub fn get(&self, name: &str) -> Option<&BoxResource> {
54 self.resources.get(name)
55 }
56
57 #[inline]
59 pub fn iter(&self) -> impl Iterator<Item = (&String, &BoxResource)> {
60 self.resources.iter()
61 }
62
63 #[inline]
65 pub fn names(&self) -> impl Iterator<Item = &String> {
66 self.resources.keys()
67 }
68
69 pub fn to_vec(&self) -> Vec<BoxResource> {
71 self.resources.values().cloned().collect()
72 }
73}
74
75#[derive(Debug, Clone)]
80pub struct Snapshot {
81 version: String,
83 resources: HashMap<TypeUrl, SnapshotResources>,
85 created_at: std::time::Instant,
87}
88
89impl Snapshot {
90 #[must_use = "builder is unused unless `.build()` is called"]
92 pub fn builder() -> SnapshotBuilder {
93 SnapshotBuilder::new()
94 }
95
96 #[inline]
98 pub fn version(&self) -> &str {
99 &self.version
100 }
101
102 #[inline]
104 pub fn created_at(&self) -> std::time::Instant {
105 self.created_at
106 }
107
108 #[inline]
110 pub fn get_resources(&self, type_url: TypeUrl) -> Option<&SnapshotResources> {
111 self.resources.get(&type_url)
112 }
113
114 #[inline]
116 pub fn get_version(&self, type_url: TypeUrl) -> Option<&str> {
117 self.resources.get(&type_url).map(|r| r.version.as_str())
118 }
119
120 #[inline]
122 pub fn contains_type(&self, type_url: TypeUrl) -> bool {
123 self.resources.contains_key(&type_url)
124 }
125
126 pub fn type_urls(&self) -> impl Iterator<Item = &TypeUrl> {
128 self.resources.keys()
129 }
130
131 pub fn total_resources(&self) -> usize {
133 self.resources.values().map(|r| r.len()).sum()
134 }
135
136 pub fn is_empty(&self) -> bool {
138 self.resources.is_empty() || self.resources.values().all(|r| r.is_empty())
139 }
140}
141
142#[must_use = "builder is unused unless `.build()` is called"]
144#[derive(Debug, Default)]
145pub struct SnapshotBuilder {
146 version: String,
147 resources: HashMap<TypeUrl, SnapshotResources>,
148}
149
150impl SnapshotBuilder {
151 pub fn new() -> Self {
153 Self::default()
154 }
155
156 pub fn version(mut self, version: impl Into<String>) -> Self {
158 self.version = version.into();
159 self
160 }
161
162 pub fn resources(
166 self,
167 type_url: TypeUrl,
168 resources: impl IntoIterator<Item = BoxResource>,
169 ) -> Self {
170 let version = self.version.clone();
171 self.resources_with_version(type_url, version, resources)
172 }
173
174 pub fn resources_with_version(
176 mut self,
177 type_url: TypeUrl,
178 version: impl Into<String>,
179 resources: impl IntoIterator<Item = BoxResource>,
180 ) -> Self {
181 let mut snapshot_resources = SnapshotResources::new(version);
182 for resource in resources {
183 snapshot_resources
184 .resources
185 .insert(resource.name().to_string(), resource);
186 }
187 self.resources.insert(type_url, snapshot_resources);
188 self
189 }
190
191 pub fn resource(mut self, type_url: TypeUrl, resource: BoxResource) -> Self {
193 let entry = self
194 .resources
195 .entry(type_url)
196 .or_insert_with(|| SnapshotResources::new(self.version.clone()));
197 entry
198 .resources
199 .insert(resource.name().to_string(), resource);
200 self
201 }
202
203 pub fn build(self) -> Snapshot {
205 Snapshot {
206 version: self.version,
207 resources: self.resources,
208 created_at: std::time::Instant::now(),
209 }
210 }
211}
212
213#[allow(dead_code)] pub type SharedSnapshot = Arc<Snapshot>;
216
217#[cfg(test)]
218mod tests {
219 use super::*;
220 use xds_core::Resource;
221
222 #[derive(Debug)]
224 struct TestResource {
225 name: String,
226 data: Vec<u8>,
227 }
228
229 impl TestResource {
230 fn new(name: impl Into<String>) -> Self {
231 Self {
232 name: name.into(),
233 data: vec![],
234 }
235 }
236 }
237
238 impl Resource for TestResource {
239 fn type_url(&self) -> &str {
240 "test.type/TestResource"
241 }
242
243 fn name(&self) -> &str {
244 &self.name
245 }
246
247 fn encode(&self) -> Result<prost_types::Any, Box<dyn std::error::Error + Send + Sync>> {
248 Ok(prost_types::Any {
249 type_url: self.type_url().to_string(),
250 value: self.data.clone(),
251 })
252 }
253
254 fn as_any(&self) -> &dyn std::any::Any {
255 self
256 }
257 }
258
259 #[test]
260 fn snapshot_builder_basic() {
261 let snapshot = Snapshot::builder().version("v1").build();
262
263 assert_eq!(snapshot.version(), "v1");
264 assert!(snapshot.is_empty());
265 }
266
267 #[test]
268 fn snapshot_builder_with_resources() {
269 let snapshot = Snapshot::builder()
270 .version("v2")
271 .resources(TypeUrl::CLUSTER.into(), vec![])
272 .build();
273
274 assert_eq!(snapshot.version(), "v2");
275 assert!(snapshot.contains_type(TypeUrl::CLUSTER.into()));
276 }
277
278 #[test]
279 fn snapshot_resources_version() {
280 let snapshot = Snapshot::builder()
281 .version("global-v1")
282 .resources(TypeUrl::CLUSTER.into(), vec![])
283 .build();
284
285 assert_eq!(
286 snapshot.get_version(TypeUrl::CLUSTER.into()),
287 Some("global-v1")
288 );
289 }
290
291 #[test]
292 fn snapshot_with_custom_version_per_type() {
293 let snapshot = Snapshot::builder()
294 .version("global-v1")
295 .resources_with_version(TypeUrl::CLUSTER.into(), "cluster-v2", vec![])
296 .resources_with_version(TypeUrl::LISTENER.into(), "listener-v3", vec![])
297 .build();
298
299 assert_eq!(
300 snapshot.get_version(TypeUrl::CLUSTER.into()),
301 Some("cluster-v2")
302 );
303 assert_eq!(
304 snapshot.get_version(TypeUrl::LISTENER.into()),
305 Some("listener-v3")
306 );
307 }
308
309 #[test]
310 fn snapshot_with_actual_resources() {
311 let resource1: BoxResource = Arc::new(TestResource::new("resource-1"));
312 let resource2: BoxResource = Arc::new(TestResource::new("resource-2"));
313
314 let type_url = TypeUrl::new("test.type/TestResource");
315
316 let snapshot = Snapshot::builder()
317 .version("v1")
318 .resources(type_url.clone(), vec![resource1, resource2])
319 .build();
320
321 assert_eq!(snapshot.total_resources(), 2);
322 assert!(!snapshot.is_empty());
323
324 let resources = snapshot.get_resources(type_url).unwrap();
325 assert_eq!(resources.len(), 2);
326 }
327
328 #[test]
329 fn snapshot_get_resource_by_name() {
330 let resource: BoxResource = Arc::new(TestResource::new("my-resource"));
331 let type_url = TypeUrl::new("test.type/TestResource");
332
333 let snapshot = Snapshot::builder()
334 .version("v1")
335 .resources(type_url.clone(), vec![resource])
336 .build();
337
338 let resources = snapshot.get_resources(type_url.clone()).unwrap();
340 let found = resources.get("my-resource");
341 assert!(found.is_some());
342 assert_eq!(found.unwrap().name(), "my-resource");
343
344 let not_found = resources.get("nonexistent");
345 assert!(not_found.is_none());
346 }
347
348 #[test]
349 fn snapshot_add_single_resource() {
350 let resource: BoxResource = Arc::new(TestResource::new("single"));
351 let type_url = TypeUrl::new("test.type/TestResource");
352
353 let snapshot = Snapshot::builder()
354 .version("v1")
355 .resource(type_url.clone(), resource)
356 .build();
357
358 assert_eq!(snapshot.total_resources(), 1);
359 }
360
361 #[test]
362 fn snapshot_multiple_types() {
363 let cluster: BoxResource = Arc::new(TestResource::new("cluster-1"));
364 let listener: BoxResource = Arc::new(TestResource::new("listener-1"));
365 let route: BoxResource = Arc::new(TestResource::new("route-1"));
366
367 let snapshot = Snapshot::builder()
368 .version("v1")
369 .resource(TypeUrl::CLUSTER.into(), cluster)
370 .resource(TypeUrl::LISTENER.into(), listener)
371 .resource(TypeUrl::ROUTE.into(), route)
372 .build();
373
374 assert_eq!(snapshot.total_resources(), 3);
375 assert!(snapshot.contains_type(TypeUrl::CLUSTER.into()));
376 assert!(snapshot.contains_type(TypeUrl::LISTENER.into()));
377 assert!(snapshot.contains_type(TypeUrl::ROUTE.into()));
378 }
379
380 #[test]
381 fn snapshot_type_urls() {
382 let cluster: BoxResource = Arc::new(TestResource::new("cluster-1"));
383 let listener: BoxResource = Arc::new(TestResource::new("listener-1"));
384
385 let snapshot = Snapshot::builder()
386 .version("v1")
387 .resource(TypeUrl::CLUSTER.into(), cluster)
388 .resource(TypeUrl::LISTENER.into(), listener)
389 .build();
390
391 let type_urls: Vec<_> = snapshot.type_urls().collect();
392 assert_eq!(type_urls.len(), 2);
393 }
394
395 #[test]
396 fn snapshot_created_at() {
397 use std::time::Instant;
398
399 let before = Instant::now();
400 let snapshot = Snapshot::builder().version("v1").build();
401 let after = Instant::now();
402
403 assert!(snapshot.created_at() >= before);
404 assert!(snapshot.created_at() <= after);
405 }
406
407 #[test]
408 fn snapshot_age_via_created_at() {
409 use std::time::Duration;
410
411 let snapshot = Snapshot::builder().version("v1").build();
412 std::thread::sleep(Duration::from_millis(10));
413
414 let elapsed = snapshot.created_at().elapsed();
416 assert!(elapsed.as_millis() >= 10);
417 }
418
419 #[test]
420 fn snapshot_resource_iteration() {
421 let resource: BoxResource = Arc::new(TestResource::new("my-resource"));
422 let type_url = TypeUrl::new("test.type/TestResource");
423
424 let snapshot = Snapshot::builder()
425 .version("v1")
426 .resources(type_url.clone(), vec![resource])
427 .build();
428
429 assert_eq!(snapshot.total_resources(), 1);
430 let resources = snapshot.get_resources(type_url).unwrap();
431 let items: Vec<_> = resources.iter().collect();
433 assert_eq!(items.len(), 1);
434 assert_eq!(items[0].1.name(), "my-resource");
435 }
436
437 #[test]
438 fn snapshot_empty_version() {
439 let snapshot = Snapshot::builder().version("").build();
440 assert_eq!(snapshot.version(), "");
441 }
442
443 #[test]
444 fn snapshot_is_empty_with_empty_resources() {
445 let snapshot = Snapshot::builder()
447 .version("v1")
448 .resources(TypeUrl::CLUSTER.into(), vec![])
449 .build();
450
451 assert!(snapshot.is_empty());
453 }
454}