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

xmtp_proto/traits/combinators/
retry.rs
