Make SpanProcessor::{on_end,force_flush,shutdown} async

Author: mendess

opentelemetry-sdk/Cargo.toml

@@ -39,6 +39,7 @@ rustdoc-args = ["--cfg", "docsrs"]
 [dev-dependencies]
 criterion = { workspace = true, features = ["html_reports"] }
 temp-env = { workspace = true }
+tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
 
 [target.'cfg(not(target_os = "windows"))'.dev-dependencies]
 pprof = { version = "0.13", features = ["flamegraph", "criterion"] }

opentelemetry-sdk/benches/batch_span_processor.rs

@@ -64,7 +64,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                             let spans = get_span_data();
                             handles.push(tokio::spawn(async move {
                                 for span in spans {
-                                    span_processor.on_end(span);
+                                    span_processor.on_end(span).await;
                                     tokio::task::yield_now().await;
                                 }
                             }));
@@ -73,7 +73,8 @@ fn criterion_benchmark(c: &mut Criterion) {
                         let _ =
                             Arc::<BatchSpanProcessor<Tokio>>::get_mut(&mut shared_span_processor)
                                 .unwrap()
-                                .shutdown();
+                                .shutdown()
+                                .await;
                     });
                 })
             },

opentelemetry-sdk/src/testing/trace/in_memory_exporter.rs

@@ -37,7 +37,7 @@ use std::sync::{Arc, Mutex};
 ///     cx.span().add_event("handling this...", Vec::new());
 ///     cx.span().end();
 ///
-///     let results = provider.force_flush();
+///     let results = provider.force_flush().await;
 ///     for result in results {
 ///         if let Err(e) = result {
 ///             println!("{:?}", e)

opentelemetry-sdk/src/trace/provider.rs

@@ -14,6 +14,7 @@ use crate::trace::{
 };
 use crate::{export::trace::SpanExporter, trace::SpanProcessor};
 use crate::{InstrumentationLibrary, Resource};
+use futures_util::StreamExt;
 use once_cell::sync::{Lazy, OnceCell};
 use opentelemetry::trace::TraceError;
 use opentelemetry::{global, trace::TraceResult};
