Fixed issues validating tracing pages which were converted to page iterators

sdk/core/azure_core/src/http/mod.rs

@@ -23,7 +23,7 @@ pub use response::{AsyncRawResponse, AsyncResponse, RawResponse, Response};
 pub use typespec_client_core::http::response;
 pub use typespec_client_core::http::{
     new_http_client, AppendToUrlQuery, Context, DeserializeWith, Format, HttpClient, JsonFormat,
-    Method, NoFormat, StatusCode, Url, UrlExt,
+    Method, NoFormat, Sanitizer, StatusCode, Url, UrlExt,
 };
 
 pub use crate::error::check_success;

sdk/core/azure_core/src/http/pager.rs

@@ -811,14 +811,36 @@ impl<P> fmt::Debug for PageIterator<P> {
     }
 }
 
-#[derive(Debug, Clone, Eq)]
-enum State<C> {
+#[derive(Clone, Eq)]
+enum State<C>
+where
+    C: AsRef<str>,
+{
     Init,
     More(C),
     Done,
 }
 
-impl<C> PartialEq for State<C> {
+impl<C> fmt::Debug for State<C>
+where
+    C: AsRef<str>,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            State::Init => write!(f, "State::Init"),
+            State::More(c) => f
+                .debug_struct("State::More")
+                .field("continuation", &c.as_ref())
+                .finish(),
+            State::Done => write!(f, "State::Done"),
+        }
+    }
+}
+
+impl<C> PartialEq for State<C>
+where
+    C: AsRef<str>,
+{
     fn eq(&self, other: &Self) -> bool {
         // Only needs to compare if both states are Init or Done; internally, we don't care about any other states.
         matches!(
@@ -829,7 +851,10 @@ impl<C> PartialEq for State<C> {
 }
 
 #[derive(Debug)]
-struct StreamState<'a, C, F> {
+struct StreamState<'a, C, F>
+where
+    C: AsRef<str>,
+{
     state: State<C>,
     make_request: F,
     continuation_token: Arc<Mutex<Option<String>>>,
@@ -863,9 +888,21 @@ where
             added_span: false,
         },
         |mut stream_state| async move {
+            // When in the "Init" state, we are either starting fresh or resuming from a continuation token. In either case,
+            // attach a span to the context for the entire paging operation.
+            if stream_state.state == State::Init {
+                tracing::debug!("establish a public API span for new pager.");
+
+                // At the very start of polling, create a span for the entire request, and attach it to the context
+                let span = create_public_api_span(&stream_state.ctx, None, None);
+                if let Some(ref s) = span {
+                    stream_state.added_span = true;
+                    stream_state.ctx = stream_state.ctx.with_value(s.clone());
+                }
+            }
+
             // Get the `continuation_token` to pick up where we left off, or None for the initial page,
             // but don't override the terminal `State::Done`.
-
             if stream_state.state != State::Done {
                 let result = match stream_state.continuation_token.lock() {
                     Ok(next_token) => match next_token.as_deref() {
@@ -895,12 +932,6 @@ where
             let result = match stream_state.state {
                 State::Init => {
                     tracing::debug!("initial page request");
-                    // At the very start of polling, create a span for the entire request, and attach it to the context
-                    let span = create_public_api_span(&stream_state.ctx, None, None);
-                    if let Some(ref s) = span {
-                        stream_state.added_span = true;
-                        stream_state.ctx = stream_state.ctx.with_value(s.clone());
-                    }
                     (stream_state.make_request)(PagerState::Initial, stream_state.ctx.clone()).await
                 }
                 State::More(n) => {

sdk/core/azure_core_test/src/tracing.rs

@@ -65,10 +65,10 @@ impl TracerProvider for MockTracingProvider {
 /// Mock Tracer - used for testing distributed tracing without involving a specific tracing implementation.
 #[derive(Debug)]
 pub struct MockTracer {
-    pub namespace: Option<&'static str>,
-    pub package_name: &'static str,
-    pub package_version: Option<&'static str>,
-    pub spans: Mutex<Vec<Arc<MockSpan>>>,
+    namespace: Option<&'static str>,
+    package_name: &'static str,
+    package_version: Option<&'static str>,
+    spans: Mutex<Vec<Arc<MockSpanInner>>>,
 }
 
 impl Tracer for MockTracer {
@@ -83,9 +83,14 @@ impl Tracer for MockTracer {
         attributes: Vec<Attribute>,
         parent: Arc<dyn crate::tracing::Span>,
     ) -> Arc<dyn crate::tracing::Span> {
-        let span = Arc::new(MockSpan::new(name, kind, attributes.clone(), Some(parent)));
+        let span = Arc::new(MockSpanInner::new(
+            name,
+            kind,
+            attributes.clone(),
+            Some(parent),
+        ));
         self.spans.lock().unwrap().push(span.clone());
-        span
+        Arc::new(MockSpan { inner: span })
     }
 
     fn start_span(
@@ -101,15 +106,15 @@ impl Tracer for MockTracer {
                 value: attr.value.clone(),
             })
             .collect();
-        let span = Arc::new(MockSpan::new(name, kind, attributes, None));
+        let span = Arc::new(MockSpanInner::new(name, kind, attributes, None));
         self.spans.lock().unwrap().push(span.clone());
-        span
+        Arc::new(MockSpan { inner: span })
     }
 }
 
 /// Mock span for testing purposes.
 #[derive(Debug)]
-pub struct MockSpan {
+struct MockSpanInner {
     pub name: Cow<'static, str>,
     pub kind: SpanKind,
     pub parent: Option<[u8; 8]>,
@@ -118,7 +123,7 @@ pub struct MockSpan {
     pub state: Mutex<SpanStatus>,
     pub is_open: Mutex<bool>,
 }
-impl MockSpan {
+impl MockSpanInner {
     fn new<C>(
         name: C,
         kind: SpanKind,
@@ -144,9 +149,22 @@ impl MockSpan {
             is_open: Mutex::new(true),
         }
     }
+
+    fn is_open(&self) -> bool {
+        let is_open = self.is_open.lock().unwrap();
+        *is_open
+    }
 }
 
-impl Span for MockSpan {
+impl AsAny for MockSpanInner {
+    fn as_any(&self) -> &dyn std::any::Any {
+        // Convert to an object that doesn't expose the lifetime parameter
+        // We're essentially erasing the lifetime here to satisfy the static requirement
+        self as &dyn std::any::Any
+    }
+}
+
+impl Span for MockSpanInner {
     fn set_attribute(&self, key: &'static str, value: AttributeValue) {
         eprintln!("{}: Setting attribute {}: {:?}", self.name, key, value);
         let mut attributes = self.attributes.lock().unwrap();
@@ -195,6 +213,19 @@ impl Span for MockSpan {
     }
 }
 
+pub struct MockSpan {
+    inner: Arc<MockSpanInner>,
+}
+
+impl Drop for MockSpan {
+    fn drop(&mut self) {
+        if self.inner.is_open() {
+            eprintln!("Warning: Dropping open span: {}", self.inner.name);
+            self.inner.end();
+        }
+    }
+}
+
 impl AsAny for MockSpan {
     fn as_any(&self) -> &dyn std::any::Any {
         // Convert to an object that doesn't expose the lifetime parameter
@@ -203,6 +234,40 @@ impl AsAny for MockSpan {
     }
 }
 
+impl Span for MockSpan {
+    fn set_attribute(&self, key: &'static str, value: AttributeValue) {
+        self.inner.set_attribute(key, value);
+    }
+
+    fn set_status(&self, status: crate::tracing::SpanStatus) {
+        self.inner.set_status(status);
+    }
+
+    fn end(&self) {
+        self.inner.end();
+    }
+
+    fn is_recording(&self) -> bool {
+        self.inner.is_recording()
+    }
+
+    fn span_id(&self) -> [u8; 8] {
+        self.inner.span_id()
+    }
+
+    fn record_error(&self, error: &dyn std::error::Error) {
+        self.inner.record_error(error);
+    }
+
+    fn set_current(&self, context: &Context) -> Box<dyn SpanGuard> {
+        self.inner.set_current(context)
+    }
+
+    fn propagate_headers(&self, request: &mut Request) {
+        self.inner.propagate_headers(request);
+    }
+}
+
 /// Expected information about a tracer.
 #[derive(Debug)]
 pub struct ExpectedTracerInformation<'a> {
@@ -293,7 +358,7 @@ pub struct ExpectedSpanInformation<'a> {
 }
 
 fn check_span_information(
-    span: &Arc<MockSpan>,
+    span: &Arc<MockSpanInner>,
     expected: &ExpectedSpanInformation<'_>,
     parent_span_map: &HashMap<Uuid, [u8; 8]>,
 ) {
@@ -405,8 +470,14 @@ impl Default for ExpectedRestApiSpan {
 #[derive(Debug, Default, Clone)]
 pub struct ExpectedInstrumentation {
     /// The package name for the service client.
+    ///
+    /// **NOTE**: Make sure that the package name comes from `env!("CARGO_PKG_NAME")` to ensure that this continues to work
+    /// if test recordings were created with a previous version of the package.
     pub package_name: String,
     /// The package version for the service client.
+    ///
+    /// **NOTE**: Make sure that the package name comes from `env!("CARGO_PKG_VERSION")` to ensure that this continues to work
+    /// if test recordings were created with a previous version of the package.
     pub package_version: String,
     /// The namespace for the service client.
     pub package_namespace: Option<&'static str>,
@@ -437,6 +508,7 @@ pub struct ExpectedInstrumentation {
 /// The `test_api` call may issue multiple service client calls, if it does, this function will verify that all expected spans were created. The caller of the `test_instrumentation_for_api` call
 /// should make sure to include all expected APIs in the call.
 ///
+///
 pub async fn assert_instrumentation_information<C, FnInit, FnTest, T>(
     create_client: FnInit,
     test_api: FnTest,

sdk/keyvault/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "rust",
   "TagPrefix": "rust/keyvault",
-  "Tag": "rust/keyvault_7e75eabce0"
+  "Tag": "rust/keyvault_5961c5368d"
 }