enhancement(buffers): improve disk buffer lifecycle visibility

changelog.d/disk_buffer_orphan_safety.enhancement.md

@@ -0,0 +1,3 @@
+Vector now detects orphaned disk buffers and provides lifecycle visibility when disk-buffered sinks continue draining in the background during configuration reloads.
+
+authors: Jansen-w

lib/vector-buffers/src/lib.rs

@@ -31,6 +31,7 @@ pub mod test;
 pub mod topology;
 
 pub(crate) mod variants;
+pub use variants::disk_v2::DiskBufferUsageSnapshot;
 
 /// `disk_v2`'s write-buffer size, re-exported under `test` so the harness can
 /// size payloads against the real value instead of hardcoding it.

lib/vector-buffers/src/topology/channel/receiver.rs

@@ -6,15 +6,15 @@ use std::{
 
 use async_recursion::async_recursion;
 use futures::Stream;
-use tokio::select;
+use tokio::{select, task::JoinHandle};
 use tokio_util::sync::ReusableBoxFuture;
 use vector_common::internal_event::emit;
 
 use super::limited_queue::LimitedReceiver;
 use crate::{
     Bufferable,
     buffer_usage_data::BufferUsageHandle,
-    variants::disk_v2::{self, ProductionFilesystem},
+    variants::disk_v2::{self, CleanupTask, ProductionFilesystem},
 };
 
 /// Adapter for papering over various receiver backends.
@@ -62,6 +62,17 @@ where
             },
         }
     }
+
+    fn take_ledger_tasks(
+        &mut self,
+        finalizers: &mut Vec<JoinHandle<()>>,
+        cleanups: &mut Vec<CleanupTask>,
+    ) {
+        if let Self::DiskV2(reader) = self {
+            finalizers.extend(reader.take_finalizer_task());
+            cleanups.extend(reader.take_cleanup_task());
+        }
+    }
 }
 
 /// A buffer receiver.
@@ -152,8 +163,22 @@ impl<T: Bufferable> BufferReceiver<T> {
         Some(item)
     }
 
-    pub fn into_stream(self) -> BufferReceiverStream<T> {
-        BufferReceiverStream::new(self)
+    pub fn into_stream(mut self) -> BufferReceiverStream<T> {
+        let mut finalizer_tasks = Vec::new();
+        let mut cleanup_tasks = Vec::new();
+        self.take_ledger_tasks(&mut finalizer_tasks, &mut cleanup_tasks);
+        BufferReceiverStream::new(self, finalizer_tasks, cleanup_tasks)
+    }
+
+    fn take_ledger_tasks(
+        &mut self,
+        finalizers: &mut Vec<JoinHandle<()>>,
+        cleanups: &mut Vec<CleanupTask>,
+    ) {
+        self.base.take_ledger_tasks(finalizers, cleanups);
+        if let Some(overflow) = self.overflow.as_mut() {
+            overflow.take_ledger_tasks(finalizers, cleanups);
+        }
     }
 }
 
@@ -167,13 +192,42 @@ enum StreamState<T: Bufferable> {
 pub struct BufferReceiverStream<T: Bufferable> {
     state: StreamState<T>,
     recv_fut: ReusableBoxFuture<'static, (Option<T>, BufferReceiver<T>)>,
+    finalizer_tasks: Vec<JoinHandle<()>>,
+    cleanup_tasks: Vec<CleanupTask>,
 }
 
 impl<T: Bufferable> BufferReceiverStream<T> {
-    pub fn new(receiver: BufferReceiver<T>) -> Self {
+    fn new(
+        receiver: BufferReceiver<T>,
+        finalizer_tasks: Vec<JoinHandle<()>>,
+        cleanup_tasks: Vec<CleanupTask>,
+    ) -> Self {
         Self {
             state: StreamState::Idle(receiver),
             recv_fut: ReusableBoxFuture::new(make_recv_future(None)),
+            finalizer_tasks,
+            cleanup_tasks,
+        }
+    }
+
+    /// Closes the receiver and waits for disk-buffer acknowledgements to release their ledgers.
+    pub async fn close(mut self) {
+        let finalizer_tasks = mem::take(&mut self.finalizer_tasks);
+        let cleanup_tasks = mem::take(&mut self.cleanup_tasks);
+        drop(self);
+
+        for (shutdown, _) in &cleanup_tasks {
+            shutdown.cancel();
+        }
+        for (_, task) in cleanup_tasks {
+            if let Err(error) = task.await {
+                error!(%error, message = "Disk buffer cleanup task failed.");
+            }
+        }
+        for task in finalizer_tasks {
+            if let Err(error) = task.await {
+                error!(%error, message = "Disk buffer finalizer task failed.");
+            }
         }
     }
 }

lib/vector-buffers/src/topology/channel/sender.rs

@@ -10,7 +10,7 @@ use crate::{
     BufferInstrumentation, Bufferable, WhenFull,
     buffer_usage_data::BufferUsageHandle,
     internal_events::BufferSendDuration,
-    variants::disk_v2::{self, ProductionFilesystem, TryWriteOutcome},
+    variants::disk_v2::{self, DiskBufferUsageSnapshot, ProductionFilesystem, TryWriteOutcome},
 };
 
 /// Adapter for papering over various sender backends.
@@ -160,6 +160,16 @@ where
             Self::DiskV2(_) => None,
         }
     }
