diff --git a/src/spanner/src/batch_read_only_transaction.rs b/src/spanner/src/batch_read_only_transaction.rs index fcfa693273..bb5449e0e9 100644 --- a/src/spanner/src/batch_read_only_transaction.rs +++ b/src/spanner/src/batch_read_only_transaction.rs @@ -406,6 +406,7 @@ impl Partition { operation: StreamOperation::Query(req.clone()), channel_hint, gax_options, + method_name: "ExecuteStreamingSql", }) .await } @@ -437,6 +438,7 @@ impl Partition { operation: StreamOperation::Read(req.clone()), channel_hint, gax_options, + method_name: "StreamingRead", }) .await } diff --git a/src/spanner/src/client.rs b/src/spanner/src/client.rs index 6c79598d31..4043812e39 100644 --- a/src/spanner/src/client.rs +++ b/src/spanner/src/client.rs @@ -126,15 +126,15 @@ macro_rules! define_idempotent_rpc { channel_hint: usize, o11y: &crate::observability::Observability, ) -> crate::Result<$response_type> { - o11y.trace_operation($canonical_name, || async move { + o11y.trace_operation( + $canonical_name, self.get_channel(channel_hint) .inner .$method() .with_request(request) .with_options(apply_request_defaults(options)) - .send() - .await - }) + .send(), + ) .await } }; diff --git a/src/spanner/src/observability/metrics.rs b/src/spanner/src/observability/metrics.rs index a6a4e84f32..cb222f707c 100644 --- a/src/spanner/src/observability/metrics.rs +++ b/src/spanner/src/observability/metrics.rs @@ -14,6 +14,8 @@ use crate::omni::InstanceType; #[cfg(feature = "_experimental-builtin-metrics")] +use std::sync::Arc; +#[cfg(feature = "_experimental-builtin-metrics")] use std::time::Duration; #[cfg(feature = "_experimental-builtin-metrics")] use std::time::Instant; @@ -73,10 +75,10 @@ impl SpannerMetrics { } #[cfg(feature = "_experimental-builtin-metrics")] -#[derive(Debug)] +#[derive(Clone, Debug)] pub(crate) struct Observability { - pub(crate) metrics: Option, - _meter_provider: Option, + pub(crate) metrics: Option>, + _meter_provider: Option>, } #[cfg(feature = "_experimental-builtin-metrics")] @@ -101,8 +103,8 @@ impl Observability { } let project_id = match project_id { - Some(id) => id, - None => return Self::disabled(), + Some(id) if !id.is_empty() => id, + _ => return Self::disabled(), }; // Create the Google Cloud Monitoring client using the same config @@ -137,22 +139,24 @@ impl Observability { let metrics = SpannerMetrics::new(meter); Self { - metrics: Some(metrics), - _meter_provider: Some(meter_provider), + metrics: Some(Arc::new(metrics)), + _meter_provider: Some(Arc::new(meter_provider)), } } - pub(crate) async fn trace_operation( + pub(crate) async fn trace_operation( &self, method: &'static str, - f: F, + fut: Fut, ) -> crate::Result where - F: FnOnce() -> Fut, Fut: std::future::Future>, { + if self.metrics.is_none() { + return fut.await; + } let start_time = Instant::now(); - let result = f().await; + let result = fut.await; let elapsed = start_time.elapsed(); self.record_operation(method, elapsed, &result); result @@ -278,7 +282,7 @@ pub(crate) fn parse_server_timing(header_val: &str) -> ServerTimings { } #[cfg(not(feature = "_experimental-builtin-metrics"))] -#[derive(Debug)] +#[derive(Clone, Debug)] pub(crate) struct Observability; #[cfg(not(feature = "_experimental-builtin-metrics"))] @@ -296,22 +300,79 @@ impl Observability { Self } - pub(crate) async fn trace_operation( + #[inline(always)] + pub(crate) async fn trace_operation( &self, _method: &'static str, - f: F, + fut: Fut, ) -> crate::Result where - F: FnOnce() -> Fut, Fut: std::future::Future>, { - f().await + fut.await } } #[cfg(all(test, feature = "_experimental-builtin-metrics"))] mod tests { use super::*; + use opentelemetry_sdk::metrics::InMemoryMetricExporter; + + #[test] + fn test_observability_disabled() { + let o11y = Observability::disabled(); + assert!(o11y.metrics.is_none()); + } + + #[test] + fn test_result_to_status_str() { + let ok_res: crate::Result<()> = Ok(()); + assert_eq!(result_to_status_str(&ok_res), "OK"); + + let status_pd = google_cloud_gax::error::rpc::Status::default() + .set_code(google_cloud_gax::error::rpc::Code::PermissionDenied); + let err_pd: crate::Result<()> = Err(crate::Error::service(status_pd)); + assert_eq!(result_to_status_str(&err_pd), "PERMISSION_DENIED"); + } + + #[test] + fn test_spanner_metrics_record_operation_and_attempt() { + let exporter = InMemoryMetricExporter::default(); + let reader = opentelemetry_sdk::metrics::PeriodicReader::builder(exporter.clone()).build(); + let provider = SdkMeterProvider::builder().with_reader(reader).build(); + let meter = provider.meter("cloud.google.com/rust"); + let metrics = SpannerMetrics::new(meter); + let o11y = Observability { + metrics: Some(Arc::new(metrics)), + _meter_provider: Some(Arc::new(provider.clone())), + }; + + let ok_res: crate::Result<()> = Ok(()); + o11y.record_operation("ExecuteSql", Duration::from_millis(50), &ok_res); + o11y.record_attempt( + "ExecuteSql", + Duration::from_millis(40), + &ok_res, + Some(12.5), + Some(5.0), + ); + + provider.force_flush().expect("force_flush failed"); + + let finished = exporter + .get_finished_metrics() + .expect("get_finished_metrics"); + assert!(!finished.is_empty()); + } + + #[tokio::test] + async fn test_trace_operation_success() { + let o11y = Observability::disabled(); + let result = o11y + .trace_operation("ExecuteSql", async { Ok::(42) }) + .await; + assert_eq!(result.expect("trace_operation result"), 42); + } #[test] fn test_parse_server_timing() { diff --git a/src/spanner/src/read_only_transaction.rs b/src/spanner/src/read_only_transaction.rs index e091d8d0f4..d67174ab57 100644 --- a/src/spanner/src/read_only_transaction.rs +++ b/src/spanner/src/read_only_transaction.rs @@ -1072,7 +1072,7 @@ fn merge_request_options( /// Helper macro to execute a streaming SQL or streaming read RPC with retry logic. macro_rules! execute_stream_with_retry { - ($self:expr, $request:ident, $gax_options:ident, $rpc_method:ident, $operation_variant:path) => {{ + ($self:expr, $request:ident, $gax_options:ident, $rpc_method:ident, $operation_variant:path, $method_name:expr) => {{ let stream = match $self .client .spanner @@ -1123,7 +1123,7 @@ macro_rules! execute_stream_with_retry { } }; - ResultSet::create(ResultSetParams { + Box::pin(ResultSet::create(ResultSetParams { stream, transaction_selector: Some($self.transaction_selector.clone()), precommit_token_tracker: $self.precommit_token_tracker.clone(), @@ -1133,7 +1133,8 @@ macro_rules! execute_stream_with_retry { operation: $operation_variant($request), channel_hint: $self.channel_hint, gax_options: $gax_options, - }) + method_name: $method_name, + })) .await }}; } @@ -1158,7 +1159,8 @@ impl ReadContext { request, gax_options, execute_streaming_sql, - StreamOperation::Query + StreamOperation::Query, + "ExecuteStreamingSql" ) } @@ -1179,7 +1181,8 @@ impl ReadContext { request, gax_options, streaming_read, - StreamOperation::Read + StreamOperation::Read, + "StreamingRead" ) } } @@ -3493,7 +3496,8 @@ pub(crate) mod tests { request, gax_options, execute_streaming_sql, - StreamOperation::Query + StreamOperation::Query, + "ExecuteStreamingSql" ) } diff --git a/src/spanner/src/result_set.rs b/src/spanner/src/result_set.rs index 20c08946ed..cc1ced1f10 100644 --- a/src/spanner/src/result_set.rs +++ b/src/spanner/src/result_set.rs @@ -98,6 +98,7 @@ pub(crate) struct ResultSetParams { pub operation: StreamOperation, pub channel_hint: usize, pub gax_options: GaxRequestOptions, + pub method_name: &'static str, } // The maximum number of PartialResultSets to buffer without a resume token. @@ -111,8 +112,11 @@ const DEFAULT_ATTEMPT_LIMIT: u32 = 10; impl ResultSet { /// Creates a new result set asynchronously, waiting for the first chunk to arrive. pub(crate) async fn create(params: ResultSetParams) -> crate::Result { + let method_name = params.method_name; let mut result_set = Self::new(params); - result_set.init_stream().await?; + let o11y = result_set.client.o11y.clone(); + let fut = Box::pin(result_set.init_stream()); + o11y.trace_operation(method_name, fut).await?; Ok(result_set) } @@ -128,6 +132,7 @@ impl ResultSet { operation, channel_hint, gax_options, + method_name: _, } = params; let gax_options = Self::apply_defaults(gax_options); @@ -1871,6 +1876,7 @@ pub(crate) mod tests { operation: StreamOperation::Query(req), channel_hint: 0, gax_options: GaxRequestOptions::default(), + method_name: "ExecuteStreamingSql", }) .await?;