xmtp_proto/traits/combinators/
retry.rs

1use std::marker::PhantomData;
2
3use xmtp_common::{
4    ExponentialBackoff, MaybeSend, MaybeSync, Retry, Strategy as RetryStrategy, retry_async,
5};
6
7use crate::api::{ApiClientError, Client, Endpoint, Query, QueryRaw};
8
9/// The concrete type of a [`crate::api::retry`] Combinators.
10/// Generally using the concrete type can be avoided with type inference
11/// or impl Trait.
12pub struct RetryQuery<E, S = ExponentialBackoff> {
13    endpoint: E,
14    pub(crate) retry: Retry<S>,
15}
16
17impl<E> RetryQuery<E> {
18    pub fn new(endpoint: E) -> Self {
19        Self {
20            endpoint,
21            retry: Default::default(),
22        }
23    }
24}
25
26#[xmtp_common::async_trait]
27impl<E, C, S> Query<C> for RetryQuery<E, S>
28where
29    E: Query<C>,
30    C: Client,
31    S: RetryStrategy,
32{
33    type Output = E::Output;
34    async fn query(&mut self, client: &C) -> Result<Self::Output, ApiClientError> {
35        retry_async!(
36            self.retry,
37            (async { Query::<C>::query(&mut self.endpoint, client).await })
38        )
39    }
40}
41
42#[xmtp_common::async_trait]
43impl<E, C, S> QueryRaw<C> for RetryQuery<E, S>
44where
45    E: Endpoint,
46    C: Client,
47    S: RetryStrategy,
48{
49    async fn query_raw(&mut self, client: &C) -> Result<bytes::Bytes, ApiClientError> {
50        retry_async!(
51            self.retry,
52            (async { QueryRaw::<C>::query_raw(&mut self.endpoint, client).await })
53        )
54    }
55}
56
57pub struct RetrySpecialized<Spec> {
58    _marker: PhantomData<Spec>,
59}
60
61impl<E, Spec> Endpoint<RetrySpecialized<Spec>> for RetryQuery<E>
62where
63    E: Endpoint<Spec>,
64    Spec: MaybeSend + MaybeSync,
65{
66    type Output = <E as Endpoint<Spec>>::Output;
67
68    fn grpc_endpoint(&self) -> std::borrow::Cow<'static, str> {
69        self.endpoint.grpc_endpoint()
70    }
71
72    fn body(&self) -> Result<bytes::Bytes, crate::api::BodyError> {
73        self.endpoint.body()
74    }
75}
76
77/// retry with the default retry strategy (ExponentialBackoff)
78pub fn retry<E>(endpoint: E) -> RetryQuery<E, ExponentialBackoff> {
79    RetryQuery::<E, _> {
80        endpoint,
81        retry: Retry::default(),
82    }
83}
84
85/// Retry the endpoint, indicating a specific strategy to retry with
86pub fn retry_with_strategy<E, S>(endpoint: E, retry: Retry<S>) -> RetryQuery<E, S> {
87    RetryQuery::<E, S> { endpoint, retry }
88}
89
90#[cfg(test)]
91mod tests {
92
93    use crate::api::{
94        EndpointExt,
95        mock::{MockError, MockNetworkClient, TestEndpoint},
96    };
97
98    use super::*;
99
100    #[xmtp_common::test]
101    async fn retries_endpoint_three_times() {
102        let mut client = MockNetworkClient::new();
103        client.expect_request().times(3).returning(|_, _, _| {
104            tracing::info!("error");
105            Err(ApiClientError::client(MockError::ARetryableError))
106        });
107        client
108            .expect_request()
109            .times(1)
110            .returning(|_, _, _| Ok(http::Response::new(vec![].into())));
111
112        let result: Result<(), _> = retry(TestEndpoint).query(&client).await;
113        assert!(result.is_ok());
114    }
115
116    #[xmtp_common::test]
117    async fn does_not_retry_non_retryable() {
118        let mut client = MockNetworkClient::new();
119        client
120            .expect_request()
121            .times(1)
122            .returning(|_, _, _| Err(ApiClientError::client(MockError::ANonRetryableError)));
123
124        let result: Result<(), _> = retry(TestEndpoint).query(&client).await;
125        assert!(result.is_err());
126    }
127
128    #[xmtp_common::test]
129    fn test_grpc_endpoint_delegates_to_wrapped_endpoint() {
130        let retry_endpoint = retry(TestEndpoint);
131        assert_eq!(retry_endpoint.grpc_endpoint(), "/test.mock/TestEndpoint");
132    }
133
134    #[xmtp_common::test]
135    fn test_body_delegates_to_wrapped_endpoint() {
136        let retry_endpoint = retry(TestEndpoint);
137        let result = retry_endpoint.body();
138        assert!(result.is_ok());
139        assert_eq!(result.unwrap(), bytes::Bytes::from(vec![]));
140    }
141
142    #[xmtp_common::test]
143    async fn retries_with_strategy() {
144        let mut client = MockNetworkClient::new();
145        client
146            .expect_request()
147            .times(2)
148            .returning(|_, _, _| Err(ApiClientError::client(MockError::ARetryableError)));
149        client
150            .expect_request()
151            .times(1)
152            .returning(|_, _, _| Ok(http::Response::new(vec![1].into())));
153
154        let result: Result<(), _> = TestEndpoint
155            .ignore_response() // ignore b/c invalid protobuf bytes
156            .retry_with_strategy(Retry::builder().retries(2).build())
157            .query(&client)
158            .await;
159        assert!(result.is_ok(), "{:?}", result.unwrap_err());
160    }
161}
xmtp_proto/traits/combinators/retry.rs

xmtp_proto/traits/combinators/
retry.rs
