From 1c91104f43dc4e97d0a67bd1a3cda8f8bed419b5 Mon Sep 17 00:00:00 2001
From: Aaron Gao <gzh@fb.com>
Date: Wed, 22 Feb 2017 10:00:25 -0800
Subject: [PATCH] posix writablefile truncate

Summary:
we occasionally missing this call so the file size will be wrong
Closes https://github.com/facebook/rocksdb/pull/1894

Differential Revision: D4598446

Pulled By: lightmark

fbshipit-source-id: 42b6ef5
---
 db/db_test2.cc    | 21 +++++++++++++++++++++
 db/db_test_util.h | 26 +++++++++++++++++++++++++-
 util/io_posix.cc  | 11 +++++++++++
 util/io_posix.h   |  6 +++---
 4 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/db/db_test2.cc b/db/db_test2.cc
index 6eb8b5951..44b7b85b9 100644
--- a/db/db_test2.cc
+++ b/db/db_test2.cc
@@ -2277,6 +2277,27 @@ TEST_F(DBTest2, GetRaceFlush2) {
   t1.join();
   rocksdb::SyncPoint::GetInstance()->DisableProcessing();
 }
+
+TEST_F(DBTest2, DirectIO) {
+  if (!IsDirectIOSupported()) {
+    return;
+  }
+  Options options = CurrentOptions();
+  options.use_direct_reads = options.use_direct_writes = true;
+  options.allow_mmap_reads = options.allow_mmap_writes = false;
+  DestroyAndReopen(options);
+
+  ASSERT_OK(Put(Key(0), "a"));
+  ASSERT_OK(Put(Key(5), "a"));
+  ASSERT_OK(Flush());
+
+  ASSERT_OK(Put(Key(10), "a"));
+  ASSERT_OK(Put(Key(15), "a"));
+  ASSERT_OK(Flush());
+
+  ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr));
+  Reopen(options);
+}
 }  // namespace rocksdb
 
 int main(int argc, char** argv) {
diff --git a/db/db_test_util.h b/db/db_test_util.h
index 6d1b872e5..550eaf326 100644
--- a/db/db_test_util.h
+++ b/db/db_test_util.h
@@ -231,6 +231,20 @@ class SpecialEnv : public EnvWrapper {
           return base_->Append(data);
         }
       }
+      Status PositionedAppend(const Slice& data, uint64_t offset) {
+        if (env_->table_write_callback_) {
+          (*env_->table_write_callback_)();
+        }
+        if (env_->drop_writes_.load(std::memory_order_acquire)) {
+          // Drop writes on the floor
+          return Status::OK();
+        } else if (env_->no_space_.load(std::memory_order_acquire)) {
+          return Status::NoSpace("No space left on device");
+        } else {
+          env_->bytes_written_ += data.size();
+          return base_->PositionedAppend(data, offset);
+        }
+      }
       Status Truncate(uint64_t size) override { return base_->Truncate(size); }
       Status Close() override {
 // SyncPoint is not supported in Released Windows Mode.
@@ -257,6 +271,9 @@ class SpecialEnv : public EnvWrapper {
       Env::IOPriority GetIOPriority() override {
         return base_->GetIOPriority();
       }
+      bool use_direct_io() const override {
+        return base_->use_direct_io();
+      }
     };
     class ManifestFile : public WritableFile {
      public:
@@ -358,7 +375,14 @@ class SpecialEnv : public EnvWrapper {
       return Status::IOError("simulated write error");
     }
 
-    Status s = target()->NewWritableFile(f, r, soptions);
+    EnvOptions optimized = soptions;
+    if (strstr(f.c_str(), "MANIFEST") != nullptr ||
+        strstr(f.c_str(), "log") != nullptr) {
+      optimized.use_mmap_writes = false;
+      optimized.use_direct_writes = false;
+    }
+
+    Status s = target()->NewWritableFile(f, r, optimized);
     if (s.ok()) {
       if (strstr(f.c_str(), ".sst") != nullptr) {
         r->reset(new SSTableFile(this, std::move(*r)));
diff --git a/util/io_posix.cc b/util/io_posix.cc
index 7f3bd5080..f46efc902 100644
--- a/util/io_posix.cc
+++ b/util/io_posix.cc
@@ -656,6 +656,17 @@ Status PosixWritableFile::PositionedAppend(const Slice& data, uint64_t offset) {
   return Status::OK();
 }
 
+Status PosixWritableFile::Truncate(uint64_t size) {
+  Status s;
+  int r = ftruncate(fd_, size);
+  if (r < 0) {
+    s = IOError(filename_, errno);
+  } else {
+    filesize_ = size;
+  }
+  return s;
+}
+
 Status PosixWritableFile::Close() {
   Status s;
 
diff --git a/util/io_posix.h b/util/io_posix.h
index 3bf155a82..40f558a19 100644
--- a/util/io_posix.h
+++ b/util/io_posix.h
@@ -93,9 +93,9 @@ class PosixWritableFile : public WritableFile {
                              const EnvOptions& options);
   virtual ~PosixWritableFile();
 
-  // Means Close() will properly take care of truncate
-  // and it does not need any additional information
-  virtual Status Truncate(uint64_t size) override { return Status::OK(); }
+  // Need to implement this so the file is truncated correctly
+  // with direct I/O
+  virtual Status Truncate(uint64_t size) override;
   virtual Status Close() override;
   virtual Status Append(const Slice& data) override;
   virtual Status PositionedAppend(const Slice& data, uint64_t offset) override;