enhancement(buffers): improve disk buffer lifecycle visibility

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::{DiskBufferUsageHandle, 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/sender.rs

@@ -14,7 +14,7 @@ use crate::{
     BufferInstrumentation, Bufferable, WhenFull,
     buffer_usage_data::BufferUsageHandle,
     internal_events::BufferSendDuration,
-    variants::disk_v2::{self, ProductionFilesystem, TryWriteOutcome},
+    variants::disk_v2::{self, DiskBufferUsageHandle, ProductionFilesystem, TryWriteOutcome},
 };
 
 /// Adapter for papering over various sender backends.
@@ -24,7 +24,10 @@ pub enum SenderAdapter<T: Bufferable> {
     InMemory(LimitedSender<T>),
 
     /// The disk v2 buffer.
-    DiskV2(Arc<Mutex<disk_v2::BufferWriter<T, ProductionFilesystem>>>),
+    DiskV2 {
+        writer: Arc<Mutex<disk_v2::BufferWriter<T, ProductionFilesystem>>>,
+        usage: DiskBufferUsageHandle,
+    },
 }
 
 impl<T: Bufferable> From<LimitedSender<T>> for SenderAdapter<T> {
@@ -35,7 +38,11 @@ impl<T: Bufferable> From<LimitedSender<T>> for SenderAdapter<T> {
 
 impl<T: Bufferable> From<disk_v2::BufferWriter<T, ProductionFilesystem>> for SenderAdapter<T> {
     fn from(v: disk_v2::BufferWriter<T, ProductionFilesystem>) -> Self {
-        Self::DiskV2(Arc::new(Mutex::new(v)))
+        let usage = v.usage_handle();
+        Self::DiskV2 {
+            writer: Arc::new(Mutex::new(v)),
+            usage,
+        }
     }
 }
 
@@ -55,7 +62,7 @@ where
     pub(crate) fn requires_encodable_items(&self) -> bool {
         match self {
             Self::InMemory(_) => false,
-            Self::DiskV2(_) => true,
+            Self::DiskV2 { .. } => true,
         }
     }
 
@@ -66,7 +73,7 @@ where
                 .await
                 .map(|()| TryWriteOutcome::Written)
                 .map_err(Into::into),
-            Self::DiskV2(writer) => {
+            Self::DiskV2 { writer, .. } => {
                 let pre_count = item.event_count() as u64;
                 let pre_size = item.size_of() as u64;
                 let mut writer = writer.lock().await;
@@ -106,7 +113,7 @@ where
                 .try_send(item)
                 .map(|()| TryWriteOutcome::Written)
                 .or_else(|e| Ok(TryWriteOutcome::Full(e.into_inner()))),
-            Self::DiskV2(writer) => {
+            Self::DiskV2 { writer, .. } => {
                 let mut writer = writer.lock().await;
 
                 // Filtering here is unconditional and independent of current occupancy.
@@ -146,7 +153,7 @@ where
     pub(crate) async fn flush(&mut self) -> crate::Result<()> {
         match self {
             Self::InMemory(_) => Ok(()),
-            Self::DiskV2(writer) => {
+            Self::DiskV2 { writer, .. } => {
                 let mut writer = writer.lock().await;
                 writer.flush().await.map_err(|e| {
                     // Errors on the I/O path, which is all that flushing touches, are never recoverable.
@@ -161,7 +168,14 @@ where
     pub fn capacity(&self) -> Option<usize> {
         match self {
             Self::InMemory(tx) => Some(tx.available_capacity()),
-            Self::DiskV2(_) => None,
+            Self::DiskV2 { .. } => None,
+        }
+    }
+
+    fn disk_usage_handle(&self) -> Option<DiskBufferUsageHandle> {
+        match self {
+            Self::InMemory(_) => None,
+            Self::DiskV2 { usage, .. } => Some(usage.clone()),
         }
     }
 }
@@ -278,6 +292,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 read-only usage handles for all disk stages in this buffer topology.
+    pub fn disk_usage_handles(&self) -> Vec<DiskBufferUsageHandle> {
+        let mut handles = self
+            .base
+            .disk_usage_handle()
+            .into_iter()
+            .collect::<Vec<_>>();
+        if let Some(overflow) = self.overflow.as_ref() {
+            handles.extend(overflow.disk_usage_handles());
+        }
+        handles
+    }
 }
 
 impl<T: Bufferable> BufferSender<T> {

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

@@ -23,7 +23,7 @@ use tokio::{
 use vector_common::finalizer::OrderedFinalizer;
 
 use super::{
-    Filesystem,
+    DiskBufferUsageHandle, Filesystem,
     backed_archive::BackedArchive,
     common::{
         DEFAULT_DATA_FILE_CLEANUP_INTERVAL, DiskBufferConfig, MAX_FILE_ID, align16,
@@ -257,6 +257,7 @@ where
     last_flush: AtomicCell<Instant>,
     // Tracks usage data about the buffer.
     usage_handle: BufferUsageHandle,
+    disk_usage_handle: DiskBufferUsageHandle,
 }
 
 impl<FS> Ledger<FS>
@@ -317,6 +318,10 @@ where
         self.total_buffer_size.load(Ordering::Acquire)
     }
 
+    pub fn disk_usage_handle(&self) -> DiskBufferUsageHandle {
+        self.disk_usage_handle.clone()
+    }
+
     /// Increments the total number of bytes for all unread records in the buffer.
     pub fn increment_total_buffer_size(&self, amount: u64) {
         let last_total_buffer_size = self.total_buffer_size.fetch_add(amount, Ordering::AcqRel);
@@ -530,6 +535,7 @@ where
         self.increment_total_buffer_size(record_size);
         self.usage_handle
             .increment_received_event_count_and_byte_size(event_count, record_size);
+        self.disk_usage_handle.increment(event_count, record_size);
         self.notify_writer_waiters();
         next_record_id
     }
@@ -553,6 +559,8 @@ where
         self.decrement_total_buffer_size(total_record_size);
         self.usage_handle
             .increment_sent_event_count_and_byte_size(event_count, total_record_size);
+        self.disk_usage_handle
+            .decrement(event_count, total_record_size);
     }
 
     /// Tracks the known byte size of a record that the reader consumed but could not decode.
@@ -562,6 +570,13 @@ where
         // provides the exact byte span that has left the buffer.
         self.usage_handle
             .increment_dropped_event_count_and_byte_size(0, total_record_size, false);
+        self.disk_usage_handle.decrement(0, total_record_size);
+    }
+
+    /// Removes an unread data file tail that could not be interpreted as a record.
+    pub fn track_abandoned_tail_bytes(&self, total_size: u64) {
+        self.decrement_total_buffer_size(total_size);
+        self.disk_usage_handle.decrement(0, total_size);
     }
 
     /// Marks the writer as finished.
@@ -689,6 +704,8 @@ where
                 initial_buffer_events,
                 initial_buffer_size,
             );
+        self.disk_usage_handle
+            .set(initial_buffer_events, initial_buffer_size);
     }
 
     pub fn track_dropped_events(&self, count: u64) {
@@ -702,6 +719,7 @@ where
         // TODO: Can we do better here?
         self.usage_handle
             .increment_dropped_event_count_and_byte_size(count, 0, false);
+        self.disk_usage_handle.decrement(count, 0);
     }
 
     /// Records that a record was dropped because it could never be written to the buffer (e.g. it
@@ -824,6 +842,7 @@ where
             data_file_cleanup: DataFileCleanupState::new(cleanup_reader_file_id),
             last_flush: AtomicCell::new(Instant::now()),
             usage_handle,
+            disk_usage_handle: DiskBufferUsageHandle::new(),
         };
 
         // NOTE: We deliberately do not seed `total_buffer_size` here. Historically, the ledger

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

@@ -174,12 +174,15 @@ use std::{
     marker::PhantomData,
     num::NonZeroU64,
     path::{Path, PathBuf},
-    sync::Arc,
+    sync::{
+        Arc,
+        atomic::{AtomicU64, Ordering},
+    },
 };
 
 use async_trait::async_trait;
 use snafu::{ResultExt, Snafu};
-use tracing::debug;
+use tracing::{debug, info};
 use vector_common::finalization::Finalizable;
 
 mod backed_archive;
@@ -212,6 +215,62 @@ use crate::{
     },
 };
 
+/// An approximate view of unread records in a disk buffer.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct DiskBufferUsageSnapshot {
+    /// Number of unread events in the buffer.
+    pub event_count: u64,
+    /// Encoded size of unread records in the buffer.
+    pub byte_size: u64,
+}
+
+/// Read-only access to the current unread usage of a disk buffer.
+#[derive(Clone, Debug)]
+pub struct DiskBufferUsageHandle {
+    event_count: Arc<AtomicU64>,
+    byte_size: Arc<AtomicU64>,
+}
+
+impl DiskBufferUsageHandle {
+    fn new() -> Self {
+        Self {
+            event_count: Arc::new(AtomicU64::new(0)),
+            byte_size: Arc::new(AtomicU64::new(0)),
+        }
+    }
+
+    /// Gets the approximate current number and encoded size of unread records.
+    pub fn snapshot(&self) -> DiskBufferUsageSnapshot {
+        DiskBufferUsageSnapshot {
+            event_count: self.event_count.load(Ordering::Relaxed),
+            byte_size: self.byte_size.load(Ordering::Relaxed),
+        }
+    }
+
+    fn set(&self, event_count: u64, byte_size: u64) {
+        self.event_count.store(event_count, Ordering::Relaxed);
+        self.byte_size.store(byte_size, Ordering::Relaxed);
+    }
+
+    fn increment(&self, event_count: u64, byte_size: u64) {
+        self.event_count.fetch_add(event_count, Ordering::Relaxed);
+        self.byte_size.fetch_add(byte_size, Ordering::Relaxed);
+    }
+
+    fn decrement(&self, event_count: u64, byte_size: u64) {
+        _ = self
+            .event_count
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
+                Some(current.saturating_sub(event_count))
+            });
+        _ = self
+            .byte_size
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
+                Some(current.saturating_sub(byte_size))
+            });
+    }
+}
+
 /// Error that occurred when creating/loading a disk buffer.
 #[derive(Debug, Snafu)]
 pub enum BufferError<T>
@@ -396,7 +455,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 +472,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 existing_usage = writer.usage_handle().snapshot();
+    if existing_usage.event_count > 0 || existing_usage.byte_size > 0 {
+        info!(
+            buffer_id = id,
+            buffer_dir = buffer_path.to_string_lossy().as_ref(),
+            unread_events = existing_usage.event_count,
+            unread_bytes = existing_usage.byte_size,
+            message = "Opened disk buffer with unread data; resuming delivery.",
+        );
+    }
+
+    Ok((writer, reader))
 }
 
 pub(crate) fn get_disk_v2_data_dir_path(base_dir: &Path, buffer_id: &str) -> PathBuf {

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

@@ -598,8 +598,7 @@ where
                 abandoned_tail_bytes,
                 "Data file tail was abandoned after bad read. Adjusting buffer size to compensate.",
             );
-            self.ledger
-                .decrement_total_buffer_size(abandoned_tail_bytes);
+            self.ledger.track_abandoned_tail_bytes(abandoned_tail_bytes);
         }
 
         Ok(())

