From 4bd7f1d6ae7506ed00a8ea9b1088f211d6104707 Mon Sep 17 00:00:00 2001 From: haphungw Date: Wed, 10 Jun 2026 15:43:28 +0000 Subject: [PATCH] feat(o11y): add error extraction to DiscoveryOperation and fix macro compatibility --- src/lro/src/internal/discovery.rs | 6 ++ src/lro/src/internal/tracing.rs | 166 ++++++++++++++++++++++++++++++ src/test-utils/src/test_layer.rs | 9 +- 3 files changed, 180 insertions(+), 1 deletion(-) diff --git a/src/lro/src/internal/discovery.rs b/src/lro/src/internal/discovery.rs index 98f65319ec..a36ca045c9 100644 --- a/src/lro/src/internal/discovery.rs +++ b/src/lro/src/internal/discovery.rs @@ -27,6 +27,7 @@ use crate::{ Poller, PollingBackoffPolicy, PollingErrorPolicy, PollingResult, Result, sealed::Poller as SealedPoller, }; +use google_cloud_gax::error::rpc::Status; use google_cloud_gax::polling_state::PollingState; use google_cloud_gax::retry_result::RetryResult; use std::sync::Arc; @@ -54,6 +55,11 @@ pub trait DiscoveryOperation { /// /// It may be `None` in which case the polling loop stops. fn name(&self) -> Option<&String>; + + /// Returns the error status of the operation, if any. + fn error(&self) -> Option { + None + } } pub fn new_discovery_poller( diff --git a/src/lro/src/internal/tracing.rs b/src/lro/src/internal/tracing.rs index 3bbd154bcb..b418294d78 100644 --- a/src/lro/src/internal/tracing.rs +++ b/src/lro/src/internal/tracing.rs @@ -67,7 +67,31 @@ impl LroRecorder { pub fn attempt_count(&self) -> Option { self.attempt_count } +} +/// Helper macro to record telemetry for Discovery LROs. +#[macro_export] +#[doc(hidden)] +macro_rules! record_discovery_polling_result { + ($span:expr, $op:expr) => { + let span = &$span; + let op = &$op; + let done = $crate::internal::DiscoveryOperation::done(op); + span.record("gcp.longrunning.done", done); + if done { + let error = $crate::internal::DiscoveryOperation::error(op); + let code = error.as_ref().map(|e| e.code as i32).unwrap_or(0); + span.record("gcp.longrunning.status_code", code); + if let Some(status) = error { + span.record("otel.status_code", "ERROR"); + span.record("otel.status_description", &status.message); + span.record("error.type", status.code.to_string()); + } + } + }; +} + +impl LroRecorder { /// Creates a new clone of `LroRecorder` carrying the specified LRO polling attempt count. /// /// Since `LroRecorder` is immutable to guarantee thread-safety, this updates the context @@ -419,4 +443,146 @@ mod tests { ); assert!(got.attributes.get("gcp.longrunning.done").is_none()); } + + #[derive(Default)] + struct MockDiscoveryOperation { + done: bool, + error: Option, + } + + impl crate::internal::DiscoveryOperation for MockDiscoveryOperation { + fn done(&self) -> bool { + self.done + } + + fn name(&self) -> Option<&String> { + None + } + + fn error(&self) -> Option { + self.error.clone() + } + } + + #[tokio::test] + async fn record_discovery_polling_result_success() { + let guard = TestLayer::initialize(); + let span = + client_request_signals!(info: &InstrumentationClientInfo::default(), method: "test"); + let op = MockDiscoveryOperation { + done: true, + error: None, + }; + + record_discovery_polling_result!(span, op); + + { + let captured = TestLayer::capture(&guard); + let got = captured + .iter() + .find(|s| s.name == "client_request") + .unwrap(); + + assert_eq!( + got.attributes + .get("gcp.longrunning.done") + .and_then(|v| v.as_bool()), + Some(true) + ); + assert_eq!( + got.attributes + .get("gcp.longrunning.status_code") + .and_then(|v| v.as_i64()), + Some(0) + ); + assert_eq!( + got.attributes + .get("otel.status_code") + .and_then(|v| v.as_string()), + Some("UNSET".to_string()) + ); + } + } + + #[tokio::test] + async fn record_discovery_polling_result_error() { + let guard = TestLayer::initialize(); + let span = + client_request_signals!(info: &InstrumentationClientInfo::default(), method: "test"); + let status = google_cloud_gax::error::rpc::Status::default() + .set_code(google_cloud_gax::error::rpc::Code::NotFound) + .set_message("not found"); + let op = MockDiscoveryOperation { + done: true, + error: Some(status), + }; + + record_discovery_polling_result!(span, op); + + { + let captured = TestLayer::capture(&guard); + let got = captured + .iter() + .find(|s| s.name == "client_request") + .unwrap(); + + assert_eq!( + got.attributes + .get("gcp.longrunning.done") + .and_then(|v| v.as_bool()), + Some(true) + ); + assert_eq!( + got.attributes + .get("gcp.longrunning.status_code") + .and_then(|v| v.as_i64()), + Some(google_cloud_gax::error::rpc::Code::NotFound as i64) + ); + assert_eq!( + got.attributes + .get("otel.status_code") + .and_then(|v| v.as_string()), + Some("ERROR".to_string()) + ); + assert_eq!( + got.attributes + .get("otel.status_description") + .and_then(|v| v.as_string()), + Some("not found".to_string()) + ); + assert_eq!( + got.attributes.get("error.type").and_then(|v| v.as_string()), + Some("NOT_FOUND".to_string()) + ); + } + } + + #[tokio::test] + async fn record_discovery_polling_result_in_progress() { + let guard = TestLayer::initialize(); + let span = + client_request_signals!(info: &InstrumentationClientInfo::default(), method: "test"); + let op = MockDiscoveryOperation { + done: false, + error: None, + }; + + record_discovery_polling_result!(span, op); + + { + let captured = TestLayer::capture(&guard); + let got = captured + .iter() + .find(|s| s.name == "client_request") + .unwrap(); + + assert_eq!( + got.attributes + .get("gcp.longrunning.done") + .and_then(|v| v.as_bool()), + Some(false) + ); + assert!(got.attributes.get("gcp.longrunning.status_code").is_none()); + } + } } diff --git a/src/test-utils/src/test_layer.rs b/src/test-utils/src/test_layer.rs index d4e9f30ffe..135114eb56 100644 --- a/src/test-utils/src/test_layer.rs +++ b/src/test-utils/src/test_layer.rs @@ -88,7 +88,14 @@ impl AttributeValue { _ => None, } } - // Add other as_ type helpers as needed + + /// Helper to get the bool value if the variant is Boolean. + pub fn as_bool(&self) -> Option { + match self { + AttributeValue::Boolean(b) => Some(*b), + _ => None, + } + } } /// Represents a captured tracing span with its attributes.