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