fix renaming encrypted directory

encryption/encryption.cc

@@ -530,23 +530,22 @@ Status KeyManagedEncryptedEnv::RenameFile(const std::string& src_fname,
   }
   s = target()->RenameFile(src_fname, dst_fname);
   if (s.ok()) {
-    s = key_manager_->DeleteFile(src_fname);
+    s = key_manager_->DeleteFileExt(src_fname, dst_fname);
   } else {
     Status delete_status __attribute__((__unused__)) =
-        key_manager_->DeleteFile(dst_fname);
+        key_manager_->DeleteFileExt(dst_fname, src_fname);
     assert(delete_status.ok());
   }
   return s;
 }
 
 Status KeyManagedEncryptedEnv::DeleteDir(const std::string& dname) {
-  Status s = target()->DeleteDir(dname);
+  // We don't guarantee atomicity. Delete keys first.
+  Status s = key_manager_->DeleteFile(dname);
   if (!s.ok()) {
     return s;
   }
-  // We don't use a dedicated `DeleteDir` function, because RocksDB already uses
-  // `RenameFile` for both file and directory.
-  return key_manager_->DeleteFile(dname);
+  return target()->DeleteDir(dname);
 }
 
 Env* NewKeyManagedEncryptedEnv(Env* base_env,

include/rocksdb/encryption.h

@@ -55,9 +55,17 @@ class KeyManager {
                          FileEncryptionInfo* file_info) = 0;
   virtual Status NewFile(const std::string& fname,
                          FileEncryptionInfo* file_info) = 0;
+  // Used with both file and directory.
   virtual Status DeleteFile(const std::string& fname) = 0;
   virtual Status LinkFile(const std::string& src_fname,
                           const std::string& dst_fname) = 0;
+  // Provide additional hint of physical file when the key name doesn't map to
+  // one. A typical use case of this is atomically deleting a directory by
+  // renaming it first.
+  virtual Status DeleteFileExt(const std::string& fname,
+                               const std::string& /*physical_fname*/) {
+    return DeleteFile(fname);
+  }
 };
 
 // An Env with underlying files being encrypted. It holds a reference to an

test_util/testutil.h

@@ -86,6 +86,18 @@ class TestKeyManager : public encryption::KeyManager {
   Status DeleteFile(const std::string& fname) override {
     std::lock_guard<std::mutex> l(mutex);
     file_set.erase(fname);
+    if (!fname.empty()) {
+      std::string copy = fname;
+      if (copy.back() != '/') {
+        copy.push_back('/');
+      }
+      auto begin = file_set.lower_bound(copy);
+      auto end = begin;
+      while (end != file_set.end() && end->compare(0, copy.size(), copy) == 0) {
+        end++;
+      }
+      file_set.erase(begin, end);
+    }
     return Status::OK();
   }