fix(buffers): recompute disk_v2 buffer size on restart to prevent underflow

changelog.d/disk_v2_buffer_size_recompute_on_restart.fix.md

@@ -0,0 +1,3 @@
+Fixed a `disk_v2` buffer accounting bug where the unread-bytes counter (`total_buffer_size`) could underflow on restart. Previously the buffer seeded the counter from the sum of on-disk data file sizes and then had the reader decrement it record-by-record while seeking to its persisted read position. Those two inputs were captured at different moments — the data files on disk versus the read position in the ledger, which are flushed on independent schedules — so a crash between their flushes could make the decrements exceed the seeded total. Because the counter is unsigned, the subtraction wrapped to a near-maximum value, making the buffer appear permanently full and wedging the writer (no further writes accepted, and recovery could stall). The reader now recomputes the unread total authoritatively at the end of its startup seek, from a single consistent snapshot — the total size of the data files still on disk minus the bytes it consumed reaching the resume position — so the counter can no longer underflow across a restart.
+
+authors: graphcareful

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

@@ -324,6 +324,42 @@ where
         );
     }
 
+    /// Overwrites the total number of bytes for all unread records in the buffer.
+    ///
+    /// Used during initialization to install an authoritatively recomputed buffer size once the
+    /// reader has established its resume position (see `BufferReader::seek_to_next_record`). Unlike
+    /// the incremental `increment`/`decrement` paths, this is an absolute store, so it is only
+    /// sound to call while no concurrent writes or acknowledgements can be mutating the counter --
+    /// i.e. during the synchronous buffer load, before the reader/writer are handed out.
+    pub fn set_total_buffer_size(&self, amount: u64) {
+        self.total_buffer_size.store(amount, Ordering::Release);
+    }
+
+    /// Sums the on-disk size of every data file currently present in the buffer's data directory.
+    ///
+    /// This is the total bytes physically on disk, across read and unread records alike; callers
+    /// that want the unread total must subtract whatever the reader has already consumed.
+    pub(super) async fn total_data_file_size(&self) -> io::Result<u64> {
+        let mut dat_reader = fs::read_dir(&self.config.data_dir).await?;
+
+        let mut total_data_file_size = 0;
+        while let Some(dir_entry) = dat_reader.next_entry().await? {
+            if let Some(file_name) = dir_entry.file_name().to_str() {
+                // I really _do_ want to only find files with a .dat extension, as that's what the
+                // code generates, and having them be .dAt or .Dat or whatever would indicate that
+                // the file is not related to our buffer.  If we had to cope with case-sensitivity
+                // of filenames from another program/OS, then it would be a different story.
+                #[allow(clippy::case_sensitive_file_extension_comparisons)]
+                if file_name.ends_with(".dat") {
+                    let metadata = dir_entry.metadata().await?;
+                    total_data_file_size += metadata.len();
+                }
+            }
+        }
+
+        Ok(total_data_file_size)
+    }
+
     /// Gets the current reader file ID.
     ///
     /// This is internally adjusted to compensate for the fact that the reader can read far past
@@ -659,7 +695,7 @@ where
         // Create the ledger object, and synchronize the buffer statistics with the buffer usage
         // handle.  This handles making sure we account for the starting size of the buffer, and
         // what not.
-        let mut ledger = Ledger {
+        let ledger = Ledger {
             config,
             lock,
             state: ledger_state,
@@ -672,69 +708,18 @@ where
             last_flush: AtomicCell::new(Instant::now()),
             usage_handle,
         };
-        ledger.update_buffer_size().await?;
 
-        Ok(ledger)
-    }
+        // NOTE: We deliberately do not seed `total_buffer_size` here. Historically, the ledger
+        // summed the on-disk data file sizes at load and then had the reader "draw that total down"
+        // record-by-record as it sought to its persisted read position. That mixed two values
+        // captured at different moments -- the directory's file sizes versus the ledger's persisted
+        // read position, which are flushed independently -- so a crash between their flushes could
+        // let the draw-down subtract past zero, wrapping the counter to a near-max value and
+        // wedging the writer. The authoritative value (the size of the *unread* records only) is
+        // now computed in a single consistent snapshot by `BufferReader::seek_to_next_record` once
+        // it has established the resume position.
 
