fix(analytics): drain queued and in-flight exports on shutdown

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<B, E>(
+        config: CollectorConfig,
+        batch_factory: B,
+        exporter: E,
+    ) -> (Self, ShutdownHandle)
+    where
+        B: BatchFactory<T>,
+        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<T, B, E> {
     batch_factory: B,
     exporter: E,
     config: CollectorConfig,
+    exports: JoinSet<Result<(), InternalError>>,
     _marker: PhantomData<T>,
 }
 
@@ -111,13 +138,29 @@ where
             batch_factory,
             exporter,
             config,
+            exports: JoinSet::new(),
             _marker: PhantomData,
         }
     }
 
-    async fn run(self, data_rx: mpsc::Receiver<T>) -> Result<(), InternalError> {
+    async fn run(mut self, data_rx: mpsc::Receiver<T>, 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<T>,
+        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,33 +223,34 @@ 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(());
         }
 
         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<Result<(), InternalError>, 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"),
+        }
+    }
 }

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},
 };

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<u8>, oneshot::Sender<Result<(), std::io::Error>>);
+
+#[derive(Clone)]
+struct ControlledExporter(mpsc::Sender<PendingExport>);
+
+#[async_trait]
+impl Exporter for ControlledExporter {
+    type Error = std::io::Error;
+
+    async fn export(self, data: Vec<u8>) -> Result<(), Self::Error> {
+        let (tx, rx) = oneshot::channel();
+        self.0.send((data, tx)).await.unwrap();
+        rx.await.unwrap()
+    }
+}
+
+fn record_ids(data: Vec<u8>) -> Vec<u32> {
+    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<DataA> for FailingBatchFactory {
+    type Batch = <ParquetBatchFactory as BatchFactory<DataA>>::Batch;
+    type Error = std::io::Error;
+
+    fn create(&self) -> Result<Self::Batch, Self::Error> {
+        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();
+}