lib/vector-buffers/src/variants/disk_v2/tests/basic.rs

@@ -238,6 +238,8 @@ async fn initial_size_correct_with_multievents() {
             // Reopen the buffer with a full reader to verify the persisted data.
             let (writer, mut reader, ledger, usage) =
                 create_default_buffer_v2_with_usage::<_, MultiEventRecord>(&data_dir).await;
+            let disk_usage_handle = writer.usage_handle();
+            let disk_usage = disk_usage_handle.snapshot();
             drop(writer);
 
             // Make sure our usage data agrees with our expected event count and byte size:
@@ -246,6 +248,8 @@ async fn initial_size_correct_with_multievents() {
             assert_eq!(expected_bytes as u64, snapshot.received_byte_size);
             assert_eq!(expected_events as u64, ledger.get_total_records());
             assert_eq!(expected_bytes, total_bytes_written);
+            assert_eq!(expected_events as u64, disk_usage.event_count);
+            assert_eq!(expected_bytes as u64, disk_usage.byte_size);
 
             // Make sure we can read all of the records we wrote, and recalculate some of these
             // values from the source:
@@ -260,6 +264,8 @@ async fn initial_size_correct_with_multievents() {
 
             assert_eq!(expected_events, total_record_events);
             assert_eq!(expected_records, total_records_read);
+            assert_eq!(0, disk_usage_handle.snapshot().event_count);
+            assert_eq!(0, disk_usage_handle.snapshot().byte_size);
         }
     })
     .await;

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

@@ -913,6 +913,10 @@ where
         }
     }
 
+    pub(crate) fn usage_handle(&self) -> super::DiskBufferUsageHandle {
+        self.ledger.disk_usage_handle()
+    }
+
     fn get_next_record_id(&mut self) -> u64 {
         self.next_record_id + self.unflushed_events
     }