forked from mooncake-track/Mooncake
Compare commits
12 Commits
main
...
copilot/in
| Author | SHA1 | Date |
|---|---|---|
copilot-swe-agent[bot]
|
c87eeb3ade | |
|
Teng Ma
|
da801bf7aa | |
|
copilot-swe-agent[bot]
|
b65e3f16c2 | |
|
copilot-swe-agent[bot]
|
85a60d8021 | |
|
copilot-swe-agent[bot]
|
ba6575937b | |
|
copilot-swe-agent[bot]
|
f1dc5cc712 | |
|
copilot-swe-agent[bot]
|
fd95e2e56e | |
|
copilot-swe-agent[bot]
|
4716cc7e78 | |
|
copilot-swe-agent[bot]
|
319c17ea77 | |
|
copilot-swe-agent[bot]
|
b1e255a948 | |
|
copilot-swe-agent[bot]
|
be44e51551 | |
|
copilot-swe-agent[bot]
|
c939d0726b |
|
|
@ -469,6 +469,9 @@ jobs:
|
|||
df -h
|
||||
shell: bash
|
||||
|
||||
- name: Install Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: Build transfer engine only
|
||||
run: |
|
||||
cd mooncake-transfer-engine
|
||||
|
|
@ -503,9 +506,8 @@ jobs:
|
|||
- name: Configure project with unit tests and examples
|
||||
run: |
|
||||
cd build
|
||||
cmake -G Ninja .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON
|
||||
cmake -G Ninja .. -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DWITH_STORE_RUST=ON -DENABLE_SCCACHE=ON
|
||||
shell: bash
|
||||
# TODO: lack WITH_RUST_EXAMPLE
|
||||
|
||||
- name: Build project with unit tests and examples
|
||||
run: |
|
||||
|
|
@ -516,6 +518,14 @@ jobs:
|
|||
sudo cmake --install .
|
||||
shell: bash
|
||||
|
||||
- name: Check Mooncake Store Rust bindings and example
|
||||
run: |
|
||||
cd mooncake-store/rust
|
||||
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
|
||||
shell: bash
|
||||
|
||||
- name: Configure project
|
||||
run: |
|
||||
cd build
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ option(WITH_TE "build mooncake transfer engine and sample code" ON)
|
|||
option(WITH_STORE "build mooncake store library and sample code" ON)
|
||||
option(WITH_P2P_STORE "build p2p store library and sample code" OFF)
|
||||
option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the transfer engine" OFF)
|
||||
option(WITH_STORE_RUST "build the Rust bindings for the Mooncake Store" ON)
|
||||
option(WITH_EP "build mooncake with expert parallelism support" OFF)
|
||||
|
||||
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11)
|
||||
|
|
@ -65,6 +66,14 @@ if (WITH_STORE)
|
|||
include_directories(mooncake-store/include)
|
||||
endif()
|
||||
|
||||
if (WITH_STORE_RUST)
|
||||
if (NOT WITH_STORE)
|
||||
message(FATAL_ERROR "WITH_STORE_RUST=ON requires WITH_STORE=ON")
|
||||
endif()
|
||||
message(STATUS "Mooncake Store Rust bindings will be built")
|
||||
add_subdirectory(mooncake-store/rust)
|
||||
endif()
|
||||
|
||||
option(EP_USE_IDE "Enable intelligent indexing for IDEs" OFF)
|
||||
if (WITH_EP)
|
||||
if (EP_USE_IDE)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,39 @@
|
|||
# CMake integration for the mooncake_store Rust crate.
|
||||
#
|
||||
# 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.
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
# 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
|
||||
)
|
||||
|
|
@ -0,0 +1,250 @@
|
|||
# This file is automatically @generated by Cargo.
|
||||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "aho-corasick"
|
||||
version = "1.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bindgen"
|
||||
version = "0.70.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f49d8fed880d473ea71efb9bf597651e77201bdd4893efe54c9e5d65ae04ce6f"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"cexpr",
|
||||
"clang-sys",
|
||||
"itertools",
|
||||
"log",
|
||||
"prettyplease",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"regex",
|
||||
"rustc-hash",
|
||||
"shlex",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "2.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af"
|
||||
|
||||
[[package]]
|
||||
name = "cexpr"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766"
|
||||
dependencies = [
|
||||
"nom",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cfg-if"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
|
||||
|
||||
[[package]]
|
||||
name = "clang-sys"
|
||||
version = "1.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4"
|
||||
dependencies = [
|
||||
"glob",
|
||||
"libc",
|
||||
"libloading",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "either"
|
||||
version = "1.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
|
||||
|
||||
[[package]]
|
||||
name = "glob"
|
||||
version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280"
|
||||
|
||||
[[package]]
|
||||
name = "itertools"
|
||||
version = "0.13.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186"
|
||||
dependencies = [
|
||||
"either",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.184"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "48f5d2a454e16a5ea0f4ced81bd44e4cfc7bd3a507b61887c99fd3538b28e4af"
|
||||
|
||||
[[package]]
|
||||
name = "libloading"
|
||||
version = "0.8.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "log"
|
||||
version = "0.4.29"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
|
||||
|
||||
[[package]]
|
||||
name = "minimal-lexical"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
|
||||
|
||||
[[package]]
|
||||
name = "mooncake_store"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"bindgen",
|
||||
"libc",
|
||||
"thiserror",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nom"
|
||||
version = "7.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
"minimal-lexical",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "prettyplease"
|
||||
version = "0.2.37"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.106"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
|
||||
dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.45"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "regex"
|
||||
version = "1.12.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276"
|
||||
dependencies = [
|
||||
"aho-corasick",
|
||||
"memchr",
|
||||
"regex-automata",
|
||||
"regex-syntax",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "regex-automata"
|
||||
version = "0.4.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f"
|
||||
dependencies = [
|
||||
"aho-corasick",
|
||||
"memchr",
|
||||
"regex-syntax",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "regex-syntax"
|
||||
version = "0.8.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a"
|
||||
|
||||
[[package]]
|
||||
name = "rustc-hash"
|
||||
version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
|
||||
|
||||
[[package]]
|
||||
name = "shlex"
|
||||
version = "1.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "2.0.117"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "2.0.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4"
|
||||
dependencies = [
|
||||
"thiserror-impl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror-impl"
|
||||
version = "2.0.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode-ident"
|
||||
version = "1.0.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
|
||||
|
||||
[[package]]
|
||||
name = "windows-link"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
[package]
|
||||
name = "mooncake_store"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
description = "Rust bindings for the Mooncake distributed KV cache store"
|
||||
license = "Apache-2.0"
|
||||
|
||||
[lib]
|
||||
name = "mooncake_store"
|
||||
# rlib for use as a Rust dependency, cdylib if a shared-library artifact is needed
|
||||
crate-type = ["rlib"]
|
||||
|
||||
[[example]]
|
||||
name = "basic_usage"
|
||||
path = "examples/basic_usage.rs"
|
||||
|
||||
[dependencies]
|
||||
thiserror = "2.0"
|
||||
libc = "0.2"
|
||||
|
||||
[build-dependencies]
|
||||
bindgen = "0.70"
|
||||
|
|
@ -0,0 +1,68 @@
|
|||
// 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::path::PathBuf;
|
||||
|
||||
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());
|
||||
|
||||
println!("cargo:rustc-link-search=native={lib_dir}");
|
||||
println!("cargo:rustc-link-lib=mooncake_store");
|
||||
|
||||
// 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");
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 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_STORE_LIB_DIR");
|
||||
println!("cargo:rerun-if-env-changed=MOONCAKE_STORE_INCLUDE_DIR");
|
||||
|
||||
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()
|
||||
.expect("Unable to generate Mooncake Store bindings");
|
||||
|
||||
let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
|
||||
bindings
|
||||
.write_to_file(out_path.join("bindings.rs"))
|
||||
.expect("Couldn't write Mooncake Store bindings");
|
||||
}
|
||||
|
|
@ -0,0 +1,159 @@
|
|||
// 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.
|
||||
|
||||
//! Basic usage example for the Mooncake Store Rust bindings.
|
||||
//!
|
||||
//! This example demonstrates the API surface of `mooncake_store`:
|
||||
//! - Creating a store handle
|
||||
//! - Setting up the client (requires a running metadata server)
|
||||
//! - Storing and retrieving values
|
||||
//! - Checking existence and sizes
|
||||
//! - Removing keys
|
||||
//!
|
||||
//! # Running
|
||||
//!
|
||||
//! The example needs a running Mooncake metadata server. Start one with:
|
||||
//!
|
||||
//! ```bash
|
||||
//! cd mooncake-transfer-engine/example/http-metadata-server-python
|
||||
//! pip install aiohttp
|
||||
//! python bootstrap_server.py &
|
||||
//! ```
|
||||
//!
|
||||
//! Then build and run via CMake (which sets the right library paths):
|
||||
//!
|
||||
//! ```bash
|
||||
//! cd build
|
||||
//! cmake -G Ninja .. -DWITH_STORE_RUST=ON
|
||||
//! cmake --build . --target build_mooncake_store_rust
|
||||
//! ```
|
||||
//!
|
||||
//! Or directly with cargo (after a CMake install):
|
||||
//!
|
||||
//! ```bash
|
||||
//! cargo run --example basic_usage
|
||||
//! ```
|
||||
|
||||
use mooncake_store::{MooncakeStore, ReplicateConfig, StoreError};
|
||||
|
||||
fn main() {
|
||||
println!("=== Mooncake Store Rust Bindings — Basic Usage ===\n");
|
||||
|
||||
// Step 1: Create a store handle.
|
||||
let store = match MooncakeStore::new() {
|
||||
Ok(s) => {
|
||||
println!("[OK] Store handle created.");
|
||||
s
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("[FAIL] Could not create store handle: {e}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
};
|
||||
|
||||
// Step 2: Connect to the metadata server and transport layer.
|
||||
//
|
||||
// In CI or environments without a running metadata server the setup
|
||||
// call is expected to fail – this is not an error in the bindings
|
||||
// themselves.
|
||||
let metadata_server =
|
||||
std::env::var("MC_METADATA_SERVER").unwrap_or_else(|_| "http://127.0.0.1:8080/metadata".to_string());
|
||||
|
||||
println!("Connecting to metadata server: {metadata_server}");
|
||||
|
||||
if let Err(e) = store.setup(
|
||||
"localhost",
|
||||
&metadata_server,
|
||||
512 << 20, // global_segment_size = 512 MiB
|
||||
128 << 20, // local_buffer_size = 128 MiB
|
||||
"tcp",
|
||||
"", // device_name (auto-select)
|
||||
"127.0.0.1:50051",
|
||||
) {
|
||||
eprintln!(
|
||||
"[SKIP] setup() returned an error (expected when no server is running): {e}\n\
|
||||
The API surface has been validated successfully."
|
||||
);
|
||||
std::process::exit(0);
|
||||
}
|
||||
|
||||
println!("[OK] Store client set up.\n");
|
||||
|
||||
// Step 3: Put a key/value pair (copies data into the store).
|
||||
let key = "example-key";
|
||||
let value = b"hello, mooncake!";
|
||||
|
||||
if let Err(e) = store.put(key, value, None) {
|
||||
eprintln!("[FAIL] put() failed: {e}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
println!("[OK] put(\"{key}\", {:?})", std::str::from_utf8(value).unwrap());
|
||||
|
||||
// Step 4: Check existence.
|
||||
match store.is_exist(key) {
|
||||
Ok(true) => println!("[OK] is_exist(\"{key}\") = true"),
|
||||
Ok(false) => println!("[WARN] is_exist(\"{key}\") = false (unexpected)"),
|
||||
Err(e) => eprintln!("[FAIL] is_exist() failed: {e}"),
|
||||
}
|
||||
|
||||
// Step 5: Get size.
|
||||
match store.get_size(key) {
|
||||
Ok(sz) => println!("[OK] get_size(\"{key}\") = {sz} bytes"),
|
||||
Err(e) => eprintln!("[FAIL] get_size() failed: {e}"),
|
||||
}
|
||||
|
||||
// Step 6: Retrieve the value.
|
||||
match store.get(key) {
|
||||
Ok(data) => {
|
||||
println!("[OK] get(\"{key}\") = {:?}", String::from_utf8_lossy(&data));
|
||||
assert_eq!(data, value, "round-trip mismatch!");
|
||||
}
|
||||
Err(e) => eprintln!("[FAIL] get() failed: {e}"),
|
||||
}
|
||||
|
||||
// Step 7: Put with explicit replication config.
|
||||
let config = ReplicateConfig {
|
||||
replica_num: 2,
|
||||
with_soft_pin: true,
|
||||
preferred_segments: vec!["seg-0".to_string()],
|
||||
};
|
||||
|
||||
if let Err(e) = store.put("replicated-key", b"replicated value", Some(&config)) {
|
||||
eprintln!("[FAIL] put() with ReplicateConfig failed: {e}");
|
||||
} else {
|
||||
println!("[OK] put(\"replicated-key\") with replica_num=2, soft_pin=true");
|
||||
}
|
||||
|
||||
// Step 8: Remove the keys.
|
||||
match store.remove(key, false) {
|
||||
Ok(()) => println!("[OK] remove(\"{key}\")"),
|
||||
Err(e) => eprintln!("[FAIL] remove() failed: {e}"),
|
||||
}
|
||||
|
||||
match store.remove("replicated-key", false) {
|
||||
Ok(()) => println!("[OK] remove(\"replicated-key\")"),
|
||||
Err(e) => eprintln!("[FAIL] remove() failed: {e}"),
|
||||
}
|
||||
|
||||
// Step 9: Demonstrate error handling for a missing key.
|
||||
match store.get_size("nonexistent") {
|
||||
Err(StoreError::OperationFailed(code)) => {
|
||||
println!("[OK] get_size(\"nonexistent\") returned OperationFailed({code}) as expected");
|
||||
}
|
||||
Ok(sz) => println!("[WARN] get_size(\"nonexistent\") unexpectedly returned {sz}"),
|
||||
Err(e) => println!("[OK] get_size(\"nonexistent\") returned error: {e}"),
|
||||
}
|
||||
|
||||
println!("\n=== Example completed successfully ===");
|
||||
}
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
// 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::ffi::NulError;
|
||||
|
||||
/// Errors returned by Mooncake Store operations.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum StoreError {
|
||||
/// A required pointer argument was null (e.g. the store handle has not
|
||||
/// been initialised yet, or an internal allocation failed).
|
||||
#[error("null handle or pointer")]
|
||||
NullHandle,
|
||||
|
||||
/// A C string could not be constructed because the input contains an
|
||||
/// interior null byte.
|
||||
#[error("string argument contains null byte: {0}")]
|
||||
InvalidString(#[from] NulError),
|
||||
|
||||
/// The C layer returned a negative status code.
|
||||
#[error("store operation failed with code {0}")]
|
||||
OperationFailed(i32),
|
||||
|
||||
/// The requested key does not exist in the store.
|
||||
///
|
||||
/// This variant is available for consumers that perform an explicit
|
||||
/// existence check (e.g. via [`crate::MooncakeStore::is_exist`]) before
|
||||
/// acting on the result. The C API itself does not expose a distinct
|
||||
/// "not found" error code, so internal methods return
|
||||
/// [`OperationFailed`](StoreError::OperationFailed) instead.
|
||||
#[error("key not found")]
|
||||
NotFound,
|
||||
}
|
||||
|
|
@ -0,0 +1,48 @@
|
|||
// 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.
|
||||
|
||||
//! # mooncake_store
|
||||
//!
|
||||
//! Safe Rust bindings for the [Mooncake](https://github.com/kvcache-ai/Mooncake)
|
||||
//! distributed KV-cache store.
|
||||
//!
|
||||
//! ## Quick start
|
||||
//!
|
||||
//! ```ignore
|
||||
//! use mooncake_store::{MooncakeStore, ReplicateConfig};
|
||||
//!
|
||||
//! let store = MooncakeStore::new()?;
|
||||
//! store.setup(
|
||||
//! "localhost",
|
||||
//! "http://127.0.0.1:8080/metadata",
|
||||
//! 512 << 20, // global_segment_size
|
||||
//! 128 << 20, // local_buffer_size
|
||||
//! "tcp",
|
||||
//! "",
|
||||
//! "127.0.0.1:50051",
|
||||
//! )?;
|
||||
//!
|
||||
//! store.put("hello", b"world", None)?;
|
||||
//!
|
||||
//! let value = store.get("hello")?;
|
||||
//! assert_eq!(value, b"world");
|
||||
//!
|
||||
//! store.remove("hello", false)?;
|
||||
//! ```
|
||||
|
||||
pub mod error;
|
||||
pub mod store;
|
||||
|
||||
pub use error::StoreError;
|
||||
pub use store::{MooncakeStore, ReplicateConfig};
|
||||
|
|
@ -0,0 +1,541 @@
|
|||
// 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.
|
||||
|
||||
//! Safe Rust wrapper around the Mooncake Store C API (`store_c.h`).
|
||||
|
||||
use std::ffi::{c_void, CString};
|
||||
|
||||
use crate::error::StoreError;
|
||||
|
||||
// Raw FFI bindings generated by build.rs / bindgen.
|
||||
#[allow(dead_code)]
|
||||
mod ffi {
|
||||
#![allow(non_snake_case)]
|
||||
#![allow(non_camel_case_types)]
|
||||
#![allow(non_upper_case_globals)]
|
||||
#![allow(clippy::useless_transmute)]
|
||||
include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ReplicateConfig
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Replication settings for a put operation.
|
||||
///
|
||||
/// This is the Rust counterpart of `mooncake_replicate_config_t` from
|
||||
/// `store_c.h`.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ReplicateConfig {
|
||||
/// Number of replicas to create (0 means use the server default).
|
||||
pub replica_num: usize,
|
||||
/// Prefer a replica on the same NUMA node / host ("soft pin").
|
||||
pub with_soft_pin: bool,
|
||||
/// Whitelist of segment names that should host a replica.
|
||||
pub preferred_segments: Vec<String>,
|
||||
}
|
||||
|
||||
impl ReplicateConfig {
|
||||
/// Converts this config into the C representation together with the
|
||||
/// owning `CString` objects that back the segment-name pointers.
|
||||
///
|
||||
/// The returned `_strings` must stay alive for as long as
|
||||
/// `ffi::mooncake_replicate_config_t` is in use.
|
||||
fn to_ffi(
|
||||
&self,
|
||||
) -> Result<
|
||||
(ffi::mooncake_replicate_config_t, Vec<CString>, Vec<*const libc::c_char>),
|
||||
StoreError,
|
||||
> {
|
||||
let strings: Vec<CString> = self
|
||||
.preferred_segments
|
||||
.iter()
|
||||
.map(|s| CString::new(s.as_str()).map_err(StoreError::from))
|
||||
.collect::<Result<_, _>>()?;
|
||||
|
||||
let ptrs: Vec<*const libc::c_char> = strings.iter().map(|s| s.as_ptr()).collect();
|
||||
|
||||
let c_config = ffi::mooncake_replicate_config_t {
|
||||
replica_num: self.replica_num,
|
||||
with_soft_pin: i32::from(self.with_soft_pin),
|
||||
preferred_segments: if ptrs.is_empty() {
|
||||
std::ptr::null_mut()
|
||||
} else {
|
||||
ptrs.as_ptr() as *mut *const libc::c_char
|
||||
},
|
||||
preferred_segments_count: ptrs.len(),
|
||||
};
|
||||
|
||||
Ok((c_config, strings, ptrs))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// MooncakeStore
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A handle to a Mooncake distributed KV-cache store.
|
||||
///
|
||||
/// # Lifecycle
|
||||
///
|
||||
/// ```ignore
|
||||
/// let store = MooncakeStore::new()?;
|
||||
/// store.setup("localhost", "http://127.0.0.1:8080/metadata",
|
||||
/// 512 << 20, 128 << 20, "tcp", "", "127.0.0.1:50051")?;
|
||||
/// store.put("my-key", b"hello world", None)?;
|
||||
/// let mut buf = vec![0u8; 11];
|
||||
/// store.get_into("my-key", buf.as_mut_ptr() as *mut c_void, buf.len())?;
|
||||
/// ```
|
||||
///
|
||||
/// The store is automatically destroyed when dropped.
|
||||
pub struct MooncakeStore {
|
||||
handle: ffi::mooncake_store_t,
|
||||
}
|
||||
|
||||
// The underlying C object is internally synchronised; sharing the handle
|
||||
// across threads is safe.
|
||||
unsafe impl Send for MooncakeStore {}
|
||||
unsafe impl Sync for MooncakeStore {}
|
||||
|
||||
impl MooncakeStore {
|
||||
// -----------------------------------------------------------------------
|
||||
// Lifecycle
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Allocate a new (uninitialised) store handle.
|
||||
///
|
||||
/// Call [`MooncakeStore::setup`] before performing any data operations.
|
||||
pub fn new() -> Result<Self, StoreError> {
|
||||
let handle = unsafe { ffi::mooncake_store_create() };
|
||||
if handle.is_null() {
|
||||
return Err(StoreError::NullHandle);
|
||||
}
|
||||
Ok(Self { handle })
|
||||
}
|
||||
|
||||
/// Initialise the store client.
|
||||
///
|
||||
/// # Parameters
|
||||
/// - `local_hostname` – IP or hostname of *this* node.
|
||||
/// - `metadata_server` – URL of the metadata server
|
||||
/// (e.g. `"http://127.0.0.1:8080/metadata"` or `"etcd://127.0.0.1:2379"`).
|
||||
/// - `global_segment_size` – size of each RDMA memory segment in bytes.
|
||||
/// - `local_buffer_size` – size of the local transfer buffer in bytes.
|
||||
/// - `protocol` – transport protocol (`"tcp"`, `"rdma"`, …).
|
||||
/// - `device_name` – network device name (empty string = auto-select).
|
||||
/// - `master_server_addr` – address of the Mooncake master service
|
||||
/// (e.g. `"127.0.0.1:50051"`).
|
||||
pub fn setup(
|
||||
&self,
|
||||
local_hostname: &str,
|
||||
metadata_server: &str,
|
||||
global_segment_size: u64,
|
||||
local_buffer_size: u64,
|
||||
protocol: &str,
|
||||
device_name: &str,
|
||||
master_server_addr: &str,
|
||||
) -> Result<(), StoreError> {
|
||||
let local_hostname_c = CString::new(local_hostname)?;
|
||||
let metadata_server_c = CString::new(metadata_server)?;
|
||||
let protocol_c = CString::new(protocol)?;
|
||||
let device_name_c = CString::new(device_name)?;
|
||||
let master_server_addr_c = CString::new(master_server_addr)?;
|
||||
|
||||
let rc = unsafe {
|
||||
ffi::mooncake_store_setup(
|
||||
self.handle,
|
||||
local_hostname_c.as_ptr(),
|
||||
metadata_server_c.as_ptr(),
|
||||
global_segment_size,
|
||||
local_buffer_size,
|
||||
protocol_c.as_ptr(),
|
||||
device_name_c.as_ptr(),
|
||||
master_server_addr_c.as_ptr(),
|
||||
)
|
||||
};
|
||||
if rc != 0 {
|
||||
return Err(StoreError::OperationFailed(rc));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Perform a connectivity health-check.
|
||||
///
|
||||
/// Returns `Ok(())` when the store is reachable, or an error otherwise.
|
||||
pub fn health_check(&self) -> Result<(), StoreError> {
|
||||
let rc = unsafe { ffi::mooncake_store_health_check(self.handle) };
|
||||
if rc != 0 {
|
||||
return Err(StoreError::OperationFailed(rc));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Put operations
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Store `value` under `key`.
|
||||
///
|
||||
/// The data is copied from `value` into the store's internal buffers, so
|
||||
/// `value` does not need to be pinned or registered.
|
||||
pub fn put(
|
||||
&self,
|
||||
key: &str,
|
||||
value: &[u8],
|
||||
config: Option<&ReplicateConfig>,
|
||||
) -> Result<(), StoreError> {
|
||||
let key_c = CString::new(key)?;
|
||||
let (_c_config, _strings, _ptrs) = Self::prepare_config(config)?;
|
||||
let cfg_ptr = _c_config
|
||||
.as_ref()
|
||||
.map_or(std::ptr::null(), |c| c as *const _);
|
||||
|
||||
let rc = unsafe {
|
||||
ffi::mooncake_store_put(
|
||||
self.handle,
|
||||
key_c.as_ptr(),
|
||||
value.as_ptr() as *const c_void,
|
||||
value.len(),
|
||||
cfg_ptr,
|
||||
)
|
||||
};
|
||||
if rc != 0 {
|
||||
return Err(StoreError::OperationFailed(rc));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Store data from a pre-registered, zero-copy buffer under `key`.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// `buffer` must be a pointer to at least `size` bytes of valid memory
|
||||
/// that has been registered with the store via
|
||||
/// [`MooncakeStore::register_buffer`].
|
||||
pub unsafe fn put_from(
|
||||
&self,
|
||||
key: &str,
|
||||
buffer: *mut c_void,
|
||||
size: usize,
|
||||
config: Option<&ReplicateConfig>,
|
||||
) -> Result<(), StoreError> {
|
||||
let key_c = CString::new(key)?;
|
||||
let (_c_config, _strings, _ptrs) = Self::prepare_config(config)?;
|
||||
let cfg_ptr = _c_config
|
||||
.as_ref()
|
||||
.map_or(std::ptr::null(), |c| c as *const _);
|
||||
|
||||
let rc = ffi::mooncake_store_put_from(self.handle, key_c.as_ptr(), buffer, size, cfg_ptr);
|
||||
if rc != 0 {
|
||||
return Err(StoreError::OperationFailed(rc));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Get operations
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Retrieve the value for `key` into a pre-allocated `buffer`.
|
||||
///
|
||||
/// Returns the number of bytes written on success.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// `buffer` must point to at least `size` bytes of writable, valid memory.
|
||||
pub unsafe fn get_into(
|
||||
&self,
|
||||
key: &str,
|
||||
buffer: *mut c_void,
|
||||
size: usize,
|
||||
) -> Result<i64, StoreError> {
|
||||
let key_c = CString::new(key)?;
|
||||
let bytes = ffi::mooncake_store_get_into(self.handle, key_c.as_ptr(), buffer, size);
|
||||
if bytes < 0 {
|
||||
return Err(StoreError::OperationFailed(bytes as i32));
|
||||
}
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
/// Retrieve the value for `key` into a Rust-owned `Vec<u8>`.
|
||||
///
|
||||
/// The buffer is allocated to the exact size of the stored value.
|
||||
pub fn get(&self, key: &str) -> Result<Vec<u8>, StoreError> {
|
||||
let size = self.get_size(key)?;
|
||||
let mut buf = vec![0u8; size as usize];
|
||||
let written =
|
||||
unsafe { self.get_into(key, buf.as_mut_ptr() as *mut c_void, buf.len())? };
|
||||
buf.truncate(written as usize);
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Existence / size / hostname
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Returns `true` when `key` exists in the store.
|
||||
pub fn is_exist(&self, key: &str) -> Result<bool, StoreError> {
|
||||
let key_c = CString::new(key)?;
|
||||
let rc = unsafe { ffi::mooncake_store_is_exist(self.handle, key_c.as_ptr()) };
|
||||
match rc {
|
||||
1 => Ok(true),
|
||||
0 => Ok(false),
|
||||
_ => Err(StoreError::OperationFailed(rc)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the size in bytes of the value stored under `key`, or an error
|
||||
/// if the key does not exist or the query fails.
|
||||
///
|
||||
/// The underlying C API returns a single negative value (`-1`) for all
|
||||
/// error conditions (null handle, missing key, internal failure), so we
|
||||
/// report `OperationFailed` with the raw code rather than attempting to
|
||||
/// distinguish `NotFound` from other failures.
|
||||
pub fn get_size(&self, key: &str) -> Result<i64, StoreError> {
|
||||
let key_c = CString::new(key)?;
|
||||
let sz = unsafe { ffi::mooncake_store_get_size(self.handle, key_c.as_ptr()) };
|
||||
if sz < 0 {
|
||||
return Err(StoreError::OperationFailed(sz as i32));
|
||||
}
|
||||
Ok(sz)
|
||||
}
|
||||
|
||||
/// Returns the hostname that the store client is registered under.
|
||||
///
|
||||
/// Uses a 1024-byte buffer, which is well above the POSIX hostname limit
|
||||
/// of 253 characters. Returns `OperationFailed` if the hostname exceeds
|
||||
/// the buffer or the call fails for another reason.
|
||||
pub fn get_hostname(&self) -> Result<String, StoreError> {
|
||||
let mut buf = vec![0u8; 1024];
|
||||
let rc = unsafe {
|
||||
ffi::mooncake_store_get_hostname(
|
||||
self.handle,
|
||||
buf.as_mut_ptr() as *mut libc::c_char,
|
||||
buf.len(),
|
||||
)
|
||||
};
|
||||
if rc != 0 {
|
||||
return Err(StoreError::OperationFailed(rc));
|
||||
}
|
||||
let cstr = unsafe { std::ffi::CStr::from_ptr(buf.as_ptr() as *const libc::c_char) };
|
||||
Ok(cstr.to_string_lossy().into_owned())
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Remove operations
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Remove the entry for `key`.
|
||||
///
|
||||
/// When `force` is `true` the key is removed even if it is currently being
|
||||
/// read by another client.
|
||||
pub fn remove(&self, key: &str, force: bool) -> Result<(), StoreError> {
|
||||
let key_c = CString::new(key)?;
|
||||
let rc = unsafe {
|
||||
ffi::mooncake_store_remove(self.handle, key_c.as_ptr(), i32::from(force))
|
||||
};
|
||||
if rc != 0 {
|
||||
return Err(StoreError::OperationFailed(rc));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove all keys matching `pattern` (a regular expression).
|
||||
///
|
||||
/// Returns the number of keys removed.
|
||||
pub fn remove_by_regex(&self, pattern: &str, force: bool) -> Result<i64, StoreError> {
|
||||
let pattern_c = CString::new(pattern)?;
|
||||
let n = unsafe {
|
||||
ffi::mooncake_store_remove_by_regex(self.handle, pattern_c.as_ptr(), i32::from(force))
|
||||
};
|
||||
if n < 0 {
|
||||
return Err(StoreError::OperationFailed(n as i32));
|
||||
}
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
/// Remove **all** keys from the store.
|
||||
///
|
||||
/// Returns the number of keys removed.
|
||||
pub fn remove_all(&self, force: bool) -> Result<i64, StoreError> {
|
||||
let n = unsafe { ffi::mooncake_store_remove_all(self.handle, i32::from(force)) };
|
||||
if n < 0 {
|
||||
return Err(StoreError::OperationFailed(n as i32));
|
||||
}
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Buffer registration (zero-copy operations)
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Register a memory region with the RDMA transport so that it can be used
|
||||
/// for zero-copy [`put_from`](MooncakeStore::put_from) /
|
||||
/// [`get_into`](MooncakeStore::get_into) operations.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// `buffer` must point to at least `size` bytes of valid, pinned memory
|
||||
/// that will remain valid until [`unregister_buffer`](MooncakeStore::unregister_buffer)
|
||||
/// is called.
|
||||
pub unsafe fn register_buffer(
|
||||
&self,
|
||||
buffer: *mut c_void,
|
||||
size: usize,
|
||||
) -> Result<(), StoreError> {
|
||||
let rc = ffi::mooncake_store_register_buffer(self.handle, buffer, size);
|
||||
if rc != 0 {
|
||||
return Err(StoreError::OperationFailed(rc));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Deregister a memory region previously registered with
|
||||
/// [`register_buffer`](MooncakeStore::register_buffer).
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// `buffer` must be the same pointer passed to `register_buffer`.
|
||||
pub unsafe fn unregister_buffer(&self, buffer: *mut c_void) -> Result<(), StoreError> {
|
||||
let rc = ffi::mooncake_store_unregister_buffer(self.handle, buffer);
|
||||
if rc != 0 {
|
||||
return Err(StoreError::OperationFailed(rc));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Internal helpers
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
/// Convert an `Option<&ReplicateConfig>` into a C-compatible value,
|
||||
/// returning the owning objects that keep the data alive.
|
||||
///
|
||||
/// The caller must bind the returned tuple to a named variable in its own
|
||||
/// stack frame *before* taking a raw pointer to the
|
||||
/// `Option<mooncake_replicate_config_t>`, so that the pointer is valid for
|
||||
/// the duration of the C call.
|
||||
#[allow(clippy::type_complexity)]
|
||||
fn prepare_config(
|
||||
config: Option<&ReplicateConfig>,
|
||||
) -> Result<
|
||||
(
|
||||
Option<ffi::mooncake_replicate_config_t>,
|
||||
Vec<CString>,
|
||||
Vec<*const libc::c_char>,
|
||||
),
|
||||
StoreError,
|
||||
> {
|
||||
match config {
|
||||
None => Ok((None, Vec::new(), Vec::new())),
|
||||
Some(cfg) => {
|
||||
let (c_config, strings, ptrs) = cfg.to_ffi()?;
|
||||
Ok((Some(c_config), strings, ptrs))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MooncakeStore {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
ffi::mooncake_store_destroy(self.handle);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::ffi::CStr;
|
||||
|
||||
#[test]
|
||||
fn replicate_config_to_ffi_empty_segments() {
|
||||
let config = ReplicateConfig {
|
||||
replica_num: 2,
|
||||
with_soft_pin: true,
|
||||
preferred_segments: Vec::new(),
|
||||
};
|
||||
|
||||
let (ffi_cfg, strings, ptrs) = config.to_ffi().expect("to_ffi should succeed");
|
||||
assert_eq!(ffi_cfg.replica_num, 2);
|
||||
assert_eq!(ffi_cfg.with_soft_pin, 1);
|
||||
assert!(ffi_cfg.preferred_segments.is_null());
|
||||
assert_eq!(ffi_cfg.preferred_segments_count, 0);
|
||||
assert!(strings.is_empty());
|
||||
assert!(ptrs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replicate_config_to_ffi_with_segments() {
|
||||
let config = ReplicateConfig {
|
||||
replica_num: 3,
|
||||
with_soft_pin: false,
|
||||
preferred_segments: vec!["seg-a".to_string(), "seg-b".to_string()],
|
||||
};
|
||||
|
||||
let (ffi_cfg, strings, ptrs) = config.to_ffi().expect("to_ffi should succeed");
|
||||
assert_eq!(ffi_cfg.replica_num, 3);
|
||||
assert_eq!(ffi_cfg.with_soft_pin, 0);
|
||||
assert!(!ffi_cfg.preferred_segments.is_null());
|
||||
assert_eq!(ffi_cfg.preferred_segments_count, 2);
|
||||
assert_eq!(strings.len(), 2);
|
||||
assert_eq!(ptrs.len(), 2);
|
||||
|
||||
let seg0 = unsafe { CStr::from_ptr(ptrs[0]) };
|
||||
let seg1 = unsafe { CStr::from_ptr(ptrs[1]) };
|
||||
assert_eq!(seg0.to_str().unwrap(), "seg-a");
|
||||
assert_eq!(seg1.to_str().unwrap(), "seg-b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replicate_config_to_ffi_rejects_interior_nul() {
|
||||
let config = ReplicateConfig {
|
||||
replica_num: 1,
|
||||
with_soft_pin: false,
|
||||
preferred_segments: vec!["bad\0segment".to_string()],
|
||||
};
|
||||
|
||||
assert!(matches!(
|
||||
config.to_ffi(),
|
||||
Err(StoreError::InvalidString(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_config_none_returns_null_config() {
|
||||
let (c_cfg, strings, ptrs) = MooncakeStore::prepare_config(None).expect("prepare should succeed");
|
||||
assert!(c_cfg.is_none());
|
||||
assert!(strings.is_empty());
|
||||
assert!(ptrs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_config_some_preserves_storage_and_values() {
|
||||
let config = ReplicateConfig {
|
||||
replica_num: 4,
|
||||
with_soft_pin: true,
|
||||
preferred_segments: vec!["x".to_string(), "y".to_string()],
|
||||
};
|
||||
|
||||
let (c_cfg, strings, ptrs) =
|
||||
MooncakeStore::prepare_config(Some(&config)).expect("prepare should succeed");
|
||||
let cfg_ref = c_cfg.as_ref().expect("expected Some config");
|
||||
assert_eq!(cfg_ref.replica_num, 4);
|
||||
assert_eq!(cfg_ref.with_soft_pin, 1);
|
||||
assert_eq!(cfg_ref.preferred_segments_count, 2);
|
||||
assert_eq!(strings.len(), 2);
|
||||
assert_eq!(ptrs.len(), 2);
|
||||
}
|
||||
}
|
||||
|
|
@ -2,12 +2,13 @@
|
|||
#include <glog/logging.h>
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <cstddef>
|
||||
#include <cstring>
|
||||
#include <chrono>
|
||||
#include <memory>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
#include <chrono>
|
||||
|
||||
#include "allocator.h"
|
||||
#include "types.h"
|
||||
|
|
@ -157,25 +158,38 @@ TEST_F(BufferAllocatorTest, ParallelAllocation) {
|
|||
const int num_threads = 4;
|
||||
const auto test_duration = std::chrono::seconds(1);
|
||||
std::vector<std::thread> threads;
|
||||
std::atomic<int> success_count{0};
|
||||
std::atomic<bool> saw_invalid_buffer{false};
|
||||
|
||||
// Create 4 threads, each performing repeated allocation and
|
||||
// deallocation for 1 second
|
||||
for (int thread_id = 0; thread_id < num_threads; ++thread_id) {
|
||||
threads.emplace_back(
|
||||
[this, &allocator, test_duration, segment_name]() {
|
||||
auto start_time = std::chrono::steady_clock::now();
|
||||
threads.emplace_back([&allocator, test_duration, segment_name,
|
||||
&success_count, &saw_invalid_buffer]() {
|
||||
auto start_time = std::chrono::steady_clock::now();
|
||||
|
||||
while (std::chrono::steady_clock::now() - start_time <
|
||||
test_duration) {
|
||||
// Allocate memory of varying sizes
|
||||
size_t alloc_size = 477;
|
||||
auto bufHandle = allocator->allocate(alloc_size);
|
||||
|
||||
ASSERT_NE(bufHandle, nullptr);
|
||||
VerifyAllocatedBuffer(*bufHandle, alloc_size,
|
||||
segment_name, segment_name);
|
||||
while (std::chrono::steady_clock::now() - start_time <
|
||||
test_duration) {
|
||||
size_t alloc_size = 477;
|
||||
auto bufHandle = allocator->allocate(alloc_size);
|
||||
if (!bufHandle) {
|
||||
std::this_thread::yield();
|
||||
continue;
|
||||
}
|
||||
});
|
||||
|
||||
auto descriptor = bufHandle->get_descriptor();
|
||||
if (bufHandle->getSegmentName() != segment_name ||
|
||||
descriptor.transport_endpoint_ != segment_name ||
|
||||
descriptor.size_ != alloc_size ||
|
||||
bufHandle->data() == nullptr) {
|
||||
saw_invalid_buffer.store(true,
|
||||
std::memory_order_relaxed);
|
||||
bufHandle.reset();
|
||||
break;
|
||||
}
|
||||
success_count.fetch_add(1, std::memory_order_relaxed);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Wait for all threads to complete
|
||||
|
|
@ -187,6 +201,8 @@ TEST_F(BufferAllocatorTest, ParallelAllocation) {
|
|||
<< (allocator_type == BufferAllocatorType::CACHELIB
|
||||
? "CACHELIB"
|
||||
: "OFFSET");
|
||||
EXPECT_FALSE(saw_invalid_buffer.load(std::memory_order_relaxed));
|
||||
EXPECT_GT(success_count.load(std::memory_order_relaxed), 0);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue