fix(store-rust): remove duplicated smoke test and benchmark blocks

Agent-Logs-Url: https://github.com/kvcache-ai/Mooncake/sessions/d2eabc20-6eb9-4b62-8fa4-a9fa594ef2a7

Co-authored-by: stmatengss <11641725+stmatengss@users.noreply.github.com>

.github/workflows/ci.yml

@@ -43,6 +43,9 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
 
+    - name: Install Rust toolchain
+      uses: dtolnay/rust-toolchain@stable
+
     - name: Install and start etcd
       run: |
         wget https://github.com/etcd-io/etcd/releases/download/v3.6.1/etcd-v3.6.1-linux-amd64.tar.gz
@@ -125,6 +128,33 @@ jobs:
         python ./bootstrap_server.py &
       shell: bash
 
+    - name: Run Mooncake Store Rust smoke test and benchmark
+      run: |
+        $GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master \
+          --eviction_high_watermark_ratio=0.95 \
+          --cluster_id=ci_rust_test_cluster \
+          --port 50051 &
+        MASTER_PID=$!
+        sleep 3
+        cd mooncake-store/rust
+        export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-asio:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd:$LD_LIBRARY_PATH
+        export MOONCAKE_BUILD_DIR=$GITHUB_WORKSPACE/build
+        export MOONCAKE_STORE_LIB_DIR=$GITHUB_WORKSPACE/build/mooncake-store/src
+        export MOONCAKE_STORE_INCLUDE_DIR=$GITHUB_WORKSPACE/mooncake-store/include
+        export MC_METADATA_SERVER=http://127.0.0.1:8080/metadata
+        export MC_RUST_STORE_RUN_INTEGRATION=true
+        export MC_RUST_STORE_MASTER_ADDR=127.0.0.1:50051
+        export MC_RUST_STORE_LOCAL_HOSTNAME=127.0.0.1
+        export MC_RUST_STORE_PROTOCOL=tcp
+        export MC_RUST_STORE_DEVICE_NAME=
+        cargo test --test minimal_smoke -- --nocapture
+        MC_RUST_BENCH_ITERATIONS=4 \
+        MC_RUST_BENCH_VALUE_SIZE=4096 \
+        MC_RUST_BENCH_WARMUP=1 \
+          cargo run --release --example store_benchmark
+        kill $MASTER_PID 2>/dev/null || true
+      shell: bash
+
     - name: Run Go store binding integration tests
       run: |
         $GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master \
@@ -547,12 +577,13 @@ jobs:
         sudo cmake --install .
       shell: bash
 
-    - name: Check Mooncake Store Rust bindings and example
+    - name: Check Mooncake Store Rust bindings, examples, and tests
       run: |
         cd mooncake-store/rust
+        export MOONCAKE_BUILD_DIR=$GITHUB_WORKSPACE/build
         MOONCAKE_STORE_LIB_DIR=$GITHUB_WORKSPACE/build/mooncake-store/src \
         MOONCAKE_STORE_INCLUDE_DIR=$GITHUB_WORKSPACE/mooncake-store/include \
-          cargo check --example basic_usage --tests
+          cargo test --examples --tests --no-run
       shell: bash
 
     - name: Configure project
@@ -670,7 +701,7 @@ jobs:
           echo "Error: code_format.sh not found or not executable"
           exit 1
         fi
-        
+
         # Determine base ref for comparison
         if [ "${{ github.event_name }}" == "pull_request" ]; then
           # For PRs: compare against the target branch

mooncake-store/rust/CMakeLists.txt

@@ -1,39 +1,51 @@
-# CMake integration for the mooncake_store Rust crate.
+# CMake integration for the mooncake_store Rust package.
 #
 # This target is only added when the parent CMakeLists.txt sets
-# WITH_STORE_RUST=ON.  It depends on the mooncake_store C++ target so that
-# the static library is available before cargo tries to link against it.
+# WITH_STORE_RUST=ON.  It depends on the mooncake_store C++ target so that the
+# static library is available before cargo tries to link against it.
 
 add_custom_target(build_mooncake_store_rust DEPENDS mooncake_store)
 
 add_custom_command(
-    TARGET build_mooncake_store_rust
-    COMMAND
-        ${CMAKE_COMMAND} -E env
-            # Where cargo should place its build artefacts.
-            CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
-            # Directory containing libmooncake_store.a (or .so).
-            MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
-            # Directory containing store_c.h.
-            MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include
-        cargo build --release
-    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-    COMMENT "Building mooncake_store Rust crate"
-    VERBATIM
-)
+  TARGET build_mooncake_store_rust
+  COMMAND
+    ${CMAKE_COMMAND} -E env
+    # Where cargo should place its build artefacts.
+    CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
+    MOONCAKE_BUILD_DIR=${PROJECT_BINARY_DIR}
+    # Directory containing libmooncake_store.a (or .so).
+    MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
+    # Directory containing store_c.h.
+    MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include cargo
+    build --release
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMENT "Building mooncake_store Rust package"
+  VERBATIM)
 
