feat(gax): retry policy decorator for 429 errors

src/gax/src/retry_policy.rs

@@ -53,13 +53,17 @@
 //!
 //! [idempotent]: https://en.wikipedia.org/wiki/Idempotence
 
+mod too_many_requests;
+
 use crate::error::Error;
 use crate::retry_result::RetryResult;
 use crate::retry_state::RetryState;
 use crate::throttle_result::ThrottleResult;
 use std::sync::Arc;
 use std::time::Duration;
 
+pub use too_many_requests::TooManyRequests;
+
 /// Determines how errors are handled in the retry loop.
 ///
 /// Implementations of this trait determine if errors are retryable, and for how
@@ -185,6 +189,42 @@ pub trait RetryPolicyExt: RetryPolicy + Sized {
     fn with_attempt_limit(self, maximum_attempts: u32) -> LimitedAttemptCount<Self> {
         LimitedAttemptCount::custom(self, maximum_attempts)
     }
+
+    /// Decorate a [RetryPolicy] to continue on certain status codes.
+    ///
+    /// This policy decorates an inner policy and retries any errors with HTTP
+    /// status code "429 - TOO_MANY_REQUESTS" **or** where the service returns
+    /// an error with code [ResourceExhausted].
+    ///
+    /// For other errors it returns the same value as the inner policy.
+    ///
+    /// Note that [ResourceExhausted] is ambiguous and may cause problems with
+    /// some services. The code is used for both "too many requests"  and for
+    /// "quota exceeded" problems. If the quota in question is some kind of rate
+    /// limit, then using this policy may be helpful. If the quota is not a rate
+    /// limit, then this retry policy may needlessly send the same RPC multiple
+    /// times.
+    ///
+    /// You should consult the documentation for the service and RPC in question
+    /// before using this policy.
+    ///
+    /// # Example
+    /// ```
+    /// use google_cloud_gax::retry_policy::{Aip194Strict, RetryPolicy, RetryPolicyExt};
+    /// use google_cloud_gax::retry_state::RetryState;
+    /// let policy = Aip194Strict;
+    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(0_u32), too_many_requests()).is_permanent());
+    /// let policy = Aip194Strict.continue_on_too_many_requests();
+    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(0_u32), too_many_requests()).is_continue());
+    ///
+    /// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
+    /// fn too_many_requests() -> Error { Error::service(Status::default().set_code(Code::ResourceExhausted)) }
+    /// ```
+    ///
+    /// [ResourceExhausted]: crate::error::rpc::Code::ResourceExhausted
+    fn continue_on_too_many_requests(self) -> TooManyRequests<Self> {
+        TooManyRequests::new(self)
+    }
 }
 
 impl<T: RetryPolicy> RetryPolicyExt for T {}
@@ -543,7 +583,7 @@ where
 }
 
 #[cfg(test)]
