Measure IODispatcher async read latency (#15001)

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

Add IODispatcher-specific histograms for async read observed completion latency, `FileSystem::Poll` wait time, and submit-to-consumer-poll prefetch lead time. Record submit timestamps on async read state, report completion latency from the callback, and report poll wait/prefetch lead when the block consumer polls the handle.

Keep Java histogram byte mappings stable by assigning the new histogram IDs after the existing range. Add a controlled async filesystem test that verifies one deferred async read records all three histograms.

Address Infer feedback by iterating `async_state->block_indices` directly when removing async IO map entries. Keep `PollAndProcessAsyncIO()` holding an owning `shared_ptr<AsyncIOState>` while erasing those entries so the coalesced request state cannot be destroyed mid-loop.

Format `util/io_dispatcher_imp.cc` and `util/io_dispatcher_test.cc` with clang-format 21.1.2 to fix the format signal.

Reviewed By: pdillinger, anand1976

Differential Revision: D113549401

fbshipit-source-id: dafc272c9252e1baccf404f681d33f76890ce0b2

include/rocksdb/io_dispatcher.h

@@ -271,9 +271,9 @@ class ReadSet {
   std::atomic<uint64_t> prefetch_bytes_ = 0;
   std::atomic<uint64_t> num_coalesced_nonadjacent_ = 0;
 
-  // Poll and process a specific async IO request
-  Status PollAndProcessAsyncIO(
-      const std::shared_ptr<AsyncIOState>& async_state);
+  // Poll and process a specific async IO request. Takes ownership by value so
+  // async_io_map_ entries can be erased without destroying the state mid-loop.
+  Status PollAndProcessAsyncIO(std::shared_ptr<AsyncIOState> async_state);
 
   // Release memory budget acquired for a prefetched block.
   void ReleasePrefetchMemory(size_t block_index);

include/rocksdb/statistics.h

@@ -796,6 +796,15 @@ enum Histograms : uint32_t {
   // Distribution of blocks prefetched per MultiScan Prepare()
   MULTISCAN_BLOCKS_PER_PREPARE,
 
+  // Time (microseconds) from IODispatcher async read submission until
+  // completion callback execution.
+  IO_DISPATCHER_ASYNC_READ_OBSERVED_COMPLETION_MICROS,
+  // Time (microseconds) IODispatcher spends waiting in FileSystem::Poll().
+  IO_DISPATCHER_ASYNC_READ_POLL_WAIT_MICROS,
+  // Time (microseconds) from async read submission until the block consumer
+  // starts polling for completion.
+  IO_DISPATCHER_ASYNC_READ_PREFETCH_LEAD_MICROS,
+
   // Coefficient of variation of key gaps in blocks, scaled by 10000
   // (e.g., CV of 0.4532 is recorded as 4532). Currently only used by index
   // blocks for uniform key distribution tracking.

java/rocksjni/portal.h

@@ -6086,6 +6086,15 @@ class HistogramTypeJni {
         return 0x4A;
       case ROCKSDB_NAMESPACE::Histograms::OPEN_AND_COMPACT_DB_OPEN_MICROS:
         return 0x4B;
+      case ROCKSDB_NAMESPACE::Histograms::
+          IO_DISPATCHER_ASYNC_READ_OBSERVED_COMPLETION_MICROS:
+        return 0x4C;
+      case ROCKSDB_NAMESPACE::Histograms::
+          IO_DISPATCHER_ASYNC_READ_POLL_WAIT_MICROS:
+        return 0x4D;
+      case ROCKSDB_NAMESPACE::Histograms::
+          IO_DISPATCHER_ASYNC_READ_PREFETCH_LEAD_MICROS:
+        return 0x4E;
       case ROCKSDB_NAMESPACE::Histograms::HISTOGRAM_ENUM_MAX:
         // 0x3E is reserved for backwards compatibility on current minor
         // version.
@@ -6259,6 +6268,15 @@ class HistogramTypeJni {
             FLUSH_WRITE_BUFFER_MANAGER_MEMTABLE_MEMORY_BYTES;
       case 0x4B:
         return ROCKSDB_NAMESPACE::Histograms::OPEN_AND_COMPACT_DB_OPEN_MICROS;
+      case 0x4C:
+        return ROCKSDB_NAMESPACE::Histograms::
+            IO_DISPATCHER_ASYNC_READ_OBSERVED_COMPLETION_MICROS;
+      case 0x4D:
+        return ROCKSDB_NAMESPACE::Histograms::
+            IO_DISPATCHER_ASYNC_READ_POLL_WAIT_MICROS;
+      case 0x4E:
+        return ROCKSDB_NAMESPACE::Histograms::
+            IO_DISPATCHER_ASYNC_READ_PREFETCH_LEAD_MICROS;
       case 0x3E:
         // 0x3E is reserved for backwards compatibility on current minor
         // version.

java/src/main/java/org/rocksdb/HistogramType.java

@@ -276,6 +276,21 @@ public enum HistogramType {
    */
   OPEN_AND_COMPACT_DB_OPEN_MICROS((byte) 0x4B),
 
+  /**
+   * Time from IODispatcher async read submission until completion callback.
+   */
+  IO_DISPATCHER_ASYNC_READ_OBSERVED_COMPLETION_MICROS((byte) 0x4C),
+
+  /**
+   * Time IODispatcher spends waiting in FileSystem::Poll().
+   */
+  IO_DISPATCHER_ASYNC_READ_POLL_WAIT_MICROS((byte) 0x4D),
+
+  /**
+   * Time from IODispatcher async read submission until the consumer polls.
+   */
+  IO_DISPATCHER_ASYNC_READ_PREFETCH_LEAD_MICROS((byte) 0x4E),
+
   // 0x3E is reserved for backwards compatibility on current minor version.
   HISTOGRAM_ENUM_MAX((byte) 0x3E);
 

monitoring/statistics.cc

@@ -410,6 +410,12 @@ const std::vector<std::pair<Histograms, std::string>> HistogramsNameMap = {
      "rocksdb.multiscan.op.prepare.iterators.micros"},
     {MULTISCAN_PREPARE_MICROS, "rocksdb.multiscan.prepare.micros"},
     {MULTISCAN_BLOCKS_PER_PREPARE, "rocksdb.multiscan.blocks.per.prepare"},
+    {IO_DISPATCHER_ASYNC_READ_OBSERVED_COMPLETION_MICROS,
+     "rocksdb.iodispatcher.async.read.observed.completion.micros"},
+    {IO_DISPATCHER_ASYNC_READ_POLL_WAIT_MICROS,
+     "rocksdb.iodispatcher.async.read.poll.wait.micros"},
+    {IO_DISPATCHER_ASYNC_READ_PREFETCH_LEAD_MICROS,
+     "rocksdb.iodispatcher.async.read.prefetch.lead.micros"},
     {BLOCK_KEY_DISTRIBUTION_CV, "rocksdb.block.key.distribution.cv"},
     {INGEST_EXTERNAL_FILE_PREPARE_TIME,
      "rocksdb.ingest.external.file.prepare.micros"},

util/io_dispatcher_imp.cc

@@ -229,8 +229,26 @@ struct AsyncIOState {
   std::vector<size_t> block_indices;
   std::vector<BlockHandle> blocks;
   FSReadRequest read_req;
+  SystemClock* clock = nullptr;
+  Statistics* statistics = nullptr;
+  uint64_t submit_time_us = 0;
 };
 
+static uint64_t GetIODispatcherNowMicros(SystemClock* clock) {
+  return clock != nullptr ? clock->NowMicros() : 0;
+}
+
+static void RecordIODispatcherTime(Statistics* statistics,
+                                   uint32_t histogram_type,
+                                   uint64_t start_time_us,
+                                   uint64_t end_time_us) {
+  if (statistics != nullptr && start_time_us != 0 &&
+      end_time_us >= start_time_us) {
+    RecordTimeToHistogram(statistics, histogram_type,
+                          end_time_us - start_time_us);
+  }
+}
+
 // ReadSet destructor - clean up IO handles
 // Must call AbortIO before deleting handles to avoid use-after-free when
 // io_uring completions arrive for deleted handles.
@@ -455,12 +473,20 @@ bool ReadSet::IsBlockAvailable(size_t block_index) const {
 
 // Poll and process async IO for a specific block
 Status ReadSet::PollAndProcessAsyncIO(
-    const std::shared_ptr<AsyncIOState>& async_state) {
+    std::shared_ptr<AsyncIOState> async_state) {
   auto* rep = job_->table->get_rep();
 
   // Poll for IO completion using FileSystem Poll API
   std::vector<void*> io_handles = {async_state->io_handle};
+  const uint64_t poll_start_time_us =
+      GetIODispatcherNowMicros(async_state->clock);
+  RecordIODispatcherTime(async_state->statistics,
+                         IO_DISPATCHER_ASYNC_READ_PREFETCH_LEAD_MICROS,
+                         async_state->submit_time_us, poll_start_time_us);
   IOStatus io_s = rep->ioptions.env->GetFileSystem()->Poll(io_handles, 1);
+  RecordIODispatcherTime(
+      async_state->statistics, IO_DISPATCHER_ASYNC_READ_POLL_WAIT_MICROS,
+      poll_start_time_us, GetIODispatcherNowMicros(async_state->clock));
   if (!io_s.ok()) {
     return io_s;
   }
@@ -496,9 +522,7 @@ Status ReadSet::PollAndProcessAsyncIO(
   DeleteAsyncIOHandle(async_state);
 
   // Remove from map - all blocks in this request have been processed
-  // Store indices in a temporary vector to avoid iterator invalidation
-  std::vector<size_t> indices_to_remove = async_state->block_indices;
-  for (const auto idx : indices_to_remove) {
+  for (const auto idx : async_state->block_indices) {
     async_io_map_.erase(idx);
   }
 
@@ -1152,6 +1176,9 @@ std::vector<size_t> IODispatcherImpl::Impl::ExecuteAsyncIO(
         read_scoped_io.read_buffer_requires_cleanup;
     async_state->block_indices = coalesced_block_indices[i];
     async_state->read_req = std::move(read_reqs[i]);
+    async_state->clock = rep->ioptions.clock;
+    async_state->statistics = rep->ioptions.stats;
+    async_state->submit_time_us = GetIODispatcherNowMicros(async_state->clock);
 
     for (const auto idx : coalesced_block_indices[i]) {
       async_state->blocks.emplace_back(job->block_handles[idx]);
@@ -1201,6 +1228,10 @@ std::vector<size_t> IODispatcherImpl::Impl::ExecuteAsyncIO(
       auto* state = static_cast<AsyncIOState*>(cb_arg);
       state->read_req.result = req.result;
       state->read_req.status = req.status;
+      RecordIODispatcherTime(
+          state->statistics,
+          IO_DISPATCHER_ASYNC_READ_OBSERVED_COMPLETION_MICROS,
+          state->submit_time_us, GetIODispatcherNowMicros(state->clock));
     };
 
     s = rep->file->ReadAsync(

util/io_dispatcher_test.cc

@@ -24,6 +24,7 @@
 #include "rocksdb/cache.h"
 #include "rocksdb/env.h"
 #include "rocksdb/options.h"
+#include "rocksdb/statistics.h"
 #include "rocksdb/table.h"
 #include "table/block_based/block_based_table_builder.h"
 #include "table/block_based/block_based_table_factory.h"
@@ -445,6 +446,7 @@ class IODispatcherTest : public DBTestBase {
   std::vector<std::unique_ptr<ImmutableOptions>> all_ioptions_;
   std::vector<std::unique_ptr<EnvOptions>> all_env_options_;
   std::vector<std::unique_ptr<InternalKeyComparator>> all_comparators_;
+  std::vector<std::shared_ptr<Statistics>> all_statistics_;
 
   class TestReadScopedBlockBufferProvider
       : public ReadScopedBlockBufferProvider {
@@ -611,11 +613,12 @@ class IODispatcherTest : public DBTestBase {
       const BlockBasedTableOptions* table_options_override = nullptr,
       bool use_direct_reads = false,
       CompressionType compression_type = kNoCompression,
-      bool allow_mmap_reads = false, Env* env_override = nullptr) {
+      bool allow_mmap_reads = false, Env* env_override = nullptr,
+      std::shared_ptr<Statistics> statistics = nullptr) {
     // Create options - store in member variables to avoid use-after-scope
     // The BlockBasedTable will keep references to these options
     Options options{};
-    options.statistics = nullptr;
+    options.statistics = statistics;
     options.allow_mmap_reads = allow_mmap_reads;
     if (env_override != nullptr) {
       options.env = env_override;
@@ -685,7 +688,7 @@ class IODispatcherTest : public DBTestBase {
     foptions.use_direct_reads = use_direct_reads;
     foptions.use_mmap_reads = allow_mmap_reads;
 
-    NewFileReader(table_name, foptions, &file, nullptr);
+    NewFileReader(table_name, foptions, &file, statistics.get());
 
     auto soptions = std::make_unique<EnvOptions>();
     soptions->use_mmap_reads = allow_mmap_reads;
@@ -708,6 +711,7 @@ class IODispatcherTest : public DBTestBase {
     all_ioptions_.push_back(std::move(ioptions));
     all_env_options_.push_back(std::move(soptions));
     all_comparators_.push_back(std::move(internal_comparator));
+    all_statistics_.push_back(std::move(statistics));
 
     table->reset(static_cast<BlockBasedTable*>(table_reader.release()));
 
@@ -1532,6 +1536,64 @@ TEST_F(IODispatcherTest, StatisticsTracking) {
   uint64_t total_reads = num_sync + num_async + num_cache;
   ASSERT_EQ(total_reads, block_handles.size());
 }
+
+TEST_F(IODispatcherTest, AsyncReadLatencyHistograms) {
+  if (!kIOUringPresent) {
+    ROCKSDB_GTEST_SKIP("Test requires io_uring support");
+    return;
+  }
+
+  auto controlled_fs = std::make_shared<ControlledAsyncFS>(base_fs_);
+  controlled_fs->SetCompleteStrayFirst(false);
+  tracking_fs_ = controlled_fs;
+  std::unique_ptr<Env> controlled_env = NewCompositeEnv(controlled_fs);
+  std::shared_ptr<Statistics> statistics = CreateDBStatistics();
+
+  BlockBasedTableOptions table_options;
+  table_options.block_cache = nullptr;
+  table_options.block_size = 16 * 1024;
+  table_options.no_block_cache = true;
+
+  std::unique_ptr<BlockBasedTable> table;
+  std::vector<BlockHandle> block_handles;
+  ASSERT_OK(CreateAndOpenSST(8, &table, &block_handles, &table_options,
+                             true /* use_direct_reads */, kNoCompression,
+                             false /* allow_mmap_reads */, controlled_env.get(),
+                             statistics));
+  ASSERT_FALSE(block_handles.empty());
+  block_handles.resize(1);
+
+  auto job = std::make_shared<IOJob>();
+  job->block_handles = block_handles;
+  job->table = table.get();
+  job->job_options.read_options.async_io = true;
+
+  std::unique_ptr<IODispatcher> dispatcher(NewIODispatcher());
+  std::shared_ptr<ReadSet> read_set;
+  ASSERT_OK(dispatcher->SubmitJob(job, &read_set));
+  ASSERT_NE(read_set, nullptr);
+  EXPECT_EQ(tracking_fs_->GetReadAsyncCount(), 1);
+
+  CachableEntry<Block> block;
+  ASSERT_OK(read_set->ReadIndex(0, &block));
+  ASSERT_NE(block.GetValue(), nullptr);
+
+  HistogramData observed_completion;
+  HistogramData poll_wait;
+  HistogramData prefetch_lead;
+  statistics->histogramData(IO_DISPATCHER_ASYNC_READ_OBSERVED_COMPLETION_MICROS,
+                            &observed_completion);
+  statistics->histogramData(IO_DISPATCHER_ASYNC_READ_POLL_WAIT_MICROS,
+                            &poll_wait);
+  statistics->histogramData(IO_DISPATCHER_ASYNC_READ_PREFETCH_LEAD_MICROS,
+                            &prefetch_lead);
+
+  EXPECT_EQ(observed_completion.count, 1);
+  EXPECT_EQ(poll_wait.count, 1);
+  EXPECT_EQ(prefetch_lead.count, 1);
+  EXPECT_GE(observed_completion.sum, prefetch_lead.sum);
+}
+
 TEST_F(IODispatcherTest, AsyncAndSyncRead) {
   // This test verifies the difference between async_io=true and async_io=false
   // by checking the statistics after reading all blocks.