[TST] Implement utility functions to instantiate blockfile provider

rust/test/Cargo.toml

@@ -16,12 +16,15 @@ async-compression = { version = "0.4.12", features = [
 ] }
 
 bincode = { workspace = true }
+criterion = { workspace = true }
 futures = { workspace = true }
 rand = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 tantivy = { workspace = true }
+tempfile = { workspace = true }
 tokio = { workspace = true }
+uuid = { workspace = true }
 
 clap = { version = "4.5.17", features = ["derive"] }
 dirs = "5.0.1"
@@ -30,4 +33,7 @@ tokio-stream = { version = "0.1.16", features = ["full"] }
 tokio-util = "0.7.12"
 bloom = "0.3.2"
 
+chroma-blockstore = { workspace = true }
+chroma-cache = { workspace = true }
+chroma-storage = { workspace = true }
 chroma-types = { workspace = true }

rust/test/src/benchmark.rs

@@ -0,0 +1,41 @@
+use std::{fmt::Debug, future::Future};
+
+use criterion::{BenchmarkId, Criterion};
+use tokio::runtime::Runtime;
+
+pub fn tokio_multi_thread() -> Runtime {
+    tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .expect("Tokio runtime should be set up correctly.")
+}
+
+pub fn bench_group<'t, Arg, Fut>(
+    name: &'t str,
+    criterion: &'t mut Criterion,
+    runtime: &'t Runtime,
+    inputs: &'t [Arg],
+    routine: impl Fn(&'t Arg) -> Fut,
+) where
+    Arg: Clone + Debug,
+    Fut: Future<Output = ()>,
+{
+    let mut input_cycle = inputs.iter().cycle();
+    let mut group = criterion.benchmark_group(name);
+    group.throughput(criterion::Throughput::Elements(1));
+    group.bench_function(
+        BenchmarkId::from_parameter(format!("{:#?}", inputs)),
+        |bencher| {
+            bencher.to_async(runtime).iter_batched(
+                || {
+                    input_cycle
+                        .next()
+                        .expect("Cycled inputs should be endless.")
+                },
+                &routine,
+                criterion::BatchSize::SmallInput,
+            );
+        },
+    );
+    group.finish();
+}

rust/test/src/lib.rs

@@ -1 +1,4 @@
+pub mod benchmark;
+pub mod storage;
+
 pub mod datasets;

rust/test/src/storage.rs

@@ -0,0 +1,58 @@
+use chroma_blockstore::{
+    arrow::{config::TEST_MAX_BLOCK_SIZE_BYTES, provider::ArrowBlockfileProvider},
+    provider::BlockfileProvider,
+};
+use chroma_cache::{
+    cache::{Cache, Cacheable},
+    config::{CacheConfig, UnboundedCacheConfig},
+};
+use chroma_storage::{local::LocalStorage, Storage};
+use chroma_types::{Segment, SegmentScope, SegmentType};
+use std::{collections::HashMap, hash::Hash};
+use tempfile::TempDir;
+use uuid::Uuid;
+
+pub fn tmp_dir() -> TempDir {
+    TempDir::new().expect("Should be able to create a temporary directory.")
+}
+
+pub fn storage() -> Storage {
+    Storage::Local(LocalStorage::new(tmp_dir().into_path().to_str().expect(
+        "Should be able to convert temporary directory path to string",
+    )))
+}
+
+pub fn unbounded_cache<K, V>() -> Cache<K, V>
+where
+    K: Send + Sync + Clone + Hash + Eq + 'static,
+    V: Send + Sync + Clone + Cacheable + 'static,
+{
+    Cache::new(&CacheConfig::Unbounded(UnboundedCacheConfig {}))
+}
+
+pub fn arrow_blockfile_provider() -> BlockfileProvider {
+    BlockfileProvider::ArrowBlockfileProvider(ArrowBlockfileProvider::new(
+        storage(),
+        TEST_MAX_BLOCK_SIZE_BYTES,
+        unbounded_cache(),
+        unbounded_cache(),
+    ))
+}
+
+pub fn segment(scope: SegmentScope) -> Segment {
+    use SegmentScope::*;
+    use SegmentType::*;
+    let r#type = match scope {
+        METADATA => BlockfileMetadata,
+        RECORD => BlockfileRecord,
+        SQLITE | VECTOR => panic!("Unsupported segment scope in testing."),
+    };
+    Segment {
+        id: Uuid::new_v4(),
+        r#type,
+        scope,
+        collection: Uuid::new_v4(),
+        metadata: None,
+        file_path: HashMap::new(),
+    }
+}