use alloc::collections::BTreeMap;
use alloc::string::String;
use crate::{ModuleMetrics, ModuleMetricsBuilder, SchemaVersion};
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "minicbor", derive(minicbor::Encode, minicbor::Decode))]
pub struct Snapshot {
#[cfg_attr(feature = "minicbor", n(0))]
pub version: SchemaVersion,
#[cfg_attr(feature = "minicbor", n(1))]
pub timestamp_ms: u64,
#[cfg_attr(feature = "minicbor", n(2))]
pub modules: BTreeMap<String, ModuleMetrics>,
}
impl Snapshot {
#[cfg(feature = "std")]
pub fn new() -> Self {
Self {
version: SchemaVersion::current(),
timestamp_ms: current_timestamp_ms(),
modules: BTreeMap::new(),
}
}
pub fn with_timestamp(timestamp_ms: u64) -> Self {
Self {
version: SchemaVersion::current(),
timestamp_ms,
modules: BTreeMap::new(),
}
}
pub fn builder() -> SnapshotBuilder {
SnapshotBuilder::new()
}
pub fn is_empty(&self) -> bool {
self.modules.is_empty()
}
pub fn len(&self) -> usize {
self.modules.len()
}
pub fn get(&self, module: &str) -> Option<&ModuleMetrics> {
self.modules.get(module)
}
pub fn iter(&self) -> impl Iterator<Item = (&String, &ModuleMetrics)> {
self.modules.iter()
}
pub fn total_reads(&self) -> u64 {
self.modules.values().map(|m| m.total_reads()).sum()
}
pub fn total_writes(&self) -> u64 {
self.modules.values().map(|m| m.total_writes()).sum()
}
}
#[cfg(feature = "std")]
impl Default for Snapshot {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub struct SnapshotBuilder {
timestamp_ms: Option<u64>,
modules: BTreeMap<String, ModuleMetrics>,
}
impl SnapshotBuilder {
pub fn new() -> Self {
Self {
timestamp_ms: None,
modules: BTreeMap::new(),
}
}
pub fn timestamp_ms(mut self, ts: u64) -> Self {
self.timestamp_ms = Some(ts);
self
}
pub fn module<F>(mut self, name: impl Into<String>, f: F) -> Self
where
F: FnOnce(ModuleMetricsBuilder) -> ModuleMetricsBuilder,
{
let metrics = f(ModuleMetricsBuilder::new()).build();
self.modules.insert(name.into(), metrics);
self
}
pub fn module_metrics(mut self, name: impl Into<String>, metrics: ModuleMetrics) -> Self {
self.modules.insert(name.into(), metrics);
self
}
#[cfg(feature = "std")]
pub fn build(self) -> Snapshot {
Snapshot {
version: SchemaVersion::current(),
timestamp_ms: self.timestamp_ms.unwrap_or_else(current_timestamp_ms),
modules: self.modules,
}
}
#[cfg(not(feature = "std"))]
pub fn build(self) -> Snapshot {
Snapshot {
version: SchemaVersion::current(),
timestamp_ms: self.timestamp_ms.unwrap_or(0),
modules: self.modules,
}
}
}
impl Default for SnapshotBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(feature = "std")]
fn current_timestamp_ms() -> u64 {
use std::time::{SystemTime, UNIX_EPOCH};
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_snapshot_builder() {
let snapshot = Snapshot::builder()
.timestamp_ms(1703160000000)
.module("producer", |m| {
m.write("events", |w| w.count(1000).rate(100.0))
})
.module("consumer", |m| {
m.read("events", |r| r.count(950).backlog(50))
})
.build();
assert_eq!(snapshot.len(), 2);
assert_eq!(snapshot.timestamp_ms, 1703160000000);
assert_eq!(snapshot.total_writes(), 1000);
assert_eq!(snapshot.total_reads(), 950);
}
#[test]
fn test_snapshot_version() {
let snapshot = Snapshot::builder().build();
assert!(snapshot.version.is_compatible());
}
#[cfg(feature = "serde")]
#[test]
fn test_serde_roundtrip() {
let snapshot = Snapshot::builder()
.timestamp_ms(1703160000000)
.module("test", |m| m.read("topic", |r| r.count(42).backlog(5)))
.build();
let json = serde_json::to_string(&snapshot).unwrap();
let parsed: Snapshot = serde_json::from_str(&json).unwrap();
assert_eq!(snapshot, parsed);
}
#[cfg(feature = "minicbor")]
#[test]
fn test_minicbor_roundtrip() {
let snapshot = Snapshot::builder()
.timestamp_ms(1703160000000)
.module("test", |m| m.read("topic", |r| r.count(42).backlog(5)))
.build();
let bytes = minicbor::to_vec(&snapshot).unwrap();
let parsed: Snapshot = minicbor::decode(&bytes).unwrap();
assert_eq!(snapshot, parsed);
}
}