From b8b4cf31617fcb6be4d1acafb4b3e5a9533d463e Mon Sep 17 00:00:00 2001
From: Congyu <52687642+Congyuwang@users.noreply.github.com>
Date: Thu, 9 Feb 2023 20:27:24 +0800
Subject: [PATCH] Add Options::load_latest() method to load latest options from
 RockDB path (#724)

---
 src/column_family.rs          |  5 +++
 src/db_options.rs             | 76 ++++++++++++++++++++++++++++++++++-
 tests/test_rocksdb_options.rs | 25 +++++++++++-
 3 files changed, 103 insertions(+), 3 deletions(-)

diff --git a/src/column_family.rs b/src/column_family.rs
index 4cd5d30..2e0c607 100644
--- a/src/column_family.rs
+++ b/src/column_family.rs
@@ -41,6 +41,11 @@ impl ColumnFamilyDescriptor {
             options,
         }
     }
+
+    /// Get the name of the ColumnFamilyDescriptor.
+    pub fn name(&self) -> &str {
+        &self.name
+    }
 }
 
 /// An opaque type used to represent a column family. Returned from some functions, and used
diff --git a/src/db_options.rs b/src/db_options.rs
index 177bd2e..d4ee073 100644
--- a/src/db_options.rs
+++ b/src/db_options.rs
@@ -14,6 +14,8 @@
 
 use std::ffi::CStr;
 use std::path::Path;
+use std::ptr::null_mut;
+use std::slice;
 use std::sync::Arc;
 
 use libc::{self, c_char, c_double, c_int, c_uchar, c_uint, c_void, size_t};
@@ -24,12 +26,12 @@ use crate::{
     comparator::{self, ComparatorCallback, CompareFn},
     db::DBAccess,
     ffi,
-    ffi_util::{to_cpath, CStrLike},
+    ffi_util::{from_cstr, to_cpath, CStrLike},
     merge_operator::{
         self, full_merge_callback, partial_merge_callback, MergeFn, MergeOperatorCallback,
     },
     slice_transform::SliceTransform,
-    Error, SnapshotWithThreadMode,
+    ColumnFamilyDescriptor, Error, SnapshotWithThreadMode,
 };
 
 fn new_cache(capacity: size_t) -> *mut ffi::rocksdb_cache_t {
@@ -875,6 +877,76 @@ pub enum LogLevel {
 }
 
 impl Options {
+    /// Constructs the DBOptions and ColumnFamilyDescriptors by loading the
+    /// latest RocksDB options file stored in the specified rocksdb database.
+    pub fn load_latest<P: AsRef<Path>>(
+        path: P,
+        env: Env,
+        ignore_unknown_options: bool,
+        cache: Cache,
+    ) -> Result<(Options, Vec<ColumnFamilyDescriptor>), Error> {
+        let path = to_cpath(path)?;
+        let mut db_options: *mut ffi::rocksdb_options_t = null_mut();
+        let mut num_column_families: usize = 0;
+        let mut column_family_names: *mut *mut c_char = null_mut();
+        let mut column_family_options: *mut *mut ffi::rocksdb_options_t = null_mut();
+        unsafe {
+            ffi_try!(ffi::rocksdb_load_latest_options(
+                path.as_ptr(),
+                env.0.inner,
+                ignore_unknown_options,
+                cache.0.inner,
+                &mut db_options,
+                &mut num_column_families,
+                &mut column_family_names,
+                &mut column_family_options,
+            ));
+        }
+        let options = Options {
+            inner: db_options,
+            outlive: OptionsMustOutliveDB::default(),
+        };
+        let column_families = unsafe {
+            Options::read_column_descriptors(
+                num_column_families,
+                column_family_names,
+                column_family_options,
+            )
+        };
+        Ok((options, column_families))
+    }
+
+    /// read column descriptors from c pointers
+    #[inline]
+    unsafe fn read_column_descriptors(
+        num_column_families: usize,
+        column_family_names: *mut *mut c_char,
+        column_family_options: *mut *mut ffi::rocksdb_options_t,
+    ) -> Vec<ColumnFamilyDescriptor> {
+        let column_family_names_iter =
+            slice::from_raw_parts(column_family_names, num_column_families)
+                .iter()
+                .map(|ptr| from_cstr(*ptr));
+        let column_family_options_iter =
+            slice::from_raw_parts(column_family_options, num_column_families)
+                .iter()
+                .map(|ptr| Options {
+                    inner: *ptr,
+                    outlive: OptionsMustOutliveDB::default(),
+                });
+        let column_descriptors = column_family_names_iter
+            .zip(column_family_options_iter)
+            .map(|(name, options)| ColumnFamilyDescriptor { name, options })
+            .collect::<Vec<_>>();
+        // free pointers
+        slice::from_raw_parts(column_family_names, num_column_families)
+            .iter()
+            .for_each(|ptr| ffi::rocksdb_free(*ptr as *mut c_void));
+        ffi::rocksdb_free(column_family_names as *mut c_void);
+        ffi::rocksdb_free(column_family_options as *mut c_void);
+        column_descriptors
+    }
+
     /// By default, RocksDB uses only one background thread for flush and
     /// compaction. Calling this function will set it up such that total of
     /// `total_threads` is used. Good value for `total_threads` is the number of
diff --git a/tests/test_rocksdb_options.rs b/tests/test_rocksdb_options.rs
index 3bc27cd..e992d8b 100644
--- a/tests/test_rocksdb_options.rs
+++ b/tests/test_rocksdb_options.rs
@@ -16,9 +16,32 @@ mod util;
 
 use std::{fs, io::Read as _};
 
-use rocksdb::{BlockBasedOptions, DBCompressionType, DataBlockIndexType, Options, ReadOptions, DB};
+use rocksdb::{
+    BlockBasedOptions, Cache, DBCompressionType, DataBlockIndexType, Env, Options, ReadOptions, DB,
+};
 use util::DBPath;
 
+#[test]
+fn test_load_latest() {
+    let n = DBPath::new("_rust_rocksdb_test_load_latest");
+    {
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+        opts.create_missing_column_families(true);
+        let _ = DB::open_cf(&opts, &n, vec!["cf0", "cf1"]).unwrap();
+    }
+    let (_, cfs) = Options::load_latest(
+        &n,
+        Env::new().unwrap(),
+        true,
+        Cache::new_lru_cache(1024 * 8).unwrap(),
+    )
+    .unwrap();
+    assert!(cfs.iter().any(|cf| cf.name() == "default"));
+    assert!(cfs.iter().any(|cf| cf.name() == "cf0"));
+    assert!(cfs.iter().any(|cf| cf.name() == "cf1"));
+}
+
 #[test]
 fn test_set_num_levels() {
     let n = DBPath::new("_rust_rocksdb_test_set_num_levels");