-mod tests {
+pub mod tests {
     use super::*;
     use http::HeaderMap;
     use std::error::Error as StdError;
@@ -795,7 +835,7 @@ mod tests {
 
     mockall::mock! {
         #[derive(Debug)]
-        Policy {}
+        pub(crate) Policy {}
         impl RetryPolicy for Policy {
             fn on_error(&self, state: &RetryState, error: Error) -> RetryResult;
             fn on_throttle(&self, state: &RetryState, error: Error) -> ThrottleResult;
@@ -1122,11 +1162,11 @@ mod tests {
         )
     }
 
-    fn idempotent_state(now: Instant) -> RetryState {
+    pub(crate) fn idempotent_state(now: Instant) -> RetryState {
         RetryState::new(true).set_start(now)
     }
 
-    fn non_idempotent_state(now: Instant) -> RetryState {
+    pub(crate) fn non_idempotent_state(now: Instant) -> RetryState {
         RetryState::new(false).set_start(now)
     }
 }

src/gax/src/retry_policy/too_many_requests.rs

@@ -0,0 +1,166 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::{Error, RetryPolicy, RetryResult, RetryState, ThrottleResult};
+use crate::error::rpc::Code;
+use std::time::Duration;
+
+/// A retry policy decorator that continues on `ResourceExhausted` or
+/// `TOO_MANY_REQUESTS`.
+///
+/// This policy returns [RetryResult::Continue] when the error is a
+/// `ResourceExhausted` (or `TOO_MANY_REQUESTS` if received from the HTTP layer).
+/// Otherwise it returns the result from the inner retry policy.
+///
+/// # Parameters
+/// * `P` - the inner retry policy.
+#[derive(Debug)]
+pub struct TooManyRequests<P>
+where
+    P: RetryPolicy,
+{
+    inner: P,
+}
+
+impl<P> TooManyRequests<P>
+where
+    P: RetryPolicy,
+{
+    /// Creates a new instance with a custom inner policy.
+    ///
+    /// # Example
+    /// ```
+    /// # use google_cloud_gax::retry_policy::{TooManyRequests, RetryPolicy};
+    /// use google_cloud_gax::retry_policy::Aip194Strict;
+    /// use google_cloud_gax::retry_state::RetryState;
+    /// let policy = TooManyRequests::new(Aip194Strict);
+    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(1_u32), too_many()).is_continue());
+    /// assert!(policy.on_error(&RetryState::new(true).set_attempt_count(2_u32), permanent()).is_permanent());
+    ///
+    /// use google_cloud_gax::error::{Error, rpc::Code, rpc::Status};
+    /// fn too_many() -> Error { Error::service(Status::default().set_code(Code::ResourceExhausted)) }
+    /// fn permanent() -> Error { Error::service(Status::default().set_code(Code::PermissionDenied)) }
+    /// ```
+    pub fn new(inner: P) -> Self {
+        Self { inner }
+    }
+
+    fn is_resource_exhausted(e: &Error) -> bool {
+        e.status()
+            .is_some_and(|s| s.code == Code::ResourceExhausted)
+    }
+
+    fn is_too_many_requests(e: &Error) -> bool {
+        e.http_status_code()
+            .is_some_and(|code| code == http::StatusCode::TOO_MANY_REQUESTS.as_u16())
+    }
+}
+
+impl<P> RetryPolicy for TooManyRequests<P>
+where
+    P: RetryPolicy,
+{
+    fn on_error(&self, state: &RetryState, error: Error) -> RetryResult {
+        if Self::is_resource_exhausted(&error) || Self::is_too_many_requests(&error) {
+            return RetryResult::Continue(error);
+        }
+        self.inner.on_error(state, error)
+    }
+
+    fn on_throttle(&self, state: &RetryState, error: Error) -> ThrottleResult {
+        self.inner.on_throttle(state, error)
+    }
+
+    fn remaining_time(&self, state: &RetryState) -> Option<Duration> {
+        self.inner.remaining_time(state)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::super::tests::{MockPolicy, idempotent_state};
+    use super::*;
+    use crate::error::rpc::Status;
+    use crate::retry_policy::NeverRetry;
+    use std::time::Instant;
+
+    fn too_many_requests() -> Error {
+        Error::service(Status::default().set_code(Code::ResourceExhausted))
+    }
+
+    fn too_many_requests_http() -> Error {
+        Error::http(
+            http::StatusCode::TOO_MANY_REQUESTS.as_u16(),
+            http::HeaderMap::new(),
+            bytes::Bytes::new(),
+        )
+    }
+
+    fn permanent() -> Error {
+        Error::service(Status::default().set_code(Code::PermissionDenied))
+    }
+
+    fn transient() -> Error {
+        Error::service(Status::default().set_code(Code::Unavailable))
+    }
+
+    #[test]
+    fn on_error() {
+        let policy = TooManyRequests::new(NeverRetry);
+        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests());
+        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
+        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests_http());
+        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
+        let result = policy.on_error(&idempotent_state(Instant::now()), permanent());
+        assert!(matches!(result, RetryResult::Exhausted(_)), "{result:?}");
+    }
+
+    #[test]
+    fn ext() {
+        use super::super::RetryPolicyExt;
+        let policy = NeverRetry.continue_on_too_many_requests();
+        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests());
+        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
+        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests_http());
+        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
+        let result = policy.on_error(&idempotent_state(Instant::now()), permanent());
+        assert!(matches!(result, RetryResult::Exhausted(_)), "{result:?}");
+    }
+
+    #[test]
+    fn forwards() {
+        let mut mock = MockPolicy::new();
+        mock.expect_on_error()
+            .times(1..)
+            .returning(|_, e| RetryResult::Permanent(e));
+        mock.expect_on_throttle()
+            .times(1..)
+            .returning(|_, e| ThrottleResult::Exhausted(e));
+        mock.expect_remaining_time().times(1).returning(|_| None);
+
+        let policy = TooManyRequests::new(mock);
+        let result = policy.on_error(&idempotent_state(Instant::now()), transient());
+        assert!(matches!(result, RetryResult::Permanent(_)), "{result:?}");
+        let result = policy.on_error(&idempotent_state(Instant::now()), too_many_requests());
+        assert!(matches!(result, RetryResult::Continue(_)), "{result:?}");
+
+        let result = policy.on_throttle(&idempotent_state(Instant::now()), transient());
+        assert!(matches!(result, ThrottleResult::Exhausted(_)));
+        let result = policy.on_throttle(&idempotent_state(Instant::now()), too_many_requests());
+        assert!(matches!(result, ThrottleResult::Exhausted(_)));
+
+        let result = policy.remaining_time(&idempotent_state(Instant::now()));
+        assert!(result.is_none(), "{result:?}");
+    }
+}