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