From 568f31274b9e20ba3e82a87a6656bc3f12162d0c Mon Sep 17 00:00:00 2001
From: Xintao <hunterlxt@live.com>
Date: Fri, 18 Dec 2020 15:16:59 +0800
Subject: [PATCH] Atomize RenameFile in KeyManagedEncryptedEnv (#222)

* adjust logic in KeyManagedEncryptedEnv::RenameFile to avoid poweroff

Signed-off-by: Xintao <hunterlxt@live.com>
---
 db/db_test_util.h                  |  3 ---
 encryption/encryption.cc           | 27 +++++++++++++++++++--------
 encryption/in_memory_key_manager.h | 11 -----------
 include/rocksdb/encryption.h       |  2 --
 4 files changed, 19 insertions(+), 24 deletions(-)

diff --git a/db/db_test_util.h b/db/db_test_util.h
index 0a67213563b..167c61dbe5e 100644
--- a/db/db_test_util.h
+++ b/db/db_test_util.h
@@ -85,9 +85,6 @@ class TestKeyManager : public encryption::KeyManager {
   Status LinkFile(const std::string&, const std::string&) override {
     return Status::OK();
   }
-  Status RenameFile(const std::string&, const std::string&) override {
-    return Status::OK();
-  }
 };
 #endif
 
diff --git a/encryption/encryption.cc b/encryption/encryption.cc
index 426958ea9e6..66e1fd775ce 100644
--- a/encryption/encryption.cc
+++ b/encryption/encryption.cc
@@ -382,9 +382,14 @@ Status KeyManagedEncryptedEnv::ReuseWritableFile(
           "Unsupported encryption method: " +
           std::to_string(static_cast<int>(file_info.method)));
   }
-  if (s.ok()) {
-    s = key_manager_->RenameFile(old_fname, fname);
+  if (!s.ok()) {
+    return s;
+  }
+  s = key_manager_->LinkFile(old_fname, fname);
+  if (!s.ok()) {
+    return s;
   }
+  s = key_manager_->DeleteFile(old_fname);
   return s;
 }
 
@@ -437,22 +442,28 @@ Status KeyManagedEncryptedEnv::LinkFile(const std::string& src_fname,
   }
   s = target()->LinkFile(src_fname, dst_fname);
   if (!s.ok()) {
-    // Ignore error
-    key_manager_->DeleteFile(dst_fname);
+    Status delete_status __attribute__((__unused__)) =
+        key_manager_->DeleteFile(dst_fname);
+    assert(delete_status.ok());
   }
   return s;
 }
 
 Status KeyManagedEncryptedEnv::RenameFile(const std::string& src_fname,
                                           const std::string& dst_fname) {
-  Status s = key_manager_->RenameFile(src_fname, dst_fname);
+  // Link(copy)File instead of RenameFile to avoid losing src_fname info when
+  // failed to rename the src_fname in the file system.
+  Status s = key_manager_->LinkFile(src_fname, dst_fname);
   if (!s.ok()) {
     return s;
   }
   s = target()->RenameFile(src_fname, dst_fname);
-  if (!s.ok()) {
-    // Ignore error
-    key_manager_->RenameFile(dst_fname, src_fname);
+  if (s.ok()) {
+    s = key_manager_->DeleteFile(src_fname);
+  } else {
+    Status delete_status __attribute__((__unused__)) =
+        key_manager_->DeleteFile(dst_fname);
+    assert(delete_status.ok());
   }
   return s;
 }
diff --git a/encryption/in_memory_key_manager.h b/encryption/in_memory_key_manager.h
index 3155d3f2d1e..6f7e216c3c3 100644
--- a/encryption/in_memory_key_manager.h
+++ b/encryption/in_memory_key_manager.h
@@ -70,17 +70,6 @@ class InMemoryKeyManager final : public KeyManager {
     return Status::OK();
   }
 
-  Status RenameFile(const std::string& src_fname,
-                    const std::string& dst_fname) override {
-    MutexLock l(&mu_);
-    if (files_.count(src_fname) == 0) {
-      return Status::Corruption("File not found: " + src_fname);
-    }
-    files_[dst_fname] = files_[src_fname];
-    files_.erase(src_fname);
-    return Status::OK();
-  }
-
  private:
   mutable port::Mutex mu_;
   Random rnd_;
diff --git a/include/rocksdb/encryption.h b/include/rocksdb/encryption.h
index 253c9c6a504..547a3498299 100644
--- a/include/rocksdb/encryption.h
+++ b/include/rocksdb/encryption.h
@@ -55,8 +55,6 @@ class KeyManager {
   virtual Status DeleteFile(const std::string& fname) = 0;
   virtual Status LinkFile(const std::string& src_fname,
                           const std::string& dst_fname) = 0;
-  virtual Status RenameFile(const std::string& src_fname,
-                            const std::string& dst_fname) = 0;
 };
 
 // An Env with underlying files being encrypted. It holds a reference to an