From 5de816e1559b361cbe0ba63f5ce539a14e0ad5d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Thu, 9 Jul 2026 08:21:17 +0200 Subject: [PATCH] chore(gax): support request attempt interception via options extensions Allow mutating request headers on every RPC attempt (including retries) by checking for a closure extension inside `RequestOptions`. If an `Arc` is present in the request options extensions, it is executed before each request attempt with the 1-based attempt number. This enables downstream clients to inject and attempt-specific request headers without modifying the public API of `google-cloud-gax[-internal]` or introducing new public traits. --- src/gax-internal/src/grpc.rs | 24 ++++++-- src/gax-internal/tests/grpc_retry_loop.rs | 49 ++++++++++++++-- src/gax-internal/tests/grpc_simple_request.rs | 52 +++++++++++++++-- .../tests/grpc_streaming_request.rs | 56 +++++++++++++++++-- 4 files changed, 160 insertions(+), 21 deletions(-) diff --git a/src/gax-internal/src/grpc.rs b/src/gax-internal/src/grpc.rs index ea8944e6f9..93dde0f7e9 100644 --- a/src/gax-internal/src/grpc.rs +++ b/src/gax-internal/src/grpc.rs @@ -33,6 +33,7 @@ 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::options::internal::RequestOptionsExt as _; use google_cloud_gax::polling_backoff_policy::PollingBackoffPolicy; use google_cloud_gax::polling_error_policy::{ Aip194Strict as PollingAip194Strict, PollingErrorPolicy, @@ -57,6 +58,8 @@ pub type GrpcService = Channel; /// The inner gRPC client type. pub type InnerClient = Grpc; +pub(crate) type AttemptInterceptor = Arc; + #[derive(Clone, Debug)] pub struct TracingAttributes { pub server_address: String, @@ -209,7 +212,10 @@ impl Client { { use ::tonic::IntoStreamingRequest; let headers = make_headers(api_client_header, request_params, &options)?; - let headers = add_auth_headers(headers, &self.credentials).await?; + let mut headers = add_auth_headers(headers, &self.credentials).await?; + if let Some(interceptor) = options.get_extension::() { + interceptor(&mut headers, 1); + } let metadata = tonic::MetadataMap::from_headers(headers); let request = ::tonic::Request::from_parts(metadata, extensions, request); let codec = tonic_prost::ProstCodec::::default(); @@ -274,7 +280,10 @@ impl Client { { use ::tonic::IntoRequest; let headers = make_headers(api_client_header, request_params, &options)?; - let headers = add_auth_headers(headers, &self.credentials).await?; + let mut headers = add_auth_headers(headers, &self.credentials).await?; + if let Some(interceptor) = options.get_extension::() { + interceptor(&mut headers, 1); + } let metadata = tonic::MetadataMap::from_headers(headers); let mut request = ::tonic::Request::from_parts(metadata, extensions, request); if let Some(timeout) = @@ -356,7 +365,7 @@ impl Client { options: &RequestOptions, remaining_time: Option, headers: HeaderMap, - _prior_attempt_count: i64, + prior_attempt_count: i64, ) -> Result> where Request: prost::Message + 'static, @@ -374,8 +383,8 @@ impl Client { } else { (None, None, None, None) }; - let resend_count = if _prior_attempt_count > 0 { - Some(_prior_attempt_count) + let resend_count = if prior_attempt_count > 0 { + Some(prior_attempt_count) } else { None }; @@ -407,6 +416,9 @@ impl Client { let mut headers = add_auth_headers(headers, &self.credentials).await?; crate::observability::propagation::inject_context(&span, &mut headers); + if let Some(interceptor) = options.get_extension::() { + interceptor(&mut headers, prior_attempt_count as u32 + 1); + } let metadata = tonic::MetadataMap::from_headers(headers); let mut request = ::tonic::Request::from_parts(metadata, extensions, request); @@ -429,7 +441,7 @@ impl Client { use crate::observability::{WithTransportLogging, WithTransportMetric, WithTransportSpan}; let pending = - WithTransportMetric::new(self.metric.clone(), pending, _prior_attempt_count as u32); + WithTransportMetric::new(self.metric.clone(), pending, prior_attempt_count as u32); let pending = WithTransportLogging::new(pending); let pending = WithTransportSpan::new(span, pending); diff --git a/src/gax-internal/tests/grpc_retry_loop.rs b/src/gax-internal/tests/grpc_retry_loop.rs index ac5ab03737..c453356f1e 100644 --- a/src/gax-internal/tests/grpc_retry_loop.rs +++ b/src/gax-internal/tests/grpc_retry_loop.rs @@ -20,8 +20,13 @@ mod tests { use google_cloud_gax::options::RequestOptions; use google_cloud_gax::retry_policy::{Aip194Strict, RetryPolicyExt}; use google_cloud_gax_internal::grpc; + use google_cloud_gax_internal::options::ClientConfig; use grpc_server::google::test::v1::EchoResponse; use grpc_server::{builder, google, start_fixed_responses}; + use http::HeaderMap; + use std::sync::{Arc, Mutex}; + + type AttemptInterceptor = Arc; fn test_credentials() -> Credentials { Anonymous::new().build() @@ -121,9 +126,45 @@ mod tests { .expect("a valid backoff policy") } + #[tokio::test] + async fn interceptor_on_retry() -> anyhow::Result<()> { + use google_cloud_gax::options::internal::RequestOptionsExt as _; + + let attempts = Arc::new(Mutex::new(Vec::new())); + let attempts_clone = attempts.clone(); + let interceptor: AttemptInterceptor = Arc::new(move |_headers, attempt| { + attempts_clone.lock().unwrap().push(attempt); + }); + + let (endpoint, _server) = + start_fixed_responses(vec![transient(), transient(), success()]).await?; + + let mut config = ClientConfig::default(); + config.cred = Some(test_credentials()); + config.endpoint = Some(endpoint); + config.backoff_policy = Some(Arc::new(test_backoff())); + + let client = grpc::Client::new(config, "https://test-only.googleapis.com").await?; + let options = RequestOptions::default().insert_extension(interceptor); + let _response = send_request_with_options(client, "interceptor_on_retry", options).await?; + + let attempts = attempts.lock().unwrap().clone(); + assert_eq!(attempts, vec![1, 2, 3]); + + Ok(()) + } + pub async fn send_request( client: grpc::Client, msg: &str, + ) -> google_cloud_gax::Result { + send_request_with_options(client, msg, RequestOptions::default()).await + } + + pub async fn send_request_with_options( + client: grpc::Client, + msg: &str, + mut request_options: RequestOptions, ) -> google_cloud_gax::Result { let extensions = { let mut e = tonic::Extensions::new(); @@ -137,11 +178,9 @@ mod tests { message: msg.into(), ..google::test::v1::EchoRequest::default() }; - let request_options = { - let mut o = RequestOptions::default(); - o.set_idempotency(true); - o - }; + if request_options.idempotent().is_none() { + request_options.set_idempotency(true); + } client .execute( extensions, diff --git a/src/gax-internal/tests/grpc_simple_request.rs b/src/gax-internal/tests/grpc_simple_request.rs index 089d3c7877..49653ca2aa 100644 --- a/src/gax-internal/tests/grpc_simple_request.rs +++ b/src/gax-internal/tests/grpc_simple_request.rs @@ -23,6 +23,10 @@ mod tests { use google_cloud_gax_internal::grpc; use google_cloud_gax_internal::options::ClientConfig; use grpc_server::{builder, google, start_echo_server, start_echo_server_with_address}; + use http::HeaderMap; + use std::sync::Arc; + + type AttemptInterceptor = Arc; fn test_credentials() -> Credentials { Anonymous::new().build() @@ -192,10 +196,50 @@ mod tests { Ok(()) } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn attempt_interceptor() -> anyhow::Result<()> { + use google_cloud_gax::options::internal::RequestOptionsExt as _; + use http::header::{HeaderName, HeaderValue}; + + let interceptor: AttemptInterceptor = Arc::new(|headers, attempt| { + headers.insert( + HeaderName::from_static("x-test-attempt"), + HeaderValue::from_str(&attempt.to_string()).expect("valid attempt number"), + ); + }); + + let (endpoint, _server) = start_echo_server().await?; + + let mut config = ClientConfig::default(); + config.cred = Some(test_credentials()); + config.endpoint = Some(endpoint); + + let client = grpc::Client::new(config, "https://test-only.googleapis.com").await?; + let options = RequestOptions::default().insert_extension(interceptor); + + let response = send_request_with_options(client, "test message", "", options).await?; + assert_eq!(&response.message, "test message"); + + assert_eq!( + response.metadata.get("x-test-attempt").map(String::as_str), + Some("1") + ); + Ok(()) + } + async fn send_request( client: grpc::Client, msg: &str, request_params: &str, + ) -> google_cloud_gax::Result { + send_request_with_options(client, msg, request_params, RequestOptions::default()).await + } + + async fn send_request_with_options( + client: grpc::Client, + msg: &str, + request_params: &str, + mut request_options: RequestOptions, ) -> google_cloud_gax::Result { let extensions = { let mut e = tonic::Extensions::new(); @@ -209,11 +253,9 @@ mod tests { message: msg.into(), ..google::test::v1::EchoRequest::default() }; - let request_options = { - let mut o = RequestOptions::default(); - o.set_retry_policy(NeverRetry); - o - }; + if request_options.retry_policy().is_none() { + request_options.set_retry_policy(NeverRetry); + } client .execute( extensions, diff --git a/src/gax-internal/tests/grpc_streaming_request.rs b/src/gax-internal/tests/grpc_streaming_request.rs index 2e84655d01..d1d3b588e9 100644 --- a/src/gax-internal/tests/grpc_streaming_request.rs +++ b/src/gax-internal/tests/grpc_streaming_request.rs @@ -22,6 +22,10 @@ mod tests { use google_cloud_gax_internal::grpc; use grpc_server::google::test::v1::{EchoRequest, EchoResponse}; use grpc_server::{builder, start_echo_server}; + use http::HeaderMap; + use std::sync::Arc; + + type AttemptInterceptor = Arc; fn test_credentials() -> Credentials { Anonymous::new().build() @@ -116,6 +120,40 @@ mod tests { Ok(()) } + #[tokio::test] + async fn attempt_interceptor() -> anyhow::Result<()> { + use google_cloud_gax::options::internal::RequestOptionsExt as _; + use http::header::{HeaderName, HeaderValue}; + + let interceptor: AttemptInterceptor = Arc::new(|headers, attempt| { + headers.insert( + HeaderName::from_static("x-test-attempt"), + HeaderValue::from_str(&attempt.to_string()).expect("valid attempt number"), + ); + }); + + let (endpoint, _server) = start_echo_server().await?; + + let mut config = google_cloud_gax_internal::options::ClientConfig::default(); + config.cred = Some(test_credentials()); + config.endpoint = Some(endpoint); + + let client = grpc::Client::new(config, "https://test-only.googleapis.com").await?; + let options = RequestOptions::default().insert_extension(interceptor); + + let (tx, rx) = tokio::sync::mpsc::channel(100); + tx.send(simple_request("msg0")).await?; + let response = send_streaming_request_with_options(client, rx, "", options).await?; + let (_, mut stream, _) = response.into_parts(); + let first_msg = stream.message().await?.unwrap(); + assert_eq!( + first_msg.metadata.get("x-test-attempt").map(String::as_str), + Some("1") + ); + + Ok(()) + } + #[tokio::test] async fn credentials_error() -> anyhow::Result<()> { let (endpoint, _server) = start_echo_server().await?; @@ -158,6 +196,16 @@ mod tests { client: grpc::Client, rx: tokio::sync::mpsc::Receiver, request_params: &str, + ) -> google_cloud_gax::Result>> { + send_streaming_request_with_options(client, rx, request_params, RequestOptions::default()) + .await + } + + async fn send_streaming_request_with_options( + client: grpc::Client, + rx: tokio::sync::mpsc::Receiver, + request_params: &str, + mut request_options: RequestOptions, ) -> google_cloud_gax::Result>> { let extensions = { let mut e = tonic::Extensions::new(); @@ -167,11 +215,9 @@ mod tests { )); e }; - let request_options = { - let mut o = RequestOptions::default(); - o.set_retry_policy(NeverRetry); - o - }; + if request_options.retry_policy().is_none() { + request_options.set_retry_policy(NeverRetry); + } client .bidi_stream::( extensions,