fix: BatchSpanProcessor implementation to match that of logs (#3116)

Co-authored-by: Lalit Kumar Bhasin <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

opentelemetry-sdk/src/logs/in_memory_exporter.rs

@@ -132,6 +132,7 @@ impl InMemoryLogExporterBuilder {
 
     /// If set, the records will not be [`InMemoryLogExporter::reset`] on shutdown.
     #[cfg(test)]
+    #[allow(dead_code)]
     pub(crate) fn keep_records_on_shutdown(self) -> Self {
         Self {
             reset_on_shutdown: false,

opentelemetry-sdk/src/trace/in_memory_exporter.rs

@@ -2,7 +2,7 @@ use crate::error::{OTelSdkError, OTelSdkResult};
 use crate::resource::Resource;
 use crate::trace::{SpanData, SpanExporter};
 use crate::InMemoryExporterError;
-use std::sync::{Arc, Mutex};
+use std::sync::{atomic::AtomicBool, Arc, Mutex};
 use std::time::Duration;
 
 /// An in-memory span exporter that stores span data in memory.
@@ -51,6 +51,8 @@ use std::time::Duration;
 pub struct InMemorySpanExporter {
     spans: Arc<Mutex<Vec<SpanData>>>,
     resource: Arc<Mutex<Resource>>,
+    should_reset_on_shutdown: bool,
+    shutdown_called: Arc<AtomicBool>,
 }
 
 impl Default for InMemorySpanExporter {
@@ -67,7 +69,9 @@ impl Default for InMemorySpanExporter {
 /// let exporter = InMemorySpanExporterBuilder::new().build();
 /// ```
 #[derive(Clone, Debug)]
-pub struct InMemorySpanExporterBuilder {}
+pub struct InMemorySpanExporterBuilder {
+    reset_on_shutdown: bool,
+}
 
 impl Default for InMemorySpanExporterBuilder {
     fn default() -> Self {
@@ -78,19 +82,38 @@ impl Default for InMemorySpanExporterBuilder {
 impl InMemorySpanExporterBuilder {
     /// Creates a new instance of the `InMemorySpanExporterBuilder`.
     pub fn new() -> Self {
-        Self {}
+        Self {
+            reset_on_shutdown: true,
+        }
     }
 
     /// Creates a new instance of the `InMemorySpanExporter`.
     pub fn build(&self) -> InMemorySpanExporter {
         InMemorySpanExporter {
             spans: Arc::new(Mutex::new(Vec::new())),
             resource: Arc::new(Mutex::new(Resource::builder().build())),
+            should_reset_on_shutdown: self.reset_on_shutdown,
+            shutdown_called: Arc::new(AtomicBool::new(false)),
+        }
+    }
+
+    /// If set, the records will not be [`InMemorySpanExporter::reset`] on shutdown.
+    #[cfg(test)]
+    #[allow(dead_code)]
+    pub(crate) fn keep_records_on_shutdown(self) -> Self {
+        Self {
+            reset_on_shutdown: false,
         }
     }
 }
 
 impl InMemorySpanExporter {
+    /// Returns true if shutdown was called.
+    pub fn is_shutdown_called(&self) -> bool {
+        self.shutdown_called
+            .load(std::sync::atomic::Ordering::Relaxed)
+    }
+
     /// Returns the finished span as a vector of `SpanData`.
     ///
     /// # Errors
@@ -140,7 +163,11 @@ impl SpanExporter for InMemorySpanExporter {
     }
 
     fn shutdown_with_timeout(&mut self, _timeout: Duration) -> OTelSdkResult {
-        self.reset();
+        self.shutdown_called
+            .store(true, std::sync::atomic::Ordering::Relaxed);
+        if self.should_reset_on_shutdown {
+            self.reset();
+        }
         Ok(())
     }