google_cloud_gax/
retry_policy.rs

1// Copyright 2024 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Defines traits for retry policies and some common implementations.
16//!
17//! The client libraries automatically retry RPCs when (1) they fail due to
18//! transient errors **and** the RPC is [idempotent], (2) or failed before an
19//! RPC was started. That is, when it is safe to attempt the RPC more than once.
20//!
21//! Applications may override the default behavior, increasing the retry
22//! attempts, or changing what errors are considered safe to retry.
23//!
24//! This module defines the traits for retry policies and some common
25//! implementations.
26//!
27//! To configure the default throttler for a client, use
28//! [ClientBuilder::with_retry_policy]. To configure the retry policy used for
29//! a specific request, use [RequestOptionsBuilder::with_retry_policy].
30//!
31//! [ClientBuilder::with_retry_policy]: crate::client_builder::ClientBuilder::with_retry_policy
32//! [RequestOptionsBuilder::with_retry_policy]: crate::options::RequestOptionsBuilder::with_retry_policy
33//!
34//! # Examples
35//!
36//! Create a policy that only retries transient errors, and retries for at
37//! most 10 seconds or at most 5 attempts: whichever limit is reached first
38//! stops the retry loop.
39//! ```
40//! # use google_cloud_gax::retry_policy::*;
41//! use std::time::Duration;
42//! let policy = Aip194Strict.with_time_limit(Duration::from_secs(10)).with_attempt_limit(5);
43//! ```
44//!
45//! Create a policy that retries on any error (even when unsafe to do so),
46//! and stops retrying after 5 attempts or 10 seconds, whichever limit is
47//! reached first stops the retry loop.
48//! ```
49//! # use google_cloud_gax::retry_policy::*;
50//! use std::time::Duration;
51//! let policy = AlwaysRetry.with_time_limit(Duration::from_secs(10)).with_attempt_limit(5);
52//! ```
53//!
54//! [idempotent]: https://en.wikipedia.org/wiki/Idempotence
55
56mod client_timeout;
57mod too_many_requests;
58
59use crate::error::Error;
60use crate::retry_result::RetryResult;
61use crate::retry_state::RetryState;
62use crate::throttle_result::ThrottleResult;
63use std::sync::Arc;
64use std::time::Duration;
65
66pub use client_timeout::ClientTimeout;
67pub use too_many_requests::TooManyRequests;
68
69/// Determines how errors are handled in the retry loop.
70///
71/// Implementations of this trait determine if errors are retryable, and for how
72/// long the retry loop may continue.
73pub trait RetryPolicy: Send + Sync + std::fmt::Debug {
74    /// Query the retry policy after an error.
75    ///
76    /// # Parameters
77    /// * `state` - the state of the retry loop.
78    /// * `error` - the last error when attempting the request.
79    #[cfg_attr(not(feature = "_internal-semver"), doc(hidden))]
80    fn on_error(&self, state: &RetryState, error: Error) -> RetryResult;
81
82    /// Query the retry policy after a retry attempt is throttled.
83    ///
84    /// Retry attempts may be throttled before they are even sent out. The retry
85    /// policy may choose to treat these as normal errors, consuming attempts,
86    /// or may prefer to ignore them and always return [RetryResult::Continue].
87    ///
88    /// # Parameters
89    /// * `_state` - the state of the retry loop.
90    /// * `error` - the previous error that caused the retry attempt. Throttling
91    ///   only applies to retry attempts, and a retry attempt implies that a
92    ///   previous attempt failed. The retry policy should preserve this error.
93    #[cfg_attr(not(feature = "_internal-semver"), doc(hidden))]
94    fn on_throttle(&self, _state: &RetryState, error: Error) -> ThrottleResult {
95        ThrottleResult::Continue(error)
96    }
97
98    /// The remaining time in the retry policy.
99    ///
100    /// For policies based on time, this returns the remaining time in the
101    /// policy. The retry loop can use this value to adjust the next RPC
102    /// timeout. For policies that are not time based this returns `None`.
103    ///
104    /// # Parameters
105    /// * `_state` - the state of the retry loop.
106    /// * `attempt_count` - the number of attempts. This method is called before
107    ///   the first attempt, so the first value is zero.
108    #[cfg_attr(not(feature = "_internal-semver"), doc(hidden))]
109    fn remaining_time(&self, _state: &RetryState) -> Option<Duration> {
110        None
111    }
112}
113
114/// A helper type to use [RetryPolicy] in client and request options.
115#[derive(Clone, Debug)]
116pub struct RetryPolicyArg(Arc<dyn RetryPolicy>);
117
118impl<T> std::convert::From<T> for RetryPolicyArg
119where
120    T: RetryPolicy + 'static,
121{
122    fn from(value: T) -> Self {
123        Self(Arc::new(value))
124    }
125}
126
127impl std::convert::From<Arc<dyn RetryPolicy>> for RetryPolicyArg {
128    fn from(value: Arc<dyn RetryPolicy>) -> Self {
129        Self(value)
130    }
131}
132
133impl From<RetryPolicyArg> for Arc<dyn RetryPolicy> {
134    fn from(value: RetryPolicyArg) -> Arc<dyn RetryPolicy> {
135        value.0
136    }
137}
138
139/// Extension trait for [`RetryPolicy`]
140pub trait RetryPolicyExt: RetryPolicy + Sized {
141    /// Decorate a [RetryPolicy] to limit the total elapsed time in the retry loop.
142    ///
143    /// While the time spent in the retry loop (including time in backoff) is
144    /// less than the prescribed duration the `on_error()` method returns the
145    /// results of the inner policy. After that time it returns
146    /// [Exhausted][RetryResult::Exhausted] if the inner policy returns
147    /// [Continue][RetryResult::Continue].
148    ///
149    /// The `remaining_time()` function returns the remaining time. This is
150    /// always [Duration::ZERO] once or after the policy's expiration time is
151    /// reached.
152    ///
153    /// # Example
154    /// ```
155    /// # use google_cloud_gax::retry_policy::*;
156    /// # use google_cloud_gax::retry_state::RetryState;
157    /// let d = std::time::Duration::from_secs(10);
158    /// let policy = Aip194Strict.with_time_limit(d);
159    /// assert!(policy.remaining_time(&RetryState::new(true)) <= Some(d));
160    /// ```
161    fn with_time_limit(self, maximum_duration: Duration) -> LimitedElapsedTime<Self> {
162        LimitedElapsedTime::custom(self, maximum_duration)
163    }
164
165    /// Decorate a [RetryPolicy] to limit the number of retry attempts.
166    ///
167    /// This policy decorates an inner policy and limits the total number of
168    /// attempts. Note that `on_error()` is not called before the initial
169    /// (non-retry) attempt. Therefore, setting the maximum number of attempts
170    /// to 0 or 1 results in no retry attempts.
171    ///
172    /// The policy passes through the results from the inner policy as long as
173    /// `attempt_count < maximum_attempts`. Once the maximum number of attempts
174    /// is reached, the policy returns [Exhausted][RetryResult::Exhausted] if the
175    /// inner policy returns [Continue][RetryResult::Continue].
176    ///
177    /// # Example
178    /// ```
179    /// # use google_cloud_gax::retry_policy::*;
180    /// # use google_cloud_gax::retry_state::RetryState;
181    /// let policy = Aip194Strict.with_attempt_limit(3);
182    /// assert_eq!(policy.remaining_time(&RetryState::new(true)), None);
183    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(0_u32), transient_error()).is_continue());
184    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(1_u32), transient_error()).is_continue());
185    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(2_u32), transient_error()).is_continue());
186    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(3_u32), transient_error()).is_exhausted());
187    ///
188    /// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
189    /// fn transient_error() -> Error { Error::service(Status::default().set_code(Code::Unavailable)) }
190    /// ```
191    fn with_attempt_limit(self, maximum_attempts: u32) -> LimitedAttemptCount<Self> {
192        LimitedAttemptCount::custom(self, maximum_attempts)
193    }
194
195    /// Decorate a [RetryPolicy] to continue on certain status codes.
196    ///
197    /// This policy decorates an inner policy and retries any errors with HTTP
198    /// status code "429 - TOO_MANY_REQUESTS" **or** where the service returns
199    /// an error with code [ResourceExhausted].
200    ///
201    /// For other errors it returns the same value as the inner policy.
202    ///
203    /// Note that [ResourceExhausted] is ambiguous and may cause problems with
204    /// some services. The code is used for both "too many requests"  and for
205    /// "quota exceeded" problems. If the quota in question is some kind of rate
206    /// limit, then using this policy may be helpful. If the quota is not a rate
207    /// limit, then this retry policy may needlessly send the same RPC multiple
208    /// times.
209    ///
210    /// You should consult the documentation for the service and RPC in question
211    /// before using this policy.
212    ///
213    /// # Example
214    /// ```
215    /// use google_cloud_gax::retry_policy::{Aip194Strict, RetryPolicy, RetryPolicyExt};
216    /// use google_cloud_gax::retry_state::RetryState;
217    /// let policy = Aip194Strict;
218    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(0_u32), too_many_requests()).is_permanent());
219    /// let policy = Aip194Strict.continue_on_too_many_requests();
220    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(0_u32), too_many_requests()).is_continue());
221    ///
222    /// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
223    /// fn too_many_requests() -> Error { Error::service(Status::default().set_code(Code::ResourceExhausted)) }
224    /// ```
225    ///
226    /// [ResourceExhausted]: crate::error::rpc::Code::ResourceExhausted
227    fn continue_on_too_many_requests(self) -> TooManyRequests<Self> {
228        TooManyRequests::new(self)
229    }
230
231    /// Decorate a [RetryPolicy] to continue on client-side timeouts.
232    ///
233    /// This policy decorates an inner policy and retries any client-side timeout errors for
234    /// idempotent requests. For other errors it returns the same value as the inner policy.
235    ///
236    /// This policy is useful if you want to ignore connection timeouts, or retry if the service is
237    /// taking too long to respond. Be aware that a client-side timeout may occur even after the
238    /// service receives the request. If your request has side-effects, such as creating or deleting
239    /// resources, it may be unsafe to retry the operation.
240    ///
241    /// # Example
242    /// ```
243    /// use google_cloud_gax::retry_policy::{Aip194Strict, RetryPolicy, RetryPolicyExt};
244    /// use google_cloud_gax::retry_state::RetryState;
245    /// let policy = Aip194Strict;
246    /// assert!(policy.on_error(&RetryState::new(false).set_attempt_count(0_u32), timeout()).is_permanent());
247    /// let policy = Aip194Strict.continue_on_client_timeout();
248    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(0_u32), timeout()).is_continue());
249    ///
250    /// # use google_cloud_gax::error::Error;
251    /// fn timeout() -> Error {
252    /// # Error::timeout("test-only")
253    /// }
254    /// ```
255    fn continue_on_client_timeout(self) -> ClientTimeout<Self> {
256        ClientTimeout::new(self)
257    }
258}
259
260impl<T: RetryPolicy> RetryPolicyExt for T {}
261
262/// A retry policy that strictly follows [AIP-194].
263///
264/// This policy must be decorated to limit the number of retry attempts or the
265/// duration of the retry loop.
266///
267/// The policy interprets AIP-194 **strictly**, the retry decision for
268/// server-side errors are based only on the status code, and the only retryable
269/// status code is "UNAVAILABLE".
270///
271/// # Example
272/// ```
273/// # use google_cloud_gax::retry_policy::*;
274/// # use google_cloud_gax::retry_state::RetryState;
275/// let policy = Aip194Strict;
276/// assert!(policy.on_error(&RetryState::new(true), transient_error()).is_continue());
277/// assert!(policy.on_error(&RetryState::new(true), permanent_error()).is_permanent());
278///
279/// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
280/// fn transient_error() -> Error { Error::service(Status::default().set_code(Code::Unavailable)) }
281/// fn permanent_error() -> Error { Error::service(Status::default().set_code(Code::PermissionDenied)) }
282/// ```
283///
284/// [AIP-194]: https://google.aip.dev/194
285#[derive(Clone, Debug)]
286pub struct Aip194Strict;
287
288impl RetryPolicy for Aip194Strict {
289    fn on_error(&self, state: &RetryState, error: Error) -> RetryResult {
290        use crate::error::rpc::Code;
291        use http::StatusCode;
292
293        if error.is_transient_and_before_rpc() {
294            return RetryResult::Continue(error);
295        }
296        if !state.idempotent {
297            return RetryResult::Permanent(error);
298        }
299        if error.is_io() {
300            return RetryResult::Continue(error);
301        }
302        if error.status().is_some_and(|s| s.code == Code::Unavailable) {
303            return RetryResult::Continue(error);
304        }
305        // Some services return a status of "Unknown" and a http status code of 503
306        // (SERVICE_UNAVAILABLE). That is not how gRPC status codes are supposed to work, but the
307        // intent is clear: we need to retry.
308        if error
309            .http_status_code()
310            .is_some_and(|code| code == StatusCode::SERVICE_UNAVAILABLE.as_u16())
311        {
312            return RetryResult::Continue(error);
313        }
314        RetryResult::Permanent(error)
315    }
316}
317
318/// A retry policy that retries all errors.
319///
320/// This policy must be decorated to limit the number of retry attempts or the
321/// duration of the retry loop.
322///
323/// The policy retries all errors. This may be useful if the service guarantees
324/// idempotency, maybe through the use of request ids.
325///
326/// # Example
327/// ```
328/// # use google_cloud_gax::retry_policy::*;
329/// # use google_cloud_gax::retry_state::RetryState;
330/// let policy = AlwaysRetry;
331/// assert!(policy.on_error(&RetryState::new(true), transient_error()).is_continue());
332/// assert!(policy.on_error(&RetryState::new(true), permanent_error()).is_continue());
333///
334/// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
335/// fn transient_error() -> Error { Error::service(Status::default().set_code(Code::Unavailable)) }
336/// fn permanent_error() -> Error { Error::service(Status::default().set_code(Code::PermissionDenied)) }
337/// ```
338#[derive(Clone, Debug)]
339pub struct AlwaysRetry;
340
341impl RetryPolicy for AlwaysRetry {
342    fn on_error(&self, _state: &RetryState, error: Error) -> RetryResult {
343        RetryResult::Continue(error)
344    }
345}
346
347/// A retry policy that never retries.
348///
349/// This policy is useful when the client already has (or may already have) a
350/// retry policy configured, and you want to avoid retrying a particular method.
351///
352/// # Example
353/// ```
354/// # use google_cloud_gax::retry_policy::*;
355/// # use google_cloud_gax::retry_state::RetryState;
356/// let policy = NeverRetry;
357/// assert!(policy.on_error(&RetryState::new(true), transient_error()).is_exhausted());
358/// assert!(policy.on_error(&RetryState::new(true), permanent_error()).is_exhausted());
359///
360/// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
361/// fn transient_error() -> Error { Error::service(Status::default().set_code(Code::Unavailable)) }
362/// fn permanent_error() -> Error { Error::service(Status::default().set_code(Code::PermissionDenied)) }
363/// ```
364#[derive(Clone, Debug)]
365pub struct NeverRetry;
366
367impl RetryPolicy for NeverRetry {
368    fn on_error(&self, _state: &RetryState, error: Error) -> RetryResult {
369        RetryResult::Exhausted(error)
370    }
371}
372
373/// Error indicating that the maximum elapsed time for retries has been exceeded.
374#[derive(thiserror::Error, Debug)]
375pub struct LimitedElapsedTimeError {
376    maximum_duration: Duration,
377    #[source]
378    source: Error,
379}
380
381impl LimitedElapsedTimeError {
382    pub(crate) fn new(maximum_duration: Duration, source: Error) -> Self {
383        Self {
384            maximum_duration,
385            source,
386        }
387    }
388
389    /// Returns the maximum number of attempts in the exhausted policy.
390    pub fn maximum_duration(&self) -> Duration {
391        self.maximum_duration
392    }
393}
394
395impl std::fmt::Display for LimitedElapsedTimeError {
396    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
397        write!(
398            f,
399            "retry policy is exhausted after {}s, the last retry attempt was throttled",
400            self.maximum_duration.as_secs_f64()
401        )
402    }
403}
404
405/// A retry policy decorator that limits the total time in the retry loop.
406///
407/// This policy decorates an inner policy and limits the duration of retry
408/// loops. While the time spent in the retry loop (including time in backoff)
409/// is less than the prescribed duration the `on_error()` method returns the
410/// results of the inner policy. After that time it returns
411/// [Exhausted][RetryResult::Exhausted] if the inner policy returns
412/// [Continue][RetryResult::Continue].
413///
414/// The `remaining_time()` function returns the remaining time. This is always
415/// [Duration::ZERO] once or after the policy's deadline is reached.
416///
417/// # Parameters
418/// * `P` - the inner retry policy, defaults to [Aip194Strict].
419#[derive(Debug)]
420pub struct LimitedElapsedTime<P = Aip194Strict>
421where
422    P: RetryPolicy,
423{
424    inner: P,
425    maximum_duration: Duration,
426}
427
428impl LimitedElapsedTime {
429    /// Creates a new instance, with the default inner policy.
430    ///
431    /// # Example
432    /// ```
433    /// # use google_cloud_gax::retry_policy::*;
434    /// # use google_cloud_gax::retry_state::RetryState;
435    /// let d = std::time::Duration::from_secs(10);
436    /// let policy = LimitedElapsedTime::new(d);
437    /// assert!(policy.remaining_time(&RetryState::new(true)) <= Some(d));
438    /// ```
439    pub fn new(maximum_duration: Duration) -> Self {
440        Self {
441            inner: Aip194Strict,
442            maximum_duration,
443        }
444    }
445}
446
447impl<P> LimitedElapsedTime<P>
448where
449    P: RetryPolicy,
450{
451    /// Creates a new instance with a custom inner policy.
452    ///
453    /// # Example
454    /// ```
455    /// # use google_cloud_gax::retry_policy::*;
456    /// # use google_cloud_gax::retry_state::RetryState;
457    /// # use google_cloud_gax::error;
458    /// use std::time::{Duration, Instant};
459    /// let d = Duration::from_secs(10);
460    /// let policy = AlwaysRetry.with_time_limit(d);
461    /// assert!(policy.remaining_time(&RetryState::new(false)) <= Some(d));
462    /// assert!(policy.on_error(&RetryState::new(false), permanent_error()).is_continue());
463    ///
464    /// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
465    /// fn transient_error() -> Error { Error::service(Status::default().set_code(Code::Unavailable)) }
466    /// fn permanent_error() -> Error { Error::service(Status::default().set_code(Code::PermissionDenied)) }
467    /// ```
468    pub fn custom(inner: P, maximum_duration: Duration) -> Self {
469        Self {
470            inner,
471            maximum_duration,
472        }
473    }
474
475    fn error_if_exhausted(&self, state: &RetryState, error: Error) -> ThrottleResult {
476        let deadline = state.start + self.maximum_duration;
477        let now = tokio::time::Instant::now().into_std();
478        if now < deadline {
479            ThrottleResult::Continue(error)
480        } else {
481            ThrottleResult::Exhausted(Error::exhausted(LimitedElapsedTimeError::new(
482                self.maximum_duration,
483                error,
484            )))
485        }
486    }
487}
488
489impl<P> RetryPolicy for LimitedElapsedTime<P>
490where
491    P: RetryPolicy + 'static,
492{
493    fn on_error(&self, state: &RetryState, error: Error) -> RetryResult {
494        match self.inner.on_error(state, error) {
495            RetryResult::Permanent(e) => RetryResult::Permanent(e),
496            RetryResult::Exhausted(e) => RetryResult::Exhausted(e),
497            RetryResult::Continue(e) => {
498                if tokio::time::Instant::now().into_std() >= state.start + self.maximum_duration {
499                    RetryResult::Exhausted(e)
500                } else {
501                    RetryResult::Continue(e)
502                }
503            }
504        }
505    }
506
507    fn on_throttle(&self, state: &RetryState, error: Error) -> ThrottleResult {
508        match self.inner.on_throttle(state, error) {
509            ThrottleResult::Continue(e) => self.error_if_exhausted(state, e),
510            ThrottleResult::Exhausted(e) => ThrottleResult::Exhausted(e),
511        }
512    }
513
514    fn remaining_time(&self, state: &RetryState) -> Option<Duration> {
515        let deadline = state.start + self.maximum_duration;
516        let remaining = deadline.saturating_duration_since(tokio::time::Instant::now().into_std());
517        if let Some(inner) = self.inner.remaining_time(state) {
518            return Some(std::cmp::min(remaining, inner));
519        }
520        Some(remaining)
521    }
522}
523
524/// A retry policy decorator that limits the number of attempts.
525///
526/// This policy decorates an inner policy and limits the total number of
527/// attempts. Note that `on_error()` is not called before the initial
528/// (non-retry) attempt. Therefore, setting the maximum number of attempts to 0
529/// or 1 results in no retry attempts.
530///
531/// The policy passes through the results from the inner policy as long as
532/// `attempt_count < maximum_attempts`. However, once the maximum number of
533/// attempts is reached, the policy replaces any [Continue][RetryResult::Continue]
534/// result with [Exhausted][RetryResult::Exhausted].
535///
536/// # Parameters
537/// * `P` - the inner retry policy.
538#[derive(Debug)]
539pub struct LimitedAttemptCount<P = Aip194Strict>
540where
541    P: RetryPolicy,
542{
543    inner: P,
544    maximum_attempts: u32,
545}
546
547impl LimitedAttemptCount {
548    /// Creates a new instance, with the default inner policy.
549    ///
550    /// # Example
551    /// ```
552    /// # use google_cloud_gax::retry_policy::*;
553    /// let policy = LimitedAttemptCount::new(5);
554    /// ```
555    pub fn new(maximum_attempts: u32) -> Self {
556        Self {
557            inner: Aip194Strict,
558            maximum_attempts,
559        }
560    }
561}
562
563impl<P> LimitedAttemptCount<P>
564where
565    P: RetryPolicy,
566{
567    /// Creates a new instance with a custom inner policy.
568    ///
569    /// # Example
570    /// ```
571    /// # use google_cloud_gax::retry_policy::*;
572    /// # use google_cloud_gax::retry_state::RetryState;
573    /// let policy = LimitedAttemptCount::custom(AlwaysRetry, 2);
574    /// assert!(policy.on_error(&RetryState::new(false).set_attempt_count(1_u32), permanent_error()).is_continue());
575    /// assert!(policy.on_error(&RetryState::new(false).set_attempt_count(2_u32), permanent_error()).is_exhausted());
576    ///
577    /// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
578    /// fn permanent_error() -> Error { Error::service(Status::default().set_code(Code::PermissionDenied)) }
579    /// ```
580    pub fn custom(inner: P, maximum_attempts: u32) -> Self {
581        Self {
582            inner,
583            maximum_attempts,
584        }
585    }
586}
587
588impl<P> RetryPolicy for LimitedAttemptCount<P>
589where
590    P: RetryPolicy,
591{
592    fn on_error(&self, state: &RetryState, error: Error) -> RetryResult {
593        match self.inner.on_error(state, error) {
594            RetryResult::Permanent(e) => RetryResult::Permanent(e),
595            RetryResult::Exhausted(e) => RetryResult::Exhausted(e),
596            RetryResult::Continue(e) => {
597                if state.attempt_count >= self.maximum_attempts {
598                    RetryResult::Exhausted(e)
599                } else {
600                    RetryResult::Continue(e)
601                }
602            }
603        }
604    }
605
606    fn on_throttle(&self, state: &RetryState, error: Error) -> ThrottleResult {
607        match self.inner.on_throttle(state, error) {
608            ThrottleResult::Continue(e) if state.attempt_count >= self.maximum_attempts => {
609                ThrottleResult::Exhausted(e)
610            }
611            res => res,
612        }
613    }
614
615    fn remaining_time(&self, state: &RetryState) -> Option<Duration> {
616        self.inner.remaining_time(state)
617    }
618}
619
620#[cfg(test)]
621pub(crate) mod tests {
622    use super::*;
623    use http::HeaderMap;
624    use std::error::Error as StdError;
625    use std::time::Instant;
626
627    // Verify `RetryPolicyArg` can be converted from the desired types.
628    #[test]
629    fn retry_policy_arg() {
630        let policy = LimitedAttemptCount::new(3);
631        let _ = RetryPolicyArg::from(policy);
632
633        let policy: Arc<dyn RetryPolicy> = Arc::new(LimitedAttemptCount::new(3));
634        let _ = RetryPolicyArg::from(policy);
635    }
636
637    #[test]
638    fn aip194_strict() {
639        let p = Aip194Strict;
640
641        let now = Instant::now();
642        assert!(
643            p.on_error(&idempotent_state(now), unavailable())
644                .is_continue()
645        );
646        assert!(
647            p.on_error(&non_idempotent_state(now), unavailable())
648                .is_permanent()
649        );
650        assert!(matches!(
651            p.on_throttle(&idempotent_state(now), unavailable()),
652            ThrottleResult::Continue(_)
653        ));
654
655        assert!(
656            p.on_error(&idempotent_state(now), unknown_and_503())
657                .is_continue()
658        );
659        assert!(
660            p.on_error(&non_idempotent_state(now), unknown_and_503())
661                .is_permanent()
662        );
663        assert!(matches!(
664            p.on_throttle(&idempotent_state(now), unknown_and_503()),
665            ThrottleResult::Continue(_)
666        ));
667
668        assert!(
669            p.on_error(&idempotent_state(now), permission_denied())
670                .is_permanent()
671        );
672        assert!(
673            p.on_error(&non_idempotent_state(now), permission_denied())
674                .is_permanent()
675        );
676
677        assert!(
678            p.on_error(&idempotent_state(now), http_unavailable())
679                .is_continue()
680        );
681        assert!(
682            p.on_error(&non_idempotent_state(now), http_unavailable())
683                .is_permanent()
684        );
685        assert!(matches!(
686            p.on_throttle(&idempotent_state(now), http_unavailable()),
687            ThrottleResult::Continue(_)
688        ));
689
690        assert!(
691            p.on_error(&idempotent_state(now), http_permission_denied())
692                .is_permanent()
693        );
694        assert!(
695            p.on_error(&non_idempotent_state(now), http_permission_denied())
696                .is_permanent()
697        );
698
699        assert!(
700            p.on_error(&idempotent_state(now), Error::io("err".to_string()))
701                .is_continue()
702        );
703        assert!(
704            p.on_error(&non_idempotent_state(now), Error::io("err".to_string()))
705                .is_permanent()
706        );
707
708        assert!(
709            p.on_error(&idempotent_state(now), pre_rpc_transient())
710                .is_continue()
711        );
712        assert!(
713            p.on_error(&non_idempotent_state(now), pre_rpc_transient())
714                .is_continue()
715        );
716
717        assert!(
718            p.on_error(&idempotent_state(now), Error::ser("err"))
719                .is_permanent()
720        );
721        assert!(
722            p.on_error(&non_idempotent_state(now), Error::ser("err"))
723                .is_permanent()
724        );
725        assert!(
726            p.on_error(&idempotent_state(now), Error::deser("err"))
727                .is_permanent()
728        );
729        assert!(
730            p.on_error(&non_idempotent_state(now), Error::deser("err"))
731                .is_permanent()
732        );
733
734        assert!(
735            p.remaining_time(&idempotent_state(now)).is_none(),
736            "p={p:?}, now={now:?}"
737        );
738    }
739
740    #[test]
741    fn always_retry() {
742        let p = AlwaysRetry;
743
744        let now = Instant::now();
745        assert!(
746            p.remaining_time(&idempotent_state(now)).is_none(),
747            "p={p:?}, now={now:?}"
748        );
749        assert!(
750            p.on_error(&idempotent_state(now), http_unavailable())
751                .is_continue()
752        );
753        assert!(
754            p.on_error(&non_idempotent_state(now), http_unavailable())
755                .is_continue()
756        );
757        assert!(matches!(
758            p.on_throttle(&idempotent_state(now), http_unavailable()),
759            ThrottleResult::Continue(_)
760        ));
761
762        assert!(
763            p.on_error(&idempotent_state(now), unavailable())
764                .is_continue()
765        );
766        assert!(
767            p.on_error(&non_idempotent_state(now), unavailable())
768                .is_continue()
769        );
770    }
771
772    #[test_case::test_case(true, Error::io("err"))]
773    #[test_case::test_case(true, pre_rpc_transient())]
774    #[test_case::test_case(true, Error::ser("err"))]
775    #[test_case::test_case(false, Error::io("err"))]
776    #[test_case::test_case(false, pre_rpc_transient())]
777    #[test_case::test_case(false, Error::ser("err"))]
778    fn always_retry_error_kind(idempotent: bool, error: Error) {
779        let p = AlwaysRetry;
780        let now = Instant::now();
781        let state = if idempotent {
782            idempotent_state(now)
783        } else {
784            non_idempotent_state(now)
785        };
786        assert!(p.on_error(&state, error).is_continue());
787    }
788
789    #[test]
790    fn never_retry() {
791        let p = NeverRetry;
792
793        let now = Instant::now();
794        assert!(
795            p.remaining_time(&idempotent_state(now)).is_none(),
796            "p={p:?}, now={now:?}"
797        );
798        assert!(
799            p.on_error(&idempotent_state(now), http_unavailable())
800                .is_exhausted()
801        );
802        assert!(
803            p.on_error(&non_idempotent_state(now), http_unavailable())
804                .is_exhausted()
805        );
806        assert!(matches!(
807            p.on_throttle(&idempotent_state(now), http_unavailable()),
808            ThrottleResult::Continue(_)
809        ));
810
811        assert!(
812            p.on_error(&idempotent_state(now), unavailable())
813                .is_exhausted()
814        );
815        assert!(
816            p.on_error(&non_idempotent_state(now), unavailable())
817                .is_exhausted()
818        );
819
820        assert!(
821            p.on_error(&idempotent_state(now), http_permission_denied())
822                .is_exhausted()
823        );
824        assert!(
825            p.on_error(&non_idempotent_state(now), http_permission_denied())
826                .is_exhausted()
827        );
828    }
829
830    #[test_case::test_case(true, Error::io("err"))]
831    #[test_case::test_case(true, pre_rpc_transient())]
832    #[test_case::test_case(true, Error::ser("err"))]
833    #[test_case::test_case(false, Error::io("err"))]
834    #[test_case::test_case(false, pre_rpc_transient())]
835    #[test_case::test_case(false, Error::ser("err"))]
836    fn never_retry_error_kind(idempotent: bool, error: Error) {
837        let p = NeverRetry;
838        let now = Instant::now();
839        let state = if idempotent {
840            idempotent_state(now)
841        } else {
842            non_idempotent_state(now)
843        };
844        assert!(p.on_error(&state, error).is_exhausted());
845    }
846
847    fn pre_rpc_transient() -> Error {
848        use crate::error::CredentialsError;
849        Error::authentication(CredentialsError::from_msg(true, "err"))
850    }
851
852    fn http_unavailable() -> Error {
853        Error::http(
854            503_u16,
855            HeaderMap::new(),
856            bytes::Bytes::from_owner("SERVICE UNAVAILABLE".to_string()),
857        )
858    }
859
860    fn http_permission_denied() -> Error {
861        Error::http(
862            403_u16,
863            HeaderMap::new(),
864            bytes::Bytes::from_owner("PERMISSION DENIED".to_string()),
865        )
866    }
867
868    fn unavailable() -> Error {
869        use crate::error::rpc::Code;
870        let status = crate::error::rpc::Status::default()
871            .set_code(Code::Unavailable)
872            .set_message("UNAVAILABLE");
873        Error::service(status)
874    }
875
876    fn unknown_and_503() -> Error {
877        use crate::error::rpc::Code;
878        let status = crate::error::rpc::Status::default()
879            .set_code(Code::Unknown)
880            .set_message("UNAVAILABLE");
881        Error::service_full(status, Some(503), None, Some("source error".into()))
882    }
883
884    fn permission_denied() -> Error {
885        use crate::error::rpc::Code;
886        let status = crate::error::rpc::Status::default()
887            .set_code(Code::PermissionDenied)
888            .set_message("PERMISSION_DENIED");
889        Error::service(status)
890    }
891
892    mockall::mock! {
893        #[derive(Debug)]
894        pub(crate) Policy {}
895        impl RetryPolicy for Policy {
896            fn on_error(&self, state: &RetryState, error: Error) -> RetryResult;
897            fn on_throttle(&self, state: &RetryState, error: Error) -> ThrottleResult;
898            fn remaining_time(&self, state: &RetryState) -> Option<Duration>;
899        }
900    }
901
902    #[test]
903    fn limited_elapsed_time_error() {
904        let limit = Duration::from_secs(123) + Duration::from_millis(567);
905        let err = LimitedElapsedTimeError::new(limit, unavailable());
906        assert_eq!(err.maximum_duration(), limit);
907        let fmt = err.to_string();
908        assert!(fmt.contains("123.567s"), "display={fmt}, debug={err:?}");
909        assert!(err.source().is_some(), "{err:?}");
910    }
911
912    #[test]
913    fn test_limited_time_forwards() {
914        let mut mock = MockPolicy::new();
915        mock.expect_on_error()
916            .times(1..)
917            .returning(|_, e| RetryResult::Continue(e));
918        mock.expect_on_throttle()
919            .times(1..)
920            .returning(|_, e| ThrottleResult::Continue(e));
921        mock.expect_remaining_time().times(1).returning(|_| None);
922
923        let now = Instant::now();
924        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
925        let rf = policy.on_error(&idempotent_state(now), transient_error());
926        assert!(rf.is_continue());
927
928        let rt = policy.remaining_time(&idempotent_state(now));
929        assert!(rt.is_some(), "policy={policy:?}, now={now:?}");
930
931        let e = policy.on_throttle(&idempotent_state(now), transient_error());
932        assert!(matches!(e, ThrottleResult::Continue(_)));
933    }
934
935    #[test]
936    fn test_limited_time_on_throttle_continue() {
937        let mut mock = MockPolicy::new();
938        mock.expect_on_throttle()
939            .times(1..)
940            .returning(|_, e| ThrottleResult::Continue(e));
941
942        let now = Instant::now();
943        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
944
945        // Before the policy expires the inner result is returned verbatim.
946        let rf = policy.on_throttle(
947            &idempotent_state(now - Duration::from_secs(50)),
948            unavailable(),
949        );
950        assert!(matches!(rf, ThrottleResult::Continue(_)), "{rf:?}");
951
952        // After the policy expires the innter result is always "exhausted".
953        let rf = policy.on_throttle(
954            &idempotent_state(now - Duration::from_secs(70)),
955            unavailable(),
956        );
957        assert!(matches!(rf, ThrottleResult::Exhausted(_)), "{rf:?}");
958    }
959
960    #[test]
961    fn test_limited_time_on_throttle_exhausted() {
962        let mut mock = MockPolicy::new();
963        mock.expect_on_throttle()
964            .times(1..)
965            .returning(|_, e| ThrottleResult::Exhausted(e));
966
967        let now = Instant::now();
968        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
969
970        // Before the policy expires the inner result is returned verbatim.
971        let rf = policy.on_throttle(
972            &idempotent_state(now - Duration::from_secs(50)),
973            unavailable(),
974        );
975        assert!(matches!(rf, ThrottleResult::Exhausted(_)), "{rf:?}");
976    }
977
978    #[test]
979    fn test_limited_time_inner_continues() {
980        let mut mock = MockPolicy::new();
981        mock.expect_on_error()
982            .times(1..)
983            .returning(|_, e| RetryResult::Continue(e));
984
985        let now = Instant::now();
986        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
987        let rf = policy.on_error(
988            &idempotent_state(now - Duration::from_secs(10)),
989            transient_error(),
990        );
991        assert!(rf.is_continue());
992
993        let rf = policy.on_error(
994            &idempotent_state(now - Duration::from_secs(70)),
995            transient_error(),
996        );
997        assert!(rf.is_exhausted());
998    }
999
1000    #[test]
1001    fn test_limited_time_inner_permanent() {
1002        let mut mock = MockPolicy::new();
1003        mock.expect_on_error()
1004            .times(2)
1005            .returning(|_, e| RetryResult::Permanent(e));
1006
1007        let now = Instant::now();
1008        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1009
1010        let rf = policy.on_error(
1011            &non_idempotent_state(now - Duration::from_secs(10)),
1012            transient_error(),
1013        );
1014        assert!(rf.is_permanent());
1015
1016        let rf = policy.on_error(
1017            &non_idempotent_state(now + Duration::from_secs(10)),
1018            transient_error(),
1019        );
1020        assert!(rf.is_permanent());
1021    }
1022
1023    #[test]
1024    fn test_limited_time_inner_exhausted() {
1025        let mut mock = MockPolicy::new();
1026        mock.expect_on_error()
1027            .times(2)
1028            .returning(|_, e| RetryResult::Exhausted(e));
1029
1030        let now = Instant::now();
1031        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1032
1033        let rf = policy.on_error(
1034            &non_idempotent_state(now - Duration::from_secs(10)),
1035            transient_error(),
1036        );
1037        assert!(rf.is_exhausted());
1038
1039        let rf = policy.on_error(
1040            &non_idempotent_state(now + Duration::from_secs(10)),
1041            transient_error(),
1042        );
1043        assert!(rf.is_exhausted());
1044    }
1045
1046    #[test]
1047    fn test_limited_time_remaining_inner_longer() {
1048        let mut mock = MockPolicy::new();
1049        mock.expect_remaining_time()
1050            .times(1)
1051            .returning(|_| Some(Duration::from_secs(30)));
1052
1053        let now = Instant::now();
1054        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1055
1056        let remaining = policy.remaining_time(&idempotent_state(now - Duration::from_secs(55)));
1057        assert!(remaining <= Some(Duration::from_secs(5)), "{remaining:?}");
1058    }
1059
1060    #[test]
1061    fn test_limited_time_remaining_inner_shorter() {
1062        let mut mock = MockPolicy::new();
1063        mock.expect_remaining_time()
1064            .times(1)
1065            .returning(|_| Some(Duration::from_secs(5)));
1066        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1067
1068        let now = Instant::now();
1069        let remaining = policy.remaining_time(&idempotent_state(now - Duration::from_secs(5)));
1070        assert!(remaining <= Some(Duration::from_secs(10)), "{remaining:?}");
1071    }
1072
1073    #[test]
1074    fn test_limited_time_remaining_inner_is_none() {
1075        let mut mock = MockPolicy::new();
1076        mock.expect_remaining_time().times(1).returning(|_| None);
1077        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1078
1079        let now = Instant::now();
1080        let remaining = policy.remaining_time(&idempotent_state(now - Duration::from_secs(50)));
1081        assert!(remaining <= Some(Duration::from_secs(10)), "{remaining:?}");
1082    }
1083
1084    #[test]
1085    fn test_limited_attempt_count_on_error() {
1086        let mut mock = MockPolicy::new();
1087        mock.expect_on_error()
1088            .times(1..)
1089            .returning(|_, e| RetryResult::Continue(e));
1090
1091        let now = Instant::now();
1092        let policy = LimitedAttemptCount::custom(mock, 3);
1093        assert!(
1094            policy
1095                .on_error(
1096                    &idempotent_state(now).set_attempt_count(1_u32),
1097                    transient_error()
1098                )
1099                .is_continue()
1100        );
1101        assert!(
1102            policy
1103                .on_error(
1104                    &idempotent_state(now).set_attempt_count(2_u32),
1105                    transient_error()
1106                )
1107                .is_continue()
1108        );
1109        assert!(
1110            policy
1111                .on_error(
1112                    &idempotent_state(now).set_attempt_count(3_u32),
1113                    transient_error()
1114                )
1115                .is_exhausted()
1116        );
1117    }
1118
1119    #[test]
1120    fn test_limited_attempt_count_on_throttle_continue() {
1121        let mut mock = MockPolicy::new();
1122        mock.expect_on_throttle()
1123            .times(1..)
1124            .returning(|_, e| ThrottleResult::Continue(e));
1125
1126        let now = Instant::now();
1127        let policy = LimitedAttemptCount::custom(mock, 3);
1128        assert!(matches!(
1129            policy.on_throttle(
1130                &idempotent_state(now).set_attempt_count(2_u32),
1131                unavailable()
1132            ),
1133            ThrottleResult::Continue(_)
1134        ));
1135    }
1136
1137    #[test]
1138    fn test_limited_attempt_count_on_throttle_error() {
1139        let mut mock = MockPolicy::new();
1140        mock.expect_on_throttle()
1141            .times(1..)
1142            .returning(|_, e| ThrottleResult::Exhausted(e));
1143
1144        let now = Instant::now();
1145        let policy = LimitedAttemptCount::custom(mock, 3);
1146        assert!(matches!(
1147            policy.on_throttle(&idempotent_state(now), unavailable()),
1148            ThrottleResult::Exhausted(_)
1149        ));
1150    }
1151
1152    #[test]
1153    fn test_limited_attempt_count_on_throttle_exhausted() {
1154        let mut mock = MockPolicy::new();
1155        mock.expect_on_throttle()
1156            .times(1..)
1157            .returning(|_, e| ThrottleResult::Continue(e));
1158
1159        let now = Instant::now();
1160        let policy = LimitedAttemptCount::custom(mock, 3);
1161        assert!(matches!(
1162            policy.on_throttle(
1163                &idempotent_state(now).set_attempt_count(3_u32),
1164                unavailable()
1165            ),
1166            ThrottleResult::Exhausted(_)
1167        ));
1168
1169        let mut mock_zero = MockPolicy::new();
1170        mock_zero
1171            .expect_on_throttle()
1172            .times(1..)
1173            .returning(|_, e| ThrottleResult::Continue(e));
1174        let policy_zero = LimitedAttemptCount::custom(mock_zero, 0);
1175        assert!(matches!(
1176            policy_zero.on_throttle(&idempotent_state(now), unavailable()),
1177            ThrottleResult::Exhausted(_)
1178        ));
1179    }
1180
1181    #[test]
1182    fn test_limited_attempt_count_remaining_none() {
1183        let mut mock = MockPolicy::new();
1184        mock.expect_remaining_time().times(1).returning(|_| None);
1185        let policy = LimitedAttemptCount::custom(mock, 3);
1186
1187        let now = Instant::now();
1188        assert!(
1189            policy.remaining_time(&idempotent_state(now)).is_none(),
1190            "policy={policy:?} now={now:?}"
1191        );
1192    }
1193
1194    #[test]
1195    fn test_limited_attempt_count_remaining_some() {
1196        let mut mock = MockPolicy::new();
1197        mock.expect_remaining_time()
1198            .times(1)
1199            .returning(|_| Some(Duration::from_secs(123)));
1200        let policy = LimitedAttemptCount::custom(mock, 3);
1201
1202        let now = Instant::now();
1203        assert_eq!(
1204            policy.remaining_time(&idempotent_state(now)),
1205            Some(Duration::from_secs(123))
1206        );
1207    }
1208
1209    #[test]
1210    fn test_limited_attempt_count_inner_permanent() {
1211        let mut mock = MockPolicy::new();
1212        mock.expect_on_error()
1213            .times(2)
1214            .returning(|_, e| RetryResult::Permanent(e));
1215        let policy = LimitedAttemptCount::custom(mock, 2);
1216        let now = Instant::now();
1217
1218        let rf = policy.on_error(&non_idempotent_state(now), transient_error());
1219        assert!(rf.is_permanent());
1220
1221        let rf = policy.on_error(&non_idempotent_state(now), transient_error());
1222        assert!(rf.is_permanent());
1223    }
1224
1225    #[test]
1226    fn test_limited_attempt_count_inner_exhausted() {
1227        let mut mock = MockPolicy::new();
1228        mock.expect_on_error()
1229            .times(2)
1230            .returning(|_, e| RetryResult::Exhausted(e));
1231        let policy = LimitedAttemptCount::custom(mock, 2);
1232        let now = Instant::now();
1233
1234        let rf = policy.on_error(&non_idempotent_state(now), transient_error());
1235        assert!(rf.is_exhausted());
1236
1237        let rf = policy.on_error(&non_idempotent_state(now), transient_error());
1238        assert!(rf.is_exhausted());
1239    }
1240
1241    fn transient_error() -> Error {
1242        use crate::error::rpc::{Code, Status};
1243        Error::service(
1244            Status::default()
1245                .set_code(Code::Unavailable)
1246                .set_message("try-again"),
1247        )
1248    }
1249
1250    pub(crate) fn idempotent_state(now: Instant) -> RetryState {
1251        RetryState::new(true).set_start(now)
1252    }
1253
1254    pub(crate) fn non_idempotent_state(now: Instant) -> RetryState {
1255        RetryState::new(false).set_start(now)
1256    }
1257}
google_cloud_gax/retry_policy.rs

google_cloud_gax/
retry_policy.rs
