Allow taking db checkpoint for state snapshot creation (MystenLabs#8041)

This PR allows taking rocksdb checkpoint which allow for taking
efficient point in time snapshot of the db (created using hard links)

crates/typed-store/src/rocks/mod.rs

@@ -11,6 +11,7 @@ use crate::{
 };
 use bincode::Options;
 use collectable::TryExtend;
+use rocksdb::checkpoint::Checkpoint;
 use rocksdb::{
     properties, AsColumnFamilyRef, CStrLike, ColumnFamilyDescriptor, DBWithThreadMode, Error,
     ErrorKind, IteratorMode, MultiThreaded, OptimisticTransactionOptions, ReadOptions, Transaction,
@@ -377,6 +378,20 @@ impl RocksDB {
         delegate_call!(self.flush())
     }
 
+    pub fn checkpoint(&self, path: &Path) -> Result<(), rocksdb::Error> {
+        match self {
+            Self::DBWithThreadMode(d) => {
+                let checkpoint = Checkpoint::new(&d.underlying)?;
+                checkpoint.create_checkpoint(path)?;
+            }
+            Self::OptimisticTransactionDB(d) => {
+                let checkpoint = Checkpoint::new(&d.underlying)?;
+                checkpoint.create_checkpoint(path)?;
+            }
+        }
+        Ok(())
+    }
+
     pub fn flush_cf(&self, cf: &impl AsColumnFamilyRef) -> Result<(), rocksdb::Error> {
         delegate_call!(self.flush_cf(cf))
     }

crates/typed-store/src/rocks/tests.rs

@@ -573,6 +573,44 @@ async fn test_multi_insert(#[values(true, false)] is_transactional: bool) {
     }
 }
 
+#[rstest]
+#[tokio::test]
+async fn test_checkpoint(#[values(true, false)] is_transactional: bool) {
+    let path_prefix = temp_dir();
+    let db_path = path_prefix.join("db");
+    let db: DBMap<i32, String> = open_map(db_path, Some("table"), is_transactional);
+    // Create kv pairs
+    let keys_vals = (0..101).map(|i| (i, i.to_string()));
+
+    db.multi_insert(keys_vals.clone())
+        .expect("Failed to multi-insert");
+    let checkpointed_path = path_prefix.join("checkpointed_db");
+    db.rocksdb
+        .checkpoint(&checkpointed_path)
+        .expect("Failed to create db checkpoint");
+    // Create more kv pairs
+    let new_keys_vals = (101..201).map(|i| (i, i.to_string()));
+    db.multi_insert(new_keys_vals.clone())
+        .expect("Failed to multi-insert");
+    // Verify checkpoint
+    let checkpointed_db: DBMap<i32, String> =
+        open_map(checkpointed_path, Some("table"), is_transactional);
+    // Ensure keys inserted before checkpoint are present in original and checkpointed db
+    for (k, v) in keys_vals {
+        let val = db.get(&k).expect("Failed to get inserted key");
+        assert_eq!(Some(v.clone()), val);
+        let val = checkpointed_db.get(&k).expect("Failed to get inserted key");
+        assert_eq!(Some(v), val);
+    }
+    // Ensure keys inserted after checkpoint are only present in original db but not in checkpointed db
+    for (k, v) in new_keys_vals {
+        let val = db.get(&k).expect("Failed to get inserted key");
+        assert_eq!(Some(v.clone()), val);
+        let val = checkpointed_db.get(&k).expect("Failed to get inserted key");
+        assert_eq!(None, val);
+    }
+}
+
 #[rstest]
 #[tokio::test]
 async fn test_multi_remove(#[values(true, false)] is_transactional: bool) {