Join worker on BatchSpanProcessor shutdown

Author: frigus02

src/sdk/trace/span_processor.rs

@@ -56,7 +56,7 @@
 //! those contexts.
 //!
 //! ```
-//! use futures::{stream};
+//! use futures::{FutureExt as _, stream};
 //! use opentelemetry::{api, sdk, global};
 //! use std::time::Duration;
 //!
@@ -68,7 +68,8 @@
 //!     // Then build a batch processor. You can use whichever executor you have available, for
 //!     // example if you are using `async-std` instead of `tokio` you can replace the spawn and
 //!     // interval functions with `async_std::task::spawn` and `async_std::stream::interval`.
-//!     let batch = sdk::BatchSpanProcessor::builder(exporter, tokio::spawn, tokio::time::interval)
+//!     let spawn = |future| tokio::spawn(future).map(|result| result.expect("task failed"));
+//!     let batch = sdk::BatchSpanProcessor::builder(exporter, spawn, tokio::time::interval)
 //!         .with_max_queue_size(4096)
 //!         .build();
 //!
@@ -96,6 +97,7 @@ use futures::{
     task::{Context, Poll},
     Future, Stream, StreamExt,
 };
+use std::fmt::Debug;
 use std::pin::Pin;
 use std::sync::{Arc, Mutex};
 use std::time;
@@ -134,9 +136,17 @@ impl api::SpanProcessor for SimpleSpanProcessor {
 /// them at a preconfigured interval.
 ///
 /// [`SpanProcessor`]: ../../../api/trace/span_processor/trait.SpanProcessor.html
-#[derive(Debug)]
 pub struct BatchSpanProcessor {
     message_sender: Mutex<mpsc::Sender<BatchMessage>>,
+    worker_handle: Mutex<Option<Pin<Box<dyn Future<Output = ()> + Send + Sync>>>>,
+}
+
+impl Debug for BatchSpanProcessor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BatchSpanProcessor")
+            .field("message_sender", &self.message_sender)
+            .finish()
+    }
 }
 
 impl api::SpanProcessor for BatchSpanProcessor {
@@ -154,6 +164,15 @@ impl api::SpanProcessor for BatchSpanProcessor {
         if let Ok(mut sender) = self.message_sender.lock() {
             let _ = sender.try_send(BatchMessage::Shutdown);
         }
+
+        if let Some(worker_handle) = self
+            .worker_handle
+            .lock()
+            .expect("worker_handle Mutex panicked")
+            .take()
+        {
+            futures::executor::block_on(worker_handle);
+        }
     }
 }
 
@@ -233,12 +252,13 @@ impl BatchSpanProcessor {
         S: Fn(BatchSpanProcessorWorker) -> SO,
         I: Fn(time::Duration) -> IS,
         IS: Stream<Item = ISI> + Send + 'static,
+        SO: Future<Output = ()> + Send + Sync + 'static,
     {
         let (message_sender, message_receiver) = mpsc::channel(config.max_queue_size);
         let ticker = interval(config.scheduled_delay).map(|_| BatchMessage::Tick);
 
         // Spawn worker process via user-defined spawn function.
-        spawn(BatchSpanProcessorWorker {
+        let worker_handle = spawn(BatchSpanProcessorWorker {
             exporter,
             messages: Box::pin(futures::stream::select(message_receiver, ticker)),
             config,
@@ -248,6 +268,7 @@ impl BatchSpanProcessor {
         // Return batch processor with link to worker
         BatchSpanProcessor {
             message_sender: Mutex::new(message_sender),
+            worker_handle: Mutex::new(Some(Box::pin(worker_handle))),
         }
     }
 
@@ -316,6 +337,7 @@ where
     S: Fn(BatchSpanProcessorWorker) -> SO,
     I: Fn(time::Duration) -> IS,
     IS: Stream<Item = ISI> + Send + 'static,
+    SO: Future<Output = ()> + Send + Sync + 'static,
 {
     /// Set max queue size for batches
     pub fn with_max_queue_size(self, size: usize) -> Self {