Merge pull request #143 from Thaumy/main

Replace std MPSC channel with crossbeam `SegQueue` for `DataExporter`

Author: saying121

Cargo.lock

@@ -49,6 +49,7 @@ tracing-subscriber = { workspace = true, features = ["env-filter"] }
 daemonize = { workspace = true }
 local-ip-address = { workspace = true }
 TinyUFO = { workspace = true }
+crossbeam = { workspace = true }
 
 [build-dependencies]
 tonic-build = { workspace = true }
@@ -89,6 +90,7 @@ which = "^7"
 num_cpus = "^1"
 local-ip-address = "^0.6"
 TinyUFO = "0.4"
+crossbeam = "0.8"
 
 [workspace.lints.rust]
 

Cargo.toml

@@ -15,13 +15,13 @@
 use std::{
     sync::{
         atomic::{AtomicUsize, Ordering},
-        mpsc::{channel, SendError, Sender, TryRecvError},
         Arc,
     },
     thread::{self, JoinHandle},
 };
 
 use chrono::{TimeZone, Utc};
+use crossbeam::queue::SegQueue;
 use profiling::data_export::common::FieldValue as WitFieldValue;
 use profiling::data_export::measurement::Point;
 use profiling::data_export::metric::Sample;
@@ -49,7 +49,7 @@ wasmtime::component::bindgen!({
 pub struct DataExporter {
     bytes_len: Arc<AtomicUsize>,
     bytes_watermark: usize,
-    data_tx: Option<Sender<Data>>,
+    data_queue: Arc<SegQueue<Option<Data>>>,
     exporter: JoinHandle<()>,
 }
 
@@ -63,22 +63,23 @@ impl DataExporter {
         // TODO: `bytes_capacity` is not used because we not have a static allocation
         // for `Data`, maybe we will remove this in the future
         let _ = bytes_capacity;
-        let (data_tx, data_rx) = channel::<Data>();
+        let data_queue = Arc::new(SegQueue::<Option<Data>>::new());
         let bytes_len = Arc::new(AtomicUsize::new(0));
 
         let exporter = thread::spawn({
+            let data_queue = Arc::clone(&data_queue);
             let bytes_len = Arc::clone(&bytes_len);
             move || {
                 let rt = Runtime::new().expect("Failed to init exporter runtime");
                 let mut data = Vec::new();
                 loop {
-                    match data_rx.try_recv() {
-                        Ok(o) => {
+                    match data_queue.pop() {
+                        Some(Some(o)) => {
                             // No critical section, relaxed ordering is fine.
                             bytes_len.fetch_sub(o.encoded_len(), Ordering::Relaxed);
                             data.push(o);
                         }
-                        Err(e) => {
+                        poped => {
                             if !data.is_empty() {
                                 let merged = ExportDataReq {
                                     task_id: task_id.clone(),
@@ -91,9 +92,10 @@ impl DataExporter {
                                 });
                                 data.clear();
                             }
-                            match e {
-                                TryRecvError::Empty => thread::park(),
-                                TryRecvError::Disconnected => break,
+                            match poped {
+                                None => thread::park(),
+                                Some(None) => break,
+                                _ => unreachable!(),
                             }
                         }
                     }
@@ -104,7 +106,7 @@ impl DataExporter {
         Self {
             bytes_len,
             bytes_watermark,
-            data_tx: Some(data_tx),
+            data_queue,
             exporter,
         }
     }
@@ -113,24 +115,21 @@ impl DataExporter {
         self.exporter.thread().unpark();
     }
 
-    pub fn schedule(&self, data: Data) -> Result<(), SendError<Data>> {
+    pub fn schedule(&self, data: Data) {
         let encoded_len = data.encoded_len();
-        // Never failes since we only take it in `Drop`.
-        let data_tx = self.data_tx.as_ref().expect("unreachable");
-        data_tx.send(data)?;
+        self.data_queue.push(Some(data));
         // No critical section, relaxed ordering is fine.
         let prev = self.bytes_len.fetch_add(encoded_len, Ordering::Relaxed);
         if prev > self.bytes_watermark {
             self.exporter.thread().unpark();
         }
-        Ok(())
     }
 }
 
 impl Drop for DataExporter {
     fn drop(&mut self) {
-        // Drop sender early to prevent the exporter thread from being parked forever.
-        drop(self.data_tx.take());
+        // Notify the consumer that there is no more data.
+        self.data_queue.push(None);
         self.exporter.thread().unpark();
     }
 }
@@ -178,7 +177,7 @@ impl profiling::data_export::file::Host for DataExportCtx {
             ty: DataType::File as _,
             bytes,
         };
-        ctx.exporter.schedule(data)?;
+        ctx.exporter.schedule(data);
 
         Ok(Ok(()))
     }
@@ -211,7 +210,7 @@ impl profiling::data_export::metric::Host for DataExportCtx {
             ty: DataType::LineProtocol as _,
             bytes,
         };
-        ctx.exporter.schedule(data)?;
+        ctx.exporter.schedule(data);
 
         Ok(Ok(()))
     }
@@ -247,7 +246,7 @@ impl profiling::data_export::measurement::Host for DataExportCtx {
             ty: DataType::LineProtocol as _,
             bytes,
         };
-        ctx.exporter.schedule(data)?;
+        ctx.exporter.schedule(data);
 
         Ok(Ok(()))
     }