google_cloud_gax/retry_policy/
too_many_requests.rs

1// Copyright 2026 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
15use super::{Error, RetryPolicy, RetryResult, RetryState, ThrottleResult};
16use crate::error::rpc::Code;
17use std::time::Duration;
18
19/// A retry policy decorator that continues on `ResourceExhausted` or
20/// `TOO_MANY_REQUESTS`.
21///
22/// This policy returns [RetryResult::Continue] when the error is a
23/// `ResourceExhausted` (or `TOO_MANY_REQUESTS` if received from the HTTP layer).
24/// Otherwise it returns the result from the inner retry policy.
25///
26/// # Parameters
27/// * `P` - the inner retry policy.
28#[derive(Debug)]
29pub struct TooManyRequests<P>
30where
31    P: RetryPolicy,
32{
33    inner: P,
34}
35
36impl<P> TooManyRequests<P>
37where
38    P: RetryPolicy,
39{
40    /// Creates a new instance with a custom inner policy.
41    ///
42    /// # Example
43    /// ```
44    /// # use google_cloud_gax::retry_policy::{TooManyRequests, RetryPolicy};
45    /// use google_cloud_gax::retry_policy::Aip194Strict;
46    /// use google_cloud_gax::retry_state::RetryState;
47    /// let policy = TooManyRequests::new(Aip194Strict);
48    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(1_u32), too_many()).is_continue());
49    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(2_u32), permanent()).is_permanent());
50    ///
51    /// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
52    /// fn too_many() -> Error { Error::service(Status::default().set_code(Code::ResourceExhausted)) }
53    /// fn permanent() -> Error { Error::service(Status::default().set_code(Code::PermissionDenied)) }
54    /// ```
55    pub fn new(inner: P) -> Self {
56        Self { inner }
57    }
58
59    fn is_resource_exhausted(e: &Error) -> bool {
60        e.status()
61            .is_some_and(|s| s.code == Code::ResourceExhausted)
62    }
63
64    fn is_too_many_requests(e: &Error) -> bool {
65        e.http_status_code()
66            .is_some_and(|code| code == http::StatusCode::TOO_MANY_REQUESTS.as_u16())
67    }
68}
69
70impl<P> RetryPolicy for TooManyRequests<P>
71where
72    P: RetryPolicy,
73{
74    fn on_error(&self, state: &RetryState, error: Error) -> RetryResult {
75        if Self::is_resource_exhausted(&error) || Self::is_too_many_requests(&error) {
76            return RetryResult::Continue(error);
77        }
78        self.inner.on_error(state, error)
79    }
80
81    fn on_throttle(&self, state: &RetryState, error: Error) -> ThrottleResult {
82        self.inner.on_throttle(state, error)
83    }
84
85    fn remaining_time(&self, state: &RetryState) -> Option<Duration> {
86        self.inner.remaining_time(state)
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::super::tests::{MockPolicy, idempotent_state};
93    use super::*;
94    use crate::error::rpc::Status;
95    use crate::retry_policy::NeverRetry;
96    use std::time::Instant;
97
98    fn too_many_requests() -> Error {
99        Error::service(Status::default().set_code(Code::ResourceExhausted))
100    }
101
102    fn too_many_requests_http() -> Error {
103        Error::http(
104            http::StatusCode::TOO_MANY_REQUESTS.as_u16(),
105            http::HeaderMap::new(),
106            bytes::Bytes::new(),
107        )
108    }
109
110    fn permanent() -> Error {
111        Error::service(Status::default().set_code(Code::PermissionDenied))
112    }
113
114    fn transient() -> Error {
115        Error::service(Status::default().set_code(Code::Unavailable))
116    }
117
118    #[test]
119    fn on_error() {
120        let policy = TooManyRequests::new(NeverRetry);
121        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests());
122        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
123        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests_http());
124        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
125        let result = policy.on_error(&idempotent_state(Instant::now()), permanent());
126        assert!(matches!(result, RetryResult::Exhausted(_)), "{result:?}");
127    }
128
129    #[test]
130    fn ext() {
131        use super::super::RetryPolicyExt;
132        let policy = NeverRetry.continue_on_too_many_requests();
133        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests());
134        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
135        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests_http());
136        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
137        let result = policy.on_error(&idempotent_state(Instant::now()), permanent());
138        assert!(matches!(result, RetryResult::Exhausted(_)), "{result:?}");
139    }
140
141    #[test]
142    fn forwards() {
143        let mut mock = MockPolicy::new();
144        mock.expect_on_error()
145            .times(1..)
146            .returning(|_, e| RetryResult::Permanent(e));
147        mock.expect_on_throttle()
148            .times(1..)
149            .returning(|_, e| ThrottleResult::Exhausted(e));
150        mock.expect_remaining_time().times(1).returning(|_| None);
151
152        let policy = TooManyRequests::new(mock);
153        let result = policy.on_error(&idempotent_state(Instant::now()), transient());
154        assert!(matches!(result, RetryResult::Permanent(_)), "{result:?}");
155        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests());
156        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
157
158        let result = policy.on_throttle(&idempotent_state(Instant::now()), transient());
159        assert!(matches!(result, ThrottleResult::Exhausted(_)));
160        let result = policy.on_throttle(&idempotent_state(Instant::now()), too_many_requests());
161        assert!(matches!(result, ThrottleResult::Exhausted(_)));
162
163        let result = policy.remaining_time(&idempotent_state(Instant::now()));
164        assert!(result.is_none(), "{result:?}");
165    }
166}
google_cloud_gax/retry_policy/too_many_requests.rs

google_cloud_gax/retry_policy/
too_many_requests.rs
