From c1fb33e1d07a30f6f1ef14a3e7c65020838c33cc Mon Sep 17 00:00:00 2001
From: Andrew Chang <andrewrchang@meta.com>
Date: Fri, 7 Feb 2025 13:25:56 -0800
Subject: [PATCH] Prefetch buffer may not contain all of requested data if EOF
 is hit (#13376)

Summary:
There was a stress test that failed at the assertion check for `IsDataBlockInBuffer`.

`IsDataBlockInBuffer` is too strict of a condition if we are trying to read past the end of the file.

This seems to be a bug from the original 2019 commit https://github.com/siying/rocksdb/commit/3737d06adc01a59e7eb29710a2a4ec64adfaa528: https://github.com/siying/rocksdb/blob/4eb51130917c260f5637731cd77baaa45dfdc5ec/file/file_prefetch_buffer.cc#L130

If the caller tries requesting more bytes than are available, then we still return `n` bytes, even if the buffer really only contains `m < n` bytes.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/13376

Test Plan: I added a unit test which caused the original `IsDataBlockInBuffer ` assertion to fail. I also updated the unit test to check for the result size, which triggered the bug (without this fix) where we return a size of `n` even if less than `n` bytes exist.

Reviewed By: anand1976

Differential Revision: D69269608

Pulled By: archang19

fbshipit-source-id: 1dc0d5930e2b73089850f6e996afbd6192cd5ac8
---
 file/file_prefetch_buffer.cc |  8 +++++---
 file/prefetch_test.cc        | 16 ++++++++++++++++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/file/file_prefetch_buffer.cc b/file/file_prefetch_buffer.cc
index d5760103534..7683db86173 100644
--- a/file/file_prefetch_buffer.cc
+++ b/file/file_prefetch_buffer.cc
@@ -865,10 +865,12 @@ bool FilePrefetchBuffer::TryReadFromCacheUntracked(
   if (copy_to_overlap_buffer) {
     buf = overlap_buf_;
   }
-  assert(buf->offset_ <= offset);
-  assert(buf->IsDataBlockInBuffer(offset, n));
+  assert(buf->IsOffsetInBuffer(offset));
   uint64_t offset_in_buffer = offset - buf->offset_;
-  *result = Slice(buf->buffer_.BufferStart() + offset_in_buffer, n);
+  assert(offset_in_buffer < buf->CurrentSize());
+  *result = Slice(
+      buf->buffer_.BufferStart() + offset_in_buffer,
+      std::min(n, buf->CurrentSize() - static_cast<size_t>(offset_in_buffer)));
   if (prefetched) {
     readahead_size_ = std::min(max_readahead_size_, readahead_size_ * 2);
   }
diff --git a/file/prefetch_test.cc b/file/prefetch_test.cc
index 726cccf29d7..2c0919ed952 100644
--- a/file/prefetch_test.cc
+++ b/file/prefetch_test.cc
@@ -3317,25 +3317,39 @@ TEST_F(FilePrefetchBufferTest, ForCompaction) {
   ASSERT_TRUE(fpb.TryReadFromCache(IOOptions(), r.get(), 0 /* offset */,
                                    3000 /* n */, &result, &s, true));
   ASSERT_EQ(s, Status::OK());
+  ASSERT_EQ(result.size(), 3000);
   ASSERT_EQ(strncmp(result.data(), content.substr(0, 3000).c_str(), 3000), 0);
 
   ASSERT_TRUE(fpb.TryReadFromCache(IOOptions(), r.get(), 3000 /* offset */,
                                    10000 /* n */, &result, &s, true));
   ASSERT_EQ(s, Status::OK());
+  ASSERT_EQ(result.size(), 10000);
   ASSERT_EQ(strncmp(result.data(), content.substr(3000, 10000).c_str(), 10000),
             0);
 
   ASSERT_TRUE(fpb.TryReadFromCache(IOOptions(), r.get(), 15000 /* offset */,
                                    4096 /* n */, &result, &s, true));
   ASSERT_EQ(s, Status::OK());
+  ASSERT_EQ(result.size(), 4096);
   ASSERT_EQ(strncmp(result.data(), content.substr(15000, 4096).c_str(), 4096),
             0);
 
   ASSERT_TRUE(fpb.TryReadFromCache(IOOptions(), r.get(), 40000 /* offset */,
                                    20000 /* n */, &result, &s, true));
   ASSERT_EQ(s, Status::OK());
+  ASSERT_EQ(result.size(), 20000);
   ASSERT_EQ(strncmp(result.data(), content.substr(40000, 20000).c_str(), 20000),
             0);
+
+  // Try reading past end of file
+  ASSERT_TRUE(fpb.TryReadFromCache(IOOptions(), r.get(), 60000 /* offset */,
+                                   10000 /* n */, &result, &s, true));
+  ASSERT_EQ(s, Status::OK());
+  ASSERT_EQ(result.size(), 64 * 1024 - 60000);
+  ASSERT_EQ(
+      strncmp(result.data(), content.substr(60000, 64 * 1024 - 60000).c_str(),
+              64 * 1024 - 60000),
+      0);
 }
 
 class FSBufferPrefetchTest
@@ -3858,6 +3872,8 @@ TEST_P(FSBufferPrefetchTest, FSBufferPrefetchRandomized) {
     }
     ASSERT_TRUE(could_read_from_cache);
     ASSERT_EQ(s, Status::OK());
+    ASSERT_EQ(result.size(),
+              std::min(len, content.size() - static_cast<size_t>(offset)));
     ASSERT_EQ(strncmp(result.data(),
                       content.substr(offset, offset + len).c_str(), len),
               0);