-# Build the basic_usage example so that CI can verify it compiles.
 add_custom_target(build_mooncake_store_rust_example DEPENDS mooncake_store)
 
 add_custom_command(
-    TARGET build_mooncake_store_rust_example
-    COMMAND
-        ${CMAKE_COMMAND} -E env
-            CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
-            MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
-            MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include
-        cargo build --example basic_usage --release
-    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-    COMMENT "Building mooncake_store Rust basic_usage example"
-    VERBATIM
-)
+  TARGET build_mooncake_store_rust_example
+  COMMAND
+    ${CMAKE_COMMAND} -E env CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
+    MOONCAKE_BUILD_DIR=${PROJECT_BINARY_DIR}
+    MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
+    MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include cargo
+    build --examples --release
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMENT "Building mooncake_store Rust examples"
+  VERBATIM)
+
+add_custom_target(build_mooncake_store_rust_tests DEPENDS mooncake_store)
+
+add_custom_command(
+  TARGET build_mooncake_store_rust_tests
+  COMMAND
+    ${CMAKE_COMMAND} -E env CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}
+    MOONCAKE_BUILD_DIR=${PROJECT_BINARY_DIR}
+    MOONCAKE_STORE_LIB_DIR=${CMAKE_CURRENT_BINARY_DIR}/../src
+    MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include cargo test
+    --tests --no-run --release
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMENT "Building mooncake_store Rust tests"
+  VERBATIM)

mooncake-store/rust/build.rs

@@ -15,47 +15,126 @@
 use std::env;
 use std::path::PathBuf;
 
