Reduce GCS QPS on ArrayRecordReader open via ReadLast.

cpp/BUILD

@@ -296,7 +296,9 @@ cc_test(
         "@abseil-cpp//absl/strings",
         "@googletest//:gtest_main",
         "@riegeli//riegeli/base:initializer",
+        "@riegeli//riegeli/base:types",
         "@riegeli//riegeli/bytes:chain_writer",
+        "@riegeli//riegeli/bytes:reader",
         "@riegeli//riegeli/bytes:string_reader",
         "@riegeli//riegeli/bytes:string_writer",
         "@riegeli//riegeli/chunk_encoding:chunk",

cpp/array_record_reader.cc

@@ -54,6 +54,9 @@ namespace array_record {
 // 64KB
 constexpr size_t kRiegeliBlockSize = (1 << 16);
 
+// 1MB tail prefetch size when reading postscript and footer from buffered tail.
+constexpr size_t kTailPrefetchSize = (1 << 20);
+
 // This number should rarely change unless there's a new great layout design
 // that wasn't backward compatible and justifies its performance and reliability
 // worth us to implement.
@@ -118,22 +121,22 @@ class ArrayChunkOffset : public ChunkOffset {
   std::vector<uint64_t> chunk_offsets_;
 };
 
+ChunkDecoder ReadChunk(Reader& reader, size_t pos, size_t len);
+
 class OffloadedChunkOffset : public ChunkOffset {
  public:
   OffloadedChunkOffset(ArrayRecordReaderBase* base_reader,
-                       uint64_t footer_offset, uint64_t num_chunks)
+                       uint64_t footer_offset, uint64_t footer_length,
+                       uint64_t num_chunks)
       : base_reader_(base_reader),
         footer_offset_(footer_offset),
+        footer_length_(footer_length),
         num_chunks_(num_chunks) {}
 
   uint64_t operator[](size_t idx) const override {
     auto backing_reader = base_reader_->get_backing_reader();
-    auto reader = backing_reader->NewReader(footer_offset_);
-    auto chunk_reader = riegeli::DefaultChunkReader<>(reader.get());
-    Chunk chunk;
-    ChunkDecoder footer_decoder;
-    chunk_reader.ReadChunk(chunk);
-    footer_decoder.Decode(chunk);
+    ChunkDecoder footer_decoder =
+        ReadChunk(*backing_reader, footer_offset_, footer_length_);
     // First item is the footer_metadata, so we want to skip.
     footer_decoder.SetIndex(idx + 1);
     ArrayRecordFooter footer;
@@ -150,6 +153,7 @@ class OffloadedChunkOffset : public ChunkOffset {
  private:
   ArrayRecordReaderBase* base_reader_;
   uint64_t footer_offset_;
+  uint64_t footer_length_;
   uint64_t num_chunks_;
 };
 
@@ -256,6 +260,19 @@ ArrayRecordReaderBase& ArrayRecordReaderBase::operator=(
   return *this;
 }
 
+class BufferedSliceReader : public Reader {
+ public:
+  BufferedSliceReader(const char* data, size_t pos, size_t len) {
+    set_buffer(data, len);
+    set_limit_pos(pos + len);
+  }
+
+ protected:
+  bool PullSlow(size_t min_length, size_t recommended_length) override {
+    return false;
+  }
+};
+
 // After the first access to the underlying `riegeli::Reader`, the lazily
 // evaluated variables for random access are all initialized. Therefore it's
 // safe to access the reader from multiple threads later on, even though the
@@ -266,6 +283,19 @@ ChunkDecoder ReadChunk(Reader& reader, size_t pos, size_t len) {
     decoder.Fail(reader.status());
     return decoder;
   }
+  if (pos >= reader.start_pos() && pos <= reader.limit_pos() &&
+      len <= reader.limit_pos() - pos) {
+    BufferedSliceReader buffered_slice(
+        reader.start() + (pos - reader.start_pos()), pos, len);
+    auto chunk_reader = riegeli::DefaultChunkReader<>(&buffered_slice);
+    Chunk chunk;
+    if (!chunk_reader.ReadChunk(chunk)) {
+      decoder.Fail(chunk_reader.status());
+      return decoder;
+    }
+    decoder.Decode(chunk);
+    return decoder;
+  }
   MaskedReader masked_reader(reader.NewReader(pos), len);
   if (!masked_reader.ok()) {
     decoder.Fail(masked_reader.status());
@@ -309,6 +339,7 @@ void ArrayRecordReaderBase::Initialize() {
   AR_ENDO_TASK("Reading ArrayRecord footer");
   RiegeliFooterMetadata footer_metadata;
   ChunkDecoder footer_decoder;
+  uint64_t footer_length = 0;
   {
     AR_ENDO_SCOPE("Reading postscript and footer chunk");
     if (!reader->SupportsRandomAccess()) {
@@ -327,6 +358,10 @@ void ArrayRecordReaderBase::Initialize() {
           InvalidArgumentError("ArrayRecord file should be at least 64KB big"));
       return;
     }
+    if (size - reader->pos() <= kTailPrefetchSize) {
+      reader->Pull(size - reader->pos());
+      reader->set_cursor(reader->limit());
+    }
     RiegeliPostscript postscript;
     auto postscript_decoder =
         ReadChunk(*reader, size - kRiegeliBlockSize, kRiegeliBlockSize);
@@ -347,8 +382,8 @@ void ArrayRecordReaderBase::Initialize() {
     }
     auto footer_offset = postscript.footer_offset();
     state_->footer_offset = footer_offset;
-    footer_decoder = ReadChunk(*reader, footer_offset,
-                               size - kRiegeliBlockSize - footer_offset);
+    footer_length = size - kRiegeliBlockSize - footer_offset;
+    footer_decoder = ReadChunk(*reader, footer_offset, footer_length);
 
     if (!footer_decoder.ReadRecord(footer_metadata)) {
       Fail(Annotate(footer_decoder.status(),
@@ -404,7 +439,7 @@ void ArrayRecordReaderBase::Initialize() {
       } break;
       case Options::IndexStorageOption::kOffloaded:
         state_->chunk_offsets = std::make_unique<OffloadedChunkOffset>(
-            this, state_->footer_offset, num_chunks);
+            this, state_->footer_offset, footer_length, num_chunks);
         break;
     }
 

cpp/array_record_reader_test.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
+#include <memory>
 #include <optional>
 #include <random>
 #include <string>
@@ -35,7 +36,9 @@ limitations under the License.
 #include "cpp/test_utils.h"
 #include "cpp/thread_pool.h"
 #include "riegeli/base/maker.h"
+#include "riegeli/base/types.h"
 #include "riegeli/bytes/chain_writer.h"
+#include "riegeli/bytes/reader.h"
 #include "riegeli/bytes/string_reader.h"
 #include "riegeli/bytes/string_writer.h"
 #include "riegeli/chunk_encoding/chunk.h"
@@ -450,5 +453,212 @@ TEST(ArrayRecordReaderTest, OutOfBoundsChunkIdx) {
   EXPECT_EQ(status.code(), absl::StatusCode::kOutOfRange);
 }
 
+// Simulates a Reader (such as riegeli::GcsReader opened with ReadLast) that
+// starts at `max(0, size - tail_prefetch_size)` with an initially empty buffer
+// that is populated on Pull().
+class TailPrefetchTestReader : public riegeli::Reader {
+ public:
+  TailPrefetchTestReader(absl::string_view data, size_t tail_prefetch_size,
+                         int* pull_count, int* new_reader_count)
+      : data_(data),
+        initial_pos_(data.size() > tail_prefetch_size
+                         ? data.size() - tail_prefetch_size
+                         : 0),
+        pull_count_(pull_count),
+        new_reader_count_(new_reader_count) {
+    set_limit_pos(initial_pos_);
+  }
+
+  bool SupportsRandomAccess() override { return true; }
+  bool SupportsNewReader() override { return true; }
+
+ protected:
+  void Done() override {
+    // Ensure available() == 0 on Close() so BufferedReader::Done() does not
+    // issue a spurious SeekBehindBuffer() RPC.
+    EXPECT_EQ(available(), 0);
+    riegeli::Reader::Done();
+  }
+
+  bool PullSlow(size_t min_length, size_t recommended_length) override {
+    if (pull_count_ != nullptr) {
+      ++(*pull_count_);
+    }
+    if (min_length > data_.size() - initial_pos_) {
+      return false;
+    }
+    set_buffer(data_.data() + initial_pos_, min_length);
+    set_limit_pos(initial_pos_ + min_length);
+    return true;
+  }
+
+  std::optional<riegeli::Position> SizeImpl() override { return data_.size(); }
+
+  std::unique_ptr<riegeli::Reader> NewReaderImpl(
+      riegeli::Position initial_pos) override {
+    if (new_reader_count_ != nullptr) {
+      ++(*new_reader_count_);
+    }
+    constexpr size_t kBlockSize = 1 << 16;
+    const size_t end_pos = initial_pos < data_.size() - kBlockSize
+                               ? data_.size() - kBlockSize
+                               : data_.size();
+    auto reader =
+        std::make_unique<ExactSliceReader>(data_.substr(0, end_pos));
+    reader->Seek(initial_pos);
+    return reader;
+  }
+
+ private:
+  class ExactSliceReader : public riegeli::StringReader<> {
+   public:
+    explicit ExactSliceReader(absl::string_view data)
+        : riegeli::StringReader<>(data) {}
+
+    ~ExactSliceReader() override { EXPECT_EQ(pos(), limit_pos()); }
+  };
+
+  absl::string_view data_;
+  size_t initial_pos_;
+  int* pull_count_;
+  int* new_reader_count_;
+};
+
+TEST(ArrayRecordReaderTest, TailPrefetchReadsPostscriptAndFooterFromBuffer) {
+  std::string encoded;
+  auto writer_options =
+      ArrayRecordWriterBase::Options().set_group_size(2).set_uncompressed();
+  auto writer = ArrayRecordWriter(
+      riegeli::Maker<riegeli::StringWriter>(&encoded), writer_options, nullptr);
+  std::vector<std::string> test_str{std::string(1 << 20, 'a'), "bbb", "ccc",
+                                    "ddd", "eee"};
+  for (const auto& s : test_str) {
+    EXPECT_TRUE(writer.WriteRecord(s));
+  }
+  ASSERT_TRUE(writer.Close());
+  ASSERT_GT(encoded.size(), 1 << 20);
+
+  int pull_count = 0;
+  int new_reader_count = 0;
+  constexpr size_t kTailPrefetch = 1 << 20;  // 1 MiB (exact boundary)
+  auto reader = ArrayRecordReader(
+      riegeli::Maker<TailPrefetchTestReader>(encoded, kTailPrefetch,
+                                             &pull_count, &new_reader_count),
+      ArrayRecordReaderBase::Options(), nullptr);
+  ASSERT_TRUE(reader.status().ok()) << reader.status().message();
+  // Both postscript and footer should be read directly from the pulled tail
+  // buffer without any NewReader() calls during Initialize().
+  EXPECT_EQ(pull_count, 1);
+  EXPECT_EQ(new_reader_count, 0);
+  EXPECT_EQ(reader.NumRecords(), test_str.size());
+  ASSERT_TRUE(reader.Close());
+}
+
+TEST(ArrayRecordReaderTest,
+     TailPrefetchDoesNotPullWhenRemainingExceedsMaxTailPrefetch) {
+  std::string encoded;
+  auto writer_options =
+      ArrayRecordWriterBase::Options().set_group_size(2).set_uncompressed();
+  auto writer = ArrayRecordWriter(
+      riegeli::Maker<riegeli::StringWriter>(&encoded), writer_options, nullptr);
+  std::vector<std::string> test_str{std::string(1 << 20, 'a'), "bbb"};
+  for (const auto& s : test_str) {
+    EXPECT_TRUE(writer.WriteRecord(s));
+  }
+  ASSERT_TRUE(writer.Close());
+  ASSERT_GT(encoded.size(), 1 << 20);
+
+  int pull_count = 0;
+  int new_reader_count = 0;
+  // Start at offset 0 on an input larger than kTailPrefetchSize (1 MiB).
+  auto reader = ArrayRecordReader(
+      riegeli::Maker<TailPrefetchTestReader>(encoded, encoded.size(),
+                                             &pull_count, &new_reader_count),
+      ArrayRecordReaderBase::Options(), nullptr);
+  ASSERT_TRUE(reader.status().ok()) << reader.status().message();
+  // Remaining bytes exceed 1 MiB, so Initialize() must not Pull() the entire
+  // file and instead reads postscript and footer via NewReader().
+  EXPECT_EQ(pull_count, 0);
+  EXPECT_EQ(new_reader_count, 2);
+  EXPECT_EQ(reader.NumRecords(), test_str.size());
+  ASSERT_TRUE(reader.Close());
+}
+
+TEST(ArrayRecordReaderTest,
+     TailPrefetchFallsBackToNewReaderWhenFooterBeforeStartPos) {
+  std::string encoded;
+  auto writer_options =
+      ArrayRecordWriterBase::Options().set_group_size(2).set_uncompressed();
+  auto writer = ArrayRecordWriter(
+      riegeli::Maker<riegeli::StringWriter>(&encoded), writer_options, nullptr);
+  std::vector<std::string> test_str{"aaa", "bbb", "ccc", "ddd", "eee"};
+  for (const auto& s : test_str) {
+    EXPECT_TRUE(writer.WriteRecord(s));
+  }
+  ASSERT_TRUE(writer.Close());
+
+  int pull_count = 0;
+  int new_reader_count = 0;
+  // Only prefetch the 64 KiB postscript block so footer_offset < start_pos().
+  constexpr size_t kPostscriptOnlyPrefetch = 1 << 16;  // 64 KiB
+  auto reader = ArrayRecordReader(
+      riegeli::Maker<TailPrefetchTestReader>(encoded, kPostscriptOnlyPrefetch,
+                                             &pull_count, &new_reader_count),
+      ArrayRecordReaderBase::Options(), nullptr);
+  ASSERT_TRUE(reader.status().ok()) << reader.status().message();
+  // Postscript is read from the buffered tail (0 NewReader calls), while footer
+  // falls back to NewReader() (1 NewReader call).
+  EXPECT_EQ(pull_count, 1);
+  EXPECT_EQ(new_reader_count, 1);
+  EXPECT_EQ(reader.NumRecords(), test_str.size());
+  ASSERT_TRUE(reader.Close());
+}
+
+TEST(ArrayRecordReaderTest, TailPrefetchReadsOffloadedIndexFromBuffer) {
+  std::string encoded;
+  auto writer_options =
+      ArrayRecordWriterBase::Options().set_group_size(2).set_uncompressed();
+  auto writer = ArrayRecordWriter(
+      riegeli::Maker<riegeli::StringWriter>(&encoded), writer_options, nullptr);
+  std::vector<std::string> test_str{"aaa", "bbb", "ccc", "ddd", "eee"};
+  for (const auto& s : test_str) {
+    EXPECT_TRUE(writer.WriteRecord(s));
+  }
+  ASSERT_TRUE(writer.Close());
+
+  int pull_count = 0;
+  int new_reader_count = 0;
+  constexpr size_t kTailPrefetch = 1 << 20;  // 1 MiB
+  auto reader_options =
+      ArrayRecordReaderBase::Options()
+          .set_index_storage_option(IndexStorageOption::kOffloaded)
+          .set_max_parallelism(0);
+  auto reader = ArrayRecordReader(
+      riegeli::Maker<TailPrefetchTestReader>(encoded, kTailPrefetch,
+                                             &pull_count, &new_reader_count),
+      reader_options, nullptr);
+  ASSERT_TRUE(reader.status().ok()) << reader.status().message();
+  // Postscript, footer, and OffloadedChunkOffset lookups during Initialize()
+  // should all be served from the pulled tail buffer without NewReader() calls.
+  EXPECT_EQ(pull_count, 1);
+  EXPECT_EQ(new_reader_count, 0);
+  EXPECT_EQ(reader.NumRecords(), test_str.size());
+  ASSERT_TRUE(reader.Close());
+}
+
+TEST(ArrayRecordReaderTest, TailPrefetchFailsOnCorruptedChunkInBuffer) {
+  std::string corrupted(1 << 16, 'x');
+  int pull_count = 0;
+  int new_reader_count = 0;
+  constexpr size_t kTailPrefetch = 1 << 20;
+  auto reader = ArrayRecordReader(
+      riegeli::Maker<TailPrefetchTestReader>(corrupted, kTailPrefetch,
+                                             &pull_count, &new_reader_count),
+      ArrayRecordReaderBase::Options(), nullptr);
+  EXPECT_FALSE(reader.status().ok());
+  EXPECT_EQ(pull_count, 1);
+  EXPECT_EQ(new_reader_count, 0);
+}
+
 }  // namespace
 }  // namespace array_record

python/array_record_module.cc

@@ -40,6 +40,14 @@ limitations under the License.
 
 namespace py = pybind11;
 
+namespace {
+
+// Prefetch the last 1 MiB when opening a GCS object so that the object size,
+// Riegeli postscript, and ArrayRecord footer can be read in a single HTTP GET.
+constexpr int64_t kGcsTailPrefetchSize = 1 << 20;
+
+}  // namespace
+
 PYBIND11_MODULE(array_record_module, m) {
   using array_record::ArrayRecordReaderBase;
   using array_record::ArrayRecordWriterBase;
@@ -103,7 +111,8 @@ PYBIND11_MODULE(array_record_module, m) {
                return new array_record::ArrayRecordReader(
                    riegeli::Maker<riegeli::GcsReader>(
                        google::cloud::storage::Client(),
-                       riegeli::GcsObject(path), std::move(gcs_reader_options)),
+                       riegeli::GcsObject(path), std::move(gcs_reader_options),
+                       google::cloud::storage::ReadLast(kGcsTailPrefetchSize)),
                    status_or_option.value(),
                    array_record::ArrayRecordGlobalPool());
              } else {