Move test assertion to DBStressFSWrapper

db_stress_tool/db_stress_env_wrapper.h

@@ -10,15 +10,48 @@
 #ifdef GFLAGS
 #pragma once
 #include "db_stress_tool/db_stress_common.h"
+#include "test_util/thread_io_activity.h"
 
 namespace ROCKSDB_NAMESPACE {
+class DbStressRandomAccessFileWrapper : public FSRandomAccessFile {
+ public:
+  explicit DbStressRandomAccessFileWrapper(
+      std::unique_ptr<FSRandomAccessFile>&& target)
+      : target_(std::move(target)) {}
+
+  IOStatus Read(uint64_t offset, size_t n, const IOOptions& options,
+                Slice* result, char* scratch,
+                IODebugContext* dbg) const override {
+    const Env::IOActivity io_activity = TEST_GetThreadIOActivity();
+    if (io_activity != Env::IOActivity::kUnknown) {
+      assert(io_activity == options.io_activity);
+    }
+    return target_->Read(offset, n, options, result, scratch, dbg);
+  }
+
+ private:
+  std::unique_ptr<FSRandomAccessFile> target_;
+};
+
 class DbStressFSWrapper : public FileSystemWrapper {
  public:
   explicit DbStressFSWrapper(const std::shared_ptr<FileSystem>& t)
       : FileSystemWrapper(t) {}
   static const char* kClassName() { return "DbStressFS"; }
   const char* Name() const override { return kClassName(); }
 
+  IOStatus NewRandomAccessFile(const std::string& f,
+                               const FileOptions& file_opts,
+                               std::unique_ptr<FSRandomAccessFile>* r,
+                               IODebugContext* dbg) override {
+    std::unique_ptr<FSRandomAccessFile> file;
+    IOStatus s = target()->NewRandomAccessFile(f, file_opts, &file, dbg);
+    if (s.ok()) {
+      r->reset(new DbStressRandomAccessFileWrapper(std::move(file)));
+    }
+    return s;
+  }
+
   IOStatus DeleteFile(const std::string& f, const IOOptions& opts,
                       IODebugContext* dbg) override {
     // We determine whether it is a manifest file by searching a strong,

db_stress_tool/db_stress_tool.cc

@@ -31,6 +31,7 @@ namespace ROCKSDB_NAMESPACE {
 namespace {
 static std::shared_ptr<ROCKSDB_NAMESPACE::Env> env_guard;
 static std::shared_ptr<ROCKSDB_NAMESPACE::Env> env_wrapper_guard;
+static std::shared_ptr<ROCKSDB_NAMESPACE::Env> legacy_env_wrapper_guard;
 static std::shared_ptr<ROCKSDB_NAMESPACE::CompositeEnvWrapper>
     dbsl_env_wrapper_guard;
 static std::shared_ptr<CompositeEnvWrapper> fault_env_guard;
@@ -104,7 +105,11 @@ int db_stress_tool(int argc, char** argv) {
     // legacy EnvWrapper. This is a workaround to prevent MultiGet and scans
     // from failing when IO uring is disabled. The EnvWrapper
     // has a default implementation of ReadAsync that redirects to Read.
-    env_wrapper_guard = std::make_shared<EnvWrapper>(env_wrapper_guard);
+    legacy_env_wrapper_guard = std::make_shared<EnvWrapper>(raw_env);
+    env_wrapper_guard = std::make_shared<CompositeEnvWrapper>(
+        legacy_env_wrapper_guard,
+        std::make_shared<DbStressFSWrapper>(
+            legacy_env_wrapper_guard->GetFileSystem()));
   }
   db_stress_env = env_wrapper_guard.get();
 

utilities/fault_injection_fs.cc

@@ -24,7 +24,6 @@
 #include "port/lang.h"
 #include "port/stack_trace.h"
 #include "test_util/sync_point.h"
-#include "test_util/thread_io_activity.h"
 #include "util/coding.h"
 #include "util/crc32c.h"
 #include "util/mutexlock.h"
@@ -399,10 +398,6 @@ IOStatus TestFSRandomAccessFile::Read(uint64_t offset, size_t n,
                                       const IOOptions& options, Slice* result,
                                       char* scratch,
                                       IODebugContext* dbg) const {
-  const Env::IOActivity io_activity = TEST_GetThreadIOActivity();
-  if (io_activity != Env::IOActivity::kUnknown) {
-    assert(io_activity == options.io_activity);
-  }
   if (!fs_->IsFilesystemActive()) {
     return fs_->GetError();
   }
@@ -453,10 +448,6 @@ IOStatus TestFSRandomAccessFile::MultiRead(FSReadRequest* reqs, size_t num_reqs,
   if (!fs_->IsFilesystemActive()) {
     return fs_->GetError();
   }
-  const Env::IOActivity io_activity = TEST_GetThreadIOActivity();
-  if (io_activity != Env::IOActivity::kUnknown) {
-    assert(io_activity == options.io_activity);
-  }
   IOStatus s = target_->MultiRead(reqs, num_reqs, options, dbg);
   bool injected_error = false;
   for (size_t i = 0; i < num_reqs; i++) {