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        // The retry loop only calls `on_throttle()` if the policy has not
608        // been exhausted.
609        assert!(state.attempt_count < self.maximum_attempts);
610        self.inner.on_throttle(state, error)
611    }
612
613    fn remaining_time(&self, state: &RetryState) -> Option<Duration> {
614        self.inner.remaining_time(state)
615    }
616}
617
618#[cfg(test)]
619pub(crate) mod tests {
620    use super::*;
621    use http::HeaderMap;
622    use std::error::Error as StdError;
623    use std::time::Instant;
624
625    // Verify `RetryPolicyArg` can be converted from the desired types.
626    #[test]
627    fn retry_policy_arg() {
628        let policy = LimitedAttemptCount::new(3);
629        let _ = RetryPolicyArg::from(policy);
630
631        let policy: Arc<dyn RetryPolicy> = Arc::new(LimitedAttemptCount::new(3));
632        let _ = RetryPolicyArg::from(policy);
633    }
634
635    #[test]
636    fn aip194_strict() {
637        let p = Aip194Strict;
638
639        let now = Instant::now();
640        assert!(
641            p.on_error(&idempotent_state(now), unavailable())
642                .is_continue()
643        );
644        assert!(
645            p.on_error(&non_idempotent_state(now), unavailable())
646                .is_permanent()
647        );
648        assert!(matches!(
649            p.on_throttle(&idempotent_state(now), unavailable()),
650            ThrottleResult::Continue(_)
651        ));
652
653        assert!(
654            p.on_error(&idempotent_state(now), unknown_and_503())
655                .is_continue()
656        );
657        assert!(
658            p.on_error(&non_idempotent_state(now), unknown_and_503())
659                .is_permanent()
660        );
661        assert!(matches!(
662            p.on_throttle(&idempotent_state(now), unknown_and_503()),
663            ThrottleResult::Continue(_)
664        ));
665
666        assert!(
667            p.on_error(&idempotent_state(now), permission_denied())
668                .is_permanent()
669        );
670        assert!(
671            p.on_error(&non_idempotent_state(now), permission_denied())
672                .is_permanent()
673        );
674
675        assert!(
676            p.on_error(&idempotent_state(now), http_unavailable())
677                .is_continue()
678        );
679        assert!(
680            p.on_error(&non_idempotent_state(now), http_unavailable())
681                .is_permanent()
682        );
683        assert!(matches!(
684            p.on_throttle(&idempotent_state(now), http_unavailable()),
685            ThrottleResult::Continue(_)
686        ));
687
688        assert!(
689            p.on_error(&idempotent_state(now), http_permission_denied())
690                .is_permanent()
691        );
692        assert!(
693            p.on_error(&non_idempotent_state(now), http_permission_denied())
694                .is_permanent()
695        );
696
697        assert!(
698            p.on_error(&idempotent_state(now), Error::io("err".to_string()))
699                .is_continue()
700        );
701        assert!(
702            p.on_error(&non_idempotent_state(now), Error::io("err".to_string()))
703                .is_permanent()
704        );
705
706        assert!(
707            p.on_error(&idempotent_state(now), pre_rpc_transient())
708                .is_continue()
709        );
710        assert!(
711            p.on_error(&non_idempotent_state(now), pre_rpc_transient())
712                .is_continue()
713        );
714
715        assert!(
716            p.on_error(&idempotent_state(now), Error::ser("err"))
717                .is_permanent()
718        );
719        assert!(
720            p.on_error(&non_idempotent_state(now), Error::ser("err"))
721                .is_permanent()
722        );
723        assert!(
724            p.on_error(&idempotent_state(now), Error::deser("err"))
725                .is_permanent()
726        );
727        assert!(
728            p.on_error(&non_idempotent_state(now), Error::deser("err"))
729                .is_permanent()
730        );
731
732        assert!(
733            p.remaining_time(&idempotent_state(now)).is_none(),
734            "p={p:?}, now={now:?}"
735        );
736    }
737
738    #[test]
739    fn always_retry() {
740        let p = AlwaysRetry;
741
742        let now = Instant::now();
743        assert!(
744            p.remaining_time(&idempotent_state(now)).is_none(),
745            "p={p:?}, now={now:?}"
746        );
747        assert!(
748            p.on_error(&idempotent_state(now), http_unavailable())
749                .is_continue()
750        );
751        assert!(
752            p.on_error(&non_idempotent_state(now), http_unavailable())
753                .is_continue()
754        );
755        assert!(matches!(
756            p.on_throttle(&idempotent_state(now), http_unavailable()),
757            ThrottleResult::Continue(_)
758        ));
759
760        assert!(
761            p.on_error(&idempotent_state(now), unavailable())
762                .is_continue()
763        );
764        assert!(
765            p.on_error(&non_idempotent_state(now), unavailable())
766                .is_continue()
767        );
768    }
769
770    #[test_case::test_case(true, Error::io("err"))]
771    #[test_case::test_case(true, pre_rpc_transient())]
772    #[test_case::test_case(true, Error::ser("err"))]
773    #[test_case::test_case(false, Error::io("err"))]
774    #[test_case::test_case(false, pre_rpc_transient())]
775    #[test_case::test_case(false, Error::ser("err"))]
776    fn always_retry_error_kind(idempotent: bool, error: Error) {
777        let p = AlwaysRetry;
778        let now = Instant::now();
779        let state = if idempotent {
780            idempotent_state(now)
781        } else {
782            non_idempotent_state(now)
783        };
784        assert!(p.on_error(&state, error).is_continue());
785    }
786
787    #[test]
788    fn never_retry() {
789        let p = NeverRetry;
790
791        let now = Instant::now();
792        assert!(
793            p.remaining_time(&idempotent_state(now)).is_none(),
794            "p={p:?}, now={now:?}"
795        );
796        assert!(
797            p.on_error(&idempotent_state(now), http_unavailable())
798                .is_exhausted()
799        );
800        assert!(
801            p.on_error(&non_idempotent_state(now), http_unavailable())
802                .is_exhausted()
803        );
804        assert!(matches!(
805            p.on_throttle(&idempotent_state(now), http_unavailable()),
806            ThrottleResult::Continue(_)
807        ));
808
809        assert!(
810            p.on_error(&idempotent_state(now), unavailable())
811                .is_exhausted()
812        );
813        assert!(
814            p.on_error(&non_idempotent_state(now), unavailable())
815                .is_exhausted()
816        );
817
818        assert!(
819            p.on_error(&idempotent_state(now), http_permission_denied())
820                .is_exhausted()
821        );
822        assert!(
823            p.on_error(&non_idempotent_state(now), http_permission_denied())
824                .is_exhausted()
825        );
826    }
827
828    #[test_case::test_case(true, Error::io("err"))]
829    #[test_case::test_case(true, pre_rpc_transient())]
830    #[test_case::test_case(true, Error::ser("err"))]
831    #[test_case::test_case(false, Error::io("err"))]
832    #[test_case::test_case(false, pre_rpc_transient())]
833    #[test_case::test_case(false, Error::ser("err"))]
834    fn never_retry_error_kind(idempotent: bool, error: Error) {
835        let p = NeverRetry;
836        let now = Instant::now();
837        let state = if idempotent {
838            idempotent_state(now)
839        } else {
840            non_idempotent_state(now)
841        };
842        assert!(p.on_error(&state, error).is_exhausted());
843    }
844
845    fn pre_rpc_transient() -> Error {
846        use crate::error::CredentialsError;
847        Error::authentication(CredentialsError::from_msg(true, "err"))
848    }
849
850    fn http_unavailable() -> Error {
851        Error::http(
852            503_u16,
853            HeaderMap::new(),
854            bytes::Bytes::from_owner("SERVICE UNAVAILABLE".to_string()),
855        )
856    }
857
858    fn http_permission_denied() -> Error {
859        Error::http(
860            403_u16,
861            HeaderMap::new(),
862            bytes::Bytes::from_owner("PERMISSION DENIED".to_string()),
863        )
864    }
865
866    fn unavailable() -> Error {
867        use crate::error::rpc::Code;
868        let status = crate::error::rpc::Status::default()
869            .set_code(Code::Unavailable)
870            .set_message("UNAVAILABLE");
871        Error::service(status)
872    }
873
874    fn unknown_and_503() -> Error {
875        use crate::error::rpc::Code;
876        let status = crate::error::rpc::Status::default()
877            .set_code(Code::Unknown)
878            .set_message("UNAVAILABLE");
879        Error::service_full(status, Some(503), None, Some("source error".into()))
880    }
881
882    fn permission_denied() -> Error {
883        use crate::error::rpc::Code;
884        let status = crate::error::rpc::Status::default()
885            .set_code(Code::PermissionDenied)
886            .set_message("PERMISSION_DENIED");
887        Error::service(status)
888    }
889
890    mockall::mock! {
891        #[derive(Debug)]
892        pub(crate) Policy {}
893        impl RetryPolicy for Policy {
894            fn on_error(&self, state: &RetryState, error: Error) -> RetryResult;
895            fn on_throttle(&self, state: &RetryState, error: Error) -> ThrottleResult;
896            fn remaining_time(&self, state: &RetryState) -> Option<Duration>;
897        }
898    }
899
900    #[test]
901    fn limited_elapsed_time_error() {
902        let limit = Duration::from_secs(123) + Duration::from_millis(567);
903        let err = LimitedElapsedTimeError::new(limit, unavailable());
904        assert_eq!(err.maximum_duration(), limit);
905        let fmt = err.to_string();
906        assert!(fmt.contains("123.567s"), "display={fmt}, debug={err:?}");
907        assert!(err.source().is_some(), "{err:?}");
908    }
909
910    #[test]
911    fn test_limited_time_forwards() {
912        let mut mock = MockPolicy::new();
913        mock.expect_on_error()
914            .times(1..)
915            .returning(|_, e| RetryResult::Continue(e));
916        mock.expect_on_throttle()
917            .times(1..)
918            .returning(|_, e| ThrottleResult::Continue(e));
919        mock.expect_remaining_time().times(1).returning(|_| None);
920
921        let now = Instant::now();
922        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
923        let rf = policy.on_error(&idempotent_state(now), transient_error());
924        assert!(rf.is_continue());
925
926        let rt = policy.remaining_time(&idempotent_state(now));
927        assert!(rt.is_some(), "policy={policy:?}, now={now:?}");
928
929        let e = policy.on_throttle(&idempotent_state(now), transient_error());
930        assert!(matches!(e, ThrottleResult::Continue(_)));
931    }
932
933    #[test]
934    fn test_limited_time_on_throttle_continue() {
935        let mut mock = MockPolicy::new();
936        mock.expect_on_throttle()
937            .times(1..)
938            .returning(|_, e| ThrottleResult::Continue(e));
939
940        let now = Instant::now();
941        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
942
943        // Before the policy expires the inner result is returned verbatim.
944        let rf = policy.on_throttle(
945            &idempotent_state(now - Duration::from_secs(50)),
946            unavailable(),
947        );
948        assert!(matches!(rf, ThrottleResult::Continue(_)), "{rf:?}");
949
950        // After the policy expires the innter result is always "exhausted".
951        let rf = policy.on_throttle(
952            &idempotent_state(now - Duration::from_secs(70)),
953            unavailable(),
954        );
955        assert!(matches!(rf, ThrottleResult::Exhausted(_)), "{rf:?}");
956    }
957
958    #[test]
959    fn test_limited_time_on_throttle_exhausted() {
960        let mut mock = MockPolicy::new();
961        mock.expect_on_throttle()
962            .times(1..)
963            .returning(|_, e| ThrottleResult::Exhausted(e));
964
965        let now = Instant::now();
966        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
967
968        // Before the policy expires the inner result is returned verbatim.
969        let rf = policy.on_throttle(
970            &idempotent_state(now - Duration::from_secs(50)),
971            unavailable(),
972        );
973        assert!(matches!(rf, ThrottleResult::Exhausted(_)), "{rf:?}");
974    }
975
976    #[test]
977    fn test_limited_time_inner_continues() {
978        let mut mock = MockPolicy::new();
979        mock.expect_on_error()
980            .times(1..)
981            .returning(|_, e| RetryResult::Continue(e));
982
983        let now = Instant::now();
984        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
985        let rf = policy.on_error(
986            &idempotent_state(now - Duration::from_secs(10)),
987            transient_error(),
988        );
989        assert!(rf.is_continue());
990
991        let rf = policy.on_error(
992            &idempotent_state(now - Duration::from_secs(70)),
993            transient_error(),
994        );
995        assert!(rf.is_exhausted());
996    }
997
998    #[test]
999    fn test_limited_time_inner_permanent() {
1000        let mut mock = MockPolicy::new();
1001        mock.expect_on_error()
1002            .times(2)
1003            .returning(|_, e| RetryResult::Permanent(e));
1004
1005        let now = Instant::now();
1006        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1007
1008        let rf = policy.on_error(
1009            &non_idempotent_state(now - Duration::from_secs(10)),
1010            transient_error(),
1011        );
1012        assert!(rf.is_permanent());
1013
1014        let rf = policy.on_error(
1015            &non_idempotent_state(now + Duration::from_secs(10)),
1016            transient_error(),
1017        );
1018        assert!(rf.is_permanent());
1019    }
1020
1021    #[test]
1022    fn test_limited_time_inner_exhausted() {
1023        let mut mock = MockPolicy::new();
1024        mock.expect_on_error()
1025            .times(2)
1026            .returning(|_, e| RetryResult::Exhausted(e));
1027
1028        let now = Instant::now();
1029        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1030
1031        let rf = policy.on_error(
1032            &non_idempotent_state(now - Duration::from_secs(10)),
1033            transient_error(),
1034        );
1035        assert!(rf.is_exhausted());
1036
1037        let rf = policy.on_error(
1038            &non_idempotent_state(now + Duration::from_secs(10)),
1039            transient_error(),
1040        );
1041        assert!(rf.is_exhausted());
1042    }
1043
1044    #[test]
1045    fn test_limited_time_remaining_inner_longer() {
1046        let mut mock = MockPolicy::new();
1047        mock.expect_remaining_time()
1048            .times(1)
1049            .returning(|_| Some(Duration::from_secs(30)));
1050
1051        let now = Instant::now();
1052        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1053
1054        let remaining = policy.remaining_time(&idempotent_state(now - Duration::from_secs(55)));
1055        assert!(remaining <= Some(Duration::from_secs(5)), "{remaining:?}");
1056    }
1057
1058    #[test]
1059    fn test_limited_time_remaining_inner_shorter() {
1060        let mut mock = MockPolicy::new();
1061        mock.expect_remaining_time()
1062            .times(1)
1063            .returning(|_| Some(Duration::from_secs(5)));
1064        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1065
1066        let now = Instant::now();
1067        let remaining = policy.remaining_time(&idempotent_state(now - Duration::from_secs(5)));
1068        assert!(remaining <= Some(Duration::from_secs(10)), "{remaining:?}");
1069    }
1070
1071    #[test]
1072    fn test_limited_time_remaining_inner_is_none() {
1073        let mut mock = MockPolicy::new();
1074        mock.expect_remaining_time().times(1).returning(|_| None);
1075        let policy = LimitedElapsedTime::custom(mock, Duration::from_secs(60));
1076
1077        let now = Instant::now();
1078        let remaining = policy.remaining_time(&idempotent_state(now - Duration::from_secs(50)));
1079        assert!(remaining <= Some(Duration::from_secs(10)), "{remaining:?}");
1080    }
1081
1082    #[test]
1083    fn test_limited_attempt_count_on_error() {
1084        let mut mock = MockPolicy::new();
1085        mock.expect_on_error()
1086            .times(1..)
1087            .returning(|_, e| RetryResult::Continue(e));
1088
1089        let now = Instant::now();
1090        let policy = LimitedAttemptCount::custom(mock, 3);
1091        assert!(
1092            policy
1093                .on_error(
1094                    &idempotent_state(now).set_attempt_count(1_u32),
1095                    transient_error()
1096                )
1097                .is_continue()
1098        );
1099        assert!(
1100            policy
1101                .on_error(
1102                    &idempotent_state(now).set_attempt_count(2_u32),
1103                    transient_error()
1104                )
1105                .is_continue()
1106        );
1107        assert!(
1108            policy
1109                .on_error(
1110                    &idempotent_state(now).set_attempt_count(3_u32),
1111                    transient_error()
1112                )
1113                .is_exhausted()
1114        );
1115    }
1116
1117    #[test]
1118    fn test_limited_attempt_count_on_throttle_continue() {
1119        let mut mock = MockPolicy::new();
1120        mock.expect_on_throttle()
1121            .times(1..)
1122            .returning(|_, e| ThrottleResult::Continue(e));
1123
1124        let now = Instant::now();
1125        let policy = LimitedAttemptCount::custom(mock, 3);
1126        assert!(matches!(
1127            policy.on_throttle(
1128                &idempotent_state(now).set_attempt_count(2_u32),
1129                unavailable()
1130            ),
1131            ThrottleResult::Continue(_)
1132        ));
1133    }
1134
1135    #[test]
1136    fn test_limited_attempt_count_on_throttle_error() {
1137        let mut mock = MockPolicy::new();
1138        mock.expect_on_throttle()
1139            .times(1..)
1140            .returning(|_, e| ThrottleResult::Exhausted(e));
1141
1142        let now = Instant::now();
1143        let policy = LimitedAttemptCount::custom(mock, 3);
1144        assert!(matches!(
1145            policy.on_throttle(&idempotent_state(now), unavailable()),
1146            ThrottleResult::Exhausted(_)
1147        ));
1148    }
1149
1150    #[test]
1151    fn test_limited_attempt_count_remaining_none() {
1152        let mut mock = MockPolicy::new();
1153        mock.expect_remaining_time().times(1).returning(|_| None);
1154        let policy = LimitedAttemptCount::custom(mock, 3);
1155
1156        let now = Instant::now();
1157        assert!(
1158            policy.remaining_time(&idempotent_state(now)).is_none(),
1159            "policy={policy:?} now={now:?}"
1160        );
1161    }
1162
1163    #[test]
1164    fn test_limited_attempt_count_remaining_some() {
1165        let mut mock = MockPolicy::new();
1166        mock.expect_remaining_time()
1167            .times(1)
1168            .returning(|_| Some(Duration::from_secs(123)));
1169        let policy = LimitedAttemptCount::custom(mock, 3);
1170
1171        let now = Instant::now();
1172        assert_eq!(
1173            policy.remaining_time(&idempotent_state(now)),
1174            Some(Duration::from_secs(123))
1175        );
1176    }
1177
1178    #[test]
1179    fn test_limited_attempt_count_inner_permanent() {
1180        let mut mock = MockPolicy::new();
1181        mock.expect_on_error()
1182            .times(2)
1183            .returning(|_, e| RetryResult::Permanent(e));
1184        let policy = LimitedAttemptCount::custom(mock, 2);
1185        let now = Instant::now();
1186
1187        let rf = policy.on_error(&non_idempotent_state(now), transient_error());
1188        assert!(rf.is_permanent());
1189
1190        let rf = policy.on_error(&non_idempotent_state(now), transient_error());
1191        assert!(rf.is_permanent());
1192    }
1193
1194    #[test]
1195    fn test_limited_attempt_count_inner_exhausted() {
1196        let mut mock = MockPolicy::new();
1197        mock.expect_on_error()
1198            .times(2)
1199            .returning(|_, e| RetryResult::Exhausted(e));
1200        let policy = LimitedAttemptCount::custom(mock, 2);
1201        let now = Instant::now();
1202
1203        let rf = policy.on_error(&non_idempotent_state(now), transient_error());
1204        assert!(rf.is_exhausted());
1205
1206        let rf = policy.on_error(&non_idempotent_state(now), transient_error());
1207        assert!(rf.is_exhausted());
1208    }
1209
1210    fn transient_error() -> Error {
1211        use crate::error::rpc::{Code, Status};
1212        Error::service(
1213            Status::default()
1214                .set_code(Code::Unavailable)
1215                .set_message("try-again"),
1216        )
1217    }
1218
1219    pub(crate) fn idempotent_state(now: Instant) -> RetryState {
1220        RetryState::new(true).set_start(now)
1221    }
1222
1223    pub(crate) fn non_idempotent_state(now: Instant) -> RetryState {
1224        RetryState::new(false).set_start(now)
1225    }
1226}
google_cloud_gax/retry_policy.rs

google_cloud_gax/
retry_policy.rs