+fn push_existing_dir(search_dirs: &mut Vec<PathBuf>, dir: PathBuf) {
+    if dir.is_dir() && !search_dirs.iter().any(|existing| existing == &dir) {
+        search_dirs.push(dir);
+    }
+}
+
+fn push_env_paths(search_dirs: &mut Vec<PathBuf>, name: &str) {
+    if let Some(value) = env::var_os(name) {
+        for dir in env::split_paths(&value) {
+            push_existing_dir(search_dirs, dir);
+        }
+    }
+}
+
+fn has_library(search_dirs: &[PathBuf], candidates: &[&str]) -> bool {
+    search_dirs.iter().any(|dir| {
+        candidates.iter().any(|candidate| {
+            ["a", "so", "dylib"]
+                .into_iter()
+                .map(|ext| dir.join(format!("lib{candidate}.{ext}")))
+                .any(|path| path.exists())
+        })
+    })
+}
+
+fn emit_link_searches(search_dirs: &[PathBuf]) {
+    for dir in search_dirs {
+        println!("cargo:rustc-link-search=native={}", dir.display());
+    }
+}
+
 fn main() {
-    // -----------------------------------------------------------------------
-    // Library search path
-    //
-    // When built via CMake (WITH_STORE_RUST=ON) the CMakeLists.txt injects
-    // MOONCAKE_STORE_LIB_DIR pointing at the directory that contains
-    // libmooncake_store.a/.so.  When cargo is invoked standalone the caller
-    // should set the variable manually or rely on the default convention of a
-    // sibling `build/` directory produced by a top-level CMake configure.
-    // -----------------------------------------------------------------------
-    let lib_dir = env::var("MOONCAKE_STORE_LIB_DIR")
-        .unwrap_or_else(|_| "../../build/mooncake-store/src".to_string());
+    let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").expect("missing CARGO_MANIFEST_DIR"));
+    let mut search_dirs = Vec::new();
 
-    println!("cargo:rustc-link-search=native={lib_dir}");
-    println!("cargo:rustc-link-lib=mooncake_store");
+    let explicit_lib_dir = env::var("MOONCAKE_STORE_LIB_DIR").ok().map(PathBuf::from);
+    if let Some(dir) = explicit_lib_dir.clone() {
+        push_existing_dir(&mut search_dirs, dir);
+    }
 
-    // Dependencies of mooncake_store that must be satisfied at link time.
-    // The list mirrors what mooncake-store/src/CMakeLists.txt links against.
-    println!("cargo:rustc-link-lib=transfer_engine");
-    println!("cargo:rustc-link-lib=stdc++");
-    println!("cargo:rustc-link-lib=glog");
-    println!("cargo:rustc-link-lib=gflags");
-    println!("cargo:rustc-link-lib=pthread");
-    println!("cargo:rustc-link-lib=xxhash");
+    if let Ok(build_dir) = env::var("MOONCAKE_BUILD_DIR") {
+        let build_dir = PathBuf::from(build_dir);
+        for dir in [
+            build_dir.join("mooncake-store/src"),
+            build_dir.join("mooncake-store/src/cachelib_memory_allocator"),
+            build_dir.join("mooncake-transfer-engine/src"),
+            build_dir.join("mooncake-transfer-engine/src/common/base"),
+            build_dir.join("mooncake-asio"),
+            build_dir.join("mooncake-common/etcd"),
+        ] {
+            push_existing_dir(&mut search_dirs, dir);
+        }
+    }
+
+    let default_build_dir = manifest_dir.join("../../build");
+    for dir in [
+        default_build_dir.join("mooncake-store/src"),
+        default_build_dir.join("mooncake-store/src/cachelib_memory_allocator"),
+        default_build_dir.join("mooncake-transfer-engine/src"),
+        default_build_dir.join("mooncake-transfer-engine/src/common/base"),
+        default_build_dir.join("mooncake-asio"),
+        default_build_dir.join("mooncake-common/etcd"),
+        PathBuf::from("/usr/local/lib"),
+        PathBuf::from("/usr/lib/x86_64-linux-gnu"),
+        PathBuf::from("/lib/x86_64-linux-gnu"),
+    ] {
+        push_existing_dir(&mut search_dirs, dir);
+    }
+
+    push_env_paths(&mut search_dirs, "LD_LIBRARY_PATH");
+    push_env_paths(&mut search_dirs, "LIBRARY_PATH");
+
+    emit_link_searches(&search_dirs);
+
+    for library in [
+        "mooncake_store",
+        "cachelib_memory_allocator",
+        "transfer_engine",
+        "base",
+        "asio",
+        "stdc++",
+        "glog",
+        "gflags",
+        "pthread",
+        "xxhash",
+        "numa",
+        "ibverbs",
+        "jsoncpp",
+        "zstd",
+        "m",
+    ] {
+        println!("cargo:rustc-link-lib={library}");
+    }
+
+    for (link_name, candidates) in [
+        ("etcd_wrapper", &["etcd_wrapper"] as &[&str]),
+        ("curl", &["curl"]),
+        ("uring", &["uring"]),
+        ("asan", &["asan"]),
+        ("gcov", &["gcov"]),
+    ] {
+        if has_library(&search_dirs, candidates) {
+            println!("cargo:rustc-link-lib={link_name}");
+        }
+    }
 
-    // -----------------------------------------------------------------------
-    // Header path for bindgen
-    // -----------------------------------------------------------------------
     let include_dir = env::var("MOONCAKE_STORE_INCLUDE_DIR")
         .unwrap_or_else(|_| "../include".to_string());
 
     let header = format!("{include_dir}/store_c.h");
 
-    // Re-run this build script if the C header changes.
     println!("cargo:rerun-if-changed={header}");
+    println!("cargo:rerun-if-env-changed=MOONCAKE_BUILD_DIR");
     println!("cargo:rerun-if-env-changed=MOONCAKE_STORE_LIB_DIR");
     println!("cargo:rerun-if-env-changed=MOONCAKE_STORE_INCLUDE_DIR");
+    println!("cargo:rerun-if-env-changed=LD_LIBRARY_PATH");
+    println!("cargo:rerun-if-env-changed=LIBRARY_PATH");
 
     let bindings = bindgen::Builder::default()
         .header(&header)
-        // Only pull in declarations from store_c.h (no transitive system headers).
         .allowlist_function("mooncake_store_.*")
         .allowlist_type("mooncake_.*")
         .generate()

mooncake-store/rust/examples/store_benchmark.rs

@@ -0,0 +1,107 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::env;
+use std::error::Error;
+use std::time::{Instant, SystemTime, UNIX_EPOCH};
+
+use mooncake_store::MooncakeStore;
+
+fn env_or(name: &str, default: &str) -> String {
+    env::var(name).unwrap_or_else(|_| default.to_string())
+}
+
+fn env_usize(name: &str, default: usize) -> usize {
+    env::var(name)
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(default)
+}
+
+fn mib_per_sec(total_bytes: usize, seconds: f64) -> f64 {
+    if seconds <= f64::EPSILON {
+        return 0.0;
+    }
+    total_bytes as f64 / (1024.0 * 1024.0) / seconds
+}
+
+fn ops_per_sec(total_ops: usize, seconds: f64) -> f64 {
+    if seconds <= f64::EPSILON {
+        return 0.0;
+    }
+    total_ops as f64 / seconds
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+    let iterations = env_usize("MC_RUST_BENCH_ITERATIONS", 32);
+    let value_size = env_usize("MC_RUST_BENCH_VALUE_SIZE", 4096);
+    let warmup = env_usize("MC_RUST_BENCH_WARMUP", 4);
+
+    let store = MooncakeStore::new()?;
+    store.setup(
+        &env_or("MC_RUST_STORE_LOCAL_HOSTNAME", "127.0.0.1"),
+        &env_or("MC_METADATA_SERVER", "http://127.0.0.1:8080/metadata"),
+        512 << 20,
+        128 << 20,
+        &env_or("MC_RUST_STORE_PROTOCOL", "tcp"),
+        &env_or("MC_RUST_STORE_DEVICE_NAME", ""),
+        &env_or("MC_RUST_STORE_MASTER_ADDR", "127.0.0.1:50051"),
+    )?;
+
+    store.health_check()?;
+
+    let unique_prefix = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos();
+    let payload = vec![b'x'; value_size];
+
+    for index in 0..warmup {
+        let key = format!("rust-bench-warmup-{unique_prefix}-{index}");
+        store.put(&key, &payload, None)?;
+        let received = store.get(&key)?;
+        assert_eq!(received.len(), payload.len());
+        store.remove(&key, false)?;
+    }
+
+    let keys: Vec<String> = (0..iterations)
+        .map(|index| format!("rust-bench-{unique_prefix}-{index}"))
+        .collect();
+
+    let put_start = Instant::now();
+    for key in &keys {
+        store.put(key, &payload, None)?;
+    }
+    let put_seconds = put_start.elapsed().as_secs_f64();
+
+    let get_start = Instant::now();
+    for key in &keys {
+        let received = store.get(key)?;
+        assert_eq!(received.len(), payload.len());
+    }
+    let get_seconds = get_start.elapsed().as_secs_f64();
+
+    let remove_start = Instant::now();
+    for key in &keys {
+        store.remove(key, false)?;
+    }
+    let remove_seconds = remove_start.elapsed().as_secs_f64();
+
+    let total_bytes = iterations * payload.len();
+    println!("put_seconds={put_seconds:.6}");
+    println!("put_mib_per_sec={:.2}", mib_per_sec(total_bytes, put_seconds));
+    println!("get_seconds={get_seconds:.6}");
+    println!("get_mib_per_sec={:.2}", mib_per_sec(total_bytes, get_seconds));
+    println!("remove_seconds={remove_seconds:.6}");
+    println!("remove_ops_per_sec={:.2}", ops_per_sec(iterations, remove_seconds));
+
+    Ok(())
+}

mooncake-store/rust/tests/minimal_smoke.rs

@@ -0,0 +1,72 @@
+// Copyright 2024 KVCache.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::env;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use mooncake_store::MooncakeStore;
+
+fn env_flag(name: &str) -> bool {
+    matches!(
+        env::var(name)
+            .unwrap_or_default()
+            .trim()
+            .to_ascii_lowercase()
+            .as_str(),
+        "1" | "true" | "yes" | "on"
+    )
+}
+
+fn env_or(name: &str, default: &str) -> String {
+    env::var(name).unwrap_or_else(|_| default.to_string())
+}
+
+#[test]
+fn smoke_round_trip() {
+    if !env_flag("MC_RUST_STORE_RUN_INTEGRATION") {
+        eprintln!("skipping Mooncake Store Rust integration smoke test");
+        return;
+    }
+
+    let store = MooncakeStore::new().expect("create store handle");
+    store
+        .setup(
+            &env_or("MC_RUST_STORE_LOCAL_HOSTNAME", "127.0.0.1"),
+            &env_or("MC_METADATA_SERVER", "http://127.0.0.1:8080/metadata"),
+            512 << 20,
+            128 << 20,
+            &env_or("MC_RUST_STORE_PROTOCOL", "tcp"),
+            &env_or("MC_RUST_STORE_DEVICE_NAME", ""),
+            &env_or("MC_RUST_STORE_MASTER_ADDR", "127.0.0.1:50051"),
+        )
+        .expect("setup store");
+
+    store.health_check().expect("health check");
+
+    let unique_suffix = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .expect("system clock before unix epoch")
+        .as_nanos();
+    let key = format!("rust-smoke-{unique_suffix}");
+    let value = b"hello from rust smoke test";
+
+    store.put(&key, value, None).expect("put value");
+    assert!(store.is_exist(&key).expect("check existence"));
+    assert_eq!(store.get_size(&key).expect("get size") as usize, value.len());
+    assert_eq!(store.get(&key).expect("get value"), value);
+    assert!(!store.get_hostname().expect("get hostname").is_empty());
+
+    store.remove(&key, false).expect("remove key");
+    assert!(!store.is_exist(&key).expect("check absence"));
+}