Fix MultiScan not respecting SupportedOps for async IO (#14510)

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

When a filesystem does not return kAsyncIO in SupportedOps(), MultiScan
still sends ReadAsync/Poll/AbortIO requests. Regular iterators check via
CheckFSFeatureSupport in ArenaWrappedDBIter::Init and ForwardIterator,
but MultiScan blindly trusted the caller-provided
MultiScanArgs::use_async_io flag.

Fix: In the MultiScan constructor, after scan_opts_ is initialized,
check CheckFSFeatureSupport and disable use_async_io if the FS doesn't
support it. Also pass the (potentially modified) scan_opts_ to Prepare()
instead of the original scan_opts parameter.

Reviewed By: mszeszko-meta

Differential Revision: D97995735

fbshipit-source-id: 331639d950fd3cb9f491feed996d5820294beb4e

db/db_iterator_test.cc

@@ -14,8 +14,10 @@
 #include "db/arena_wrapped_db_iter.h"
 #include "db/db_iter.h"
 #include "db/db_test_util.h"
+#include "env/composite_env_wrapper.h"
 #include "port/port.h"
 #include "port/stack_trace.h"
+#include "rocksdb/file_system.h"
 #include "rocksdb/io_dispatcher.h"
 #include "rocksdb/iostats_context.h"
 #include "rocksdb/perf_context.h"
@@ -4806,6 +4808,124 @@ TEST_P(DBMultiScanIteratorTest, AsyncPrefetchAcrossMultipleFiles) {
   iter.reset();
 }
 
+// Wrapper filesystem that does not support async IO.
+// Used to verify that MultiScan gracefully falls back to sync IO.
+class NoAsyncIOFS : public FileSystemWrapper {
+ public:
+  explicit NoAsyncIOFS(const std::shared_ptr<FileSystem>& target)
+      : FileSystemWrapper(target) {}
+  static const char* kClassName() { return "NoAsyncIOFS"; }
+  const char* Name() const override { return kClassName(); }
+
+  void SupportedOps(int64_t& supported_ops) override {
+    target()->SupportedOps(supported_ops);
+    supported_ops &= ~(1 << FSSupportedOps::kAsyncIO);
+  }
+
+  IOStatus Poll(std::vector<void*>& /*io_handles*/,
+                size_t /*min_completions*/) override {
+    ADD_FAILURE() << "Poll should not be called when kAsyncIO is unsupported";
+    return IOStatus::NotSupported();
+  }
+
+  IOStatus AbortIO(std::vector<void*>& /*io_handles*/) override {
+    ADD_FAILURE()
+        << "AbortIO should not be called when kAsyncIO is unsupported";
+    return IOStatus::NotSupported();
+  }
+
+  IOStatus NewRandomAccessFile(const std::string& fname,
+                               const FileOptions& opts,
+                               std::unique_ptr<FSRandomAccessFile>* result,
+                               IODebugContext* dbg) override {
+    IOStatus s = target()->NewRandomAccessFile(fname, opts, result, dbg);
+    if (s.ok()) {
+      *result = std::make_unique<NoAsyncIOFile>(std::move(*result));
+    }
+    return s;
+  }
+
+ private:
+  class NoAsyncIOFile : public FSRandomAccessFileOwnerWrapper {
+   public:
+    using FSRandomAccessFileOwnerWrapper::FSRandomAccessFileOwnerWrapper;
+    IOStatus ReadAsync(FSReadRequest& /*req*/, const IOOptions& /*opts*/,
+                       std::function<void(FSReadRequest&, void*)> /*cb*/,
+                       void* /*cb_arg*/, void** /*io_handle*/,
+                       IOHandleDeleter* /*del_fn*/,
+                       IODebugContext* /*dbg*/) override {
+      ADD_FAILURE()
+          << "ReadAsync should not be called when kAsyncIO is unsupported";
+      return IOStatus::NotSupported();
+    }
+  };
+};
+
+TEST_P(DBMultiScanIteratorTest, AsyncIOFallbackWithoutFSSupport) {
+  // Verify that MultiScan works correctly when use_async_io = true but the
+  // filesystem does NOT support kAsyncIO. The constructor should silently
+  // disable async IO, and no async operations should be attempted.
+  auto options = CurrentOptions();
+  options.target_file_size_base = 1 << 15;  // 32KiB
+  options.compaction_style = kCompactionStyleUniversal;
+  options.num_levels = 50;
+  options.compression = kNoCompression;
+
+  auto no_async_fs = std::make_shared<NoAsyncIOFS>(env_->GetFileSystem());
+  std::unique_ptr<Env> wrapped_env(new CompositeEnvWrapper(env_, no_async_fs));
+  options.env = wrapped_env.get();
+
+  DestroyAndReopen(options);
+
+  Random rnd(305);
+
+  // Create data across multiple files
+  for (int i = 0; i < 1000; ++i) {
+    std::stringstream ss;
+    ss << "k" << std::setw(5) << std::setfill('0') << i;
+    ASSERT_OK(Put(ss.str(), rnd.RandomString(1 << 10)));
+  }
+  ASSERT_OK(Flush());
+  ASSERT_OK(db_->CompactRange({}, nullptr, nullptr));
+
+  // Set up multiple non-overlapping ranges
+  std::vector<std::string> key_ranges(
+      {"k00000", "k00100", "k00500", "k00600", "k00800", "k00900"});
+  ReadOptions ro;
+  ro.fill_cache = GetParam();
+
+  ColumnFamilyHandle* cfh = dbfull()->DefaultColumnFamily();
+
+  MultiScanArgs scan_options(BytewiseComparator());
+  scan_options.use_async_io = true;  // Deliberately request async IO
+  scan_options.insert(key_ranges[0], key_ranges[1]);
+  scan_options.insert(key_ranges[2], key_ranges[3]);
+  scan_options.insert(key_ranges[4], key_ranges[5]);
+
+  // The MultiScan constructor should disable async IO since the FS
+  // doesn't support it. If it doesn't, the NoAsyncIOFS wrapper will
+  // cause ADD_FAILURE() when ReadAsync/Poll/AbortIO are called.
+  std::unique_ptr<MultiScan> iter =
+      dbfull()->NewMultiScan(ro, cfh, scan_options);
+  ASSERT_NE(iter, nullptr);
+
+  // Iterate all ranges and verify keys are returned correctly
+  try {
+    for (auto range : *iter) {
+      for (auto it : range) {
+        it.first.ToString();
+      }
+    }
+  } catch (MultiScanException& ex) {
+    FAIL() << "Iterator returned status " << ex.what();
+  } catch (std::logic_error& ex) {
+    FAIL() << "Iterator returned logic error " << ex.what();
+  }
+
+  iter.reset();
+  Close();
+}
+
 TEST_P(DBMultiScanIteratorTest, AsyncPrefetchMultipleLevels) {
   // Test async prefetch with files in L0 and non-L0 levels
   // Similar setup to AsyncPrefetchAcrossMultipleFiles but with L0 files

db/multi_scan.cc

@@ -3,7 +3,9 @@
 //  COPYING file in the root directory) and Apache 2.0 License
 //  (found in the LICENSE.Apache file in the root directory).
 
+#include "file/file_util.h"
 #include "rocksdb/db.h"
+#include "rocksdb/file_system.h"
 
 namespace ROCKSDB_NAMESPACE {
 
@@ -12,30 +14,43 @@ using MultiScanIterator = MultiScan::MultiScanIterator;
 MultiScan::MultiScan(const ReadOptions& read_options,
                      const MultiScanArgs& scan_opts, DB* db,
                      ColumnFamilyHandle* cfh)
-    : read_options_(read_options), scan_opts_(scan_opts), db_(db), cfh_(cfh) {
+    : read_options_(read_options),
+      scan_opts_([&] {
+        // Disable async IO if the filesystem does not support it, consistent
+        // with how ArenaWrappedDBIter::Init() handles ReadOptions::async_io.
+        auto opts = scan_opts;
+        if (opts.use_async_io &&
+            !CheckFSFeatureSupport(db->GetFileSystem(),
+                                   FSSupportedOps::kAsyncIO)) {
+          opts.use_async_io = false;
+        }
+        return opts;
+      }()),
+      db_(db),
+      cfh_(cfh) {
   bool slow_path = false;
   // Setup read_options with iterate_uuper_bound based on the first scan.
   // Subsequent scans will update and allocate a new DB iterator as necessary
-  if (scan_opts.GetScanRanges()[0].range.limit) {
-    upper_bound_ = *scan_opts.GetScanRanges()[0].range.limit;
+  if (scan_opts_.GetScanRanges()[0].range.limit) {
+    upper_bound_ = *scan_opts_.GetScanRanges()[0].range.limit;
     read_options_.iterate_upper_bound = &upper_bound_;
   } else {
     read_options_.iterate_upper_bound = nullptr;
   }
-  for (const auto& opts : scan_opts.GetScanRanges()) {
+  for (const auto& opts : scan_opts_.GetScanRanges()) {
     // Check that all the ScanOptions either specify an upper bound or not. If
     // its mixed we take the slow path which avoids calling Prepare: we have to
     // reallocate the Iterator with updated read_options everytime we switch
     // between upper bound or no upper bound, which complicates Prepare.
     if (opts.range.limit.has_value() !=
-        scan_opts.GetScanRanges()[0].range.limit.has_value()) {
+        scan_opts_.GetScanRanges()[0].range.limit.has_value()) {
       slow_path = true;
       break;
     }
   }
   db_iter_.reset(db->NewIterator(read_options_, cfh));
   if (!slow_path) {
-    db_iter_->Prepare(scan_opts);
+    db_iter_->Prepare(scan_opts_);
   }
 }