-    async fn update_buffer_size(&mut self) -> Result<(), LedgerLoadCreateError> {
-        // Under normal operation, the reader and writer maintain a consistent state within the
-        // ledger.  However, due to the nature of how we update the ledger, process crashes could
-        // lead to missed updates as we execute reads and writes as non-atomic units of execution:
-        // update a field, do the read/write, update some more fields depending on success or
-        // failure, etc.
-        //
-        // This is an issue because we depend on knowing the total buffer size (the total size of
-        // unread records, specifically) so that we can correctly limit writes when we've reached
-        // the configured maximum buffer size.
-        //
-        // While it's not terribly efficient, and I'd like to eventually formulate a better design,
-        // this approach is absolutely correct: get the file size of every data file on disk,
-        // and set the "total buffer size" to the sum of all of those file sizes.
-        //
-        // When the reader does any necessary seeking to get to the record it left off on, it will
-        // adjust the "total buffer size" downwards for each record it runs through, leaving "total
-        // buffer size" at the correct value.
-        let mut dat_reader = fs::read_dir(&self.config.data_dir).await.context(IoSnafu)?;
-
-        let mut total_buffer_size = 0;
-        while let Some(dir_entry) = dat_reader.next_entry().await.context(IoSnafu)? {
-            if let Some(file_name) = dir_entry.file_name().to_str() {
-                // I really _do_ want to only find files with a .dat extension, as that's what the
-                // code generates, and having them be .dAt or .Dat or whatever would indicate that
-                // the file is not related to our buffer.  If we had to cope with case-sensitivity
-                // of filenames from another program/OS, then it would be a different story.
-                #[allow(clippy::case_sensitive_file_extension_comparisons)]
-                if file_name.ends_with(".dat") {
-                    let metadata = dir_entry.metadata().await.context(IoSnafu)?;
-                    total_buffer_size += metadata.len();
-
-                    debug!(
-                        data_file = file_name,
-                        file_size = metadata.len(),
-                        total_buffer_size,
-                        "Found existing data file."
-                    );
-                }
-            }
-        }
-
-        // A non-zero sum means the buffer reopened on top of records left on disk by
-        // a previous run, the reseed path whose value the reader later draws down and
-        // the one most exposed to the buffer-size underflow.
-        #[cfg(feature = "antithesis-disk-asserts")]
-        {
-            #![allow(clippy::disallowed_types)] // once_cell::Lazy
-            antithesis_sdk::assert_sometimes!(
-                total_buffer_size > 0,
-                "the buffer reopens with pre-existing on-disk records",
-                &serde_json::json!({ "total_buffer_size": total_buffer_size })
-            );
-        }
-
-        self.increment_total_buffer_size(total_buffer_size);
-
-        Ok(())
+        Ok(ledger)
     }
 
     #[must_use]

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

@@ -461,12 +461,19 @@ where
             self.data_file_start_record_id = Some(record_id);
         }
 
-        // Track the amount of data we read.  If we're still loading the buffer, then the only thing
-        // other we need to do is update the total buffer size.  Everything else below only matters
-        // when we're doing real record reads.
+        // Track the amount of data we read.  `bytes_read` accumulates the on-disk size of the
+        // records consumed within the current data file; at the end of initialization it is exactly
+        // the size of the already-read prefix of the file the reader resumes on, which
+        // `seek_to_next_record` uses to install the authoritative buffer size.
         self.bytes_read += record_bytes;
         if !self.ready_to_read {
-            self.ledger.decrement_total_buffer_size(record_bytes);
+            // During initialization we are only replaying already-read records to reposition the
+            // reader; we deliberately do NOT adjust `total_buffer_size` here. The authoritative
+            // unread total is installed once, at the end of `seek_to_next_record`, from a single
+            // consistent snapshot. Decrementing per-record during this replay was the historical
+            // approach, and -- because the seeded total and the replayed records came from
+            // snapshots taken at different times across a crash -- it is what allowed the
+            // buffer-size counter to underflow and wrap.
             return;
         }
 
