raftstore-v2: support renaming encrypted dir (inefficiently) and batc…

…h importing data keys (tikv#14583) ref tikv#12842, ref tikv#14095, ref tikv#14097 support renaming encrypted dir (inefficiently) and batch importing data keys Signed-off-by: tabokie <[email protected]> Signed-off-by: lidezhu <[email protected]>
lidezhu · Apr 27, 2023 · 1eaea68 · 1eaea68
1 parent 28dfff5
commit 1eaea68
Show file tree

Hide file tree

Showing 34 changed files with 828 additions and 194 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/components/encryption/Cargo.toml b/components/encryption/Cargo.toml
@@ -37,6 +37,7 @@ thiserror = "1.0"
 tikv_alloc = { workspace = true }
 tikv_util = { workspace = true }
 tokio = { version = "1.5", features = ["time", "rt"] }
+walkdir = "2"
 
 [dev-dependencies]
 matches = "0.1.8"

diff --git a/components/encryption/export/src/lib.rs b/components/encryption/export/src/lib.rs
@@ -14,8 +14,9 @@ use derive_more::Deref;
 #[cfg(feature = "cloud-aws")]
 pub use encryption::KmsBackend;
 pub use encryption::{
-    from_engine_encryption_method, Backend, DataKeyManager, DataKeyManagerArgs, DecrypterReader,
-    EncryptionConfig, Error, FileConfig, Iv, KmsConfig, MasterKeyConfig, Result,
+    clean_up_dir, clean_up_trash, from_engine_encryption_method, trash_dir_all, Backend,
+    DataKeyManager, DataKeyManagerArgs, DecrypterReader, EncryptionConfig, Error, FileConfig, Iv,
+    KmsConfig, MasterKeyConfig, Result,
 };
 use encryption::{
     DataKeyPair, EncryptedKey, FileBackend, KmsProvider, PlainKey, PlaintextBackend,

diff --git a/components/encryption/src/file_dict_file.rs b/components/encryption/src/file_dict_file.rs
@@ -217,10 +217,11 @@ impl FileDictionaryFile {
         Ok(file_dict)
     }
 
-    /// Append an insert operation to the log file.
+    /// Append an insert operation to the log file. The record is guaranteed to
+    /// be persisted if `sync` is set.
     ///
     /// Warning: `self.write(file_dict)` must be called before.
-    pub fn insert(&mut self, name: &str, info: &FileInfo) -> Result<()> {
+    pub fn insert(&mut self, name: &str, info: &FileInfo, sync: bool) -> Result<()> {
         self.file_dict.files.insert(name.to_owned(), info.clone());
         if self.enable_log {
             let file = self.append_file.as_mut().unwrap();
@@ -231,12 +232,16 @@ impl FileDictionaryFile {
                 let truncate_num: usize = truncate_num.map_or(0, |c| c.parse().unwrap());
                 bytes.truncate(truncate_num);
                 file.write_all(&bytes)?;
-                file.sync_all()?;
+                if sync {
+                    file.sync_all()?;
+                }
                 Ok(())
             });
 
             file.write_all(&bytes)?;
-            file.sync_all()?;
+            if sync {
+                file.sync_all()?;
+            }
 
             self.file_size += bytes.len();
             self.check_compact()?;
@@ -250,13 +255,15 @@ impl FileDictionaryFile {
     /// Append a remove operation to the log file.
     ///
     /// Warning: `self.write(file_dict)` must be called before.
-    pub fn remove(&mut self, name: &str) -> Result<()> {
+    pub fn remove(&mut self, name: &str, sync: bool) -> Result<()> {
         self.file_dict.files.remove(name);
         if self.enable_log {
             let file = self.append_file.as_mut().unwrap();
             let bytes = Self::convert_record_to_bytes(name, LogRecord::Remove)?;
             file.write_all(&bytes)?;
-            file.sync_all()?;
+            if sync {
+                file.sync_all()?;
+            }
 
             self.removed += 1;
             self.file_size += bytes.len();
@@ -268,6 +275,13 @@ impl FileDictionaryFile {
         Ok(())
     }
 
+    pub fn sync(&mut self) -> Result<()> {
+        if self.enable_log {
+            self.append_file.as_mut().unwrap().sync_all()?;
+        }
+        Ok(())
+    }
+
     /// This function needs to be called after each append operation to check
     /// if compact is needed.
     fn check_compact(&mut self) -> Result<()> {
@@ -407,9 +421,9 @@ mod tests {
         let info4 = create_file_info(4, EncryptionMethod::Aes128Ctr);
         let info5 = create_file_info(3, EncryptionMethod::Aes128Ctr);
 
-        file_dict_file.insert("info1", &info1).unwrap();
-        file_dict_file.insert("info2", &info2).unwrap();
-        file_dict_file.insert("info3", &info3).unwrap();
+        file_dict_file.insert("info1", &info1, true).unwrap();
+        file_dict_file.insert("info2", &info2, true).unwrap();
+        file_dict_file.insert("info3", &info3, true).unwrap();
 
         let file_dict = file_dict_file.recovery().unwrap();
 
@@ -418,18 +432,18 @@ mod tests {
         assert_eq!(*file_dict.files.get("info3").unwrap(), info3);
         assert_eq!(file_dict.files.len(), 3);
 
-        file_dict_file.remove("info2").unwrap();
-        file_dict_file.remove("info1").unwrap();
-        file_dict_file.insert("info2", &info4).unwrap();
+        file_dict_file.remove("info2", true).unwrap();
+        file_dict_file.remove("info1", true).unwrap();
+        file_dict_file.insert("info2", &info4, true).unwrap();
 
         let file_dict = file_dict_file.recovery().unwrap();
         assert_eq!(file_dict.files.get("info1"), None);
         assert_eq!(*file_dict.files.get("info2").unwrap(), info4);
         assert_eq!(*file_dict.files.get("info3").unwrap(), info3);
         assert_eq!(file_dict.files.len(), 2);
 
-        file_dict_file.insert("info5", &info5).unwrap();
-        file_dict_file.remove("info3").unwrap();
+        file_dict_file.insert("info5", &info5, true).unwrap();
+        file_dict_file.remove("info3", true).unwrap();
 
         let file_dict = file_dict_file.recovery().unwrap();
         assert_eq!(file_dict.files.get("info1"), None);
@@ -460,7 +474,7 @@ mod tests {
         .unwrap();
 
         let info = create_file_info(1, EncryptionMethod::Aes256Ctr);
-        file_dict_file.insert("info", &info).unwrap();
+        file_dict_file.insert("info", &info, true).unwrap();
 
         let (_, file_dict) = FileDictionaryFile::open(
             tempdir.path(),
@@ -550,14 +564,14 @@ mod tests {
             )
             .unwrap();
 
-            file_dict.insert("f1", &info1).unwrap();
-            file_dict.insert("f2", &info2).unwrap();
-            file_dict.insert("f3", &info3).unwrap();
+            file_dict.insert("f1", &info1, true).unwrap();
+            file_dict.insert("f2", &info2, true).unwrap();
+            file_dict.insert("f3", &info3, true).unwrap();
 
-            file_dict.insert("f4", &info4).unwrap();
-            file_dict.remove("f3").unwrap();
+            file_dict.insert("f4", &info4, true).unwrap();
+            file_dict.remove("f3", true).unwrap();
 
-            file_dict.remove("f2").unwrap();
+            file_dict.remove("f2", true).unwrap();
         }
         // Try open as v1 file. Should fail.
         {

diff --git a/components/encryption/src/lib.rs b/components/encryption/src/lib.rs
@@ -10,6 +10,8 @@ mod manager;
 mod master_key;
 mod metrics;
 
+use std::{io::ErrorKind, path::Path};
+
 pub use self::{
     config::*,
     crypter::{
@@ -27,3 +29,120 @@ pub use self::{
         Backend, DataKeyPair, EncryptedKey, FileBackend, KmsBackend, KmsProvider, PlaintextBackend,
     },
 };
+
+const TRASH_PREFIX: &str = "TRASH-";
+
+/// Remove a directory.
+///
+/// Rename it before actually removal.
+#[inline]
+pub fn trash_dir_all(
+    path: impl AsRef<Path>,
+    key_manager: Option<&DataKeyManager>,
+) -> std::io::Result<()> {
+    let path = path.as_ref();
+    let name = match path.file_name() {
+        Some(n) => n,
+        None => {
+            return Err(std::io::Error::new(
+                ErrorKind::InvalidInput,
+                "path is invalid",
+            ));
+        }
+    };
+    let trash_path = path.with_file_name(format!("{}{}", TRASH_PREFIX, name.to_string_lossy()));
+    if let Err(e) = file_system::rename(path, &trash_path) {
+        if e.kind() == ErrorKind::NotFound {
+            return Ok(());
+        }
+        return Err(e);
+    } else if let Some(m) = key_manager {
+        m.remove_dir(path, Some(&trash_path))?;
+    }
+    file_system::remove_dir_all(trash_path)
+}
+
+/// When using `trash_dir_all`, it's possible the directory is marked as trash
+/// but not being actually deleted after a restart. This function can be used
+/// to resume all those removal in the given directory.
+#[inline]
+pub fn clean_up_trash(
+    path: impl AsRef<Path>,
+    key_manager: Option<&DataKeyManager>,
+) -> std::io::Result<()> {
+    for e in file_system::read_dir(path)? {
+        let e = e?;
+        let os_fname = e.file_name();
+        let fname = os_fname.to_str().unwrap();
+        if let Some(original) = fname.strip_prefix(TRASH_PREFIX) {
+            let original = e.path().with_file_name(original);
+            if let Some(m) = &key_manager {
+                m.remove_dir(&original, Some(&e.path()))?;
+            }
+            file_system::remove_dir_all(e.path())?;
+        }
+    }
+    Ok(())
+}
+
+/// Removes all directories with the given prefix.
+#[inline]
+pub fn clean_up_dir(
+    path: impl AsRef<Path>,
+    prefix: &str,
+    key_manager: Option<&DataKeyManager>,
+) -> std::io::Result<()> {
+    for e in file_system::read_dir(path)? {
+        let e = e?;
+        let fname = e.file_name().to_str().unwrap().to_owned();
+        if fname.starts_with(prefix) {
+            if let Some(m) = &key_manager {
+                m.remove_dir(&e.path(), None)?;
+            }
+            file_system::remove_dir_all(e.path())?;
+        }
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use tempfile::Builder;
+
+    use super::*;
+
+    #[test]
+    fn test_trash_dir_all() {
+        let tmp_dir = Builder::new()
+            .prefix("test_reserve_space_for_recover")
+            .tempdir()
+            .unwrap();
+        let data_path = tmp_dir.path();
+        let sub_dir0 = data_path.join("sub_dir0");
+        let trash_sub_dir0 = data_path.join(format!("{}sub_dir0", TRASH_PREFIX));
+        file_system::create_dir_all(&sub_dir0).unwrap();
+        assert!(sub_dir0.exists());
+
+        trash_dir_all(&sub_dir0, None).unwrap();
+        assert!(!sub_dir0.exists());
+        assert!(!trash_sub_dir0.exists());
+
+        file_system::create_dir_all(&sub_dir0).unwrap();
+        file_system::create_dir_all(&trash_sub_dir0).unwrap();
+        trash_dir_all(&sub_dir0, None).unwrap();
+        assert!(!sub_dir0.exists());
+        assert!(!trash_sub_dir0.exists());
+
+        clean_up_trash(data_path, None).unwrap();
+
+        file_system::create_dir_all(&trash_sub_dir0).unwrap();
+        assert!(trash_sub_dir0.exists());
+        clean_up_trash(data_path, None).unwrap();
+        assert!(!trash_sub_dir0.exists());
+
+        file_system::create_dir_all(&sub_dir0).unwrap();
+        assert!(sub_dir0.exists());
+        clean_up_dir(data_path, "sub", None).unwrap();
+        assert!(!sub_dir0.exists());
+    }
+}