vtubestudio/service/
retry.rs

1use crate::data::{RequestEnvelope, ResponseEnvelope};
2use crate::error::{Error, ErrorKind};
3
4use futures_util::future;
5use tower::retry::{Policy, Retry};
6use tower::Layer;
7use tracing::debug;
8
9/// Determines whether a request should be retried.
10///
11/// This can be used as either a [`Layer`] or a [`Policy`].
12#[derive(Debug, Clone)]
13pub struct RetryPolicy {
14    retry_on_disconnect: bool,
15    retry_on_auth_error: bool,
16}
17
18impl RetryPolicy {
19    /// Creates a new [`RetryPolicy`] with default values.
20    pub fn new() -> Self {
21        RetryPolicy {
22            retry_on_disconnect: true,
23            retry_on_auth_error: true,
24        }
25    }
26
27    /// Whether requests should be retried on disconnect. Default `true`.
28    pub fn on_disconnect(mut self, value: bool) -> Self {
29        self.retry_on_disconnect = value;
30        self
31    }
32
33    /// Whether requests should be retried on auth error. Default `true`.
34    pub fn on_auth_error(mut self, value: bool) -> Self {
35        self.retry_on_auth_error = value;
36        self
37    }
38}
39
40impl<S> Layer<S> for RetryPolicy {
41    type Service = Retry<Self, S>;
42
43    fn layer(&self, service: S) -> Self::Service {
44        let policy = self.clone();
45        Retry::new(policy, service)
46    }
47}
48
49impl Policy<RequestEnvelope, ResponseEnvelope, Error> for RetryPolicy {
50    type Future = future::Ready<Self>;
51
52    fn retry(
53        &self,
54        req: &RequestEnvelope,
55        result: Result<&ResponseEnvelope, &Error>,
56    ) -> Option<Self::Future> {
57        Some(future::ready(match result {
58            Ok(resp) if resp.is_unauthenticated_error() && self.retry_on_auth_error => {
59                self.clone().on_auth_error(false)
60            }
61
62            Err(e) => {
63                if self.retry_on_auth_error && e.is_unauthenticated_error() {
64                    debug!(
65                        message_type = req.message_type.as_str(),
66                        "Retrying request due to API auth error"
67                    );
68                    self.clone().on_auth_error(false)
69                } else if self.retry_on_disconnect && e.has_kind(ErrorKind::ConnectionDropped) {
70                    debug!(
71                        message_type = req.message_type.as_str(),
72                        "Retrying request due to disconnection"
73                    );
74                    self.clone().on_disconnect(false)
75                } else {
76                    return None;
77                }
78            }
79
80            _ => return None,
81        }))
82    }
83
84    fn clone_request(&self, req: &RequestEnvelope) -> Option<RequestEnvelope> {
85        Some(req.clone())
86    }
87}
vtubestudio/service/retry.rs

vtubestudio/service/
retry.rs
