From 2ecac9f96d677099ed9054ff3921ad94161e0b4f Mon Sep 17 00:00:00 2001
From: Mike Kolupaev <kolmike@fb.com>
Date: Mon, 1 Jun 2015 11:22:36 -0700
Subject: [PATCH] add rocksdb::WritableFileWrapper similar to
 rocksdb::EnvWrapper

Summary: It used to be no good (known to me) non-intrusive way to wrap WritableFile - you can't call protected virtual methods of the wrapped pointer to WritableFile. This diff adds a convenience class WritableFileWrapper that makes wrapping WritableFile both possible and easy.

Test Plan: `make clean; make -j release`, `make clean; OPT=-DROCKSDB_LITE make release`, `make clean; USE_CLANG=1 make -j all`.

Reviewers: sdong, yhchiang, rven

Reviewed By: rven

Subscribers: dhruba, tnovak, march

Differential Revision: https://reviews.facebook.net/D39147
---
 include/rocksdb/env.h | 43 +++++++++++++++++++++++
 util/env_test.cc      | 81 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 124 insertions(+)

diff --git a/include/rocksdb/env.h b/include/rocksdb/env.h
index 2fb924214..ceb598f80 100644
--- a/include/rocksdb/env.h
+++ b/include/rocksdb/env.h
@@ -543,6 +543,8 @@ class WritableFile {
   void operator=(const WritableFile&);
 
  protected:
+  friend class WritableFileWrapper;
+
   Env::IOPriority io_priority_;
 };
 
@@ -878,6 +880,47 @@ class EnvWrapper : public Env {
   Env* target_;
 };
 
+// An implementation of WritableFile that forwards all calls to another
+// WritableFile. May be useful to clients who wish to override just part of the
+// functionality of another WritableFile.
+// It's declared as friend of WritableFile to allow forwarding calls to
+// protected virtual methods.
+class WritableFileWrapper : public WritableFile {
+ public:
+  explicit WritableFileWrapper(WritableFile* t) : target_(t) { }
+
+  Status Append(const Slice& data) override { return target_->Append(data); }
+  Status Close() override { return target_->Close(); }
+  Status Flush() override { return target_->Flush(); }
+  Status Sync() override { return target_->Sync(); }
+  Status Fsync() override { return target_->Fsync(); }
+  void SetIOPriority(Env::IOPriority pri) override {
+    target_->SetIOPriority(pri);
+  }
+  uint64_t GetFileSize() override { return target_->GetFileSize(); }
+  void GetPreallocationStatus(size_t* block_size,
+                              size_t* last_allocated_block) override {
+    target_->GetPreallocationStatus(block_size, last_allocated_block);
+  }
+  size_t GetUniqueId(char* id, size_t max_size) const override {
+    return target_->GetUniqueId(id, max_size);
+  }
+  Status InvalidateCache(size_t offset, size_t length) override {
+    return target_->InvalidateCache(offset, length);
+  }
+
+ protected:
+  Status Allocate(off_t offset, off_t len) override {
+    return target_->Allocate(offset, len);
+  }
+  Status RangeSync(off_t offset, off_t nbytes) override {
+    return target_->RangeSync(offset, nbytes);
+  }
+
+ private:
+  WritableFile* target_;
+};
+
 // Returns a new environment that stores its data in memory and delegates
 // all non-file-storage tasks to base_env. The caller must delete the result
 // when it is no longer needed.
diff --git a/util/env_test.cc b/util/env_test.cc
index 081a10f59..552175191 100644
--- a/util/env_test.cc
+++ b/util/env_test.cc
@@ -982,6 +982,87 @@ TEST_F(EnvPosixTest, Preallocation) {
   ASSERT_EQ(last_allocated_block, 7UL);
 }
 
+// Test that all WritableFileWrapper forwards all calls to WritableFile.
+TEST_F(EnvPosixTest, WritableFileWrapper) {
+  class Base : public WritableFile {
+   public:
+    mutable int *step_;
+
+    void inc(int x) const {
+      EXPECT_EQ(x, (*step_)++);
+    }
+
+    explicit Base(int* step) : step_(step) {
+      inc(0);
+    }
+
+    Status Append(const Slice& data) override { inc(1); return Status::OK(); }
+    Status Close() override { inc(2); return Status::OK(); }
+    Status Flush() override { inc(3); return Status::OK(); }
+    Status Sync() override { inc(4); return Status::OK(); }
+    Status Fsync() override { inc(5); return Status::OK(); }
+    void SetIOPriority(Env::IOPriority pri) override { inc(6); }
+    uint64_t GetFileSize() override { inc(7); return 0; }
+    void GetPreallocationStatus(size_t* block_size,
+                                size_t* last_allocated_block) override {
+      inc(8);
+    }
+    size_t GetUniqueId(char* id, size_t max_size) const override {
+      inc(9);
+      return 0;
+    }
+    Status InvalidateCache(size_t offset, size_t length) override {
+      inc(10);
+      return Status::OK();
+    }
+
+   protected:
+    Status Allocate(off_t offset, off_t len) override {
+      inc(11);
+      return Status::OK();
+    }
+    Status RangeSync(off_t offset, off_t nbytes) override {
+      inc(12);
+      return Status::OK();
+    }
+
+   public:
+    ~Base() {
+      inc(13);
+    }
+  };
+
+  class Wrapper : public WritableFileWrapper {
+   public:
+    explicit Wrapper(WritableFile* target) : WritableFileWrapper(target) {}
+
+    void CallProtectedMethods() {
+      Allocate(0, 0);
+      RangeSync(0, 0);
+    }
+  };
+
+  int step = 0;
+
+  {
+    Base b(&step);
+    Wrapper w(&b);
+    w.Append(Slice());
+    w.Close();
+    w.Flush();
+    w.Sync();
+    w.Fsync();
+    w.SetIOPriority(Env::IOPriority::IO_HIGH);
+    w.GetFileSize();
+    w.GetPreallocationStatus(nullptr, nullptr);
+    w.GetUniqueId(nullptr, 0);
+    w.InvalidateCache(0, 0);
+    w.CallProtectedMethods();
+  }
+
+  EXPECT_EQ(14, step);
+}
+
 }  // namespace rocksdb
 
 int main(int argc, char** argv) {
-- 
GitLab