Further cleanup, docs.

Signed-off-by: Alex Feinberg <alex@strlen.net>

components/hybrid_engine/Cargo.toml

@@ -14,6 +14,7 @@ txn_types = { workspace = true }
 tikv_util = { workspace = true }
 engine_rocks = { workspace = true }
 region_cache_memory_engine = { workspace = true }
+tempfile = "3.0"
 
 [dev-dependencies]
 tempfile = "3.0"

components/hybrid_engine/src/lib.rs

@@ -21,6 +21,7 @@ mod snapshot;
 mod sst;
 mod table_properties;
 mod ttl_properties;
+pub mod util;
 mod write_batch;
 
 pub use engine::HybridEngine;

components/hybrid_engine/src/snapshot.rs

@@ -75,6 +75,7 @@ where
     }
 }
 
+/// TODO: May be possible to replace this with an Either.
 pub struct HybridDbVector(Box<dyn DbVector>);
 
 impl DbVector for HybridDbVector {}

components/hybrid_engine/src/util.rs

@@ -0,0 +1,47 @@
+// Copyright 2024 TiKV Project Authors. Licensed under Apache-2.0.
+
+use std::sync::Arc;
+
+use engine_rocks::{util::new_engine, RocksEngine};
+use engine_traits::{Result, CF_DEFAULT, CF_LOCK, CF_WRITE};
+use region_cache_memory_engine::RangeCacheMemoryEngine;
+use tempfile::{Builder, TempDir};
+
+use crate::HybridEngine;
+
+/// Create a [`HybridEngine`] using temporary storage in `prefix`.
+/// Once the memory engine is created, runs `configure_memory_engine_fn`.
+/// Returns the handle to temporary directory and HybridEngine.
+///
+/// # Example
+///
+/// ```
+/// use hybrid_engine::util::hybrid_engine_for_tests;
+/// let (_path, _hybrid_engine) = hybrid_engine_for_tests("temp", |memory_engine| {
+///     let range = engine_traits::CacheRange::new(b"k00".to_vec(), b"k10".to_vec());
+///     memory_engine.new_range(range.clone());
+///     {
+///         let mut core = memory_engine.core().lock().unwrap();
+///         core.mut_range_manager().set_range_readable(&range, true);
+///         core.mut_range_manager().set_safe_ts(&range, 10);
+///     }
+/// })
+/// .unwrap();
+/// ```
+pub fn hybrid_engine_for_tests<F>(
+    prefix: &str,
+    configure_memory_engine_fn: F,
+) -> Result<(TempDir, HybridEngine<RocksEngine, RangeCacheMemoryEngine>)>
+where
+    F: FnOnce(&RangeCacheMemoryEngine),
+{
+    let path = Builder::new().prefix(prefix).tempdir()?;
+    let disk_engine = new_engine(
+        path.path().to_str().unwrap(),
+        &[CF_DEFAULT, CF_LOCK, CF_WRITE],
+    )?;
+    let memory_engine = RangeCacheMemoryEngine::new(Arc::default());
+    configure_memory_engine_fn(&memory_engine);
+    let hybrid_engine = HybridEngine::new(disk_engine, memory_engine);
+    Ok((path, hybrid_engine))
+}

components/hybrid_engine/src/write_batch.rs

@@ -124,36 +124,11 @@ impl<EK: KvEngine> Mutable for HybridEngineWriteBatch<EK> {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
-
-    use engine_rocks::{util::new_engine, RocksEngine};
     use engine_traits::{
         CacheRange, KvEngine, Mutable, Peekable, SnapshotContext, WriteBatch, WriteBatchExt,
-        CF_DEFAULT, CF_LOCK, CF_WRITE,
     };
-    use region_cache_memory_engine::RangeCacheMemoryEngine;
-    use tempfile::{Builder, TempDir};
-
-    use super::Result;
-    use crate::HybridEngine;
-
-    fn hybrid_engine_for_tests<F>(
-        prefix: &str,
-        configure_memory_engine_fn: F,
-    ) -> Result<(TempDir, HybridEngine<RocksEngine, RangeCacheMemoryEngine>)>
-    where
-        F: FnOnce(&RangeCacheMemoryEngine) -> (),
-    {
-        let path = Builder::new().prefix(prefix).tempdir()?;
-        let disk_engine = new_engine(
-            path.path().to_str().unwrap(),
-            &[CF_DEFAULT, CF_LOCK, CF_WRITE],
-        )?;
-        let memory_engine = RangeCacheMemoryEngine::new(Arc::default());
-        configure_memory_engine_fn(&memory_engine);
-        let hybrid_engine = HybridEngine::new(disk_engine, memory_engine);
-        Ok((path, hybrid_engine))
-    }
+
+    use crate::util::hybrid_engine_for_tests;
 
     #[test]
     fn test_write_to_both_engines() {