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