From 4c6f2028548d78ce39e3ed9545467c76ab2d92a5 Mon Sep 17 00:00:00 2001 From: Joshua Tan Date: Fri, 10 Jul 2026 05:36:01 +0000 Subject: [PATCH] refactor(gax): extract transport policies to grpc_helpers --- src/gax-internal/src/grpc.rs | 92 ++------ src/gax-internal/src/grpc/grpc_helpers.rs | 264 +++++++++++++++++++++- 2 files changed, 279 insertions(+), 77 deletions(-) diff --git a/src/gax-internal/src/grpc.rs b/src/gax-internal/src/grpc.rs index 21a510b946..a90e2cccc8 100644 --- a/src/gax-internal/src/grpc.rs +++ b/src/gax-internal/src/grpc.rs @@ -31,23 +31,15 @@ use from_status::to_gax_error; use futures::TryFutureExt; use google_cloud_auth::credentials::Credentials; use google_cloud_gax::Result; -use google_cloud_gax::backoff_policy::BackoffPolicy; use google_cloud_gax::client_builder::Error as BuilderError; use google_cloud_gax::client_builder::Result as ClientBuilderResult; use google_cloud_gax::error::Error; -use google_cloud_gax::exponential_backoff::ExponentialBackoff; use google_cloud_gax::options::RequestOptions; use google_cloud_gax::polling_backoff_policy::PollingBackoffPolicy; -use google_cloud_gax::polling_error_policy::{ - Aip194Strict as PollingAip194Strict, PollingErrorPolicy, -}; +use google_cloud_gax::polling_error_policy::PollingErrorPolicy; use google_cloud_gax::response::{Parts, Response}; use google_cloud_gax::retry_loop_internal::retry_loop; -use google_cloud_gax::retry_policy::{ - Aip194Strict as RetryAip194Strict, RetryPolicy, RetryPolicyExt as _, -}; -use google_cloud_gax::retry_throttler::SharedRetryThrottler; -use grpc_helpers::{add_auth_headers, make_credentials, make_headers}; +use grpc_helpers::{TransportPolicies, add_auth_headers, make_credentials, make_headers}; use http::HeaderMap; use opentelemetry_semantic_conventions::{attribute as otel_attr, trace as otel_trace}; use std::sync::Arc; @@ -75,12 +67,7 @@ pub struct Client { metric: crate::observability::TransportMetric, tracing_attributes: Option, credentials: Credentials, - retry_policy: Arc, - backoff_policy: Arc, - retry_throttler: SharedRetryThrottler, - polling_error_policy: Arc, - polling_backoff_policy: Arc, - attempt_timeout: Option, + transport_policies: TransportPolicies, } impl Client { @@ -126,25 +113,7 @@ impl Client { metric: crate::observability::TransportMetric::new(instrumentation), tracing_attributes, credentials, - retry_policy: config.retry_policy.clone().unwrap_or_else(|| { - Arc::new( - RetryAip194Strict - .with_attempt_limit(10) - .with_time_limit(Duration::from_secs(60)), - ) - }), - backoff_policy: config - .backoff_policy - .clone() - .unwrap_or_else(|| Arc::new(ExponentialBackoff::default())), - retry_throttler: config.retry_throttler, - polling_error_policy: config - .polling_error_policy - .unwrap_or_else(|| Arc::new(PollingAip194Strict)), - polling_backoff_policy: config - .polling_backoff_policy - .unwrap_or_else(|| Arc::new(ExponentialBackoff::default())), - attempt_timeout: config.attempt_timeout, + transport_policies: TransportPolicies::from_config(&config), }) } @@ -281,9 +250,11 @@ impl Client { let headers = add_auth_headers(headers, &self.credentials).await?; let metadata = tonic::MetadataMap::from_headers(headers); let mut request = ::tonic::Request::from_parts(metadata, extensions, request); - if let Some(timeout) = - crate::options::resolve_effective_timeout(&options, self.attempt_timeout, None) - { + if let Some(timeout) = crate::options::resolve_effective_timeout( + &options, + self.transport_policies.attempt_timeout(), + None, + ) { request.set_timeout(timeout); } let codec = tonic_prost::ProstCodec::::default(); @@ -318,9 +289,9 @@ impl Client { Response: prost::Message + Default + 'static, { let idempotent = options.idempotent().unwrap_or(false); - let retry_throttler = self.get_retry_throttler(&options); - let retry_policy = self.get_retry_policy(&options); - let backoff_policy = self.get_backoff_policy(&options); + let retry_throttler = self.transport_policies.get_retry_throttler(&options); + let retry_policy = self.transport_policies.get_retry_policy(&options); + let backoff_policy = self.transport_policies.get_backoff_policy(&options); let this = self.clone(); let mut prior_attempt_count: i64 = 0; let inner = async move |remaining_time: Option| { @@ -415,9 +386,11 @@ impl Client { let metadata = tonic::MetadataMap::from_headers(headers); let mut request = ::tonic::Request::from_parts(metadata, extensions, request); - if let Some(timeout) = - crate::options::resolve_effective_timeout(options, self.attempt_timeout, remaining_time) - { + if let Some(timeout) = crate::options::resolve_effective_timeout( + options, + self.transport_policies.attempt_timeout(), + remaining_time, + ) { request.set_timeout(timeout); } let codec = tonic_prost::ProstCodec::::default(); @@ -528,45 +501,18 @@ impl Client { Ok(endpoint) } - fn get_retry_policy(&self, options: &RequestOptions) -> Arc { - options - .retry_policy() - .clone() - .unwrap_or_else(|| self.retry_policy.clone()) - } - - pub(crate) fn get_backoff_policy(&self, options: &RequestOptions) -> Arc { - options - .backoff_policy() - .clone() - .unwrap_or_else(|| self.backoff_policy.clone()) - } - - pub(crate) fn get_retry_throttler(&self, options: &RequestOptions) -> SharedRetryThrottler { - options - .retry_throttler() - .clone() - .unwrap_or_else(|| self.retry_throttler.clone()) - } - pub fn get_polling_error_policy( &self, options: &RequestOptions, ) -> Arc { - options - .polling_error_policy() - .clone() - .unwrap_or_else(|| self.polling_error_policy.clone()) + self.transport_policies.get_polling_error_policy(options) } pub fn get_polling_backoff_policy( &self, options: &RequestOptions, ) -> Arc { - options - .polling_backoff_policy() - .clone() - .unwrap_or_else(|| self.polling_backoff_policy.clone()) + self.transport_policies.get_polling_backoff_policy(options) } } diff --git a/src/gax-internal/src/grpc/grpc_helpers.rs b/src/gax-internal/src/grpc/grpc_helpers.rs index 8b80345ecf..03145a5588 100644 --- a/src/gax-internal/src/grpc/grpc_helpers.rs +++ b/src/gax-internal/src/grpc/grpc_helpers.rs @@ -12,15 +12,28 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::options::ClientConfig; use google_cloud_auth::credentials::{ Builder as CredentialsBuilder, CacheableResource, Credentials, }; use google_cloud_gax::Result; +use google_cloud_gax::backoff_policy::BackoffPolicy; use google_cloud_gax::client_builder::{Error as BuilderError, Result as ClientBuilderResult}; use google_cloud_gax::error::Error; +use google_cloud_gax::exponential_backoff::ExponentialBackoff; use google_cloud_gax::options::RequestOptions; use google_cloud_gax::options::internal::RequestOptionsExt as _; +use google_cloud_gax::polling_backoff_policy::PollingBackoffPolicy; +use google_cloud_gax::polling_error_policy::{ + Aip194Strict as PollingAip194Strict, PollingErrorPolicy, +}; +use google_cloud_gax::retry_policy::{ + Aip194Strict as RetryAip194Strict, RetryPolicy, RetryPolicyExt as _, +}; +use google_cloud_gax::retry_throttler::SharedRetryThrottler; use http::{HeaderMap, header::HeaderName}; +use std::sync::Arc; +use std::time::Duration; const X_GOOG_API_CLIENT: HeaderName = HeaderName::from_static("x-goog-api-client"); const X_GOOG_REQUEST_PARAMS: HeaderName = HeaderName::from_static("x-goog-request-params"); @@ -49,9 +62,7 @@ pub(crate) async fn add_auth_headers( /// Returns a clone of `Credentials` if already present in `config`; /// otherwise, returns a new default `Credentials` object. -pub(crate) fn make_credentials( - config: &crate::options::ClientConfig, -) -> ClientBuilderResult { +pub(crate) fn make_credentials(config: &ClientConfig) -> ClientBuilderResult { if let Some(c) = config.cred.clone() { return Ok(c); } @@ -110,13 +121,103 @@ pub(crate) fn make_headers( Ok(headers) } +/// Policies governing gRPC client transport behaviors. +/// These policies are initialized from a [`ClientConfig`] and can +/// be overridden on a per-request basis using [`RequestOptions`]. +#[derive(Clone, Debug)] +pub(crate) struct TransportPolicies { + retry_policy: Arc, + backoff_policy: Arc, + retry_throttler: SharedRetryThrottler, + polling_error_policy: Arc, + polling_backoff_policy: Arc, + attempt_timeout: Option, +} + +impl TransportPolicies { + /// Creates a new `TransportPolicies` from the given [`ClientConfig`]. + /// Missing policies are populated with default values. + pub(crate) fn from_config(config: &ClientConfig) -> Self { + Self { + retry_policy: config.retry_policy.clone().unwrap_or_else(|| { + Arc::new( + RetryAip194Strict + .with_attempt_limit(10) + .with_time_limit(Duration::from_secs(60)), + ) + }), + backoff_policy: config + .backoff_policy + .clone() + .unwrap_or_else(|| Arc::new(ExponentialBackoff::default())), + retry_throttler: config.retry_throttler.clone(), + polling_error_policy: config + .polling_error_policy + .clone() + .unwrap_or_else(|| Arc::new(PollingAip194Strict)), + polling_backoff_policy: config + .polling_backoff_policy + .clone() + .unwrap_or_else(|| Arc::new(ExponentialBackoff::default())), + attempt_timeout: config.attempt_timeout, + } + } + + pub(crate) fn get_retry_policy(&self, options: &RequestOptions) -> Arc { + options + .retry_policy() + .clone() + .unwrap_or_else(|| self.retry_policy.clone()) + } + + pub(crate) fn get_backoff_policy(&self, options: &RequestOptions) -> Arc { + options + .backoff_policy() + .clone() + .unwrap_or_else(|| self.backoff_policy.clone()) + } + + pub(crate) fn get_retry_throttler(&self, options: &RequestOptions) -> SharedRetryThrottler { + options + .retry_throttler() + .clone() + .unwrap_or_else(|| self.retry_throttler.clone()) + } + + pub(crate) fn get_polling_error_policy( + &self, + options: &RequestOptions, + ) -> Arc { + options + .polling_error_policy() + .clone() + .unwrap_or_else(|| self.polling_error_policy.clone()) + } + + pub(crate) fn get_polling_backoff_policy( + &self, + options: &RequestOptions, + ) -> Arc { + options + .polling_backoff_policy() + .clone() + .unwrap_or_else(|| self.polling_backoff_policy.clone()) + } + + pub(crate) fn attempt_timeout(&self) -> Option { + self.attempt_timeout + } +} + #[cfg(test)] mod tests { use super::*; use google_cloud_auth::credentials::{CacheableResource, CredentialsProvider, EntityTag}; use google_cloud_auth::errors::CredentialsError; + use google_cloud_gax::retry_throttler::AdaptiveThrottler; use http::{Extensions, header::HeaderName, header::HeaderValue}; use pretty_assertions::assert_eq; + use std::sync::Mutex; type AuthResult = std::result::Result; type TestResult = anyhow::Result<()>; @@ -210,7 +311,7 @@ mod tests { .return_once(|| Some(expected_domain.to_string())); let credentials = Credentials::from(provider); - let mut config = crate::options::ClientConfig::default(); + let mut config = ClientConfig::default(); config.cred = Some(credentials); // Act @@ -306,4 +407,159 @@ mod tests { let res = make_headers(API_CLIENT_HEADER, "invalid\nparams", &options); assert!(res.is_err(), "{res:?}"); } + + #[test] + fn transport_policies_from_config_defaults() { + // Arrange + let config = ClientConfig::default(); + + // Act + let policies = TransportPolicies::from_config(&config); + + // Assert + let expected_retry_policy = RetryAip194Strict + .with_attempt_limit(10) + .with_time_limit(Duration::from_secs(60)); + let expected_backoff_policy = ExponentialBackoff::default(); + let expected_polling_error_policy = PollingAip194Strict; + let expected_polling_backoff_policy = ExponentialBackoff::default(); + + // Some of these are dyn, so we compare their string representations. + let options = RequestOptions::default(); + assert_eq!( + format!("{:?}", policies.get_retry_policy(&options)), + format!("{:?}", expected_retry_policy) + ); + assert_eq!( + format!("{:?}", policies.get_backoff_policy(&options)), + format!("{:?}", expected_backoff_policy) + ); + assert!(Arc::ptr_eq( + &policies.get_retry_throttler(&options), + &config.retry_throttler + )); + assert_eq!( + format!("{:?}", policies.get_polling_error_policy(&options)), + format!("{:?}", expected_polling_error_policy) + ); + assert_eq!( + format!("{:?}", policies.get_polling_backoff_policy(&options)), + format!("{:?}", expected_polling_backoff_policy) + ); + assert_eq!(policies.attempt_timeout(), None); + } + + #[test] + fn transport_policies_from_config() { + // Arrange + let config = create_test_config(); + + // Act + let policies = TransportPolicies::from_config(&config); + + // Assert + let options = RequestOptions::default(); + assert!(Arc::ptr_eq( + &policies.get_retry_policy(&options), + config + .retry_policy + .as_ref() + .expect("retry policy should be set") + )); + assert!(Arc::ptr_eq( + &policies.get_backoff_policy(&options), + config + .backoff_policy + .as_ref() + .expect("backoff policy should be set") + )); + assert!(Arc::ptr_eq( + &policies.get_retry_throttler(&options), + &config.retry_throttler + )); + assert!(Arc::ptr_eq( + &policies.get_polling_error_policy(&options), + config + .polling_error_policy + .as_ref() + .expect("polling error policy should be set") + )); + assert!(Arc::ptr_eq( + &policies.get_polling_backoff_policy(&options), + config + .polling_backoff_policy + .as_ref() + .expect("polling backoff policy should be set") + )); + assert_eq!(policies.attempt_timeout(), config.attempt_timeout); + } + + #[test] + fn transport_policies_request_options_overrides() { + // Arrange + let config = create_test_config(); + let policies = TransportPolicies::from_config(&config); + + // Overrides + let mut options = RequestOptions::default(); + let override_retry_policy: Arc = + Arc::new(RetryAip194Strict.with_attempt_limit(3)); + let override_backoff_policy: Arc = + Arc::new(ExponentialBackoff::default()); + let override_retry_throttler: SharedRetryThrottler = + Arc::new(Mutex::new(AdaptiveThrottler::default())); + let override_polling_error_policy: Arc = + Arc::new(PollingAip194Strict); + let override_polling_backoff_policy: Arc = + Arc::new(ExponentialBackoff::default()); + + options.set_retry_policy(override_retry_policy.clone()); + options.set_backoff_policy(override_backoff_policy.clone()); + options.set_retry_throttler(override_retry_throttler.clone()); + options.set_polling_error_policy(override_polling_error_policy.clone()); + options.set_polling_backoff_policy(override_polling_backoff_policy.clone()); + + // Act & Assert + assert!(Arc::ptr_eq( + &policies.get_retry_policy(&options), + &override_retry_policy + )); + assert!(Arc::ptr_eq( + &policies.get_backoff_policy(&options), + &override_backoff_policy + )); + assert!(Arc::ptr_eq( + &policies.get_retry_throttler(&options), + &override_retry_throttler + )); + assert!(Arc::ptr_eq( + &policies.get_polling_error_policy(&options), + &override_polling_error_policy + )); + assert!(Arc::ptr_eq( + &policies.get_polling_backoff_policy(&options), + &override_polling_backoff_policy + )); + } + + fn create_test_config() -> ClientConfig { + let mut config = ClientConfig::default(); + let retry_policy: Arc = Arc::new(RetryAip194Strict.with_attempt_limit(5)); + let backoff_policy: Arc = Arc::new(ExponentialBackoff::default()); + let retry_throttler: SharedRetryThrottler = + Arc::new(Mutex::new(AdaptiveThrottler::default())); + let polling_error_policy: Arc = Arc::new(PollingAip194Strict); + let polling_backoff_policy: Arc = + Arc::new(ExponentialBackoff::default()); + let attempt_timeout = Some(Duration::from_secs(42)); + + config.retry_policy = Some(retry_policy); + config.backoff_policy = Some(backoff_policy); + config.retry_throttler = retry_throttler; + config.polling_error_policy = Some(polling_error_policy); + config.polling_backoff_policy = Some(polling_backoff_policy); + config.attempt_timeout = attempt_timeout; + + config + } }