Add support for TiKV IO rate limiter (#383)

 

Signed-off-by: Yang Zhang <[email protected]>

Co-authored-by: Xinye Tao <[email protected]>

CMakeLists.txt

@@ -725,6 +725,7 @@ set(SOURCES
         env/env.cc
         env/env_chroot.cc
         env/env_encryption.cc
+        env/env_inspected.cc
         env/file_system.cc
         env/file_system_tracer.cc
         env/fs_remap.cc

TARGETS

@@ -112,6 +112,7 @@ cpp_library_wrapper(name="rocksdb_lib", srcs=[
         "env/env.cc",
         "env/env_chroot.cc",
         "env/env_encryption.cc",
+        "env/env_inspected.cc",
         "env/env_posix.cc",
         "env/file_system.cc",
         "env/file_system_tracer.cc",

env/env_basic_test.cc

@@ -14,6 +14,7 @@
 #include "rocksdb/convenience.h"
 #include "rocksdb/env.h"
 #include "rocksdb/env_encryption.h"
+#include "rocksdb/env_inspected.h"
 #include "test_util/testharness.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -81,6 +82,41 @@ static Env* GetTestFS() {
   return fs_env;
 }
 
+class DummyFileSystemInspector : public FileSystemInspector {
+ public:
+  DummyFileSystemInspector(size_t refill_bytes = 0)
+      : refill_bytes_(refill_bytes) {}
+
+  Status Read(size_t len, size_t* allowed) override {
+    assert(allowed);
+    if (refill_bytes_ == 0) {
+      *allowed = len;
+    } else {
+      *allowed = std::min(refill_bytes_, len);
+    }
+    return Status::OK();
+  }
+
+  Status Write(size_t len, size_t* allowed) override {
+    assert(allowed);
+    if (refill_bytes_ == 0) {
+      *allowed = len;
+    } else {
+      *allowed = std::min(refill_bytes_, len);
+    }
+    return Status::OK();
+  }
+
+ private:
+  size_t refill_bytes_;
+};
+
+static Env* GetInspectedEnv() {
+  static std::unique_ptr<Env> inspected_env(NewFileSystemInspectedEnv(
+      Env::Default(), std::make_shared<DummyFileSystemInspector>(1)));
+  return inspected_env.get();
+}
+
 }  // namespace
 class EnvBasicTestWithParam
     : public testing::Test,
@@ -118,6 +154,9 @@ INSTANTIATE_TEST_CASE_P(EncryptedEnv, EnvMoreTestWithParam,
 INSTANTIATE_TEST_CASE_P(MemEnv, EnvBasicTestWithParam,
                         ::testing::Values(&GetMemoryEnv));
 
+INSTANTIATE_TEST_CASE_P(InspectedEnv, EnvBasicTestWithParam,
+                        ::testing::Values(&GetInspectedEnv));
+
 namespace {
 
 // Returns a vector of 0 or 1 Env*, depending whether an Env is registered for