@@ -49,11 +50,14 @@ pub(crate) struct TracerProviderInner {
 
 impl Drop for TracerProviderInner {
     fn drop(&mut self) {
-        for processor in &mut self.processors {
-            if let Err(err) = processor.shutdown() {
-                global::handle_error(err);
+        let processors = std::mem::take(&mut self.processors);
+        crate::util::spawn_future(async move {
+            for processor in processors {
+                if let Err(err) = processor.shutdown().await {
+                    global::handle_error(err);
+                }
             }
-        }
+        })
     }
 }
 
@@ -121,13 +125,14 @@ impl TracerProvider {
     ///     provider
     /// }
     ///
-    /// fn main() {
+    /// #[tokio::main]
+    /// async fn main() {
     ///     let provider = init_tracing();
     ///
     ///     // create spans..
     ///
     ///     // force all spans to flush
-    ///     for result in provider.force_flush() {
+    ///     for result in provider.force_flush().await {
     ///         if let Err(err) = result {
     ///             // .. handle flush error
     ///         }
@@ -141,17 +146,17 @@ impl TracerProvider {
     ///     global::shutdown_tracer_provider();
     /// }
     /// ```
-    pub fn force_flush(&self) -> Vec<TraceResult<()>> {
-        self.span_processors()
-            .iter()
-            .map(|processor| processor.force_flush())
+    pub async fn force_flush(&self) -> Vec<TraceResult<()>> {
+        futures_util::stream::iter(self.span_processors())
+            .then(|processor| processor.force_flush())
             .collect()
+            .await
     }
 
     /// Shuts down the current `TracerProvider`.
     ///
     /// Note that shut down doesn't means the TracerProvider has dropped
-    pub fn shutdown(&self) -> TraceResult<()> {
+    pub async fn shutdown(&self) -> TraceResult<()> {
         if self
             .is_shutdown
             .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
@@ -161,7 +166,7 @@ impl TracerProvider {
             // it's up to the processor to properly block new spans after shutdown
             let mut errs = vec![];
             for processor in &self.inner.processors {
-                if let Err(err) = processor.shutdown() {
+                if let Err(err) = processor.shutdown().await {
                     errs.push(err);
                 }
             }
@@ -348,6 +353,7 @@ mod tests {
         }
     }
 
+    #[async_trait::async_trait]
     impl SpanProcessor for TestSpanProcessor {
         fn on_start(&self, _span: &mut Span, _cx: &Context) {
             self.assert_info
@@ -356,19 +362,19 @@ mod tests {
                 .fetch_add(1, Ordering::SeqCst);
         }
 
-        fn on_end(&self, _span: SpanData) {
+        async fn on_end(&self, _span: SpanData) {
             // ignore
         }
 
-        fn force_flush(&self) -> TraceResult<()> {
+        async fn force_flush(&self) -> TraceResult<()> {
             if self.success {
                 Ok(())
             } else {
                 Err(TraceError::from("cannot export"))
             }
         }
 
-        fn shutdown(&self) -> TraceResult<()> {
+        async fn shutdown(&self) -> TraceResult<()> {
             if self.assert_info.0.is_shutdown.load(Ordering::SeqCst) {
                 Ok(())
             } else {
@@ -378,13 +384,13 @@ mod tests {
                     Ordering::SeqCst,
                     Ordering::SeqCst,
                 );
-                self.force_flush()
+                self.force_flush().await
             }
         }
     }
 
-    #[test]
-    fn test_force_flush() {
+    #[tokio::test]
+    async fn test_force_flush() {
         let tracer_provider = super::TracerProvider::new(TracerProviderInner {
             processors: vec![
                 Box::from(TestSpanProcessor::new(true)),
@@ -393,7 +399,7 @@ mod tests {
             config: Default::default(),
         });
 
-        let results = tracer_provider.force_flush();
+        let results = tracer_provider.force_flush().await;
         assert_eq!(results.len(), 2);
     }
 
@@ -527,8 +533,8 @@ mod tests {
         assert_eq!(no_service_name.config().resource.len(), 0)
     }
 
-    #[test]
-    fn test_shutdown_noops() {
+    #[tokio::test]
+    async fn test_shutdown_noops() {
         let processor = TestSpanProcessor::new(false);
         let assert_handle = processor.assert_info();
         let tracer_provider = super::TracerProvider::new(TracerProviderInner {
@@ -545,12 +551,12 @@ mod tests {
 
         assert!(assert_handle.started_span_count(2));
 
-        let shutdown = |tracer_provider: super::TracerProvider| {
-            let _ = tracer_provider.shutdown(); // shutdown once
+        let shutdown = |tracer_provider: super::TracerProvider| async move {
+            let _ = tracer_provider.shutdown().await; // shutdown once
         };
 
         // assert tracer provider can be shutdown using on a cloned version
-        shutdown(tracer_provider.clone());
+        shutdown(tracer_provider.clone()).await;
 
         // after shutdown we should get noop tracer
         let noop_tracer = tracer_provider.tracer("noop");

opentelemetry-sdk/src/trace/span.rs

@@ -14,6 +14,8 @@ use opentelemetry::KeyValue;
 use std::borrow::Cow;
 use std::time::SystemTime;
 
+use super::TracerProvider;
+
 /// Single operation within a trace.
 #[derive(Debug)]
 pub struct Span {
@@ -192,24 +194,45 @@ impl opentelemetry::trace::Span for Span {
 
     /// Finishes the span with given timestamp.
     fn end_with_timestamp(&mut self, timestamp: SystemTime) {
-        self.ensure_ended_and_exported(Some(timestamp));
+        if let Some(fut) = self.ensure_ended_and_exported(Some(timestamp)) {
+            crate::util::spawn_future(fut)
+        }
     }
 }
 
 impl Span {
-    fn ensure_ended_and_exported(&mut self, timestamp: Option<SystemTime>) {
+    fn ensure_ended_and_exported(
+        &mut self,
+        timestamp: Option<SystemTime>,
+    ) -> Option<impl std::future::Future<Output = ()>> {
         // skip if data has already been exported
-        let mut data = match self.data.take() {
+        let data = match self.data.take() {
             Some(data) => data,
-            None => return,
+            None => return None,
         };
 
         let provider = self.tracer.provider();
         // skip if provider has been shut down
         if provider.is_shutdown() {
-            return;
+            return None;
         }
 
+        Some(Self::ensure_ended_and_exported_impl(
+            data,
+            provider.clone(),
+            self.tracer.clone(),
+            self.span_context.clone(),
+            timestamp,
+        ))
+    }
+
+    async fn ensure_ended_and_exported_impl(
+        mut data: SpanData,
+        provider: TracerProvider,
+        tracer: crate::trace::Tracer,
+        span_context: SpanContext,
+        timestamp: Option<SystemTime>,
+    ) {
         // ensure end time is set via explicit end or implicitly on drop
         if let Some(timestamp) = timestamp {
             data.end_time = timestamp;
@@ -220,19 +243,19 @@ impl Span {
         match provider.span_processors() {
             [] => {}
             [processor] => {
-                processor.on_end(build_export_data(
-                    data,
-                    self.span_context.clone(),
-                    &self.tracer,
-                ));
+                processor
+                    .on_end(build_export_data(data, span_context, &tracer))
+                    .await;
             }
             processors => {
                 for processor in processors {
-                    processor.on_end(build_export_data(
-                        data.clone(),
-                        self.span_context.clone(),
-                        &self.tracer,
-                    ));
+                    processor
+                        .on_end(build_export_data(
+                            data.clone(),
+                            span_context.clone(),
+                            &tracer,
+                        ))
+                        .await;
                 }
             }
         }
@@ -242,7 +265,9 @@ impl Span {
 impl Drop for Span {
     /// Report span on inner drop
     fn drop(&mut self) {
-        self.ensure_ended_and_exported(None);
+        if let Some(fut) = self.ensure_ended_and_exported(None) {
+            crate::util::spawn_future(fut)
+        }
     }
 }
 
@@ -705,8 +730,8 @@ mod tests {
         assert_eq!(event_vec.len(), DEFAULT_MAX_EVENT_PER_SPAN as usize);
     }
 
-    #[test]
-    fn test_span_exported_data() {
+    #[tokio::test]
+    async fn test_span_exported_data() {
         let provider = crate::trace::TracerProvider::builder()
             .with_simple_exporter(NoopSpanExporter::new())
             .build();
@@ -719,7 +744,7 @@ mod tests {
         let exported_data = span.exported_data();
         assert!(exported_data.is_some());
 
-        provider.shutdown().expect("shutdown panicked");
+        provider.shutdown().await.expect("shutdown panicked");
         let dropped_span = tracer.start("span_with_dropped_provider");
         // return none if the provider has already been dropped
         assert!(dropped_span.exported_data().is_none());