@@ -534,7 +541,11 @@ where
         let metadata = data_file.metadata().await?;
 
         let decrease_amount = bytes_read.map_or_else(
-            || metadata.len(),
+            // `None` is only passed from `seek_to_next_record`'s fast path while it deletes
+            // fully-read straggler files during initialization. Initialization does not mutate
+            // `total_buffer_size` incrementally -- the authoritative value is installed once at the
+            // end of the seek -- so deleting a file here must not decrement it.
+            || 0,
             |bytes_read| {
                 // A file shorter than bytes_read makes the delta below underflow
                 // and feed a wrapped value into decrement_total_buffer_size.
@@ -968,6 +979,44 @@ where
             "Synchronized with ledger. Reader ready."
         );
 
+        // Install the authoritative buffer size now that the reader is positioned at the first
+        // unread record.
+        //
+        // At this point the seek has deleted every data file that was fully read (the fast path
+        // above), so the files still on disk are exactly: the file the reader resumed on, plus any
+        // later, entirely-unread files. `self.bytes_read` is the on-disk size of the already-read
+        // prefix of that resumed-on file. The unread total is therefore simply the total size of
+        // the remaining files minus that consumed prefix -- a single, internally-consistent
+        // snapshot, rather than a file-size seed drawn down against a separately-persisted read
+        // position. This is what makes the result immune to the crash-window underflow.
+        //
+        // `saturating_sub` is belt-and-suspenders: `bytes_read` is a prefix of the remaining files
+        // and so can never exceed their total, but clamping guarantees we can never wrap even if a
+        // corrupted/short file slips through.
+        let total_data_file_size = self.ledger.total_data_file_size().await.context(IoSnafu)?;
+
+        // A non-zero on-disk total means we reopened on top of records left behind by a previous
+        // run -- the recovery path most exposed to the historical buffer-size underflow.
+        #[cfg(feature = "antithesis-disk-asserts")]
+        {
+            #![allow(clippy::disallowed_types)] // once_cell::Lazy
+            antithesis_sdk::assert_sometimes!(
+                total_data_file_size > 0,
+                "the buffer reopens with pre-existing on-disk records",
+                &serde_json::json!({ "total_buffer_size": total_data_file_size })
+            );
+        }
+
+        let unread_buffer_size = total_data_file_size.saturating_sub(self.bytes_read);
+        self.ledger.set_total_buffer_size(unread_buffer_size);
+
+        debug!(
+            total_data_file_size,
+            consumed_prefix = self.bytes_read,
+            unread_buffer_size,
+            "Recalculated buffer size from on-disk data files."
+        );
+
         self.ready_to_read = true;
 
         Ok(())

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

@@ -927,3 +927,87 @@ async fn reader_writer_positions_aligned_through_multiple_files_and_records() {
     let parent = trace_span!("reader_writer_positions_aligned_through_multiple_files_and_records");
     fut.instrument(parent.or_current()).await;
 }
+
+/// Regression test for the buffer-size underflow on restart.
+///
+/// Historically, reopening a buffer seeded `total_buffer_size` from the sum of the on-disk data
+/// file sizes and then had the reader "draw it down" record-by-record as it sought to its persisted
+/// read position. Those two inputs were captured at different moments (the directory's file sizes
+/// vs. the ledger's persisted position), so a crash between their respective flushes could make the
+/// draw-down subtract more than the seed held -- wrapping the unsigned counter to a near-maximum
+/// value, making the buffer look permanently full, and wedging the writer.
+///
+/// The reader now recomputes the unread total authoritatively at the end of `seek_to_next_record`,
+/// from a single consistent snapshot: `(total on-disk file size) - (bytes consumed reaching the
+/// resume position)`. This test exercises the case that actually stresses that arithmetic -- a
+/// single data file that is *partially* read and acknowledged, so the reader resumes mid-file with
+/// a non-empty already-read prefix -- and asserts the reopened buffer reports exactly the bytes of
+/// the still-unread records.
+#[tokio::test]
+async fn buffer_size_recalculated_correctly_after_partial_read_reload() {
+    let _a = install_tracing_helpers();
+    with_temp_dir(|dir| {
+        let data_dir = dir.to_path_buf();
+
+        async move {
+            // Write several small records that all land in a single data file.
+            let total_records = 8;
+            let mut record_sizes = Vec::new();
+            let (mut writer, mut reader, ledger) =
+                create_default_buffer_v2::<_, SizedRecord>(data_dir.clone()).await;
+            for _ in 0..total_records {
+                let bytes_written = writer
+                    .write_record(SizedRecord::new(64))
+                    .await
+                    .expect("write should not fail");
+                record_sizes.push(bytes_written as u64);
+            }
+            writer.flush().await.expect("flush should not fail");
+
+            let total_bytes: u64 = record_sizes.iter().sum();
+            assert_buffer_size!(ledger, total_records, total_bytes);
+
+            // Read and acknowledge a strict prefix of the records.
+            let acked_records = 3usize;
+            for _ in 0..acked_records {
+                let record = read_next_some(&mut reader).await;
+                acknowledge(record).await;
+            }
+
+            // Let the spawned finalizer deliver the acknowledgements, then drive one more read so
+            // the reader consumes them, advancing its persisted record position past the prefix.
+            tokio::task::yield_now().await;
+            let _ = read_next(&mut reader).await;
+
+            let expected_unread: u64 = record_sizes[acked_records..].iter().sum();
+
+            // Checkpoint on the live buffer: the acknowledged prefix has been drawn down, so the
+            // running size already reflects only the unread records. (This also guards the test
+            // against the acknowledgements not having been processed before we reload.)
+            assert_eq!(
+                expected_unread,
+                ledger.get_total_buffer_size(),
+                "live buffer size should reflect only the unread records after acking the prefix",
+            );
+
+            // Persist the advanced read position so the reopened buffer resumes mid-file rather
+            // than replaying from the start.
+            ledger.flush().expect("ledger flush should not fail");
+
+            drop(writer);
+            drop(reader);
+            drop(ledger);
+
+            // Reopen. The authoritatively recomputed buffer size must equal exactly the bytes of
+            // the unread records -- not the full file, and never an underflowed/wrapped value.
+            let (_writer, _reader, ledger) =
+                create_default_buffer_v2::<_, SizedRecord>(data_dir).await;
+            assert_eq!(
+                expected_unread,
+                ledger.get_total_buffer_size(),
+                "reopened buffer size should equal the unread records' on-disk bytes",
+            );
+        }
+    })
+    .await;
+}