From 325003655cf771a6470a1f251ec13493eb25180f Mon Sep 17 00:00:00 2001 From: Sam Stoelinga Date: Fri, 2 Oct 2026 17:24:19 -0700 Subject: [PATCH] Reduce GCS QPS on ArrayRecordReader open via ReadLast. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, opening a `gs://` ArrayRecord shard issued 3 separate GCS HTTP GET requests (5 when `index_storage_option` is `offloaded`): 1. An unbounded GET at offset 0 during `GcsReader` construction solely to obtain the object size (and then aborted mid-stream). 2. A GET via `reader->NewReader(size - 64KB)` to read the Riegeli postscript. 3. A GET via `reader->NewReader(footer_offset)` to read the ArrayRecord footer (and additional GETs in `OffloadedChunkOffset::operator[]` when `index_storage_option` is `offloaded`). Open `riegeli::GcsReader` with `google::cloud::storage::ReadLast(1 MiB)` (using a per-reader `google::cloud::storage::Client` instance to avoid cross-thread `libcurl` handle pool contention and `Curl_peer` refcount data races) so the initial HTTP GET retrieves the total object size (`Content-Range`) and prefetches the last 1 MiB of the shard. In `ArrayRecordReaderBase::Initialize()`, pull the prefetched tail into `reader`'s buffer (resetting the cursor to `limit()` so `available() == 0` on close) and read the postscript and footer (including offloaded chunk index lookups in `OffloadedChunkOffset::operator[]`) directly from `reader`'s buffered tail via `ReadChunk` when `footer_offset >= reader->start_pos()`, without spawning subordinate `NewReader` HTTP streams. Real-GCS integration benchmarks show 3.0×–5.0× fewer HTTP GETs on open (23.0× fewer on shards < 1 MiB), 2.0×–3.0× fewer TLS handshakes on open, 2.2×–4.7× faster shard initialization, and 1.13×–1.65× higher bulk read throughput with zero data races. PiperOrigin-RevId: 992624347 --- cpp/BUILD | 2 + cpp/array_record_reader.cc | 55 ++++++-- cpp/array_record_reader_test.cc | 210 +++++++++++++++++++++++++++++ python/array_record_module.cc | 11 +- python/array_record_module_test.py | 200 +++++++++++++++++++++++++++ 5 files changed, 467 insertions(+), 11 deletions(-) diff --git a/cpp/BUILD b/cpp/BUILD index 64796ce..3c5ef4a 100644 --- a/cpp/BUILD +++ b/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", diff --git a/cpp/array_record_reader.cc b/cpp/array_record_reader.cc index 484385f..367db94 100644 --- a/cpp/array_record_reader.cc +++ b/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 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( - this, state_->footer_offset, num_chunks); + this, state_->footer_offset, footer_length, num_chunks); break; } diff --git a/cpp/array_record_reader_test.cc b/cpp/array_record_reader_test.cc index 55a9a56..e9d8a76 100644 --- a/cpp/array_record_reader_test.cc +++ b/cpp/array_record_reader_test.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -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 SizeImpl() override { return data_.size(); } + + std::unique_ptr 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(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(&encoded), writer_options, nullptr); + std::vector 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(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(&encoded), writer_options, nullptr); + std::vector 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(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(&encoded), writer_options, nullptr); + std::vector 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(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(&encoded), writer_options, nullptr); + std::vector 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(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(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 diff --git a/python/array_record_module.cc b/python/array_record_module.cc index 88b8961..b52c004 100644 --- a/python/array_record_module.cc +++ b/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( 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 { diff --git a/python/array_record_module_test.py b/python/array_record_module_test.py index 452e517..50d6bd2 100644 --- a/python/array_record_module_test.py +++ b/python/array_record_module_test.py @@ -14,7 +14,11 @@ """Tests for array_record_module.""" +from concurrent import futures +from http import server import os +import re +import threading from absl.testing import absltest @@ -22,6 +26,80 @@ from python.array_record_module import ArrayRecordWriter +class _MockGcsServer: + """Minimal HTTP server emulating the GCS JSON API media download endpoint.""" + + def __init__(self, payload: bytes): + self.payload = payload + self.requests = [] + self.tcp_connections = 0 + self.aborted_streams = 0 + self._lock = threading.Lock() + parent = self + + class Handler(server.BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def handle(self): + with parent._lock: + parent.tcp_connections += 1 + try: + super().handle() + except ConnectionResetError: + with parent._lock: + parent.aborted_streams += 1 + + def do_GET(self): # pylint: disable=invalid-name + range_header = self.headers.get("Range") + with parent._lock: + parent.requests.append({ + "path": self.path, + "range": range_header, + }) + total = len(parent.payload) + start = 0 + end = total - 1 + status = 200 + if range_header: + m_last = re.fullmatch(r"bytes=-(\d+)", range_header) + m_range = re.fullmatch(r"bytes=(\d+)-(\d*)", range_header) + if m_last: + last_n = int(m_last.group(1)) + start = max(0, total - last_n) + status = 206 + elif m_range: + start = int(m_range.group(1)) + if m_range.group(2): + end = min(total - 1, int(m_range.group(2))) + status = 206 + data = parent.payload[start : end + 1] + self.send_response(status) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Range", f"bytes {start}-{end}/{total}") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + try: + self.wfile.write(data) + except (BrokenPipeError, ConnectionResetError): + with parent._lock: + parent.aborted_streams += 1 + + def log_message(self, fmt, *args): + del fmt, args + + self._httpd = server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self.port = self._httpd.server_address[1] + self._thread = threading.Thread( + target=self._httpd.serve_forever, daemon=True + ) + self._thread.start() + + def close(self): + self._httpd.shutdown() + self._httpd.server_close() + self._thread.join(timeout=5.0) + + class ArrayRecordModuleTest(absltest.TestCase): def setUp(self): @@ -142,6 +220,128 @@ def test_writer_options(self): "group_size:42,transpose:false,pad_to_block_boundary:false,zstd:3," "window_log:20,max_parallelism:1") + def test_gcs_read_single_open_rpc(self): + # Write a >2 MiB uncompressed ArrayRecord file so that the file size exceeds + # the 1 MiB tail prefetch window. + writer = ArrayRecordWriter(self.test_file, "group_size:1,uncompressed") + records = [f"record-{i:04d}-".encode() + (b"x" * 65536) for i in range(36)] + for r in records: + writer.write(r) + writer.close() + with open(self.test_file, "rb") as f: + payload = f.read() + self.assertGreater(len(payload), 2 * 1024 * 1024) + + mock_gcs = _MockGcsServer(payload) + old_endpoint = os.environ.get("CLOUD_STORAGE_EMULATOR_ENDPOINT") + os.environ["CLOUD_STORAGE_EMULATOR_ENDPOINT"] = ( + f"http://127.0.0.1:{mock_gcs.port}" + ) + try: + # Opening and closing a gs:// shard should issue exactly 1 HTTP GET + # (with Range: bytes=-1048576) instead of 3 HTTP GETs. + num_records = len(records) + reader = ArrayRecordReader( + "gs://test-bucket/shard-00000-of-00001", + "readahead_buffer_size:0,max_parallelism:0", + file_reader_buffer_size=32768, + ) + self.assertEqual(reader.num_records(), num_records) + reader.close() + self.assertLen( + mock_gcs.requests, + 1, + msg=( + f"ranges={[r['range'] for r in mock_gcs.requests]}, " + f"tcp_connections={mock_gcs.tcp_connections}, " + f"aborted_streams={mock_gcs.aborted_streams}" + ), + ) + self.assertEqual(mock_gcs.requests[0]["range"], "bytes=-1048576") + self.assertEqual(mock_gcs.tcp_connections, 1) + self.assertEqual(mock_gcs.aborted_streams, 0) + + # Opening a gs:// shard with index_storage_option:offloaded should also + # issue only 1 HTTP GET by serving OffloadedChunkOffset lookups from the + # prefetched tail buffer. + mock_gcs.requests.clear() + conns_before = mock_gcs.tcp_connections + reader = ArrayRecordReader( + "gs://test-bucket/shard-00000-of-00001", + "index_storage_option:offloaded,readahead_buffer_size:0,max_parallelism:0", + file_reader_buffer_size=32768, + ) + self.assertEqual(reader.num_records(), num_records) + reader.close() + self.assertLen(mock_gcs.requests, 1) + self.assertEqual(mock_gcs.requests[0]["range"], "bytes=-1048576") + self.assertEqual(mock_gcs.tcp_connections - conns_before, 1) + self.assertEqual(mock_gcs.aborted_streams, 0) + + # Verify record payload reading works over gs:// and reuses the same + # TCP connection within the ArrayRecordReader instance. + mock_gcs.requests.clear() + conns_before = mock_gcs.tcp_connections + reader = ArrayRecordReader( + "gs://test-bucket/shard-00000-of-00001", + "readahead_buffer_size:16M,max_parallelism:1", + file_reader_buffer_size=32768, + ) + self.assertEqual( + reader.read([0, 17, 35]), [records[0], records[17], records[35]] + ) + reader.close() + self.assertEqual(mock_gcs.tcp_connections - conns_before, 1) + finally: + if old_endpoint is None: + os.environ.pop("CLOUD_STORAGE_EMULATOR_ENDPOINT", None) + else: + os.environ["CLOUD_STORAGE_EMULATOR_ENDPOINT"] = old_endpoint + mock_gcs.close() + + def test_gcs_concurrent_multithreaded_reads(self): + writer = ArrayRecordWriter(self.test_file, "group_size:1,uncompressed") + records = [f"record-{i:04d}-".encode() + (b"x" * 65536) for i in range(36)] + for r in records: + writer.write(r) + writer.close() + with open(self.test_file, "rb") as f: + payload = f.read() + + mock_gcs = _MockGcsServer(payload) + old_endpoint = os.environ.get("CLOUD_STORAGE_EMULATOR_ENDPOINT") + os.environ["CLOUD_STORAGE_EMULATOR_ENDPOINT"] = ( + f"http://127.0.0.1:{mock_gcs.port}" + ) + try: + def worker(idx: int): + options = ( + "index_storage_option:offloaded,readahead_buffer_size:16M,max_parallelism:1" + if idx % 2 == 1 + else "readahead_buffer_size:16M,max_parallelism:1" + ) + reader = ArrayRecordReader(f"gs://test-bucket/shard-{idx:05d}", options) + self.assertLen(records, reader.num_records()) + fetched = reader.read([0, idx % len(records), len(records) - 1]) + reader.close() + return fetched + + with futures.ThreadPoolExecutor(max_workers=16) as executor: + results = list(executor.map(worker, range(32))) + + for idx, fetched in enumerate(results): + self.assertEqual( + fetched, + [records[0], records[idx % len(records)], records[-1]], + ) + self.assertLen(mock_gcs.requests, 64) + finally: + if old_endpoint is None: + os.environ.pop("CLOUD_STORAGE_EMULATOR_ENDPOINT", None) + else: + os.environ["CLOUD_STORAGE_EMULATOR_ENDPOINT"] = old_endpoint + mock_gcs.close() + if __name__ == "__main__": absltest.main()