Skip to main content

xmtp_api_backend/queries/
api_stats.rs

1use xmtp_proto::{
2    api::HasStats,
3    api_client::{AggregateStats, ApiStats, IdentityStats, XmtpBackendClient, XmtpMlsStreams},
4    backend_v1::*,
5    types::{GroupId, InstallationId, TopicCursor},
6};
7#[derive(Clone, Debug)]
8pub struct TrackedStatsClient<C> {
9    inner: C,
10    stats: ApiStats,
11    identity_stats: IdentityStats,
12}
13impl<C> TrackedStatsClient<C> {
14    pub fn new(inner: C) -> Self {
15        Self {
16            inner,
17            stats: Default::default(),
18            identity_stats: Default::default(),
19        }
20    }
21    pub fn inner(&self) -> &C {
22        &self.inner
23    }
24}
25#[xmtp_common::async_trait]
26impl<C: XmtpBackendClient> XmtpBackendClient for TrackedStatsClient<C> {
27    type Error = C::Error;
28    fn register_client_event_writer(
29        &self,
30        writer: &std::sync::Arc<dyn xmtp_events::EventWriter<()>>,
31    ) {
32        self.inner.register_client_event_writer(writer);
33    }
34    fn unregister_client_event_writer(
35        &self,
36        writer: &std::sync::Arc<dyn xmtp_events::EventWriter<()>>,
37    ) {
38        self.inner.unregister_client_event_writer(writer);
39    }
40    async fn publish(&self, request: PublishRequest) -> Result<PublishResponse, Self::Error> {
41        self.stats.publish.count_request();
42        self.inner.publish(request).await
43    }
44    async fn create_upload(
45        &self,
46        request: CreateUploadRequest,
47    ) -> Result<CreateUploadResponse, Self::Error> {
48        self.stats.create_upload.count_request();
49        self.inner.create_upload(request).await
50    }
51    async fn query(&self, request: QueryRequest) -> Result<QueryResponse, Self::Error> {
52        self.stats.query.count_request();
53        self.inner.query(request).await
54    }
55    async fn query_newest(
56        &self,
57        request: QueryNewestRequest,
58    ) -> Result<QueryNewestResponse, Self::Error> {
59        self.stats.query_newest.count_request();
60        self.inner.query_newest(request).await
61    }
62    async fn get_inbox_ids(
63        &self,
64        request: GetInboxIdsRequest,
65    ) -> Result<GetInboxIdsResponse, Self::Error> {
66        self.identity_stats.get_inbox_ids.count_request();
67        self.inner.get_inbox_ids(request).await
68    }
69    // Not counted: the configuration read is hourly and carries no identity,
70    // and `AggregateStats` is public SDK surface that spec 006 does not change.
71    async fn get_configuration(
72        &self,
73        request: GetConfigurationRequest,
74    ) -> Result<GetConfigurationResponse, Self::Error> {
75        self.inner.get_configuration(request).await
76    }
77    fn backend_url(&self) -> Option<&str> {
78        self.inner.backend_url()
79    }
80
81    fn has_credential_source(&self) -> bool {
82        self.inner.has_credential_source()
83    }
84
85    fn set_limits(&self, limits: std::sync::Arc<xmtp_configuration::LimitsConfiguration>) {
86        self.inner.set_limits(limits)
87    }
88    async fn verify_smart_contract_wallet_signatures(
89        &self,
90        request: VerifySmartContractWalletSignaturesRequest,
91    ) -> Result<VerifySmartContractWalletSignaturesResponse, Self::Error> {
92        self.identity_stats
93            .verify_smart_contract_wallet_signatures
94            .count_request();
95        self.inner
96            .verify_smart_contract_wallet_signatures(request)
97            .await
98    }
99    async fn register(&self, request: RegisterRequest) -> Result<RecipientState, Self::Error> {
100        self.inner.register(request).await
101    }
102    async fn unregister(
103        &self,
104        request: UnregisterRequest,
105    ) -> Result<UnregisterResponse, Self::Error> {
106        self.inner.unregister(request).await
107    }
108    async fn update_subscriptions(
109        &self,
110        request: UpdateSubscriptionsRequest,
111    ) -> Result<RecipientState, Self::Error> {
112        self.inner.update_subscriptions(request).await
113    }
114}
115#[xmtp_common::async_trait]
116impl<C: XmtpMlsStreams> XmtpMlsStreams for TrackedStatsClient<C> {
117    type Error = C::Error;
118    type GroupMessageStream = C::GroupMessageStream;
119    type WelcomeMessageStream = C::WelcomeMessageStream;
120    async fn subscribe_envelopes_with_cursors(
121        &self,
122        cursors: &TopicCursor,
123        limits: xmtp_proto::types::IncomingBatchLimits,
124    ) -> Result<xmtp_proto::types::IncomingSubscription<Self::Error>, Self::Error> {
125        self.stats.subscribe_static.count_request();
126        self.inner
127            .subscribe_envelopes_with_cursors(cursors, limits)
128            .await
129    }
130    async fn subscribe_group_messages(
131        &self,
132        groups: &[&GroupId],
133    ) -> Result<Self::GroupMessageStream, Self::Error> {
134        self.stats.subscribe_static.count_request();
135        self.inner.subscribe_group_messages(groups).await
136    }
137    async fn subscribe_group_messages_with_cursors(
138        &self,
139        groups: &TopicCursor,
140    ) -> Result<Self::GroupMessageStream, Self::Error> {
141        self.stats.subscribe_static.count_request();
142        self.inner
143            .subscribe_group_messages_with_cursors(groups)
144            .await
145    }
146    async fn subscribe_welcome_messages(
147        &self,
148        installations: &[&InstallationId],
149    ) -> Result<Self::WelcomeMessageStream, Self::Error> {
150        self.stats.subscribe_static.count_request();
151        self.inner.subscribe_welcome_messages(installations).await
152    }
153    async fn subscribe_welcome_messages_with_cursors(
154        &self,
155        installations: &TopicCursor,
156    ) -> Result<Self::WelcomeMessageStream, Self::Error> {
157        self.stats.subscribe_static.count_request();
158        self.inner
159            .subscribe_welcome_messages_with_cursors(installations)
160            .await
161    }
162}
163xmtp_common::if_native! {
164    use xmtp_proto::api_client::XmtpMlsBidiStreams;
165
166    // `XmtpMlsBidiStreams` is native-only, so this forward is gated like the
167    // trait. It carries no per-call stat (the bidi stream is opened once and
168    // mutated in place, not counted per RPC like the unary/stream calls above);
169    // it exists so the stats wrapper is bidi-capable, letting the standard
170    // feature-switched test client open a `BidiConnection`.
171    #[xmtp_common::async_trait]
172    impl<C> XmtpMlsBidiStreams for TrackedStatsClient<C>
173    where
174        C: XmtpMlsBidiStreams,
175    {
176        type SubscribeStream = <C as XmtpMlsBidiStreams>::SubscribeStream;
177        type Error = <C as XmtpMlsBidiStreams>::Error;
178
179        fn host(&self) -> &str {
180            self.inner.host()
181        }
182
183        fn bidi_limits(&self) -> std::sync::Arc<xmtp_configuration::LimitsConfiguration> {
184            self.inner.bidi_limits()
185        }
186
187        async fn subscribe_bidi(
188            &self,
189            requests: futures::stream::BoxStream<'static, xmtp_proto::backend_v1::SubscribeRequest>,
190        ) -> Result<Self::SubscribeStream, Self::Error> {
191            self.stats.subscribe.count_request();
192            self.inner.subscribe_bidi(requests).await
193        }
194    }
195}
196
197impl<C> HasStats for TrackedStatsClient<C> {
198    fn aggregate_stats(&self) -> AggregateStats {
199        AggregateStats {
200            identity: self.identity_stats.clone(),
201            mls: self.stats.clone(),
202        }
203    }
204    fn mls_stats(&self) -> ApiStats {
205        self.stats.clone()
206    }
207    fn identity_stats(&self) -> IdentityStats {
208        self.identity_stats.clone()
209    }
210}
211
212#[xmtp_common::async_trait]
213impl<C: xmtp_proto::api::IsConnectedCheck> xmtp_proto::api::IsConnectedCheck
214    for TrackedStatsClient<C>
215{
216    async fn is_connected(&self) -> bool {
217        self.inner.is_connected().await
218    }
219}
220
221#[cfg(all(test, not(target_arch = "wasm32")))]
222mod tests;