+
+    fn usage_snapshot(&self) -> Option<DiskBufferUsageSnapshot> {
+        match self {
+            Self::InMemory(_) => Some(DiskBufferUsageSnapshot {
+                event_count: 0,
+                byte_size: 0,
+            }),
+            Self::DiskV2(writer) => writer.try_lock().ok().map(|writer| writer.usage_snapshot()),
+        }
+    }
 }
 
 enum UsageAccounting {
@@ -277,6 +287,19 @@ impl<T: Bufferable> BufferSender<T> {
     pub fn with_custom_instrumentation(&mut self, instrumentation: impl BufferInstrumentation<T>) {
         self.custom_instrumentation = Some(Arc::new(instrumentation));
     }
+
+    /// Returns the aggregate approximate usage of all disk stages in this buffer topology.
+    ///
+    /// Returns `None` rather than waiting if any disk writer is currently busy.
+    pub fn disk_usage_snapshot(&self) -> Option<DiskBufferUsageSnapshot> {
+        let mut usage = self.base.usage_snapshot()?;
+        if let Some(overflow) = self.overflow.as_ref() {
+            let overflow_usage = overflow.disk_usage_snapshot()?;
+            usage.event_count = usage.event_count.saturating_add(overflow_usage.event_count);
+            usage.byte_size = usage.byte_size.saturating_add(overflow_usage.byte_size);
+        }
+        Some(usage)
+    }
 }
 
 impl<T: Bufferable> BufferSender<T> {
@@ -380,3 +403,45 @@ impl<T: Bufferable> BufferSender<T> {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod usage_snapshot_tests {
+    use std::num::{NonZeroU64, NonZeroUsize};
+
+    use tempfile::tempdir;
+    use tracing::Span;
+
+    use super::SenderAdapter;
+    use crate::{BufferConfig, BufferType, MemoryBufferSize, WhenFull, test::SizedRecord};
+
+    #[tokio::test]
+    async fn busy_disk_stage_makes_diagnostic_snapshot_unavailable() {
+        let config = BufferConfig::Chained(vec![
+            BufferType::Memory {
+                size: MemoryBufferSize::MaxEvents(NonZeroUsize::new(1).unwrap()),
+                when_full: WhenFull::Overflow,
+            },
+            BufferType::DiskV2 {
+                max_size: NonZeroU64::new(268_435_488).unwrap(),
+                when_full: WhenFull::Block,
+            },
+        ]);
+        let data_dir = tempdir().unwrap();
+        let (sender, _receiver) = config
+            .build::<SizedRecord>(
+                Some(data_dir.path().to_path_buf()),
+                "busy".into(),
+                Span::none(),
+            )
+            .await
+            .unwrap();
+        let overflow = sender.overflow.as_ref().unwrap();
+        let writer = match &overflow.base {
+            SenderAdapter::DiskV2(writer) => writer,
+            SenderAdapter::InMemory(_) => unreachable!(),
+        };
+        let _guard = writer.lock().await;
+
+        assert_eq!(sender.disk_usage_snapshot(), None);
+    }
+}

lib/vector-buffers/src/variants/disk_v2/checkpoint_recovery.rs

@@ -775,7 +775,7 @@ mod tests {
                 .expect("ledger should not fail to load/create");
             let ledger = Arc::new(ledger);
             let (finalizer, _finalizer_stream) = OrderedFinalizer::new(None);
-            let reader = BufferReader::new(Arc::clone(&ledger), finalizer);
+            let reader = BufferReader::new(Arc::clone(&ledger), finalizer, None);
 
             Self { reader, ledger }
         }

lib/vector-buffers/src/variants/disk_v2/ledger.rs

@@ -20,6 +20,7 @@ use tokio::{
     io::AsyncWriteExt,
     sync::{Mutex, MutexGuard, Notify},
 };
+use tokio_util::sync::CancellationToken;
 use vector_common::finalizer::OrderedFinalizer;
 
 use super::{
@@ -34,6 +35,8 @@ use super::{
 };
 use crate::buffer_usage_data::BufferUsageHandle;
 
+pub(crate) type CleanupTask = (CancellationToken, tokio::task::JoinHandle<()>);
+
 pub const LEDGER_LEN: usize = align16(mem::size_of::<ArchivedLedgerState>());
 
 /// Error that occurred during calls to [`Ledger`].
@@ -857,25 +860,32 @@ where
     }
 
     #[must_use]
-    pub(super) fn spawn_finalizer(self: Arc<Self>) -> OrderedFinalizer<u64> {
+    pub(super) fn spawn_finalizer(
+        self: Arc<Self>,
+    ) -> (OrderedFinalizer<u64>, tokio::task::JoinHandle<()>) {
         let (finalizer, mut stream) = OrderedFinalizer::new(None);
-        vector_common::spawn_in_current_span(async move {
+        let task = vector_common::spawn_in_current_span(async move {
             while let Some((_status, amount)) = stream.next().await {
                 self.increment_pending_acks(amount);
                 self.notify_writer_waiters();
             }
         });
-        finalizer
+        (finalizer, task)
     }
 
-    pub(super) fn spawn_data_file_cleanup(self: &Arc<Self>)
+    pub(super) fn spawn_data_file_cleanup(self: &Arc<Self>) -> CleanupTask
     where
         FS: 'static,
     {
         let ledger = Arc::downgrade(self);
-        vector_common::spawn_in_current_span(async move {
+        let shutdown = CancellationToken::new();
+        let task_shutdown = shutdown.clone();
+        let task = vector_common::spawn_in_current_span(async move {
             loop {
-                tokio::time::sleep(DEFAULT_DATA_FILE_CLEANUP_INTERVAL).await;
+                tokio::select! {
+                    () = task_shutdown.cancelled() => break,
+                    () = tokio::time::sleep(DEFAULT_DATA_FILE_CLEANUP_INTERVAL) => {}
+                }
 
                 let Some(ledger) = ledger.upgrade() else {
                     break;
@@ -889,8 +899,13 @@ where
                 }
 
                 drop(ledger);
+
+                if task_shutdown.is_cancelled() {
+                    break;
+                }
             }
         });
+        (shutdown, task)
     }
 }
 

lib/vector-buffers/src/variants/disk_v2/mod.rs

@@ -179,14 +179,15 @@ use std::{
 
 use async_trait::async_trait;
 use snafu::{ResultExt, Snafu};
-use tracing::debug;
+use tracing::{debug, info};
 use vector_common::finalization::Finalizable;
 
 mod backed_archive;
 mod checkpoint_recovery;
 pub(crate) mod common;
 mod io;
 mod ledger;
+pub(crate) use ledger::CleanupTask;
 mod reader;
 mod record;
 mod ser;
@@ -212,6 +213,15 @@ use crate::{
     },
 };
 
+/// An approximate view of unacknowledged records in a disk buffer.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct DiskBufferUsageSnapshot {
+    /// Number of events that have not been acknowledged.
+    pub event_count: u64,
+    /// Encoded size of records that have not been acknowledged.
+    pub byte_size: u64,
+}
+
 /// Error that occurred when creating/loading a disk buffer.
 #[derive(Debug, Snafu)]
 pub enum BufferError<T>
@@ -260,9 +270,9 @@ where
             .await
             .context(WriterSeekFailedSnafu)?;
 
-        let finalizer = Arc::clone(&ledger).spawn_finalizer();
+        let (finalizer, finalizer_task) = Arc::clone(&ledger).spawn_finalizer();
 
-        let mut reader = BufferReader::new(Arc::clone(&ledger), finalizer);
+        let mut reader = BufferReader::new(Arc::clone(&ledger), finalizer, Some(finalizer_task));
         let mut unread_buffer_size = reader
             .reconcile_checkpoint_window()
             .await
@@ -306,7 +316,7 @@ where
 
         ledger.synchronize_buffer_usage();
         if ledger.filesystem().supports_background_cleanup() {
-            ledger.spawn_data_file_cleanup();
+            reader.set_cleanup_task(ledger.spawn_data_file_cleanup());
         }
 
         Ok((writer, reader, ledger))
@@ -396,7 +406,8 @@ where
     usage_handle.set_buffer_limits(Some(max_size.get()), None);
 
     let buffer_path = get_disk_v2_data_dir_path(data_dir, id);
-    let builder = DiskBufferConfigBuilder::from_path(buffer_path).max_buffer_size(max_size.get());
+    let builder =
+        DiskBufferConfigBuilder::from_path(buffer_path.clone()).max_buffer_size(max_size.get());
     // Shrink the data-file size (and the matching record size) so files fill and
     // rotate constantly. That is what reaches the rare recovery paths the bug hides
     // in: reopening a file whose last write was cut short, and reusing a file number
@@ -412,9 +423,21 @@ where
         None => builder,
     };
     let config = builder.build()?;
-    Buffer::from_config(config, usage_handle)
+    let (writer, reader) = Buffer::from_config(config, usage_handle)
         .await
-        .map_err(Into::into)
+        .map_err(Box::<dyn Error + Send + Sync>::from)?;
+    let usage = writer.usage_snapshot();
+    if usage.event_count > 0 || usage.byte_size > 0 {
+        info!(
+            buffer_id = id,
+            buffer_dir = buffer_path.to_string_lossy().as_ref(),
+            accounted_events = usage.event_count,
+            accounted_bytes = usage.byte_size,
+            message = "Opened disk buffer with remaining accounted data; resuming processing.",
+        );
+    }
+
+    Ok((writer, reader))
 }
 
 pub(crate) fn get_disk_v2_data_dir_path(base_dir: &Path, buffer_id: &str) -> PathBuf {