Skip to main content

xmtp_proto/api_client/
stats.rs

1use std::sync::{
2    Arc,
3    atomic::{AtomicUsize, Ordering},
4};
5
6#[derive(Clone, Default, Debug)]
7pub struct ApiStats {
8    pub publish: Arc<EndpointStats>,
9    pub create_upload: Arc<EndpointStats>,
10    pub query: Arc<EndpointStats>,
11    pub query_newest: Arc<EndpointStats>,
12    pub subscribe: Arc<EndpointStats>,
13    pub subscribe_static: Arc<EndpointStats>,
14}
15
16impl ApiStats {
17    pub fn clear(&self) {
18        self.publish.clear();
19        self.create_upload.clear();
20        self.query.clear();
21        self.query_newest.clear();
22        self.subscribe.clear();
23        self.subscribe_static.clear();
24    }
25}
26
27#[derive(Clone, Default, Debug)]
28pub struct IdentityStats {
29    pub get_inbox_ids: Arc<EndpointStats>,
30    pub verify_smart_contract_wallet_signatures: Arc<EndpointStats>,
31}
32
33impl IdentityStats {
34    pub fn clear(&self) {
35        self.get_inbox_ids.clear();
36        self.verify_smart_contract_wallet_signatures.clear();
37    }
38}
39
40#[derive(Debug)]
41pub struct AggregateStats {
42    pub mls: ApiStats,
43    pub identity: IdentityStats,
44}
45
46#[derive(Default, Debug)]
47pub struct EndpointStats {
48    request_count: AtomicUsize,
49}
50
51impl std::fmt::Display for EndpointStats {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        write!(f, "{}", self.request_count.load(Ordering::Relaxed))
54    }
55}
56
57impl EndpointStats {
58    pub fn count_request(&self) {
59        self.request_count.fetch_add(1, Ordering::Relaxed);
60    }
61
62    pub fn get_count(&self) -> usize {
63        self.request_count.load(Ordering::Relaxed)
64    }
65    pub fn clear(&self) {
66        self.request_count.store(0, Ordering::Relaxed)
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[xmtp_common::test]
75    fn test_endpoint_stats_clear() {
76        let stats = EndpointStats::default();
77
78        // Add some counts
79        stats.count_request();
80        stats.count_request();
81        assert_eq!(stats.get_count(), 2);
82
83        // Clear should reset to 0
84        stats.clear();
85        assert_eq!(stats.get_count(), 0);
86
87        // Should be able to count again after clear
88        stats.count_request();
89        assert_eq!(stats.get_count(), 1);
90    }
91
92    #[xmtp_common::test]
93    fn test_endpoint_stats_display() {
94        let stats = EndpointStats::default();
95
96        // Display shows current count
97        assert_eq!(format!("{}", stats), "0");
98
99        stats.count_request();
100        stats.count_request();
101        assert_eq!(format!("{}", stats), "2");
102
103        stats.clear();
104        assert_eq!(format!("{}", stats), "0");
105    }
106}