Skip to content

chore(gax): support request attempt interception via options extensions #6014

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
olavloite wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

chore(gax): support request attempt interception via options extensions #6014

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 18 additions & 6 deletions src/gax-internal/src/grpc.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -57,6 +58,8 @@ pub type GrpcService = Channel;
/// The inner gRPC client type.
pub type InnerClient = Grpc<GrpcService>;

pub(crate) type AttemptInterceptor = Arc<dyn Fn(&mut http::HeaderMap, u32) + Send + Sync>;

#[derive(Clone, Debug)]
pub struct TracingAttributes {
pub server_address: String,
Expand Down Expand Up @@ -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::<AttemptInterceptor>() {
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::<Request, Response>::default();
Expand Down Expand Up @@ -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::<AttemptInterceptor>() {
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) =
Expand Down Expand Up @@ -356,7 +365,7 @@ impl Client {
options: &RequestOptions,
remaining_time: Option<std::time::Duration>,
headers: HeaderMap,
_prior_attempt_count: i64,
prior_attempt_count: i64,
) -> Result<tonic::Response<Response>>
where
Request: prost::Message + 'static,
Expand All @@ -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
};
Expand Down Expand Up @@ -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::<AttemptInterceptor>() {
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);
Expand All @@ -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);

Expand Down
49 changes: 44 additions & 5 deletions src/gax-internal/tests/grpc_retry_loop.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<dyn Fn(&mut HeaderMap, u32) + Send + Sync>;

fn test_credentials() -> Credentials {
Anonymous::new().build()
Expand Down Expand Up @@ -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<EchoResponse> {
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<EchoResponse> {
let extensions = {
let mut e = tonic::Extensions::new();
Expand All @@ -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,
Expand Down
52 changes: 47 additions & 5 deletions src/gax-internal/tests/grpc_simple_request.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<dyn Fn(&mut HeaderMap, u32) + Send + Sync>;

fn test_credentials() -> Credentials {
Anonymous::new().build()
Expand Down Expand Up @@ -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<google::test::v1::EchoResponse> {
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<google::test::v1::EchoResponse> {
let extensions = {
let mut e = tonic::Extensions::new();
Expand All @@ -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,
Expand Down
56 changes: 51 additions & 5 deletions src/gax-internal/tests/grpc_streaming_request.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<dyn Fn(&mut HeaderMap, u32) + Send + Sync>;

fn test_credentials() -> Credentials {
Anonymous::new().build()
Expand Down Expand Up @@ -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?;
Expand Down Expand Up @@ -158,6 +196,16 @@ mod tests {
client: grpc::Client,
rx: tokio::sync::mpsc::Receiver<EchoRequest>,
request_params: &str,
) -> google_cloud_gax::Result<tonic::Response<tonic::codec::Streaming<EchoResponse>>> {
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<EchoRequest>,
request_params: &str,
mut request_options: RequestOptions,
) -> google_cloud_gax::Result<tonic::Response<tonic::codec::Streaming<EchoResponse>>> {
let extensions = {
let mut e = tonic::Extensions::new();
Expand All @@ -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::<EchoRequest, EchoResponse>(
extensions,
Expand Down
Loading