diff --git a/crates/analytics/src/collectors.rs b/crates/analytics/src/collectors.rs index 6f36204..1fb18b5 100644 --- a/crates/analytics/src/collectors.rs +++ b/crates/analytics/src/collectors.rs @@ -1,7 +1,10 @@ use { crate::{AnalyticsEvent, Batch, BatchFactory, Collector, Exporter}, std::{marker::PhantomData, pin::pin, time::Duration}, - tokio::sync::{mpsc, mpsc::error::TrySendError}, + tokio::{ + sync::{mpsc, mpsc::error::TrySendError, oneshot}, + task::{JoinError, JoinHandle, JoinSet}, + }, }; #[derive(Debug, thiserror::Error)] @@ -16,6 +19,20 @@ enum InternalError { Serialization, } +pub struct ShutdownHandle { + shutdown_tx: oneshot::Sender<()>, + task: JoinHandle<()>, +} + +impl ShutdownHandle { + /// Closes the queue and waits for exports. Batch/export failures are + /// logged. + pub async fn shutdown(self) -> Result<(), JoinError> { + let _ = self.shutdown_tx.send(()); + self.task.await + } +} + #[derive(Debug, thiserror::Error)] pub enum CollectionError { #[error("Data channel overflow")] @@ -66,17 +83,26 @@ where B::Error: std::error::Error, E: Exporter, { - let (data_tx, data_rx) = mpsc::channel(config.data_queue_capacity); - - tokio::spawn(async move { - let event_loop = EventLoop::new(batch_factory, exporter, config); + Self::new_with_shutdown(config, batch_factory, exporter).0 + } - if let Err(err) = event_loop.run(data_rx).await { - tracing::warn!(?err, "analytics event loop failed"); - } - }); + /// Dropping the shutdown handle leaves the collector running. + pub fn new_with_shutdown( + config: CollectorConfig, + batch_factory: B, + exporter: E, + ) -> (Self, ShutdownHandle) + where + B: BatchFactory, + B::Error: std::error::Error, + E: Exporter, + { + let (data_tx, data_rx) = mpsc::channel(config.data_queue_capacity); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let event_loop = EventLoop::new(batch_factory, exporter, config); + let task = tokio::spawn(event_loop.run(data_rx, shutdown_rx)); - Self { data_tx } + (Self { data_tx }, ShutdownHandle { shutdown_tx, task }) } } @@ -95,6 +121,7 @@ struct EventLoop { batch_factory: B, exporter: E, config: CollectorConfig, + exports: JoinSet>, _marker: PhantomData, } @@ -111,13 +138,29 @@ where batch_factory, exporter, config, + exports: JoinSet::new(), _marker: PhantomData, } } - async fn run(self, data_rx: mpsc::Receiver) -> Result<(), InternalError> { + async fn run(mut self, data_rx: mpsc::Receiver, shutdown_rx: oneshot::Receiver<()>) { + if let Err(err) = self.collect(data_rx, shutdown_rx).await { + tracing::warn!(?err, "analytics event loop failed"); + } + + while let Some(result) = self.exports.join_next().await { + Self::finish_export(result); + } + } + + async fn collect( + &mut self, + data_rx: mpsc::Receiver, + mut shutdown_rx: oneshot::Receiver<()>, + ) -> Result<(), InternalError> { let mut data_rx = pin!(data_rx); let mut export_interval = pin!(tokio::time::interval(self.config.export_interval)); + let mut shutdown_pending = true; let mut current_batch = self .batch_factory @@ -126,6 +169,17 @@ where loop { tokio::select! { + result = &mut shutdown_rx, if shutdown_pending => { + shutdown_pending = false; + if result.is_ok() { + data_rx.close(); + } + }, + + Some(result) = self.exports.join_next(), if !self.exports.is_empty() => { + Self::finish_export(result); + }, + data = data_rx.recv() => match data { Some(data) => { if let Err(err) = current_batch.push(data) { @@ -147,7 +201,7 @@ where } }, - // The transmitter has been dropped. Export current batch and shutdown. + // The channel is drained. Export the final partial batch. None => { return self.export_batch(&mut current_batch); }, @@ -169,7 +223,7 @@ where Ok(std::mem::replace(current_batch, next_batch)) } - fn export_batch(&self, current_batch: &mut B::Batch) -> Result<(), InternalError> { + fn export_batch(&mut self, current_batch: &mut B::Batch) -> Result<(), InternalError> { if current_batch.is_empty() { return Ok(()); } @@ -177,25 +231,26 @@ where let current_batch = self.replace_batch(current_batch)?; let exporter = self.exporter.clone(); - tokio::spawn(async move { - let result = async { - let data = tokio::task::spawn_blocking(move || current_batch.serialize()) - .await - .map_err(|_| InternalError::Serialization)? - .map_err(|err| InternalError::Batch(err.to_string()))?; - - exporter - .export(data) - .await - .map_err(|err| InternalError::Export(err.to_string())) - } - .await; + self.exports.spawn(async move { + let data = tokio::task::spawn_blocking(move || current_batch.serialize()) + .await + .map_err(|_| InternalError::Serialization)? + .map_err(|err| InternalError::Batch(err.to_string()))?; - if let Err(err) = result { - tracing::warn!(?err, "failed to export batch data"); - } + exporter + .export(data) + .await + .map_err(|err| InternalError::Export(err.to_string())) }); Ok(()) } + + fn finish_export(result: Result, JoinError>) { + match result { + Ok(Ok(())) => {} + Ok(Err(err)) => tracing::warn!(?err, "failed to export batch data"), + Err(err) => tracing::warn!(?err, "analytics export task failed"), + } + } } diff --git a/crates/analytics/src/lib.rs b/crates/analytics/src/lib.rs index 1e49f86..30e1d7f 100644 --- a/crates/analytics/src/lib.rs +++ b/crates/analytics/src/lib.rs @@ -7,7 +7,7 @@ use { tap::Tap, }; pub use { - collectors::{BatchCollector, CollectionError, CollectorConfig}, + collectors::{BatchCollector, CollectionError, CollectorConfig, ShutdownHandle}, exporters::{AwsConfig, AwsError, AwsExporter, NoopExporter}, serializers::{NoopBatchFactory, ParquetBatchFactory, ParquetConfig, ParquetError}, }; diff --git a/crates/analytics/tests/integration.rs b/crates/analytics/tests/integration.rs index 5e0834b..8e09259 100644 --- a/crates/analytics/tests/integration.rs +++ b/crates/analytics/tests/integration.rs @@ -2,6 +2,7 @@ use { analytics::{ AnalyticsExt, BatchCollector, + BatchFactory, BatchObserver, CollectionObserver, Collector, @@ -12,6 +13,10 @@ use { ParquetConfig, }, async_trait::async_trait, + parquet::{ + file::reader::{FileReader, SerializedFileReader}, + record::RowAccessor, + }, parquet_derive::ParquetRecordWriter, std::{ sync::{ @@ -20,7 +25,7 @@ use { }, time::Duration, }, - tokio::sync::{mpsc, mpsc::error::TrySendError}, + tokio::sync::{mpsc, mpsc::error::TrySendError, oneshot}, }; #[derive(Clone)] @@ -200,3 +205,168 @@ async fn observability() { assert_eq!(observer.batch_serialization.load(Ordering::SeqCst), 1); assert_eq!(observer.collection.load(Ordering::SeqCst), 2); } + +type PendingExport = (Vec, oneshot::Sender>); + +#[derive(Clone)] +struct ControlledExporter(mpsc::Sender); + +#[async_trait] +impl Exporter for ControlledExporter { + type Error = std::io::Error; + + async fn export(self, data: Vec) -> Result<(), Self::Error> { + let (tx, rx) = oneshot::channel(); + self.0.send((data, tx)).await.unwrap(); + rx.await.unwrap() + } +} + +fn record_ids(data: Vec) -> Vec { + SerializedFileReader::new(bytes::Bytes::from(data)) + .unwrap() + .get_row_iter(None) + .unwrap() + .map(|row| row.unwrap().get_uint(0).unwrap()) + .collect() +} + +#[tokio::test] +async fn shutdown_drains_records_and_waits_for_exports() { + tokio::time::timeout(Duration::from_secs(5), async { + let (tx, mut rx) = mpsc::channel(8); + let (collector, shutdown) = BatchCollector::new_with_shutdown( + CollectorConfig { + data_queue_capacity: 3, + ..Default::default() + }, + ParquetBatchFactory::new(ParquetConfig { + batch_capacity: 2, + alloc_buffer_size: 8192, + }), + ControlledExporter(tx), + ); + + for a in 0..2 { + collector + .collect(DataA { + a, + b: "foo", + c: true, + }) + .unwrap(); + } + let (data, first_export) = rx.recv().await.unwrap(); + let mut ids = record_ids(data); + for a in 2..5 { + collector + .collect(DataA { + a, + b: "foo", + c: true, + }) + .unwrap(); + } + + let mut shutdown = tokio::spawn(shutdown.shutdown()); + while ids.len() < 5 { + let (data, complete) = rx.recv().await.unwrap(); + ids.extend(record_ids(data)); + complete.send(Ok(())).unwrap(); + } + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut shutdown) + .await + .is_err() + ); + first_export.send(Ok(())).unwrap(); + shutdown.await.unwrap().unwrap(); + + ids.sort_unstable(); + assert_eq!(ids, [0, 1, 2, 3, 4]); + assert!(matches!( + collector.collect(DataA { + a: 5, + b: "foo", + c: true + }), + Err(analytics::CollectionError::DataChannelClosed) + )); + }) + .await + .unwrap(); +} + +struct FailingBatchFactory { + inner: ParquetBatchFactory, + creates: AtomicUsize, + fail_after: usize, +} + +impl BatchFactory for FailingBatchFactory { + type Batch = >::Batch; + type Error = std::io::Error; + + fn create(&self) -> Result { + if self.creates.fetch_add(1, Ordering::Relaxed) == self.fail_after { + return Err(std::io::Error::other("batch failed")); + } + self.inner.create().map_err(std::io::Error::other) + } +} + +#[tokio::test] +async fn shutdown_waits_for_exports_after_an_error() { + tokio::time::timeout(Duration::from_secs(5), async { + for fail_after in [2, usize::MAX] { + let (tx, mut rx) = mpsc::channel(8); + let (collector, shutdown) = BatchCollector::new_with_shutdown( + CollectorConfig::default(), + FailingBatchFactory { + inner: ParquetBatchFactory::new(ParquetConfig { + batch_capacity: 1, + alloc_buffer_size: 8192, + }), + creates: AtomicUsize::new(0), + fail_after, + }, + ControlledExporter(tx), + ); + collector + .collect(DataA { + a: 0, + b: "foo", + c: true, + }) + .unwrap(); + let (_, first_export) = rx.recv().await.unwrap(); + collector + .collect(DataA { + a: 1, + b: "foo", + c: true, + }) + .unwrap(); + let blocked_export = if fail_after == 2 { + first_export + } else { + let (_, second_export) = rx.recv().await.unwrap(); + first_export + .send(Err(std::io::Error::other("upload failed"))) + .unwrap(); + second_export + }; + + let mut shutdown = tokio::spawn(shutdown.shutdown()); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut shutdown) + .await + .is_err() + ); + blocked_export.send(Ok(())).unwrap(); + shutdown.await.unwrap().unwrap(); + } + }) + .await + .unwrap(); +}