fix(test): await tracing span closure before assertions

Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>

Author: matthewgrossman

crates/openshell-server/src/compute/mod.rs

@@ -14118,7 +14118,6 @@ mod tests {
     /// Driver watch events arrive on a background stream, so the store writes
     /// they trigger land outside the request that caused them.
     #[tokio::test]
-    #[ignore = "flaky under concurrent test execution"]
     async fn driver_watch_events_are_roots_and_store_operations_have_parents() {
         use crate::otel_tracing::test_exporter;
 
@@ -14133,20 +14132,12 @@ mod tests {
             .await
             .unwrap();
 
+        let root = traced.wait_for_span("driver_watch.sandbox_deleted").await;
         let spans = traced.finished_spans();
-        let root = spans
-            .iter()
-            .find(|s| s.name == "driver_watch.sandbox_deleted")
-            .unwrap_or_else(|| {
-                panic!(
-                    "the event records a span of its own, got {:?}",
-                    spans.iter().map(|s| &s.name).collect::<Vec<_>>()
-                )
-            });
 
-        test_exporter::assert_is_root(root);
+        test_exporter::assert_is_root(&root);
         assert_eq!(
-            test_exporter::attribute(root, "sandbox.id").as_deref(),
+            test_exporter::attribute(&root, "sandbox.id").as_deref(),
             Some("sb-1"),
             "the span names which sandbox the driver reported on"
         );
@@ -14164,7 +14155,6 @@ mod tests {
     /// The reconciler runs on a timer with no inbound request, so without a
     /// span of its own each store call becomes its own anonymous trace.
     #[tokio::test]
-    #[ignore = "flaky under concurrent test execution"]
     async fn reconcile_sweeps_are_roots_and_operations_have_parents() {
         use crate::otel_tracing::test_exporter;
 
@@ -14179,9 +14169,8 @@ mod tests {
             .await
             .unwrap();
 
-        // Other tests drive their own reconcile loops into the shared
-        // exporter, so match on the shape of a sweep rather than assuming
-        // there is exactly one.
+        // A closed sweep has no remaining worker-held child spans.
+        traced.wait_for_span("reconcile.sandboxes").await;
         let spans = traced.finished_spans();
         let roots = traced.spans_named("reconcile.sandboxes");
         assert!(

crates/openshell-server/src/grpc/sandbox.rs

@@ -4543,7 +4543,6 @@ mod tests {
     }
 
     #[tokio::test]
-    #[ignore = "flaky under concurrent test execution"]
     async fn watch_producer_releases_request_span_when_client_disconnects() {
         use crate::otel_tracing::test_exporter;
         use tokio_stream::StreamExt as _;
@@ -4582,6 +4581,7 @@ mod tests {
 
         drop(request_span);
         stream.disconnect_and_wait().await;
+        traced.wait_for_span("disconnected_watch_request").await;
 
         assert_eq!(
             traced.spans_named("disconnected_watch_request").len(),

crates/openshell-server/src/otel_tracing.rs

@@ -146,7 +146,7 @@ where
 /// Isolated in-memory span exporters for tracing tests.
 #[cfg(test)]
 pub mod test_exporter {
-    use std::sync::OnceLock;
+    use std::sync::{Arc, OnceLock};
 
     use tracing::{Dispatch, Subscriber, dispatcher::WeakDispatch};
 
@@ -159,6 +159,7 @@ pub mod test_exporter {
     struct CloseWithDispatch<S> {
         inner: S,
         dispatch: OnceLock<WeakDispatch>,
+        closed: Arc<tokio::sync::Notify>,
     }
 
     impl<S: Subscriber> Subscriber for CloseWithDispatch<S> {
@@ -224,7 +225,14 @@ pub mod test_exporter {
                 .get()
                 .and_then(WeakDispatch::upgrade)
                 .expect("a live span keeps its test dispatcher alive");
-            tracing::dispatcher::with_default(&dispatch, || self.inner.try_close(id))
+            let closed = tracing::dispatcher::with_default(&dispatch, || self.inner.try_close(id));
+            if closed {
+                // The simple exporter has finished before try_close returns.
+                // Wake assertions only after the owning registry and layers
+                // have completed cleanup, including recursive parent closure.
+                self.closed.notify_waiters();
+            }
+            closed
         }
 
         fn current_span(&self) -> tracing_core::span::Current {
@@ -273,14 +281,17 @@ pub mod test_exporter {
             .with_simple_exporter(exporter.clone())
             .build();
         let subscriber = tracing_subscriber::registry().with(super::layer(&provider, None));
+        let closed = Arc::new(tokio::sync::Notify::new());
         let dispatch = Dispatch::new(CloseWithDispatch {
             inner: subscriber,
             dispatch: OnceLock::new(),
+            closed: Arc::clone(&closed),
         });
         TracingTestGuard {
             _default: tracing::dispatcher::set_default(&dispatch),
             _provider: provider,
             exporter,
+            closed,
             _lock: lock,
         }
     }
@@ -291,6 +302,52 @@ pub mod test_exporter {
             self.exporter.get_finished_spans().expect("in-memory spans")
         }
 
+        /// Wait for expected spans to finish before taking an assertion snapshot.
+        ///
+        /// A completed `SQLx` query may still have its span held by a `SQLite`
+        /// worker. Export is synchronous once the span closes, but flushing
+        /// cannot close that live span. Await closure notifications instead of
+        /// assuming the query result also means tracing cleanup has completed.
+        pub async fn wait_for_spans(
+            &self,
+            predicate: impl Fn(&[opentelemetry_sdk::trace::SpanData]) -> bool + Send + Sync,
+        ) -> Vec<opentelemetry_sdk::trace::SpanData> {
+            tokio::time::timeout(std::time::Duration::from_secs(5), async {
+                loop {
+                    // notify_waiters wakes futures created before notification,
+                    // even before polling. Subscribe before reading so closure
+                    // between the snapshot and await cannot lose a wakeup.
+                    let notified = self.closed.notified();
+                    let spans = self.finished_spans();
+                    if predicate(&spans) {
+                        return spans;
+                    }
+                    notified.await;
+                }
+            })
+            .await
+            .unwrap_or_else(|_| {
+                panic!(
+                    "timed out waiting for expected spans, got {:?}",
+                    self.finished_spans()
+                        .iter()
+                        .map(|span| &span.name)
+                        .collect::<Vec<_>>()
+                )
+            })
+        }
+
+        /// Wait for the completed span named `name`.
+        pub async fn wait_for_span(&self, name: &str) -> opentelemetry_sdk::trace::SpanData {
+            let spans = self
+                .wait_for_spans(|spans| spans.iter().any(|span| span.name == name))
+                .await;
+            spans
+                .into_iter()
+                .find(|span| span.name == name)
+                .expect("the awaited snapshot contains the expected span")
+        }
+
         /// Spans named `name`.
         pub fn spans_named(&self, name: &str) -> Vec<opentelemetry_sdk::trace::SpanData> {
             self.finished_spans()
@@ -378,6 +435,7 @@ pub mod test_exporter {
         _default: tracing::dispatcher::DefaultGuard,
         _provider: opentelemetry_sdk::trace::SdkTracerProvider,
         exporter: opentelemetry_sdk::trace::InMemorySpanExporter,
+        closed: Arc<tokio::sync::Notify>,
         _lock: std::sync::MutexGuard<'static, ()>,
     }
 
@@ -690,6 +748,56 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn tracing_waits_for_worker_held_spans_to_close() {
+        let traced = test_exporter::install_traced();
+        let parent = tracing::info_span!("awaited_parent");
+        let child = tracing::info_span!(parent: &parent, "awaited_child");
+        drop(parent);
+
+        let (release, released) = std::sync::mpsc::channel();
+        let worker = std::thread::spawn(move || {
+            released.recv().expect("test releases the worker's span");
+            drop(child);
+        });
+
+        assert!(traced.finished_spans().is_empty());
+        let waiting = traced.wait_for_span("awaited_parent");
+        tokio::pin!(waiting);
+        assert!(futures::poll!(waiting.as_mut()).is_pending());
+
+        release.send(()).unwrap();
+        let parent = waiting.await;
+        worker.join().expect("worker closes child and parent");
+        let child = traced.span_named("awaited_child");
+        test_exporter::assert_is_root(&parent);
+        assert_eq!(child.parent_span_id, parent.span_context.span_id());
+    }
+
+    #[tokio::test]
+    async fn tracing_wait_does_not_lose_closure_between_snapshot_and_await() {
+        let traced = test_exporter::install_traced();
+        let span = std::sync::Mutex::new(Some(tracing::info_span!("close_before_await")));
+        let spans = traced
+            .wait_for_spans(|spans| {
+                // The first snapshot is empty. Close its span before polling
+                // the notification future, as a worker could do concurrently.
+                drop(span.lock().unwrap().take());
+                spans.iter().any(|span| span.name == "close_before_await")
+            })
+            .await;
+        assert_eq!(spans.len(), 1);
+        test_exporter::assert_is_root(&spans[0]);
+    }
+
+    #[tokio::test(start_paused = true)]
+    #[should_panic(expected = "timed out waiting for expected spans")]
+    async fn tracing_wait_times_out_when_a_span_never_closes() {
+        let traced = test_exporter::install_traced();
+        let _span = tracing::info_span!("still_open");
+        traced.wait_for_span("still_open").await;
+    }
+
     #[tokio::test]
     async fn tracing_events_are_not_exported() {
         let traced = test_exporter::install_traced();

crates/openshell-server/src/persistence/tests.rs

@@ -35,7 +35,6 @@ async fn failed_store_calls_are_marked_on_the_span() {
 /// the span must stay clean — otherwise every lease a replica does not win, and
 /// every gateway restart, exports as a failure.
 #[tokio::test]
-#[ignore = "flaky under concurrent test execution"]
 async fn expected_conflicts_leave_the_span_unmarked() {
     use crate::otel_tracing::test_exporter;
 
@@ -66,6 +65,7 @@ async fn expected_conflicts_leave_the_span_unmarked() {
         .await
         .expect_err("the name is already taken");
 
+    traced.wait_for_span("store.put_if").await;
     let span = traced.span_with("store.put_if", "object.id", "expected-conflict-second");
 
     assert_eq!(
@@ -79,7 +79,6 @@ async fn expected_conflicts_leave_the_span_unmarked() {
 /// Span names stay low-cardinality so they group across object types; what
 /// each call touched is carried as attributes.
 #[tokio::test]
-#[ignore = "flaky under concurrent test execution"]
 async fn store_spans_record_what_they_touched_as_attributes() {
     use crate::otel_tracing::test_exporter;
 
@@ -97,6 +96,13 @@ async fn store_spans_record_what_they_touched_as_attributes() {
         .unwrap();
     store.list("sandbox", "default", 10, 0).await.unwrap();
 
+    traced
+        .wait_for_spans(|spans| {
+            ["store.get", "store.get_by_name", "store.list"]
+                .iter()
+                .all(|name| spans.iter().any(|span| &span.name == name))
+        })
+        .await;
     let by_name = traced.span_with("store.get_by_name", "object.name", "my-sandbox");
     assert_eq!(
         test_exporter::attribute(&by_name, "object_type").as_deref(),
@@ -2807,7 +2813,6 @@ async fn membership_selector_escapes_adversarial_label_key() {
 /// so a trace decomposes an RPC into the storage work it did rather than
 /// bottoming out at the request boundary.
 #[tokio::test]
-#[ignore = "flaky under concurrent test execution"]
 async fn store_operations_export_spans_with_parents() {
     use tracing::Instrument as _;
 
@@ -2827,7 +2832,7 @@ async fn store_operations_export_spans_with_parents() {
     .await;
     drop(request_span);
 
-    let root = traced.span_named("request");
+    let root = traced.wait_for_span("request").await;
     let spans = traced.finished_spans();
     let child = spans
         .iter()

crates/openshell-server/src/provider_refresh.rs

@@ -3580,7 +3580,6 @@ mod tests {
     }
 
     #[tokio::test]
-    #[ignore = "flaky under concurrent test execution"]
     async fn refresh_worker_records_a_root_span_only_when_a_state_has_work() {
         use crate::otel_tracing::test_exporter;
 
@@ -3621,7 +3620,7 @@ mod tests {
         Box::pin(run_refresh_worker_tick(&store, None, None))
             .await
             .unwrap();
-        test_exporter::assert_is_root(&traced.span_named("refresh.provider_credentials"));
+        test_exporter::assert_is_root(&traced.wait_for_span("refresh.provider_credentials").await);
     }
 
     #[test]