├── .cargo └── config.toml ├── img └── hero.png ├── src ├── allocator.rs ├── terminal.rs ├── profiling.rs ├── memory.rs ├── system.rs ├── storage.rs ├── dry.rs ├── docker.rs ├── dragonfly.rs ├── redis.rs ├── keydb.rs ├── scylladb.rs ├── keyprovider.rs ├── map.rs ├── arangodb.rs ├── surrealds.rs ├── dialect.rs ├── neo4j.rs ├── surrealmx.rs ├── lmdb.rs ├── mdbx.rs ├── surrealkv.rs ├── valueprovider.rs ├── fjall.rs ├── redb.rs └── mongodb.rs ├── Dockerfile ├── .rustfmt.toml ├── .editorconfig ├── Makefile ├── .gitignore ├── schema.surql ├── Cargo.toml └── LICENSE /.cargo/config.toml: -------------------------------------------------------------------------------- 1 | [net] 2 | git-fetch-with-cli = false 3 | 4 | -------------------------------------------------------------------------------- /img/hero.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/surrealdb/crud-bench/HEAD/img/hero.png -------------------------------------------------------------------------------- /src/allocator.rs: -------------------------------------------------------------------------------- 1 | #[global_allocator] 2 | static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; 3 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM cgr.dev/chainguard/glibc-dynamic:latest 2 | 3 | ARG TARGETARCH 4 | COPY artifacts/crud-bench-${TARGETARCH}/crud-bench /crud-bench 5 | 6 | ENTRYPOINT ["/crud-bench"] 7 | -------------------------------------------------------------------------------- /.rustfmt.toml: -------------------------------------------------------------------------------- 1 | edition = "2021" 2 | hard_tabs = true 3 | merge_derives = true 4 | reorder_imports = true 5 | reorder_modules = true 6 | use_field_init_shorthand = true 7 | use_small_heuristics = "Off" 8 | 9 | # ----------------------------------- 10 | # Unstable options we would like to use in future 11 | # ----------------------------------- 12 | 13 | #blank_lines_lower_bound = 1 14 | #group_imports = "One" 15 | #indent_style = "Block" 16 | #match_arm_blocks = true 17 | #reorder_impl_items = true 18 | #wrap_comments = true 19 | -------------------------------------------------------------------------------- /.editorconfig: -------------------------------------------------------------------------------- 1 | # EditorConfig coding styles definitions. For more information about the 2 | # properties used in this file, please see the EditorConfig documentation: 3 | # http://editorconfig.org/ 4 | 5 | root = true 6 | 7 | [*] 8 | charset = utf-8 9 | end_of_line = LF 10 | indent_size = 4 11 | indent_style = tab 12 | insert_final_newline = true 13 | trim_trailing_whitespace = true 14 | 15 | [*.toml] 16 | indent_size = 4 17 | indent_style = space 18 | 19 | [*.{yml,json}] 20 | indent_size = 2 21 | indent_style = space 22 | 23 | [*.{md,diff}] 24 | trim_trailing_whitespace = false 25 | -------------------------------------------------------------------------------- /src/terminal.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use std::fmt::Display; 3 | use std::io::Write; 4 | use std::io::{IsTerminal, Stdout, stdout}; 5 | 6 | pub(crate) struct Terminal(Option); 7 | 8 | impl Default for Terminal { 9 | fn default() -> Self { 10 | let stdout = stdout(); 11 | if stdout.is_terminal() { 12 | Self(Some(stdout)) 13 | } else { 14 | Self(None) 15 | } 16 | } 17 | } 18 | 19 | impl Clone for Terminal { 20 | fn clone(&self) -> Self { 21 | Self(self.0.as_ref().map(|_| stdout())) 22 | } 23 | } 24 | 25 | impl Terminal { 26 | /// Write a line to this Terminal via a callback 27 | pub(crate) fn write(&mut self, mut f: F) -> Result<()> 28 | where 29 | F: FnMut() -> Option, 30 | S: Display, 31 | { 32 | if let Some(ref mut o) = self.0 33 | && let Some(s) = f() 34 | { 35 | write!(o, "{s}")?; 36 | o.flush()?; 37 | } 38 | Ok(()) 39 | } 40 | } 41 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | define CRUD_BENCH_VALUE 2 | { 3 | "text": "text:50", 4 | "integer": "int", 5 | "number": "int:1..5000", 6 | "nested": { 7 | "text": "text:1000", 8 | "array": [ 9 | "string:50", 10 | "string:50", 11 | "string:50", 12 | "string:50", 13 | "string:50" 14 | ] 15 | } 16 | } 17 | endef 18 | 19 | database ?= surrealdb 20 | 21 | .PHONY: default 22 | default: 23 | @echo "Choose a Makefile target:" 24 | @$(MAKE) -pRrq -f $(lastword $(MAKEFILE_LIST)) : 2>/dev/null | awk -v RS= -F: '/^# File/,/^# Finished Make data base/ {if ($$1 !~ "^[#.]") {print " - " $$1}}' | sort 25 | 26 | .PHONY: build 27 | build: 28 | cargo build -r 29 | 30 | export CRUD_BENCH_VALUE 31 | .PHONY: dev 32 | dev: 33 | cargo run -- -d $(database) -s 100000 -c 128 -t 48 -k string26 -r 34 | 35 | export CRUD_BENCH_VALUE 36 | .PHONY: test 37 | test: 38 | target/release/crud-bench -d $(database) -s 5000000 -c 128 -t 48 -k string26 -r 39 | -------------------------------------------------------------------------------- /src/profiling.rs: -------------------------------------------------------------------------------- 1 | use pprof::ProfilerGuard; 2 | use pprof::ProfilerGuardBuilder; 3 | use pprof::protos::Message; 4 | use std::io::Write; 5 | use std::sync::OnceLock; 6 | 7 | static PROFILER: OnceLock> = OnceLock::new(); 8 | 9 | pub(crate) fn initialise() { 10 | PROFILER.get_or_init(|| { 11 | ProfilerGuardBuilder::default() 12 | .frequency(1000) 13 | .blocklist(&["libc", "libgcc", "pthread", "vdso"]) 14 | .build() 15 | .unwrap() 16 | }); 17 | } 18 | 19 | pub(crate) fn process() { 20 | if let Some(guard) = PROFILER.get() 21 | && let Ok(report) = guard.report().build() 22 | { 23 | // Output a flamegraph 24 | let file = std::fs::File::create("flamegraph.svg").unwrap(); 25 | report.flamegraph(file).unwrap(); 26 | // Output a pprof 27 | let mut file = std::fs::File::create("profile.pb").unwrap(); 28 | let profile = report.pprof().unwrap(); 29 | let mut content = Vec::new(); 30 | profile.encode(&mut content).unwrap(); 31 | file.write_all(&content).unwrap(); 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /src/memory.rs: -------------------------------------------------------------------------------- 1 | use sysinfo::System; 2 | 3 | /// System memory information for database optimization 4 | pub(crate) struct Config { 5 | /// Recommended memory allocation for database cache/buffer pools in GB 6 | pub cache_gb: u64, 7 | } 8 | 9 | impl Config { 10 | /// Get system memory information and calculate recommended database memory allocation 11 | pub fn new() -> Self { 12 | // Load the system attributed 13 | let system = System::new_all(); 14 | // Get the total system memory 15 | let total_memory = system.total_memory(); 16 | // Convert to GB for easier calculations 17 | let total_gb = total_memory / (1024 * 1024 * 1024); 18 | // Use ~75% of total memory for database cache 19 | let cache_gb = if total_gb <= 8 { 20 | // Small systems: use ~50% of memory 21 | (total_gb / 2).max(1) 22 | } else if total_gb <= 32 { 23 | // Medium systems: use ~60% of memory 24 | (total_gb * 3 / 5).max(4) 25 | } else { 26 | // Large systems: use ~75% of memory, but leave at least 8GB for OS 27 | (total_gb * 3 / 4).clamp(8, total_gb - 8) 28 | }; 29 | // Return configuration 30 | Self { 31 | cache_gb, 32 | } 33 | } 34 | } 35 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # ----------------------------------- 2 | # OS X 3 | # ----------------------------------- 4 | 5 | # Directory files 6 | .DS_Store 7 | .AppleDouble 8 | .LSOverride 9 | 10 | # Thumbnail files 11 | ._* 12 | 13 | # Files that might appear on external disk 14 | .Spotlight-V100 15 | .Trashes 16 | 17 | # Directories potentially created on remote AFP share 18 | .AppleDB 19 | .AppleDesktop 20 | Network Trash Folder 21 | Temporary Items 22 | .apdisk 23 | 24 | # ----------------------------------- 25 | # Files 26 | # ----------------------------------- 27 | 28 | **/*.rs.bk 29 | *.db 30 | *.skv 31 | *.sw? 32 | *.skv 33 | 34 | # ----------------------------------- 35 | # Folders 36 | # ----------------------------------- 37 | 38 | /debug/ 39 | /target/ 40 | /lib/cache/ 41 | /lib/store/ 42 | /lib/target/ 43 | /tests/sdk/*/Cargo.lock 44 | /tests/sdk/*/target 45 | 46 | .idea/ 47 | .vscode/ 48 | /result 49 | /bin/ 50 | /.direnv/ 51 | 52 | # --- Ignore benchmark database files 53 | data/ 54 | redb 55 | lmdb/ 56 | mdbx/ 57 | rocksdb/ 58 | surrealkv/ 59 | surrealmx/ 60 | 61 | # --- Ignore generated result files 62 | result.csv 63 | result.json 64 | result.html 65 | result-*.csv 66 | result-*.json 67 | result-*.html 68 | -------------------------------------------------------------------------------- /src/system.rs: -------------------------------------------------------------------------------- 1 | use serde::{Deserialize, Serialize}; 2 | use sysinfo::System; 3 | 4 | #[derive(Debug, Clone, Serialize, Deserialize)] 5 | pub struct SystemInfo { 6 | pub timestamp: i64, 7 | pub hostname: String, 8 | pub os_name: String, 9 | pub os_version: String, 10 | pub kernel_version: String, 11 | pub cpu_cores: usize, 12 | pub cpu_physical_cores: usize, 13 | pub cpu_arch: String, 14 | pub total_memory: u64, 15 | pub available_memory: u64, 16 | } 17 | 18 | pub fn collect() -> SystemInfo { 19 | SystemInfo::collect() 20 | } 21 | 22 | impl SystemInfo { 23 | pub fn collect() -> Self { 24 | // Create a new system instance 25 | let mut sys = System::new_all(); 26 | // Refresh the system information 27 | sys.refresh_all(); 28 | // Get the system details 29 | let timestamp = chrono::Utc::now().timestamp(); 30 | let hostname = System::host_name().unwrap_or_else(|| "unknown".to_string()); 31 | let os_name = System::name().unwrap_or_else(|| "unknown".to_string()); 32 | let os_version = System::os_version().unwrap_or_else(|| "unknown".to_string()); 33 | let kernel_version = System::kernel_version().unwrap_or_else(|| "unknown".to_string()); 34 | // Get the CPU details 35 | let cpu_arch = System::cpu_arch(); 36 | let cpu_cores = num_cpus::get(); 37 | let cpu_physical_cores = num_cpus::get_physical(); 38 | // Get the memory details 39 | let total_memory = sys.total_memory(); 40 | let available_memory = sys.available_memory(); 41 | // Return the system information 42 | Self { 43 | timestamp, 44 | hostname, 45 | os_name, 46 | os_version, 47 | kernel_version, 48 | cpu_cores, 49 | cpu_physical_cores, 50 | cpu_arch, 51 | total_memory, 52 | available_memory, 53 | } 54 | } 55 | } 56 | -------------------------------------------------------------------------------- /src/storage.rs: -------------------------------------------------------------------------------- 1 | use crate::result::BenchmarkResult; 2 | use anyhow::Result; 3 | use surrealdb::Surreal; 4 | use surrealdb::engine::remote::ws::Client; 5 | use surrealdb::engine::remote::ws::Ws; 6 | use surrealdb::opt::auth::Root; 7 | 8 | pub struct StorageClient { 9 | db: Surreal, 10 | } 11 | 12 | impl StorageClient { 13 | pub async fn connect(endpoint: &str) -> Result { 14 | // Create a new client 15 | let db = Surreal::new::(endpoint).await?; 16 | // Sign in as root user 17 | db.signin(Root { 18 | username: "root".to_string(), 19 | password: "root".to_string(), 20 | }) 21 | .await?; 22 | // Use namespace and database 23 | db.use_ns("surrealdb").use_db("crud-bench").await?; 24 | Ok(Self { 25 | db, 26 | }) 27 | } 28 | 29 | pub async fn store_result(&self, result: &BenchmarkResult) -> Result<()> { 30 | // Create the schema if it doesn't exist 31 | self.db 32 | .query( 33 | r#" 34 | DEFINE TABLE IF NOT EXISTS result SCHEMAFULL; 35 | DEFINE FIELD IF NOT EXISTS database ON result TYPE option; 36 | DEFINE FIELD IF NOT EXISTS system_info ON result TYPE option; 37 | DEFINE FIELD IF NOT EXISTS benchmark_metadata ON result TYPE option; 38 | DEFINE FIELD IF NOT EXISTS creates ON result TYPE option; 39 | DEFINE FIELD IF NOT EXISTS reads ON result TYPE option; 40 | DEFINE FIELD IF NOT EXISTS updates ON result TYPE option; 41 | DEFINE FIELD IF NOT EXISTS deletes ON result TYPE option; 42 | DEFINE FIELD IF NOT EXISTS scans ON result TYPE array; 43 | DEFINE FIELD IF NOT EXISTS batches ON result TYPE array; 44 | DEFINE FIELD IF NOT EXISTS sample ON result TYPE object; 45 | DEFINE FIELD IF NOT EXISTS timestamp ON result TYPE datetime DEFAULT time::now(); 46 | DEFINE INDEX IF NOT EXISTS idx_database ON result FIELDS database; 47 | DEFINE INDEX IF NOT EXISTS idx_timestamp ON result FIELDS timestamp; 48 | "#, 49 | ) 50 | .await?; 51 | // Convert to serde_json::Value for insertion 52 | let result = serde_json::to_value(result)?; 53 | // Insert the result using a query 54 | self.db.query("CREATE result CONTENT $result").bind(("result", result)).await?.check()?; 55 | // All ok 56 | Ok(()) 57 | } 58 | } 59 | -------------------------------------------------------------------------------- /schema.surql: -------------------------------------------------------------------------------- 1 | -- ---------------------------------------- 2 | -- SurrealDB schema for crud-bench result storage 3 | -- This file defines the schema used to store benchmark 4 | -- results in SurrealDB. This schema is automatically 5 | -- applied to the database when crud-bench is run. 6 | -- ---------------------------------------- 7 | 8 | -- Define the result table 9 | DEFINE TABLE IF NOT EXISTS result SCHEMAFULL; 10 | 11 | -- Define fields 12 | DEFINE FIELD IF NOT EXISTS database ON result TYPE option; 13 | DEFINE FIELD IF NOT EXISTS system_info ON result TYPE option; 14 | DEFINE FIELD IF NOT EXISTS benchmark_metadata ON result TYPE option; 15 | DEFINE FIELD IF NOT EXISTS creates ON result TYPE option; 16 | DEFINE FIELD IF NOT EXISTS reads ON result TYPE option; 17 | DEFINE FIELD IF NOT EXISTS updates ON result TYPE option; 18 | DEFINE FIELD IF NOT EXISTS deletes ON result TYPE option; 19 | DEFINE FIELD IF NOT EXISTS scans ON result TYPE array; 20 | DEFINE FIELD IF NOT EXISTS batches ON result TYPE array; 21 | DEFINE FIELD IF NOT EXISTS sample ON result TYPE object; 22 | DEFINE FIELD IF NOT EXISTS timestamp ON result TYPE datetime DEFAULT time::now(); 23 | 24 | -- Define indexes for efficient querying 25 | DEFINE INDEX IF NOT EXISTS idx_database ON result FIELDS database; 26 | DEFINE INDEX IF NOT EXISTS idx_timestamp ON result FIELDS timestamp; 27 | 28 | -- ---------------------------------------- 29 | -- Get all results for a specific database 30 | -- ---------------------------------------- 31 | 32 | -- SELECT * FROM result WHERE database = 'Redis' ORDER BY timestamp DESC; 33 | 34 | -- ---------------------------------------- 35 | -- Get latest results for each database 36 | -- ---------------------------------------- 37 | 38 | -- SELECT database, creates.ops as create_ops, reads.ops as read_ops 39 | -- FROM result 40 | -- ORDER BY timestamp DESC 41 | -- GROUP BY database; 42 | 43 | -- ---------------------------------------- 44 | -- Compare results by system specs 45 | -- ---------------------------------------- 46 | 47 | -- SELECT * FROM result 48 | -- WHERE system_info.cpu_cores = 8 49 | -- ORDER BY creates.ops DESC; 50 | 51 | -- ---------------------------------------- 52 | -- Get historical trend for a database 53 | -- ---------------------------------------- 54 | 55 | -- SELECT timestamp, creates.ops, reads.ops, updates.ops, deletes.ops 56 | -- FROM result 57 | -- WHERE database = 'SurrealDB (RocksDB)' 58 | -- ORDER BY timestamp ASC; 59 | -------------------------------------------------------------------------------- /src/dry.rs: -------------------------------------------------------------------------------- 1 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 2 | use crate::valueprovider::Columns; 3 | use crate::{Benchmark, KeyType, Scan}; 4 | use anyhow::Result; 5 | use serde_json::Value; 6 | use std::hint::black_box; 7 | use std::time::Duration; 8 | 9 | pub(crate) struct DryClientProvider {} 10 | 11 | impl BenchmarkEngine for DryClientProvider { 12 | /// The number of seconds to wait before connecting 13 | fn wait_timeout(&self) -> Option { 14 | None 15 | } 16 | /// Initiates a new datastore benchmarking engine 17 | async fn setup(_kt: KeyType, _columns: Columns, _options: &Benchmark) -> Result { 18 | Ok(Self {}) 19 | } 20 | /// Creates a new client for this benchmarking engine 21 | async fn create_client(&self) -> Result { 22 | Ok(DryClient {}) 23 | } 24 | } 25 | 26 | pub(crate) struct DryClient {} 27 | 28 | impl BenchmarkClient for DryClient { 29 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 30 | black_box((key, val)); 31 | Ok(()) 32 | } 33 | 34 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 35 | black_box((key, val)); 36 | Ok(()) 37 | } 38 | 39 | async fn read_u32(&self, key: u32) -> Result<()> { 40 | black_box(key); 41 | Ok(()) 42 | } 43 | 44 | async fn read_string(&self, key: String) -> Result<()> { 45 | black_box(key); 46 | Ok(()) 47 | } 48 | 49 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 50 | black_box((key, val)); 51 | Ok(()) 52 | } 53 | 54 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 55 | black_box((key, val)); 56 | Ok(()) 57 | } 58 | 59 | async fn delete_u32(&self, key: u32) -> Result<()> { 60 | black_box(key); 61 | Ok(()) 62 | } 63 | 64 | async fn delete_string(&self, key: String) -> Result<()> { 65 | black_box(key); 66 | Ok(()) 67 | } 68 | 69 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 70 | black_box(scan); 71 | Ok(scan.expect.unwrap_or(0)) 72 | } 73 | 74 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 75 | black_box(scan); 76 | Ok(scan.expect.unwrap_or(0)) 77 | } 78 | 79 | async fn batch_create_u32( 80 | &self, 81 | key_vals: impl Iterator + Send, 82 | ) -> Result<()> { 83 | for (key, val) in key_vals { 84 | black_box((key, val)); 85 | } 86 | Ok(()) 87 | } 88 | 89 | async fn batch_create_string( 90 | &self, 91 | key_vals: impl Iterator + Send, 92 | ) -> Result<()> { 93 | for (key, val) in key_vals { 94 | black_box((key, val)); 95 | } 96 | Ok(()) 97 | } 98 | 99 | async fn batch_read_u32(&self, keys: impl Iterator + Send) -> Result<()> { 100 | for key in keys { 101 | black_box(key); 102 | } 103 | Ok(()) 104 | } 105 | 106 | async fn batch_read_string(&self, keys: impl Iterator + Send) -> Result<()> { 107 | for key in keys { 108 | black_box(key); 109 | } 110 | Ok(()) 111 | } 112 | 113 | async fn batch_update_u32( 114 | &self, 115 | key_vals: impl Iterator + Send, 116 | ) -> Result<()> { 117 | for (key, val) in key_vals { 118 | black_box((key, val)); 119 | } 120 | Ok(()) 121 | } 122 | 123 | async fn batch_update_string( 124 | &self, 125 | key_vals: impl Iterator + Send, 126 | ) -> Result<()> { 127 | for (key, val) in key_vals { 128 | black_box((key, val)); 129 | } 130 | Ok(()) 131 | } 132 | 133 | async fn batch_delete_u32(&self, keys: impl Iterator + Send) -> Result<()> { 134 | for key in keys { 135 | black_box(key); 136 | } 137 | Ok(()) 138 | } 139 | 140 | async fn batch_delete_string(&self, keys: impl Iterator + Send) -> Result<()> { 141 | for key in keys { 142 | black_box(key); 143 | } 144 | Ok(()) 145 | } 146 | } 147 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "crud-bench" 3 | edition = "2024" 4 | version = "0.1.0" 5 | license = "Apache-2.0" 6 | readme = "README.md" 7 | rust-version = "1.89.0" 8 | 9 | [features] 10 | default = [ 11 | "arangodb", 12 | "dragonfly", 13 | "fjall", 14 | "keydb", 15 | "mdbx", 16 | "lmdb", 17 | "mariadb", 18 | "mongodb", 19 | "mysql", 20 | "neo4j", 21 | "postgres", 22 | "redb", 23 | "redis", 24 | "rocksdb", 25 | "scylladb", 26 | "sqlite", 27 | "surrealdb", 28 | "surrealkv", 29 | "surrealmx", 30 | ] 31 | arangodb = ["dep:arangors"] 32 | dragonfly = ["dep:redis"] 33 | keydb = ["dep:redis"] 34 | fjall = ["dep:fjall"] 35 | mdbx = ["dep:libmdbx"] 36 | lmdb = ["dep:heed"] 37 | mariadb = ["dep:mysql_async"] 38 | mongodb = ["dep:mongodb"] 39 | mysql = ["dep:mysql_async"] 40 | neo4j = ["dep:neo4rs"] 41 | postgres = ["dep:tokio-postgres"] 42 | redb = ["dep:redb"] 43 | redis = ["dep:redis"] 44 | rocksdb = ["dep:rocksdb"] 45 | scylladb = ["dep:scylla"] 46 | sqlite = ["dep:tokio-rusqlite"] 47 | surrealdb = [ 48 | "dep:surrealdb", 49 | "dep:surrealdb-types", 50 | "surrealdb/kv-mem", 51 | "surrealdb/kv-rocksdb", 52 | "surrealdb/kv-surrealkv", 53 | "surrealdb/protocol-http", 54 | "surrealdb/protocol-ws", 55 | "surrealdb/rustls", 56 | ] 57 | surrealkv = ["dep:surrealkv"] 58 | surrealmx = ["dep:surrealmx"] 59 | 60 | [profile.release] 61 | lto = true 62 | strip = "debuginfo" 63 | opt-level = 3 64 | debug = false 65 | panic = 'abort' 66 | codegen-units = 1 67 | overflow-checks = false 68 | 69 | [dependencies] 70 | affinitypool = "0.4.0" 71 | anyhow = "1.0.100" 72 | arangors = { version = "0.6.0", optional = true } 73 | bincode = { version = "2.0.1", features = ["serde"] } 74 | bytesize = "2.3.1" 75 | comfy-table = "7.2.1" 76 | chrono = "0.4.42" 77 | clap = { version = "4.5.53", features = ["derive", "string", "env", "color"] } 78 | csv = "1.4.0" 79 | dashmap = "6.1.0" 80 | env_logger = "0.11.8" 81 | fjall = { version = "2.11.2", optional = true } 82 | flatten-json-object = "0.6.1" 83 | futures = "0.3.31" 84 | hdrhistogram = "7.5.4" 85 | heed = { version = "0.22.0", optional = true } 86 | libmdbx = { version = "0.6.4", optional = true } 87 | log = "0.4.28" 88 | mimalloc = "0.1.48" 89 | mongodb = { version = "3.4.1", optional = true } 90 | mysql_async = { version = "0.36.1", default-features = false, features = [ 91 | "bigdecimal", 92 | "binlog", 93 | "derive", 94 | "frunk", 95 | "rust_decimal", 96 | "time", 97 | ], optional = true } 98 | neo4rs = { version = "0.8.0", optional = true } 99 | num_cpus = "1.17.0" 100 | pprof = { version = "0.15.0", features = ["flamegraph", "prost-codec"] } 101 | rand = { version = "0.9.2", features = ["small_rng"] } 102 | redb = { version = "3.1.0", optional = true } 103 | redis = { version = "0.32.7", features = ["tokio-comp"], optional = true } 104 | rocksdb = { version = "0.24.0-surreal.1", package = "surrealdb-rocksdb", features = ["lz4", "snappy"], optional = true } 105 | scylla = { version = "1.4.1", optional = true } 106 | serde = { version = "1.0.228", features = ["derive"] } 107 | serde_json = "1.0.145" 108 | serial_test = "3.2.0" 109 | surrealdb = { version = "3.0.0-alpha.18", default-features = false, optional = true } 110 | surrealdb-types = { version = "3.0.0-alpha.18", default-features = false, optional = true } 111 | surrealkv = { version = "0.16.0", optional = true } 112 | surrealmx = { version = "0.15.0", optional = true } 113 | sysinfo = { version = "0.37.2", features = ["serde"] } 114 | tokio = { version = "1.48.0", features = ["macros", "time", "rt-multi-thread"] } 115 | tokio-postgres = { version = "0.7.15", optional = true, features = [ 116 | "with-serde_json-1", 117 | "with-uuid-1", 118 | ] } 119 | tokio-rusqlite = { version = "0.7.0", optional = true, features = ["bundled"] } 120 | twox-hash = "2.1.2" 121 | uuid = { version = "1.18.1", features = ["v4"] } 122 | 123 | [profile.profiling] 124 | inherits = "release" # start from the production settings 125 | debug = true # keep debug info so symbols show up in flamegraph 126 | strip = "none" 127 | -------------------------------------------------------------------------------- /src/docker.rs: -------------------------------------------------------------------------------- 1 | use crate::benchmark::Benchmark; 2 | use log::{debug, error, info}; 3 | use std::fmt; 4 | use std::process::{Command, exit}; 5 | use std::time::Duration; 6 | 7 | const RETRIES: i32 = 10; 8 | 9 | const TIMEOUT: Duration = Duration::from_secs(6); 10 | 11 | pub(crate) struct DockerParams { 12 | pub(crate) image: &'static str, 13 | pub(crate) pre_args: String, 14 | pub(crate) post_args: String, 15 | } 16 | 17 | pub(crate) struct Container { 18 | image: String, 19 | } 20 | 21 | impl Drop for Container { 22 | fn drop(&mut self) { 23 | let _ = Self::stop(); 24 | } 25 | } 26 | 27 | impl Container { 28 | /// Get the name of the Docker image 29 | pub(crate) fn image(&self) -> &str { 30 | &self.image 31 | } 32 | 33 | /// Start the Docker container 34 | pub(crate) fn start(image: String, pre: &str, post: &str, options: &Benchmark) -> Self { 35 | // Output debug information to the logs 36 | info!("Starting Docker image '{image}'"); 37 | // Attempt to start Docker 10 times 38 | for i in 1..=RETRIES { 39 | // Configure the Docker command arguments 40 | let mut args = Arguments::new(["run"]); 41 | // Configure the default pre arguments 42 | args.append(pre); 43 | // Configure any custom pre arguments 44 | if let Ok(v) = std::env::var("DOCKER_PRE_ARGS") { 45 | args.append(&v); 46 | } 47 | // Run in privileged mode if specified 48 | if options.privileged { 49 | args.add(["--privileged"]); 50 | } 51 | // Configure the Docker container options 52 | args.add(["--rm"]); 53 | args.add(["--quiet"]); 54 | args.add(["--pull", "always"]); 55 | args.add(["--name", "crud-bench"]); 56 | args.add(["--net", "host"]); 57 | args.add(["-d", &image]); 58 | // Configure the default post arguments 59 | args.append(post); 60 | // Configure any custom post arguments 61 | if let Ok(v) = std::env::var("DOCKER_POST_ARGS") { 62 | args.append(&v); 63 | } 64 | // Execute the Docker run command 65 | match Self::execute(args.clone()) { 66 | // The command executed successfully 67 | Ok(_) => break, 68 | // There was an error with the command 69 | Err(e) => match i { 70 | // This is the last attempt so exit fully 71 | RETRIES => { 72 | error!("Docker command failure: `docker {args}`"); 73 | error!("{e}"); 74 | exit(1); 75 | } 76 | // Let's log the output and retry the command 77 | _ => { 78 | debug!("Docker command failure: `docker {args}`"); 79 | debug!("{e}"); 80 | std::thread::sleep(TIMEOUT); 81 | } 82 | }, 83 | } 84 | } 85 | // Return the container name 86 | Self { 87 | image, 88 | } 89 | } 90 | 91 | /// Stop the Docker container 92 | pub(crate) fn stop() -> Result { 93 | info!("Stopping Docker container 'crud-bench'"); 94 | let args = ["container", "stop", "--time", "300", "crud-bench"]; 95 | Self::execute(Arguments::new(args)) 96 | } 97 | 98 | /// Output the container logs 99 | pub(crate) fn logs() -> Result { 100 | info!("Logging Docker container 'crud-bench'"); 101 | let args = ["container", "logs", "crud-bench"]; 102 | Self::execute(Arguments::new(args)) 103 | } 104 | 105 | fn execute(args: Arguments) -> Result { 106 | // Output debug information to the logs 107 | println!("Running command: `docker {args}`"); 108 | // Create a new process command 109 | let mut command = Command::new("docker"); 110 | // Set the arguments on the command 111 | let command = command.args(args.0.clone()); 112 | // Catch all output from the command 113 | let output = command.output().expect("Failed to execute process"); 114 | // Output command failure if errored 115 | match output.status.success() { 116 | // Get the stderr out from the command 117 | false => Err(String::from_utf8(output.stderr).unwrap().trim().to_string()), 118 | // Get the stdout out from the command 119 | true => Ok(String::from_utf8(output.stdout).unwrap().trim().to_string()), 120 | } 121 | } 122 | } 123 | 124 | #[derive(Clone)] 125 | pub(crate) struct Arguments(Vec); 126 | 127 | impl fmt::Display for Arguments { 128 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 129 | write!(f, "{}", self.0.join(" ")) 130 | } 131 | } 132 | 133 | impl Arguments { 134 | fn new(args: I) -> Self 135 | where 136 | I: IntoIterator, 137 | S: Into, 138 | { 139 | let mut a = Self(vec![]); 140 | a.add(args); 141 | a 142 | } 143 | 144 | fn add(&mut self, args: I) 145 | where 146 | I: IntoIterator, 147 | S: Into, 148 | { 149 | for arg in args { 150 | self.0.push(arg.into()); 151 | } 152 | } 153 | 154 | fn append(&mut self, args: &str) { 155 | let split: Vec<&str> = args.split(' ').filter(|a| !a.is_empty()).collect(); 156 | self.add(split); 157 | } 158 | } 159 | -------------------------------------------------------------------------------- /src/dragonfly.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "dragonfly")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::docker::DockerParams; 5 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 6 | use crate::valueprovider::Columns; 7 | use crate::{Benchmark, KeyType, Projection, Scan}; 8 | use anyhow::{Result, bail}; 9 | use futures::StreamExt; 10 | use redis::aio::MultiplexedConnection; 11 | use redis::{AsyncCommands, Client, ScanOptions}; 12 | use serde_json::Value; 13 | use std::hint::black_box; 14 | use tokio::sync::Mutex; 15 | 16 | pub const DEFAULT: &str = "redis://:root@127.0.0.1:6379/"; 17 | 18 | pub(crate) fn docker(_: &Benchmark) -> DockerParams { 19 | DockerParams { 20 | image: "docker.dragonflydb.io/dragonflydb/dragonfly", 21 | pre_args: "-p 127.0.0.1:6379:6379 --ulimit memlock=-1".to_string(), 22 | post_args: "--requirepass root".to_string(), 23 | } 24 | } 25 | 26 | pub(crate) struct DragonflyClientProvider { 27 | url: String, 28 | } 29 | 30 | impl BenchmarkEngine for DragonflyClientProvider { 31 | /// Initiates a new datastore benchmarking engine 32 | async fn setup(_kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 33 | Ok(Self { 34 | url: options.endpoint.as_deref().unwrap_or(DEFAULT).to_owned(), 35 | }) 36 | } 37 | /// Creates a new client for this benchmarking engine 38 | async fn create_client(&self) -> Result { 39 | let client = Client::open(self.url.as_str())?; 40 | let conn_record = Mutex::new(client.get_multiplexed_async_connection().await?); 41 | let conn_iter = Mutex::new(client.get_multiplexed_async_connection().await?); 42 | Ok(DragonflyClient { 43 | conn_record, 44 | conn_iter, 45 | }) 46 | } 47 | } 48 | 49 | pub(crate) struct DragonflyClient { 50 | conn_iter: Mutex, 51 | conn_record: Mutex, 52 | } 53 | 54 | impl BenchmarkClient for DragonflyClient { 55 | #[allow(dependency_on_unit_never_type_fallback)] 56 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 57 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 58 | let _: () = self.conn_record.lock().await.set(key, val).await?; 59 | Ok(()) 60 | } 61 | 62 | #[allow(dependency_on_unit_never_type_fallback)] 63 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 64 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 65 | let _: () = self.conn_record.lock().await.set(key, val).await?; 66 | Ok(()) 67 | } 68 | 69 | #[allow(dependency_on_unit_never_type_fallback)] 70 | async fn read_u32(&self, key: u32) -> Result<()> { 71 | let val: Vec = self.conn_record.lock().await.get(key).await?; 72 | assert!(!val.is_empty()); 73 | black_box(val); 74 | Ok(()) 75 | } 76 | 77 | #[allow(dependency_on_unit_never_type_fallback)] 78 | async fn read_string(&self, key: String) -> Result<()> { 79 | let val: Vec = self.conn_record.lock().await.get(key).await?; 80 | assert!(!val.is_empty()); 81 | black_box(val); 82 | Ok(()) 83 | } 84 | 85 | #[allow(dependency_on_unit_never_type_fallback)] 86 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 87 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 88 | let _: () = self.conn_record.lock().await.set(key, val).await?; 89 | Ok(()) 90 | } 91 | 92 | #[allow(dependency_on_unit_never_type_fallback)] 93 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 94 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 95 | let _: () = self.conn_record.lock().await.set(key, val).await?; 96 | Ok(()) 97 | } 98 | 99 | #[allow(dependency_on_unit_never_type_fallback)] 100 | async fn delete_u32(&self, key: u32) -> Result<()> { 101 | let _: () = self.conn_record.lock().await.del(key).await?; 102 | Ok(()) 103 | } 104 | 105 | #[allow(dependency_on_unit_never_type_fallback)] 106 | async fn delete_string(&self, key: String) -> Result<()> { 107 | let _: () = self.conn_record.lock().await.del(key).await?; 108 | Ok(()) 109 | } 110 | 111 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 112 | self.scan_bytes(scan).await 113 | } 114 | 115 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 116 | self.scan_bytes(scan).await 117 | } 118 | } 119 | 120 | impl DragonflyClient { 121 | async fn scan_bytes(&self, scan: &Scan) -> Result { 122 | // Conditional scans are not supported 123 | if scan.condition.is_some() { 124 | bail!(NOT_SUPPORTED_ERROR); 125 | } 126 | // Extract parameters 127 | let s = scan.start.unwrap_or(0); 128 | let l = scan.limit.unwrap_or(usize::MAX); 129 | let p = scan.projection()?; 130 | // Get the two connection types 131 | let mut conn_iter = self.conn_iter.lock().await; 132 | let mut conn_record = self.conn_record.lock().await; 133 | // Configure the scan options for improve iteration 134 | let opts = ScanOptions::default().with_count(5000); 135 | // Create an iterator starting at the beginning 136 | let mut iter = conn_iter.scan_options::(opts).await?.skip(s); 137 | // Perform the relevant projection scan type 138 | match p { 139 | Projection::Id => { 140 | // We use a for loop to iterate over the results, while 141 | // calling black_box internally. This is necessary as 142 | // an iterator with `filter_map` or `map` is optimised 143 | // out by the compiler when calling `count` at the end. 144 | let mut count = 0; 145 | for _ in 0..l { 146 | if let Some(k) = iter.next().await { 147 | black_box(k); 148 | count += 1; 149 | } else { 150 | break; 151 | } 152 | } 153 | Ok(count) 154 | } 155 | Projection::Full => { 156 | // We use a for loop to iterate over the results, while 157 | // calling black_box internally. This is necessary as 158 | // an iterator with `filter_map` or `map` is optimised 159 | // out by the compiler when calling `count` at the end. 160 | let mut count = 0; 161 | while let Some(k) = iter.next().await { 162 | let v: Vec = conn_record.get(k).await?; 163 | black_box(v); 164 | count += 1; 165 | if count >= l { 166 | break; 167 | } 168 | } 169 | Ok(count) 170 | } 171 | Projection::Count => match scan.limit { 172 | // Full count queries are too slow 173 | None => bail!(NOT_SUPPORTED_ERROR), 174 | Some(l) => Ok(iter.take(l).count().await), 175 | }, 176 | } 177 | } 178 | } 179 | -------------------------------------------------------------------------------- /src/redis.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "redis")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::docker::DockerParams; 5 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 6 | use crate::valueprovider::Columns; 7 | use crate::{Benchmark, KeyType, Projection, Scan}; 8 | use anyhow::{Result, bail}; 9 | use futures::StreamExt; 10 | use redis::aio::MultiplexedConnection; 11 | use redis::{AsyncCommands, Client, ScanOptions}; 12 | use serde_json::Value; 13 | use std::hint::black_box; 14 | use tokio::sync::Mutex; 15 | 16 | pub const DEFAULT: &str = "redis://:root@127.0.0.1:6379/"; 17 | 18 | pub(crate) fn docker(options: &Benchmark) -> DockerParams { 19 | DockerParams { 20 | image: "redis", 21 | pre_args: "-p 127.0.0.1:6379:6379".to_string(), 22 | post_args: match (options.persisted, options.sync) { 23 | (false, _) => "redis-server --requirepass root --appendonly no --save ''".to_string(), 24 | (true, false) => { 25 | "redis-server --requirepass root --appendonly yes --appendfsync no".to_string() 26 | } 27 | (true, true) => { 28 | "redis-server --requirepass root --appendonly yes --appendfsync always".to_string() 29 | } 30 | }, 31 | } 32 | } 33 | 34 | pub(crate) struct RedisClientProvider { 35 | url: String, 36 | } 37 | 38 | impl BenchmarkEngine for RedisClientProvider { 39 | /// Initiates a new datastore benchmarking engine 40 | async fn setup(_kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 41 | Ok(Self { 42 | url: options.endpoint.as_deref().unwrap_or(DEFAULT).to_owned(), 43 | }) 44 | } 45 | /// Creates a new client for this benchmarking engine 46 | async fn create_client(&self) -> Result { 47 | let client = Client::open(self.url.as_str())?; 48 | Ok(RedisClient { 49 | conn_iter: Mutex::new(client.get_multiplexed_async_connection().await?), 50 | conn_record: Mutex::new(client.get_multiplexed_async_connection().await?), 51 | }) 52 | } 53 | } 54 | 55 | pub(crate) struct RedisClient { 56 | conn_iter: Mutex, 57 | conn_record: Mutex, 58 | } 59 | 60 | impl BenchmarkClient for RedisClient { 61 | #[allow(dependency_on_unit_never_type_fallback)] 62 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 63 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 64 | let _: () = self.conn_record.lock().await.set(key, val).await?; 65 | Ok(()) 66 | } 67 | 68 | #[allow(dependency_on_unit_never_type_fallback)] 69 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 70 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 71 | let _: () = self.conn_record.lock().await.set(key, val).await?; 72 | Ok(()) 73 | } 74 | 75 | async fn read_u32(&self, key: u32) -> Result<()> { 76 | let val: Vec = self.conn_record.lock().await.get(key).await?; 77 | assert!(!val.is_empty()); 78 | black_box(val); 79 | Ok(()) 80 | } 81 | 82 | #[allow(dependency_on_unit_never_type_fallback)] 83 | async fn read_string(&self, key: String) -> Result<()> { 84 | let val: Vec = self.conn_record.lock().await.get(key).await?; 85 | assert!(!val.is_empty()); 86 | black_box(val); 87 | Ok(()) 88 | } 89 | 90 | #[allow(dependency_on_unit_never_type_fallback)] 91 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 92 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 93 | let _: () = self.conn_record.lock().await.set(key, val).await?; 94 | Ok(()) 95 | } 96 | 97 | #[allow(dependency_on_unit_never_type_fallback)] 98 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 99 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 100 | let _: () = self.conn_record.lock().await.set(key, val).await?; 101 | Ok(()) 102 | } 103 | 104 | #[allow(dependency_on_unit_never_type_fallback)] 105 | async fn delete_u32(&self, key: u32) -> Result<()> { 106 | let _: () = self.conn_record.lock().await.del(key).await?; 107 | Ok(()) 108 | } 109 | 110 | #[allow(dependency_on_unit_never_type_fallback)] 111 | async fn delete_string(&self, key: String) -> Result<()> { 112 | let _: () = self.conn_record.lock().await.del(key).await?; 113 | Ok(()) 114 | } 115 | 116 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 117 | self.scan_bytes(scan).await 118 | } 119 | 120 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 121 | self.scan_bytes(scan).await 122 | } 123 | } 124 | 125 | impl RedisClient { 126 | async fn scan_bytes(&self, scan: &Scan) -> Result { 127 | // Conditional scans are not supported 128 | if scan.condition.is_some() { 129 | bail!(NOT_SUPPORTED_ERROR); 130 | } 131 | // Extract parameters 132 | let s = scan.start.unwrap_or(0); 133 | let l = scan.limit.unwrap_or(usize::MAX); 134 | let p = scan.projection()?; 135 | // Get the two connection types 136 | let mut conn_iter = self.conn_iter.lock().await; 137 | let mut conn_record = self.conn_record.lock().await; 138 | // Configure the scan options for improve iteration 139 | let opts = ScanOptions::default().with_count(5000); 140 | // Create an iterator starting at the beginning 141 | let mut iter = conn_iter.scan_options::(opts).await?.skip(s); 142 | // Perform the relevant projection scan type 143 | match p { 144 | Projection::Id => { 145 | // We use a for loop to iterate over the results, while 146 | // calling black_box internally. This is necessary as 147 | // an iterator with `filter_map` or `map` is optimised 148 | // out by the compiler when calling `count` at the end. 149 | let mut count = 0; 150 | for _ in 0..l { 151 | if let Some(k) = iter.next().await { 152 | black_box(k); 153 | count += 1; 154 | } else { 155 | break; 156 | } 157 | } 158 | Ok(count) 159 | } 160 | Projection::Full => { 161 | // We use a for loop to iterate over the results, while 162 | // calling black_box internally. This is necessary as 163 | // an iterator with `filter_map` or `map` is optimised 164 | // out by the compiler when calling `count` at the end. 165 | let mut count = 0; 166 | while let Some(k) = iter.next().await { 167 | let v: Vec = conn_record.get(k).await?; 168 | black_box(v); 169 | count += 1; 170 | if count >= l { 171 | break; 172 | } 173 | } 174 | Ok(count) 175 | } 176 | Projection::Count => match scan.limit { 177 | // Full count queries are too slow 178 | None => bail!(NOT_SUPPORTED_ERROR), 179 | Some(l) => Ok(iter.take(l).count().await), 180 | }, 181 | } 182 | } 183 | } 184 | -------------------------------------------------------------------------------- /src/keydb.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "keydb")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::docker::DockerParams; 5 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 6 | use crate::valueprovider::Columns; 7 | use crate::{Benchmark, KeyType, Projection, Scan}; 8 | use anyhow::{Result, bail}; 9 | use futures::StreamExt; 10 | use redis::aio::MultiplexedConnection; 11 | use redis::{AsyncCommands, Client, ScanOptions}; 12 | use serde_json::Value; 13 | use std::hint::black_box; 14 | use tokio::sync::Mutex; 15 | 16 | pub const DEFAULT: &str = "redis://:root@127.0.0.1:6379/"; 17 | 18 | pub(crate) fn docker(options: &Benchmark) -> DockerParams { 19 | DockerParams { 20 | image: "eqalpha/keydb", 21 | pre_args: "-p 127.0.0.1:6379:6379".to_string(), 22 | post_args: match (options.persisted, options.sync) { 23 | (false, _) => "keydb-server --requirepass root --appendonly no --save ''".to_string(), 24 | (true, false) => { 25 | "keydb-server --requirepass root --appendonly yes --appendfsync no".to_string() 26 | } 27 | (true, true) => { 28 | "keydb-server --requirepass root --appendonly yes --appendfsync always".to_string() 29 | } 30 | }, 31 | } 32 | } 33 | 34 | pub(crate) struct KeydbClientProvider { 35 | url: String, 36 | } 37 | 38 | impl BenchmarkEngine for KeydbClientProvider { 39 | /// Initiates a new datastore benchmarking engine 40 | async fn setup(_kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 41 | Ok(KeydbClientProvider { 42 | url: options.endpoint.as_deref().unwrap_or(DEFAULT).to_owned(), 43 | }) 44 | } 45 | /// Creates a new client for this benchmarking engine 46 | async fn create_client(&self) -> Result { 47 | let client = Client::open(self.url.as_str())?; 48 | let conn_record = Mutex::new(client.get_multiplexed_async_connection().await?); 49 | let conn_iter = Mutex::new(client.get_multiplexed_async_connection().await?); 50 | Ok(KeydbClient { 51 | conn_record, 52 | conn_iter, 53 | }) 54 | } 55 | } 56 | 57 | pub(crate) struct KeydbClient { 58 | conn_record: Mutex, 59 | conn_iter: Mutex, 60 | } 61 | 62 | impl BenchmarkClient for KeydbClient { 63 | #[allow(dependency_on_unit_never_type_fallback)] 64 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 65 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 66 | let _: () = self.conn_record.lock().await.set(key, val).await?; 67 | Ok(()) 68 | } 69 | 70 | #[allow(dependency_on_unit_never_type_fallback)] 71 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 72 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 73 | let _: () = self.conn_record.lock().await.set(key, val).await?; 74 | Ok(()) 75 | } 76 | 77 | #[allow(dependency_on_unit_never_type_fallback)] 78 | async fn read_u32(&self, key: u32) -> Result<()> { 79 | let val: Vec = self.conn_record.lock().await.get(key).await?; 80 | assert!(!val.is_empty()); 81 | black_box(val); 82 | Ok(()) 83 | } 84 | 85 | #[allow(dependency_on_unit_never_type_fallback)] 86 | async fn read_string(&self, key: String) -> Result<()> { 87 | let val: Vec = self.conn_record.lock().await.get(key).await?; 88 | assert!(!val.is_empty()); 89 | black_box(val); 90 | Ok(()) 91 | } 92 | 93 | #[allow(dependency_on_unit_never_type_fallback)] 94 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 95 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 96 | let _: () = self.conn_record.lock().await.set(key, val).await?; 97 | Ok(()) 98 | } 99 | 100 | #[allow(dependency_on_unit_never_type_fallback)] 101 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 102 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 103 | let _: () = self.conn_record.lock().await.set(key, val).await?; 104 | Ok(()) 105 | } 106 | 107 | #[allow(dependency_on_unit_never_type_fallback)] 108 | async fn delete_u32(&self, key: u32) -> Result<()> { 109 | let _: () = self.conn_record.lock().await.del(key).await?; 110 | Ok(()) 111 | } 112 | 113 | #[allow(dependency_on_unit_never_type_fallback)] 114 | async fn delete_string(&self, key: String) -> Result<()> { 115 | let _: () = self.conn_record.lock().await.del(key).await?; 116 | Ok(()) 117 | } 118 | 119 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 120 | self.scan_bytes(scan).await 121 | } 122 | 123 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 124 | self.scan_bytes(scan).await 125 | } 126 | } 127 | 128 | impl KeydbClient { 129 | async fn scan_bytes(&self, scan: &Scan) -> Result { 130 | // Conditional scans are not supported 131 | if scan.condition.is_some() { 132 | bail!(NOT_SUPPORTED_ERROR); 133 | } 134 | // Extract parameters 135 | let s = scan.start.unwrap_or(0); 136 | let l = scan.limit.unwrap_or(usize::MAX); 137 | let p = scan.projection()?; 138 | // Get the two connection types 139 | let mut conn_iter = self.conn_iter.lock().await; 140 | let mut conn_record = self.conn_record.lock().await; 141 | // Configure the scan options for improve iteration 142 | let opts = ScanOptions::default().with_count(5000); 143 | // Create an iterator starting at the beginning 144 | let mut iter = conn_iter.scan_options::(opts).await?.skip(s); 145 | // Perform the relevant projection scan type 146 | match p { 147 | Projection::Id => { 148 | // We use a for loop to iterate over the results, while 149 | // calling black_box internally. This is necessary as 150 | // an iterator with `filter_map` or `map` is optimised 151 | // out by the compiler when calling `count` at the end. 152 | let mut count = 0; 153 | for _ in 0..l { 154 | if let Some(k) = iter.next().await { 155 | black_box(k); 156 | count += 1; 157 | } else { 158 | break; 159 | } 160 | } 161 | Ok(count) 162 | } 163 | Projection::Full => { 164 | // We use a for loop to iterate over the results, while 165 | // calling black_box internally. This is necessary as 166 | // an iterator with `filter_map` or `map` is optimised 167 | // out by the compiler when calling `count` at the end. 168 | let mut count = 0; 169 | while let Some(k) = iter.next().await { 170 | let v: Vec = conn_record.get(k).await?; 171 | black_box(v); 172 | count += 1; 173 | if count >= l { 174 | break; 175 | } 176 | } 177 | Ok(count) 178 | } 179 | Projection::Count => match scan.limit { 180 | // Full count queries are too slow 181 | None => bail!(NOT_SUPPORTED_ERROR), 182 | Some(l) => Ok(iter.take(l).count().await), 183 | }, 184 | } 185 | } 186 | } 187 | -------------------------------------------------------------------------------- /src/scylladb.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "scylladb")] 2 | 3 | use crate::dialect::AnsiSqlDialect; 4 | use crate::docker::DockerParams; 5 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 6 | use crate::valueprovider::{ColumnType, Columns}; 7 | use crate::{Benchmark, KeyType, Projection, Scan}; 8 | use anyhow::Result; 9 | use futures::StreamExt; 10 | use scylla::_macro_internal::SerializeValue; 11 | use scylla::client::{PoolSize, session::Session, session_builder::SessionBuilder}; 12 | use serde_json::Value; 13 | use std::hint::black_box; 14 | use std::num::NonZeroUsize; 15 | 16 | pub const DEFAULT: &str = "127.0.0.1:9042"; 17 | 18 | pub(crate) fn docker(options: &Benchmark) -> DockerParams { 19 | DockerParams { 20 | image: "scylladb/scylla", 21 | pre_args: match options.sync { 22 | true => { 23 | "-p 9042:9042 -e SCYLLA_ARGS='--commitlog-sync=batch --commitlog-sync-batch-window-in-ms=1'".to_string() 24 | } 25 | false => { 26 | "-p 9042:9042 -e SCYLLA_ARGS='--commitlog-sync=periodic --commitlog-sync-period-in-ms=1000'".to_string() 27 | } 28 | }, 29 | post_args: "".to_string(), 30 | } 31 | } 32 | 33 | pub(crate) struct ScyllaDBClientProvider(KeyType, Columns, String); 34 | 35 | impl BenchmarkEngine for ScyllaDBClientProvider { 36 | /// Initiates a new datastore benchmarking engine 37 | async fn setup(kt: KeyType, columns: Columns, options: &Benchmark) -> Result { 38 | let url = options.endpoint.as_deref().unwrap_or(DEFAULT).to_owned(); 39 | Ok(ScyllaDBClientProvider(kt, columns, url)) 40 | } 41 | /// Creates a new client for this benchmarking engine 42 | async fn create_client(&self) -> Result { 43 | let session = SessionBuilder::new() 44 | .pool_size(PoolSize::PerHost(NonZeroUsize::new(1).unwrap())) 45 | .known_node(&self.2) 46 | .tcp_nodelay(true) 47 | .build() 48 | .await?; 49 | Ok(ScylladbClient { 50 | session, 51 | kt: self.0, 52 | columns: self.1.clone(), 53 | }) 54 | } 55 | } 56 | 57 | pub(crate) struct ScylladbClient { 58 | session: Session, 59 | kt: KeyType, 60 | columns: Columns, 61 | } 62 | 63 | impl BenchmarkClient for ScylladbClient { 64 | async fn startup(&self) -> Result<()> { 65 | self.session 66 | .query_unpaged( 67 | " 68 | CREATE KEYSPACE bench 69 | WITH replication = { 'class': 'SimpleStrategy', 'replication_factor' : 1 } 70 | AND durable_writes = true 71 | ", 72 | (), 73 | ) 74 | .await?; 75 | let id_type = match self.kt { 76 | KeyType::Integer => "INT", 77 | KeyType::String26 | KeyType::String90 | KeyType::String250 | KeyType::String506 => { 78 | "TEXT" 79 | } 80 | KeyType::Uuid => { 81 | todo!() 82 | } 83 | }; 84 | let fields: Vec = self 85 | .columns 86 | .0 87 | .iter() 88 | .map(|(n, t)| match t { 89 | ColumnType::String => format!("{n} TEXT"), 90 | ColumnType::Integer => format!("{n} INT"), 91 | ColumnType::Object => format!("{n} TEXT"), 92 | ColumnType::Float => format!("{n} FLOAT"), 93 | ColumnType::DateTime => format!("{n} TIMESTAMP"), 94 | ColumnType::Uuid => format!("{n} UUID"), 95 | ColumnType::Bool => format!("{n} BOOLEAN"), 96 | }) 97 | .collect(); 98 | let fields = fields.join(","); 99 | self.session 100 | .query_unpaged( 101 | format!("CREATE TABLE bench.record ( id {id_type} PRIMARY KEY, {fields})"), 102 | (), 103 | ) 104 | .await?; 105 | Ok(()) 106 | } 107 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 108 | self.create(key as i32, val).await 109 | } 110 | 111 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 112 | self.create(key, val).await 113 | } 114 | 115 | async fn read_u32(&self, key: u32) -> Result<()> { 116 | self.read(key as i32).await 117 | } 118 | 119 | async fn read_string(&self, key: String) -> Result<()> { 120 | self.read(key).await 121 | } 122 | 123 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 124 | self.scan(scan).await 125 | } 126 | 127 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 128 | self.scan(scan).await 129 | } 130 | 131 | #[allow(dependency_on_unit_never_type_fallback)] 132 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 133 | self.update(key as i32, val).await 134 | } 135 | 136 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 137 | self.update(key, val).await 138 | } 139 | 140 | #[allow(dependency_on_unit_never_type_fallback)] 141 | async fn delete_u32(&self, key: u32) -> Result<()> { 142 | self.delete(key as i32).await 143 | } 144 | 145 | async fn delete_string(&self, key: String) -> Result<()> { 146 | self.delete(key).await 147 | } 148 | } 149 | 150 | impl ScylladbClient { 151 | async fn create(&self, key: T, val: Value) -> Result<()> 152 | where 153 | T: SerializeValue, 154 | { 155 | let (fields, values) = AnsiSqlDialect::create_clause(&self.columns, val); 156 | let stm = format!("INSERT INTO bench.record (id, {fields}) VALUES (?, {values})"); 157 | self.session.query_unpaged(stm, (&key,)).await?; 158 | Ok(()) 159 | } 160 | 161 | async fn read(&self, key: T) -> Result<()> 162 | where 163 | T: SerializeValue, 164 | { 165 | let stm = "SELECT * FROM bench.record WHERE id=?"; 166 | let res = self.session.query_unpaged(stm, (&key,)).await?; 167 | assert_eq!(res.into_rows_result()?.rows_num(), 1); 168 | Ok(()) 169 | } 170 | 171 | async fn update(&self, key: T, val: Value) -> Result<()> 172 | where 173 | T: SerializeValue, 174 | { 175 | let fields = AnsiSqlDialect::update_clause(&self.columns, val); 176 | let stm = format!("UPDATE bench.record SET {fields} WHERE id=?"); 177 | self.session.query_unpaged(stm, (&key,)).await?; 178 | Ok(()) 179 | } 180 | 181 | async fn delete(&self, key: T) -> Result<()> 182 | where 183 | T: SerializeValue, 184 | { 185 | let stm = "DELETE FROM bench.record WHERE id=?"; 186 | self.session.query_unpaged(stm, (&key,)).await?; 187 | Ok(()) 188 | } 189 | 190 | async fn scan(&self, scan: &Scan) -> Result { 191 | // Extract parameters 192 | let s = scan.start.unwrap_or_default(); 193 | let l = scan.limit.map(|l| format!("LIMIT {}", l + s)).unwrap_or_default(); 194 | let c = AnsiSqlDialect::filter_clause(scan)?; 195 | let p = scan.projection()?; 196 | // Perform the relevant projection scan type 197 | match p { 198 | Projection::Id => { 199 | let stm = format!("SELECT id FROM bench.record {c} {l}"); 200 | let mut res = self.session.query_iter(stm, ()).await?.rows_stream()?.skip(s); 201 | let mut count = 0; 202 | while let Some(v) = res.next().await { 203 | let v: (String,) = v?; 204 | black_box(v); 205 | count += 1; 206 | } 207 | Ok(count) 208 | } 209 | Projection::Full => { 210 | let stm = format!("SELECT id FROM bench.record {c} {l}"); 211 | let mut res = self.session.query_iter(stm, ()).await?.rows_stream()?.skip(s); 212 | let mut count = 0; 213 | while let Some(v) = res.next().await { 214 | let v: (String,) = v?; 215 | black_box(v); 216 | count += 1; 217 | } 218 | Ok(count) 219 | } 220 | Projection::Count => { 221 | let stm = format!("SELECT count(*) FROM bench.record {c} {l}"); 222 | let mut res = self.session.query_iter(stm, ()).await?.rows_stream()?.skip(s); 223 | let count: (String,) = res.next().await.unwrap()?; 224 | let count: usize = count.0.parse()?; 225 | Ok(count) 226 | } 227 | } 228 | } 229 | } 230 | -------------------------------------------------------------------------------- /src/keyprovider.rs: -------------------------------------------------------------------------------- 1 | use crate::KeyType; 2 | use twox_hash::XxHash64; 3 | 4 | #[derive(Clone, Copy)] 5 | pub(crate) enum KeyProvider { 6 | OrderedInteger(OrderedInteger), 7 | UnorderedInteger(UnorderedInteger), 8 | OrderedString(OrderedString), 9 | UnorderedString(UnorderedString), 10 | } 11 | 12 | impl KeyProvider { 13 | pub(crate) fn new(key_type: KeyType, random: bool) -> Self { 14 | match key_type { 15 | KeyType::Integer => { 16 | if random { 17 | Self::UnorderedInteger(UnorderedInteger::default()) 18 | } else { 19 | Self::OrderedInteger(OrderedInteger::default()) 20 | } 21 | } 22 | KeyType::String26 => { 23 | if random { 24 | Self::UnorderedString(UnorderedString::new(1)) 25 | } else { 26 | Self::OrderedString(OrderedString::new(1)) 27 | } 28 | } 29 | KeyType::String90 => { 30 | if random { 31 | Self::UnorderedString(UnorderedString::new(5)) 32 | } else { 33 | Self::OrderedString(OrderedString::new(5)) 34 | } 35 | } 36 | KeyType::String250 => { 37 | if random { 38 | Self::UnorderedString(UnorderedString::new(15)) 39 | } else { 40 | Self::OrderedString(OrderedString::new(15)) 41 | } 42 | } 43 | KeyType::String506 => { 44 | if random { 45 | Self::UnorderedString(UnorderedString::new(31)) 46 | } else { 47 | Self::OrderedString(OrderedString::new(31)) 48 | } 49 | } 50 | KeyType::Uuid => { 51 | todo!() 52 | } 53 | } 54 | } 55 | } 56 | 57 | pub(crate) trait IntegerKeyProvider: Send { 58 | fn key(&mut self, n: u32) -> u32; 59 | } 60 | 61 | pub(crate) trait StringKeyProvider: Send { 62 | fn key(&mut self, n: u32) -> String; 63 | } 64 | 65 | #[derive(Default, Clone, Copy)] 66 | pub(crate) struct OrderedInteger(); 67 | 68 | impl IntegerKeyProvider for OrderedInteger { 69 | fn key(&mut self, n: u32) -> u32 { 70 | // We need to increment by 1 71 | // because MySQL PRIMARY IDs 72 | // can not be 0, resulting in 73 | // duplicate ID errors. 74 | n + 1 75 | } 76 | } 77 | #[derive(Default, Clone, Copy)] 78 | pub(crate) struct UnorderedInteger(); 79 | 80 | impl IntegerKeyProvider for UnorderedInteger { 81 | fn key(&mut self, n: u32) -> u32 { 82 | Self::feistel_transform(n) 83 | } 84 | } 85 | 86 | impl UnorderedInteger { 87 | // A very simple round function: XOR the input with the key and shift 88 | fn feistel_round_function(value: u32, key: u32) -> u32 { 89 | (value ^ key).rotate_left(5).wrapping_add(key) 90 | } 91 | 92 | // Perform one round of the Feistel network 93 | fn feistel_round(left: u16, right: u16, round_key: u32) -> (u16, u16) { 94 | let new_left = right; 95 | let new_right = left ^ (Self::feistel_round_function(right as u32, round_key) as u16); 96 | (new_left, new_right) 97 | } 98 | 99 | fn feistel_transform(input: u32) -> u32 { 100 | let mut left = (input >> 16) as u16; 101 | let mut right = (input & 0xFFFF) as u16; 102 | 103 | // Hard-coded keys for simplicity 104 | let keys = [0xA5A5A5A5, 0x5A5A5A5A, 0x3C3C3C3C]; 105 | 106 | for &key in &keys { 107 | let (new_left, new_right) = Self::feistel_round(left, right, key); 108 | left = new_left; 109 | right = new_right; 110 | } 111 | 112 | // Combine left and right halves back into a single u32 113 | ((left as u32) << 16) | (right as u32) 114 | } 115 | } 116 | 117 | fn hash_string(n: u32, repeat: usize) -> String { 118 | let mut hex_string = String::with_capacity(repeat * 16 + 10); 119 | for s in 0..repeat as u64 { 120 | let hash_result = XxHash64::oneshot(s, &n.to_be_bytes()); 121 | hex_string.push_str(&format!("{hash_result:x}")); 122 | } 123 | hex_string 124 | } 125 | 126 | #[derive(Clone, Copy)] 127 | pub(crate) struct OrderedString(usize); 128 | 129 | impl OrderedString { 130 | fn new(repeat: usize) -> Self { 131 | Self(repeat) 132 | } 133 | } 134 | 135 | impl StringKeyProvider for OrderedString { 136 | fn key(&mut self, n: u32) -> String { 137 | let hex_string = hash_string(n, self.0); 138 | format!("{n:010}{hex_string}") 139 | } 140 | } 141 | 142 | #[derive(Default, Clone, Copy)] 143 | pub(crate) struct UnorderedString(usize); 144 | 145 | impl UnorderedString { 146 | fn new(repeat: usize) -> Self { 147 | Self(repeat) 148 | } 149 | } 150 | 151 | impl StringKeyProvider for UnorderedString { 152 | fn key(&mut self, n: u32) -> String { 153 | let hex_string = hash_string(n, self.0); 154 | format!("{hex_string}{n:010}") 155 | } 156 | } 157 | 158 | #[cfg(test)] 159 | mod test { 160 | use crate::keyprovider::{OrderedString, StringKeyProvider, UnorderedString}; 161 | 162 | #[test] 163 | fn ordered_string_26() { 164 | let mut o = OrderedString::new(1); 165 | let s = o.key(12345678); 166 | assert_eq!(s.len(), 26); 167 | assert_eq!(s, "0012345678d79235c904e704c6"); 168 | } 169 | 170 | #[test] 171 | fn unordered_string_26() { 172 | let mut o = UnorderedString::new(1); 173 | let s = o.key(12345678); 174 | assert_eq!(s.len(), 26); 175 | assert_eq!(s, "d79235c904e704c60012345678"); 176 | } 177 | 178 | #[test] 179 | fn ordered_string_90() { 180 | let mut o = OrderedString::new(5); 181 | let s = o.key(12345678); 182 | assert_eq!(s.len(), 90); 183 | assert_eq!( 184 | s, 185 | "0012345678d79235c904e704c6c379c25fea98cd11b4d0f71900f91df2ecc87c25d7fff4b03be1bd13590485d3" 186 | ); 187 | } 188 | 189 | #[test] 190 | fn unordered_string_90() { 191 | let mut o = UnorderedString::new(5); 192 | let s = o.key(12345678); 193 | assert_eq!(s.len(), 90); 194 | assert_eq!( 195 | s, 196 | "d79235c904e704c6c379c25fea98cd11b4d0f71900f91df2ecc87c25d7fff4b03be1bd13590485d30012345678" 197 | ); 198 | } 199 | 200 | #[test] 201 | fn ordered_string_250() { 202 | let mut o = OrderedString::new(15); 203 | let s = o.key(12345678); 204 | assert_eq!(s.len(), 250); 205 | assert_eq!( 206 | s, 207 | "0012345678d79235c904e704c6c379c25fea98cd11b4d0f71900f91df2ecc87c25d7fff4b03be1bd13590485d31bc0feb2815d5c908f5a4633b8a9d5d6ec1c074d5d64ab296c6495f784f8294ac42b828a9c4ef45d3decc0a8dff00062adfb547fea6132f38afda36acf629cc15413acfe35a50fecbec285e9ee42b136" 208 | ); 209 | } 210 | 211 | #[test] 212 | fn unordered_string_250() { 213 | let mut o = UnorderedString::new(15); 214 | let s = o.key(12345678); 215 | assert_eq!(s.len(), 250); 216 | assert_eq!( 217 | s, 218 | "d79235c904e704c6c379c25fea98cd11b4d0f71900f91df2ecc87c25d7fff4b03be1bd13590485d31bc0feb2815d5c908f5a4633b8a9d5d6ec1c074d5d64ab296c6495f784f8294ac42b828a9c4ef45d3decc0a8dff00062adfb547fea6132f38afda36acf629cc15413acfe35a50fecbec285e9ee42b1360012345678" 219 | ); 220 | } 221 | 222 | #[test] 223 | fn ordered_string_506() { 224 | let mut o = OrderedString::new(31); 225 | let s = o.key(12345678); 226 | assert_eq!(s.len(), 506); 227 | assert_eq!( 228 | s, 229 | "0012345678d79235c904e704c6c379c25fea98cd11b4d0f71900f91df2ecc87c25d7fff4b03be1bd13590485d31bc0feb2815d5c908f5a4633b8a9d5d6ec1c074d5d64ab296c6495f784f8294ac42b828a9c4ef45d3decc0a8dff00062adfb547fea6132f38afda36acf629cc15413acfe35a50fecbec285e9ee42b13691088df6c3740c87c3d003e3addf1888a582ac5cb408feec138fe9a43c9fda574006e770bb0b5e84edcbeecc6f723960ed7d02591a7b2487bb317f83bfd95e44a69d957deb6b10e22d895a375acfa54143137feeb53921625bc9d582166477e562454fecc90f130662338c070bd709c27d8478abaa825dc69bc3aa89dc7ce076" 230 | ); 231 | } 232 | 233 | #[test] 234 | fn unordered_string_506() { 235 | let mut o = UnorderedString::new(31); 236 | let s = o.key(12345678); 237 | assert_eq!(s.len(), 506); 238 | assert_eq!( 239 | s, 240 | "d79235c904e704c6c379c25fea98cd11b4d0f71900f91df2ecc87c25d7fff4b03be1bd13590485d31bc0feb2815d5c908f5a4633b8a9d5d6ec1c074d5d64ab296c6495f784f8294ac42b828a9c4ef45d3decc0a8dff00062adfb547fea6132f38afda36acf629cc15413acfe35a50fecbec285e9ee42b13691088df6c3740c87c3d003e3addf1888a582ac5cb408feec138fe9a43c9fda574006e770bb0b5e84edcbeecc6f723960ed7d02591a7b2487bb317f83bfd95e44a69d957deb6b10e22d895a375acfa54143137feeb53921625bc9d582166477e562454fecc90f130662338c070bd709c27d8478abaa825dc69bc3aa89dc7ce0760012345678" 241 | ); 242 | } 243 | } 244 | -------------------------------------------------------------------------------- /src/map.rs: -------------------------------------------------------------------------------- 1 | use crate::benchmark::NOT_SUPPORTED_ERROR; 2 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 3 | use crate::valueprovider::Columns; 4 | use crate::{Benchmark, KeyType, Projection, Scan}; 5 | use anyhow::{Result, bail}; 6 | use dashmap::DashMap; 7 | use serde_json::Value; 8 | use std::hash::Hash; 9 | use std::hint::black_box; 10 | use std::sync::Arc; 11 | use std::time::Duration; 12 | 13 | #[derive(Clone)] 14 | pub(crate) enum MapDatabase { 15 | Integer(Arc>), 16 | String(Arc>), 17 | } 18 | 19 | impl From for MapDatabase { 20 | fn from(t: KeyType) -> Self { 21 | match t { 22 | KeyType::Integer => Self::Integer(DashMap::new().into()), 23 | KeyType::String26 | KeyType::String90 | KeyType::String250 | KeyType::String506 => { 24 | Self::String(DashMap::new().into()) 25 | } 26 | KeyType::Uuid => todo!(), 27 | } 28 | } 29 | } 30 | 31 | pub(crate) struct MapClientProvider(MapDatabase); 32 | 33 | impl BenchmarkEngine for MapClientProvider { 34 | /// The number of seconds to wait before connecting 35 | fn wait_timeout(&self) -> Option { 36 | None 37 | } 38 | /// Initiates a new datastore benchmarking engine 39 | async fn setup(kt: KeyType, _columns: Columns, _options: &Benchmark) -> Result { 40 | Ok(Self(kt.into())) 41 | } 42 | /// Creates a new client for this benchmarking engine 43 | async fn create_client(&self) -> Result { 44 | Ok(MapClient(self.0.clone())) 45 | } 46 | } 47 | 48 | pub(crate) struct MapClient(MapDatabase); 49 | 50 | impl BenchmarkClient for MapClient { 51 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 52 | if let MapDatabase::Integer(m) = &self.0 { 53 | assert!(m.insert(key, val).is_none()); 54 | } else { 55 | bail!("Invalid MapDatabase variant"); 56 | } 57 | Ok(()) 58 | } 59 | 60 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 61 | if let MapDatabase::String(m) = &self.0 { 62 | assert!(m.insert(key, val).is_none()); 63 | } else { 64 | bail!("Invalid MapDatabase variant"); 65 | } 66 | Ok(()) 67 | } 68 | 69 | async fn read_u32(&self, key: u32) -> Result<()> { 70 | if let MapDatabase::Integer(m) = &self.0 { 71 | assert!(m.get(&key).is_some()); 72 | } else { 73 | bail!("Invalid MapDatabase variant"); 74 | } 75 | Ok(()) 76 | } 77 | 78 | async fn read_string(&self, key: String) -> Result<()> { 79 | if let MapDatabase::String(m) = &self.0 { 80 | assert!(m.get(&key).is_some()); 81 | } else { 82 | bail!("Invalid MapDatabase variant"); 83 | } 84 | Ok(()) 85 | } 86 | 87 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 88 | if let MapDatabase::Integer(m) = &self.0 { 89 | assert!(m.insert(key, val).is_some()); 90 | } else { 91 | bail!("Invalid MapDatabase variant"); 92 | } 93 | Ok(()) 94 | } 95 | 96 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 97 | if let MapDatabase::String(m) = &self.0 { 98 | assert!(m.insert(key, val).is_some()); 99 | } else { 100 | bail!("Invalid MapDatabase variant"); 101 | } 102 | Ok(()) 103 | } 104 | 105 | async fn delete_u32(&self, key: u32) -> Result<()> { 106 | if let MapDatabase::Integer(m) = &self.0 { 107 | assert!(m.remove(&key).is_some()); 108 | } else { 109 | bail!("Invalid MapDatabase variant"); 110 | } 111 | Ok(()) 112 | } 113 | 114 | async fn delete_string(&self, key: String) -> Result<()> { 115 | if let MapDatabase::String(m) = &self.0 { 116 | assert!(m.remove(&key).is_some()); 117 | } else { 118 | bail!("Invalid MapDatabase variant"); 119 | } 120 | Ok(()) 121 | } 122 | 123 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 124 | if let MapDatabase::Integer(m) = &self.0 { 125 | Self::scan(m, scan).await 126 | } else { 127 | bail!("Invalid MapDatabase variant"); 128 | } 129 | } 130 | 131 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 132 | if let MapDatabase::String(m) = &self.0 { 133 | Self::scan(m, scan).await 134 | } else { 135 | bail!("Invalid MapDatabase variant"); 136 | } 137 | } 138 | 139 | async fn batch_create_u32( 140 | &self, 141 | key_vals: impl Iterator + Send, 142 | ) -> Result<()> { 143 | if let MapDatabase::Integer(m) = &self.0 { 144 | for (key, val) in key_vals { 145 | assert!(m.insert(key, val).is_none()); 146 | } 147 | } else { 148 | bail!("Invalid MapDatabase variant"); 149 | } 150 | Ok(()) 151 | } 152 | 153 | async fn batch_create_string( 154 | &self, 155 | key_vals: impl Iterator + Send, 156 | ) -> Result<()> { 157 | if let MapDatabase::String(m) = &self.0 { 158 | for (key, val) in key_vals { 159 | assert!(m.insert(key, val).is_none()); 160 | } 161 | } else { 162 | bail!("Invalid MapDatabase variant"); 163 | } 164 | Ok(()) 165 | } 166 | 167 | async fn batch_read_u32(&self, keys: impl Iterator + Send) -> Result<()> { 168 | if let MapDatabase::Integer(m) = &self.0 { 169 | for key in keys { 170 | assert!(m.get(&key).is_some()); 171 | } 172 | } else { 173 | bail!("Invalid MapDatabase variant"); 174 | } 175 | Ok(()) 176 | } 177 | 178 | async fn batch_read_string(&self, keys: impl Iterator + Send) -> Result<()> { 179 | if let MapDatabase::String(m) = &self.0 { 180 | for key in keys { 181 | assert!(m.get(&key).is_some()); 182 | } 183 | } else { 184 | bail!("Invalid MapDatabase variant"); 185 | } 186 | Ok(()) 187 | } 188 | 189 | async fn batch_update_u32( 190 | &self, 191 | key_vals: impl Iterator + Send, 192 | ) -> Result<()> { 193 | if let MapDatabase::Integer(m) = &self.0 { 194 | for (key, val) in key_vals { 195 | assert!(m.insert(key, val).is_some()); 196 | } 197 | } else { 198 | bail!("Invalid MapDatabase variant"); 199 | } 200 | Ok(()) 201 | } 202 | 203 | async fn batch_update_string( 204 | &self, 205 | key_vals: impl Iterator + Send, 206 | ) -> Result<()> { 207 | if let MapDatabase::String(m) = &self.0 { 208 | for (key, val) in key_vals { 209 | assert!(m.insert(key, val).is_some()); 210 | } 211 | } else { 212 | bail!("Invalid MapDatabase variant"); 213 | } 214 | Ok(()) 215 | } 216 | 217 | async fn batch_delete_u32(&self, keys: impl Iterator + Send) -> Result<()> { 218 | if let MapDatabase::Integer(m) = &self.0 { 219 | for key in keys { 220 | assert!(m.remove(&key).is_some()); 221 | } 222 | } else { 223 | bail!("Invalid MapDatabase variant"); 224 | } 225 | Ok(()) 226 | } 227 | 228 | async fn batch_delete_string(&self, keys: impl Iterator + Send) -> Result<()> { 229 | if let MapDatabase::String(m) = &self.0 { 230 | for key in keys { 231 | assert!(m.remove(&key).is_some()); 232 | } 233 | } else { 234 | bail!("Invalid MapDatabase variant"); 235 | } 236 | Ok(()) 237 | } 238 | } 239 | 240 | impl MapClient { 241 | async fn scan(m: &DashMap, scan: &Scan) -> Result 242 | where 243 | T: Eq + Hash, 244 | { 245 | // Contional scans are not supported 246 | if scan.condition.is_some() { 247 | bail!(NOT_SUPPORTED_ERROR); 248 | } 249 | // Extract parameters 250 | let s = scan.start.unwrap_or(0); 251 | let l = scan.limit.unwrap_or(usize::MAX); 252 | let p = scan.projection()?; 253 | // Perform the relevant projection scan type 254 | match p { 255 | Projection::Id => { 256 | // We use a for loop to iterate over the results, while 257 | // calling black_box internally. This is necessary as 258 | // an iterator with `filter_map` or `map` is optimised 259 | // out by the compiler when calling `count` at the end. 260 | let mut count = 0; 261 | for v in m.iter().skip(s).take(l) { 262 | black_box(v); 263 | count += 1; 264 | } 265 | Ok(count) 266 | } 267 | Projection::Full => { 268 | // We use a for loop to iterate over the results, while 269 | // calling black_box internally. This is necessary as 270 | // an iterator with `filter_map` or `map` is optimised 271 | // out by the compiler when calling `count` at the end. 272 | let mut count = 0; 273 | for v in m.iter().skip(s).take(l) { 274 | black_box(v); 275 | count += 1; 276 | } 277 | Ok(count) 278 | } 279 | Projection::Count => Ok(m 280 | .iter() 281 | .skip(s) // Skip the first `offset` entries 282 | .take(l) // Take the next `limit` entries 283 | .count()), 284 | } 285 | } 286 | } 287 | -------------------------------------------------------------------------------- /src/arangodb.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "arangodb")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::dialect::ArangoDBDialect; 5 | use crate::docker::DockerParams; 6 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 7 | use crate::valueprovider::Columns; 8 | use crate::{Benchmark, KeyType, Projection, Scan}; 9 | use anyhow::{Result, bail}; 10 | use arangors::client::reqwest::ReqwestClient; 11 | use arangors::document::options::InsertOptions; 12 | use arangors::document::options::RemoveOptions; 13 | use arangors::{Collection, Connection, Database, GenericConnection}; 14 | use serde_json::Value; 15 | use std::hint::black_box; 16 | use std::time::Duration; 17 | use tokio::sync::Mutex; 18 | 19 | pub const DEFAULT: &str = "http://127.0.0.1:8529"; 20 | 21 | pub(crate) fn docker(options: &Benchmark) -> DockerParams { 22 | DockerParams { 23 | image: "arangodb", 24 | pre_args: "--ulimit nofile=65536:65536 -p 127.0.0.1:8529:8529 -e ARANGO_NO_AUTH=1".to_string(), 25 | post_args: match options.optimised { 26 | true => "--server.scheduler-queue-size 8192 --server.prio1-size 8192 --server.prio2-size 8192 --server.maximal-queue-size 8192".to_string(), 27 | false => "".to_string(), 28 | }, 29 | } 30 | } 31 | 32 | pub(crate) struct ArangoDBClientProvider { 33 | sync: bool, 34 | key: KeyType, 35 | url: String, 36 | } 37 | 38 | impl BenchmarkEngine for ArangoDBClientProvider { 39 | /// Initiates a new datastore benchmarking engine 40 | async fn setup(kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 41 | Ok(Self { 42 | sync: options.sync, 43 | key: kt, 44 | url: options.endpoint.as_deref().unwrap_or(DEFAULT).to_owned(), 45 | }) 46 | } 47 | /// Creates a new client for this benchmarking engine 48 | async fn create_client(&self) -> Result { 49 | let (conn, db, co) = create_arango_client(&self.url).await?; 50 | Ok(ArangoDBClient { 51 | sync: self.sync, 52 | keytype: self.key, 53 | connection: conn, 54 | database: Mutex::new(db), 55 | collection: Mutex::new(co), 56 | }) 57 | } 58 | /// The number of seconds to wait before connecting 59 | fn wait_timeout(&self) -> Option { 60 | Some(Duration::from_secs(15)) 61 | } 62 | } 63 | 64 | pub(crate) struct ArangoDBClient { 65 | sync: bool, 66 | keytype: KeyType, 67 | connection: GenericConnection, 68 | database: Mutex>, 69 | collection: Mutex>, 70 | } 71 | 72 | async fn create_arango_client( 73 | url: &str, 74 | ) -> Result<(GenericConnection, Database, Collection)> 75 | { 76 | // Create the connection to the database 77 | let conn = Connection::establish_without_auth(url).await.unwrap(); 78 | // Create the benchmarking database 79 | let db = match conn.create_database("crud-bench").await { 80 | Err(_) => conn.db("crud-bench").await.unwrap(), 81 | Ok(db) => db, 82 | }; 83 | // Create the becnhmark record collection 84 | let co = match db.create_collection("record").await { 85 | Err(_) => db.collection("record").await.unwrap(), 86 | Ok(db) => db, 87 | }; 88 | Ok((conn, db, co)) 89 | } 90 | 91 | impl BenchmarkClient for ArangoDBClient { 92 | async fn startup(&self) -> Result<()> { 93 | // Ensure we drop the database first. 94 | // We can drop the database initially 95 | // because the other clients will be 96 | // created subsequently, and will then 97 | // create the database as necessary. 98 | self.connection.drop_database("crud-bench").await?; 99 | // Everything ok 100 | Ok(()) 101 | } 102 | 103 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 104 | match self.keytype { 105 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 106 | _ => self.create(key.to_string(), val).await, 107 | } 108 | } 109 | 110 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 111 | match self.keytype { 112 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 113 | _ => self.create(key, val).await, 114 | } 115 | } 116 | 117 | async fn read_u32(&self, key: u32) -> Result<()> { 118 | match self.keytype { 119 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 120 | _ => self.read(key.to_string()).await, 121 | } 122 | } 123 | 124 | async fn read_string(&self, key: String) -> Result<()> { 125 | match self.keytype { 126 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 127 | _ => self.read(key).await, 128 | } 129 | } 130 | 131 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 132 | match self.keytype { 133 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 134 | _ => self.update(key.to_string(), val).await, 135 | } 136 | } 137 | 138 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 139 | match self.keytype { 140 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 141 | _ => self.update(key, val).await, 142 | } 143 | } 144 | 145 | async fn delete_u32(&self, key: u32) -> Result<()> { 146 | match self.keytype { 147 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 148 | _ => self.delete(key.to_string()).await, 149 | } 150 | } 151 | 152 | async fn delete_string(&self, key: String) -> Result<()> { 153 | match self.keytype { 154 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 155 | _ => self.delete(key).await, 156 | } 157 | } 158 | 159 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 160 | match self.keytype { 161 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 162 | _ => self.scan(scan).await, 163 | } 164 | } 165 | 166 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 167 | match self.keytype { 168 | KeyType::String506 => bail!(NOT_SUPPORTED_ERROR), 169 | _ => self.scan(scan).await, 170 | } 171 | } 172 | } 173 | 174 | impl ArangoDBClient { 175 | fn to_doc(key: String, mut val: Value) -> Result { 176 | let obj = val.as_object_mut().unwrap(); 177 | obj.insert("_key".to_string(), key.into()); 178 | Ok(val) 179 | } 180 | 181 | async fn create(&self, key: String, val: Value) -> Result<()> { 182 | let val = Self::to_doc(key, val)?; 183 | let opt = InsertOptions::builder() 184 | .wait_for_sync(self.sync) 185 | .return_new(true) 186 | .overwrite(false) 187 | .build(); 188 | let res = { self.collection.lock().await.create_document(val, opt).await? }; 189 | assert!(res.new_doc().is_some()); 190 | Ok(()) 191 | } 192 | 193 | async fn read(&self, key: String) -> Result<()> { 194 | let doc = { self.collection.lock().await.document::(&key).await? }; 195 | assert!(doc.is_object()); 196 | assert_eq!(doc.get("_key").unwrap().as_str().unwrap(), key); 197 | Ok(()) 198 | } 199 | 200 | async fn update(&self, key: String, val: Value) -> Result<()> { 201 | let val = Self::to_doc(key, val)?; 202 | let opt = InsertOptions::builder() 203 | .wait_for_sync(self.sync) 204 | .return_new(true) 205 | .overwrite(true) 206 | .build(); 207 | let res = { self.collection.lock().await.create_document(val, opt).await? }; 208 | assert!(res.new_doc().is_some()); 209 | Ok(()) 210 | } 211 | 212 | async fn delete(&self, key: String) -> Result<()> { 213 | let opt = RemoveOptions::builder().wait_for_sync(self.sync).build(); 214 | let res = { self.collection.lock().await.remove_document::(&key, opt, None).await? }; 215 | assert!(res.has_response()); 216 | Ok(()) 217 | } 218 | 219 | async fn scan(&self, scan: &Scan) -> Result { 220 | // Extract parameters 221 | let l = match (scan.start, scan.limit) { 222 | (Some(s), Some(l)) => format!("LIMIT {s}, {l}"), 223 | (Some(s), None) => format!("LIMIT {s}, 1000000000"), 224 | (None, Some(l)) => format!("LIMIT {l}"), 225 | (None, None) => "".to_string(), 226 | }; 227 | let c = ArangoDBDialect::filter_clause(scan)?; 228 | let p = scan.projection()?; 229 | // Perform the relevant projection scan type 230 | match p { 231 | Projection::Id => { 232 | let stm = format!("FOR r IN record {c} {l} RETURN {{ _id: r._id }}"); 233 | let res: Vec = { self.database.lock().await.aql_str(&stm).await.unwrap() }; 234 | // We use a for loop to iterate over the results, while 235 | // calling black_box internally. This is necessary as 236 | // an iterator with `filter_map` or `map` is optimised 237 | // out by the compiler when calling `count` at the end. 238 | let mut count = 0; 239 | for v in res { 240 | black_box(v); 241 | count += 1; 242 | } 243 | Ok(count) 244 | } 245 | Projection::Full => { 246 | let stm = format!("FOR r IN record {c} {l} RETURN r"); 247 | let res: Vec = { self.database.lock().await.aql_str(&stm).await.unwrap() }; 248 | // We use a for loop to iterate over the results, while 249 | // calling black_box internally. This is necessary as 250 | // an iterator with `filter_map` or `map` is optimised 251 | // out by the compiler when calling `count` at the end. 252 | let mut count = 0; 253 | for v in res { 254 | black_box(v); 255 | count += 1; 256 | } 257 | Ok(count) 258 | } 259 | Projection::Count => { 260 | let stm = 261 | format!("FOR r IN record {c} {l} COLLECT WITH COUNT INTO count RETURN count"); 262 | let res: Vec = { self.database.lock().await.aql_str(&stm).await.unwrap() }; 263 | let count = res.first().unwrap().as_i64().unwrap(); 264 | Ok(count as usize) 265 | } 266 | } 267 | } 268 | } 269 | -------------------------------------------------------------------------------- /src/surrealds.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "surrealdb")] 2 | 3 | //! # SurrealDS - Multi-Instance SurrealDB Client Provider 4 | //! 5 | //! This module provides the `SurrealDBClientsProvider`, which enables benchmarking against 6 | //! multiple SurrealDB instances simultaneously. This is particularly useful for testing 7 | //! distributed SurrealDB deployments, load balancing scenarios, and assessing the performance 8 | //! characteristics of multi-node SurrealDB clusters. 9 | //! 10 | //! ## Features 11 | //! 12 | //! - **Multi-Instance Support**: Connect to multiple SurrealDB instances using a single endpoint configuration 13 | //! - **Round-Robin Load Balancing**: Automatically distributes client connections across available instances 14 | //! - **Networked Connections Only**: Designed for remote SurrealDB instances (ws://, wss://, http://, https://) 15 | //! 16 | //! ## Usage 17 | //! 18 | //! The SurrealDS feature is enabled through the `surrealdb` feature flag and requires specifying 19 | //! multiple endpoints in the endpoint configuration, separated by semicolons: 20 | //! 21 | //! ```bash 22 | //! cargo run -r -- -d surrealdb -e "ws://127.0.0.1:8001;ws://127.0.0.1:8002;ws://127.0.0.1:8003" -s 100000 23 | //! ``` 24 | //! 25 | //! ## Architecture 26 | //! 27 | //! The `SurrealDBClientsProvider` implements the `BenchmarkEngine` trait and creates 28 | //! `SurrealDBClient` instances on demand. Each client creation request uses the round-robin 29 | //! algorithm to select the next available endpoint, ensuring even distribution of connections 30 | //! across all configured SurrealDB instances. 31 | 32 | use crate::engine::BenchmarkEngine; 33 | use crate::surrealdb::{SurrealDBClient, initialise_db, surrealdb_password, surrealdb_username}; 34 | use crate::valueprovider::Columns; 35 | use crate::{Benchmark, KeyType}; 36 | use anyhow::{Result, bail}; 37 | use std::sync::atomic::{AtomicUsize, Ordering}; 38 | use surrealdb::opt::auth::Root; 39 | 40 | /// Default endpoints for SurrealDS when no custom endpoint is specified. 41 | /// 42 | /// This configuration assumes three local SurrealDB instances running on ports 8001, 8002, and 8003. 43 | /// In production scenarios, these would typically point to different hosts or distributed nodes. 44 | const DEFAULT: &str = "ws://127.0.0.1:8001;ws://127.0.0.1:8002;ws://127.0.0.1:8003"; 45 | 46 | /// A benchmark engine provider that manages connections to multiple SurrealDB instances. 47 | /// 48 | /// `SurrealDBClientsProvider` enables distributed benchmarking by maintaining a pool of 49 | /// endpoints and creating client connections using a round-robin selection algorithm. 50 | /// This ensures that benchmark workloads are evenly distributed across all configured 51 | /// SurrealDB instances. 52 | /// 53 | /// ## Connection Strategy 54 | /// 55 | /// - Connections are created on-demand for each concurrent client 56 | /// - Round-robin selection ensures even distribution of load 57 | /// - Each endpoint must be a remote connection (ws://, wss://, http://, https://) 58 | /// - All instances must be accessible and properly configured with the same credentials 59 | /// 60 | /// ## Example Configuration 61 | /// 62 | /// ```text 63 | /// endpoints: ["ws://node1:8000", "ws://node2:8000", "ws://node3:8000"] 64 | /// ``` 65 | /// 66 | /// When creating clients, the provider will cycle through endpoints in order: 67 | /// - Client 0 → node1 68 | /// - Client 1 → node2 69 | /// - Client 2 → node3 70 | /// - Client 3 → node1 (wraps around) 71 | pub(crate) struct SurrealDBClientsProvider { 72 | /// Atomic counter for round-robin endpoint selection. 73 | /// 74 | /// This counter is incremented atomically on each client creation to determine 75 | /// the next endpoint to use. The value is taken modulo the number of endpoints 76 | /// to cycle through the available instances. 77 | round_robin: AtomicUsize, 78 | 79 | /// List of SurrealDB endpoints to distribute connections across. 80 | /// 81 | /// Each endpoint should be a valid remote connection string (e.g., "ws://host:port"). 82 | /// The endpoints are parsed from the configuration string by splitting on semicolons. 83 | endpoints: Vec, 84 | 85 | /// Root user authentication credentials. 86 | /// 87 | /// These credentials are used to authenticate with all configured SurrealDB instances. 88 | /// All instances in the cluster must accept the same root credentials. 89 | root: Root, 90 | } 91 | 92 | impl BenchmarkEngine for SurrealDBClientsProvider { 93 | /// Initializes a new multi-instance SurrealDB benchmarking engine. 94 | /// 95 | /// This method sets up the provider by parsing the endpoint configuration string, 96 | /// validating that all endpoints are remote connections, and preparing the 97 | /// round-robin counter for distributing client connections. 98 | /// 99 | /// # Arguments 100 | /// 101 | /// * `_` - Key type (unused, inherited from trait) 102 | /// * `_columns` - Column configuration (unused, inherited from trait) 103 | /// * `options` - Benchmark configuration containing the endpoint string 104 | /// 105 | /// # Endpoint Format 106 | /// 107 | /// The endpoint string should contain one or more SurrealDB endpoints separated 108 | /// by semicolons. Each endpoint must use a remote protocol (ws, wss, http, or https). 109 | /// 110 | /// Example: `"ws://127.0.0.1:8001;ws://127.0.0.1:8002;ws://127.0.0.1:8003"` 111 | /// 112 | /// # Returns 113 | /// 114 | /// Returns a configured `SurrealDBClientsProvider` ready to create client connections, 115 | /// or an error if: 116 | /// - The endpoint configuration is empty 117 | /// - Any endpoint uses an unsupported protocol (e.g., embedded databases) 118 | /// 119 | /// # Errors 120 | /// 121 | /// This function will return an error if: 122 | /// - The endpoint string is empty or contains no valid endpoints 123 | /// - Any endpoint uses a non-remote protocol (mem://, file://, rocksdb://, surrealkv://) 124 | async fn setup(_: KeyType, _columns: Columns, options: &Benchmark) -> Result { 125 | // Get the custom endpoint configuration from options, or use the default 126 | // three-instance local setup if no endpoint is specified 127 | let endpoint = options.endpoint.as_deref().unwrap_or(DEFAULT); 128 | 129 | // Define root user authentication credentials from environment variables or use defaults 130 | // All configured SurrealDB instances must accept these credentials 131 | let root = Root { 132 | username: surrealdb_username(), 133 | password: surrealdb_password(), 134 | }; 135 | 136 | // Parse and validate endpoints from the semicolon-separated string 137 | let mut endpoints = Vec::new(); 138 | for e in endpoint.split(';') { 139 | // Skip empty segments that result from trailing or double semicolons 140 | if e.is_empty() { 141 | continue; 142 | } 143 | 144 | // Validate that each endpoint uses a remote connection protocol 145 | // SurrealDS only supports networked instances, not embedded databases 146 | let scheme = e.split_once(':').map(|(scheme, _)| scheme).unwrap_or(""); 147 | 148 | match scheme { 149 | "ws" | "wss" | "http" | "https" => endpoints.push(e.to_string()), 150 | _ => bail!("A remote connection is expected: {e}"), 151 | }; 152 | } 153 | 154 | // Ensure at least one valid endpoint was provided 155 | if endpoints.is_empty() { 156 | bail!("Invalid endpoint: {endpoint}") 157 | } 158 | 159 | // Create the provider with initialized round-robin counter starting at 0 160 | Ok(Self { 161 | round_robin: AtomicUsize::new(0), 162 | endpoints, 163 | root, 164 | }) 165 | } 166 | 167 | /// Creates a new SurrealDB client connected to one of the configured instances. 168 | /// 169 | /// This method implements the round-robin load balancing strategy by: 170 | /// 1. Atomically incrementing the round-robin counter 171 | /// 2. Using modulo arithmetic to select the next endpoint 172 | /// 3. Establishing a connection to the selected endpoint 173 | /// 4. Returning a fully initialized and authenticated client 174 | /// 175 | /// # Load Distribution 176 | /// 177 | /// With 3 endpoints and 6 concurrent clients, the distribution would be: 178 | /// - Client 0 → endpoint 0 179 | /// - Client 1 → endpoint 1 180 | /// - Client 2 → endpoint 2 181 | /// - Client 3 → endpoint 0 182 | /// - Client 4 → endpoint 1 183 | /// - Client 5 → endpoint 2 184 | /// 185 | /// # Returns 186 | /// 187 | /// Returns a `SurrealDBClient` connected and authenticated to one of the 188 | /// configured SurrealDB instances, or an error if the connection fails. 189 | /// 190 | /// # Errors 191 | /// 192 | /// This function will return an error if: 193 | /// - The connection to the selected endpoint fails 194 | /// - Authentication with the root credentials fails 195 | /// - The namespace or database selection fails 196 | async fn create_client(&self) -> Result { 197 | // Atomically fetch the current counter value and increment it for the next call 198 | // Using Relaxed ordering is sufficient here as we don't need synchronization 199 | // beyond the atomic increment itself 200 | let next = self.round_robin.fetch_add(1, Ordering::Relaxed); 201 | 202 | // Select the endpoint using modulo to cycle through the available instances 203 | // This ensures even distribution: counter % endpoint_count 204 | let endpoint = &self.endpoints[next % self.endpoints.len()]; 205 | 206 | // Initialize a new connection to the selected endpoint with root credentials 207 | // This function handles connection, authentication, and namespace/database selection 208 | let client = initialise_db(endpoint, self.root.clone()).await?; 209 | 210 | // Wrap the connected client in a SurrealDBClient for the benchmark interface 211 | Ok(SurrealDBClient::new(client)) 212 | } 213 | } 214 | -------------------------------------------------------------------------------- /src/dialect.rs: -------------------------------------------------------------------------------- 1 | use crate::Scan; 2 | use crate::benchmark::NOT_SUPPORTED_ERROR; 3 | use crate::valueprovider::Columns; 4 | use anyhow::{Result, bail}; 5 | use chrono::{DateTime, TimeZone, Utc}; 6 | use flatten_json_object::ArrayFormatting; 7 | use flatten_json_object::Flattener; 8 | #[cfg(feature = "mongodb")] 9 | use mongodb::bson::{Document, doc, to_document}; 10 | use serde_json::Value; 11 | use uuid::Uuid; 12 | 13 | /// Help converting generated values to the right database representation 14 | pub(crate) trait Dialect { 15 | fn uuid(u: Uuid) -> Value { 16 | Value::String(u.to_string()) 17 | } 18 | fn date_time(secs_from_epoch: i64) -> Value { 19 | // Get the current UTC time 20 | let datetime: DateTime = Utc.timestamp_opt(secs_from_epoch, 0).unwrap(); 21 | // Format it to the SQL-friendly ISO 8601 format 22 | let formatted = datetime.to_rfc3339(); 23 | Value::String(formatted) 24 | } 25 | fn escape_field(field: String) -> String { 26 | field 27 | } 28 | fn arg_string(val: Value) -> String { 29 | val.to_string() 30 | } 31 | } 32 | 33 | // -------------------------------------------------- 34 | // Default 35 | // -------------------------------------------------- 36 | 37 | pub(crate) struct DefaultDialect(); 38 | 39 | impl Dialect for DefaultDialect {} 40 | 41 | // -------------------------------------------------- 42 | // SQL 43 | // -------------------------------------------------- 44 | 45 | pub(crate) struct AnsiSqlDialect(); 46 | 47 | impl Dialect for AnsiSqlDialect { 48 | fn escape_field(field: String) -> String { 49 | format!("\"{field}\"") 50 | } 51 | 52 | fn arg_string(val: Value) -> String { 53 | match val { 54 | Value::Null => "null".to_string(), 55 | Value::Bool(b) => b.to_string(), 56 | Value::Number(n) => n.to_string(), 57 | Value::String(s) => format!("'{s}'"), 58 | Value::Array(a) => serde_json::to_string(&a).unwrap(), 59 | Value::Object(o) => format!("'{}'", serde_json::to_string(&o).unwrap()), 60 | } 61 | } 62 | } 63 | 64 | impl AnsiSqlDialect { 65 | /// Constructs the field clauses for the [C]reate tests 66 | pub fn create_clause(cols: &Columns, val: Value) -> (String, String) { 67 | let mut fields = Vec::with_capacity(cols.0.len()); 68 | let mut values = Vec::with_capacity(cols.0.len()); 69 | if let Value::Object(map) = val { 70 | for (f, v) in map { 71 | fields.push(Self::escape_field(f)); 72 | values.push(Self::arg_string(v)); 73 | } 74 | } 75 | let fields = fields.join(", "); 76 | let values = values.join(", "); 77 | (fields, values) 78 | } 79 | /// Constructs the field clauses for the [U]pdate tests 80 | pub fn update_clause(cols: &Columns, val: Value) -> String { 81 | let mut updates = Vec::with_capacity(cols.0.len()); 82 | if let Value::Object(map) = val { 83 | for (f, v) in map { 84 | let f = Self::escape_field(f); 85 | let v = Self::arg_string(v); 86 | updates.push(format!("{f} = {v}")); 87 | } 88 | } 89 | updates.join(", ") 90 | } 91 | /// Constructs the WHERE clause for [S]elect tests 92 | pub fn filter_clause(scan: &Scan) -> Result { 93 | if let Some(ref c) = scan.condition { 94 | if let Some(ref c) = c.sql { 95 | return Ok(format!("WHERE {c}")); 96 | } else { 97 | bail!(NOT_SUPPORTED_ERROR); 98 | } 99 | } 100 | Ok(String::new()) 101 | } 102 | } 103 | 104 | // -------------------------------------------------- 105 | // MySQL 106 | // -------------------------------------------------- 107 | 108 | pub(crate) struct MySqlDialect(); 109 | 110 | impl Dialect for MySqlDialect { 111 | fn escape_field(field: String) -> String { 112 | format!("`{field}`") 113 | } 114 | 115 | fn arg_string(val: Value) -> String { 116 | match val { 117 | Value::Null => "null".to_string(), 118 | Value::Bool(b) => b.to_string(), 119 | Value::Number(n) => n.to_string(), 120 | Value::String(s) => format!("'{s}'"), 121 | Value::Array(a) => serde_json::to_string(&a).unwrap(), 122 | Value::Object(o) => format!("'{}'", serde_json::to_string(&o).unwrap()), 123 | } 124 | } 125 | } 126 | 127 | impl MySqlDialect { 128 | /// Constructs the field clauses for the [C]reate tests 129 | pub fn create_clause(cols: &Columns, val: Value) -> (String, String) { 130 | let mut fields = Vec::with_capacity(cols.0.len()); 131 | let mut values = Vec::with_capacity(cols.0.len()); 132 | if let Value::Object(map) = val { 133 | for (f, v) in map { 134 | fields.push(Self::escape_field(f)); 135 | values.push(Self::arg_string(v)); 136 | } 137 | } 138 | let fields = fields.join(", "); 139 | let values = values.join(", "); 140 | (fields, values) 141 | } 142 | /// Constructs the field clauses for the [U]pdate tests 143 | pub fn update_clause(cols: &Columns, val: Value) -> String { 144 | let mut updates = Vec::with_capacity(cols.0.len()); 145 | if let Value::Object(map) = val { 146 | for (f, v) in map { 147 | let f = Self::escape_field(f); 148 | let v = Self::arg_string(v); 149 | updates.push(format!("{f} = {v}")); 150 | } 151 | } 152 | updates.join(", ") 153 | } 154 | /// Constructs the WHERE clause for [S]elect tests 155 | pub fn filter_clause(scan: &Scan) -> Result { 156 | if let Some(ref c) = scan.condition { 157 | if let Some(ref c) = c.mysql { 158 | return Ok(format!("WHERE {c}")); 159 | } else { 160 | bail!(NOT_SUPPORTED_ERROR); 161 | } 162 | } 163 | Ok(String::new()) 164 | } 165 | } 166 | 167 | // -------------------------------------------------- 168 | // Neo4j 169 | // -------------------------------------------------- 170 | 171 | pub(crate) struct Neo4jDialect(); 172 | 173 | impl Dialect for Neo4jDialect {} 174 | 175 | impl Neo4jDialect { 176 | /// Constructs the field clauses for the [C]reate tests 177 | pub fn create_clause(val: Value) -> Result { 178 | let val = Flattener::new() 179 | .set_key_separator("_") 180 | .set_array_formatting(ArrayFormatting::Surrounded { 181 | start: "_".to_string(), 182 | end: "".to_string(), 183 | }) 184 | .set_preserve_empty_arrays(false) 185 | .set_preserve_empty_objects(false) 186 | .flatten(&val)?; 187 | let obj = val.as_object().unwrap(); 188 | let mut fields = Vec::with_capacity(obj.len()); 189 | if let Value::Object(map) = val { 190 | for (f, v) in map { 191 | let f = Self::escape_field(f); 192 | let v = Self::arg_string(v); 193 | fields.push(format!("{f}: {v}")); 194 | } 195 | } 196 | Ok(fields.join(", ")) 197 | } 198 | /// Constructs the field clauses for the [U]pdate tests 199 | pub fn update_clause(val: Value) -> Result { 200 | let val = Flattener::new() 201 | .set_key_separator("_") 202 | .set_array_formatting(ArrayFormatting::Surrounded { 203 | start: "_".to_string(), 204 | end: "".to_string(), 205 | }) 206 | .set_preserve_empty_arrays(false) 207 | .set_preserve_empty_objects(false) 208 | .flatten(&val)?; 209 | let obj = val.as_object().unwrap(); 210 | let mut fields = Vec::with_capacity(obj.len()); 211 | if let Value::Object(map) = val { 212 | for (f, v) in map { 213 | let f = Self::escape_field(f); 214 | let v = Self::arg_string(v); 215 | fields.push(format!("r.{f} = {v}")); 216 | } 217 | } 218 | Ok(fields.join(", ")) 219 | } 220 | /// Constructs the WHERE clause for [S]elect tests 221 | pub fn filter_clause(scan: &Scan) -> Result { 222 | if let Some(ref c) = scan.condition { 223 | if let Some(ref c) = c.neo4j { 224 | if let Some(index) = &scan.index 225 | && let Some(kind) = &index.index_type 226 | && kind == "fulltext" 227 | { 228 | return Ok(c.to_string()); 229 | } else { 230 | return Ok(format!("WHERE {c}")); 231 | } 232 | } else { 233 | bail!(NOT_SUPPORTED_ERROR); 234 | } 235 | } 236 | Ok(String::new()) 237 | } 238 | } 239 | 240 | // -------------------------------------------------- 241 | // SurrealDB 242 | // -------------------------------------------------- 243 | 244 | pub(crate) struct SurrealDBDialect(); 245 | 246 | impl Dialect for SurrealDBDialect {} 247 | 248 | impl SurrealDBDialect { 249 | /// Constructs the WHERE clause for [S]elect tests 250 | pub fn filter_clause(scan: &Scan) -> Result { 251 | if let Some(ref c) = scan.condition { 252 | if let Some(ref c) = c.surrealdb { 253 | return Ok(format!("WHERE {c}")); 254 | } else { 255 | bail!(NOT_SUPPORTED_ERROR); 256 | } 257 | } 258 | Ok(String::new()) 259 | } 260 | } 261 | 262 | // -------------------------------------------------- 263 | // ArangoDB 264 | // -------------------------------------------------- 265 | 266 | pub(crate) struct ArangoDBDialect(); 267 | 268 | impl Dialect for ArangoDBDialect {} 269 | 270 | impl ArangoDBDialect { 271 | /// Constructs the WHERE clause for [S]elect tests 272 | pub fn filter_clause(scan: &Scan) -> Result { 273 | if let Some(ref c) = scan.condition { 274 | if let Some(ref c) = c.arangodb { 275 | return Ok(format!("FILTER {c}")); 276 | } else { 277 | bail!(NOT_SUPPORTED_ERROR); 278 | } 279 | } 280 | Ok(String::new()) 281 | } 282 | } 283 | 284 | // -------------------------------------------------- 285 | // MongoDB 286 | // -------------------------------------------------- 287 | 288 | pub(crate) struct MongoDBDialect(); 289 | 290 | #[cfg(feature = "mongodb")] 291 | impl Dialect for MongoDBDialect {} 292 | 293 | #[cfg(feature = "mongodb")] 294 | impl MongoDBDialect { 295 | /// Constructs the filter document for [S]elect tests 296 | pub fn filter_clause(scan: &Scan) -> Result { 297 | if let Some(ref c) = scan.condition { 298 | if let Some(ref c) = c.mongodb { 299 | return Ok(to_document(c)?); 300 | } else { 301 | bail!(NOT_SUPPORTED_ERROR); 302 | } 303 | } 304 | Ok(doc! {}) 305 | } 306 | } 307 | -------------------------------------------------------------------------------- /src/neo4j.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "neo4j")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::dialect::Neo4jDialect; 5 | use crate::docker::DockerParams; 6 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 7 | use crate::valueprovider::Columns; 8 | use crate::{Benchmark, Index, KeyType, Projection, Scan}; 9 | use anyhow::{Result, bail}; 10 | use neo4rs::BoltType; 11 | use neo4rs::ConfigBuilder; 12 | use neo4rs::Graph; 13 | use neo4rs::query; 14 | use serde_json::Value; 15 | use std::hint::black_box; 16 | 17 | pub const DEFAULT: &str = "127.0.0.1:7687"; 18 | 19 | pub(crate) fn docker(options: &Benchmark) -> DockerParams { 20 | DockerParams { 21 | image: "neo4j", 22 | pre_args: match options.sync { 23 | true => { 24 | // Neo4j does not have the ability to configure 25 | // per-transaction on-disk sync control, so the 26 | // closest option when sync is true, is to 27 | // checkpoint after every transaction, and to 28 | // checkpoint in the background every second 29 | "--ulimit nofile=65536:65536 -p 127.0.0.1:7474:7474 -p 127.0.0.1:7687:7687 -e NEO4J_AUTH=none -e NEO4J_dbms_checkpoint_interval_time=1s -e NEO4J_dbms_checkpoint_interval_tx=1".to_string() 30 | } 31 | false => { 32 | // Neo4j does not have the ability to configure 33 | // per-transaction on-disk sync control, so the 34 | // closest option when sync is false, is to 35 | // checkpoint in the background every second, 36 | // and to checkpoint every 10,000 transactions 37 | "--ulimit nofile=65536:65536 -p 127.0.0.1:7474:7474 -p 127.0.0.1:7687:7687 -e NEO4J_AUTH=none -e NEO4J_dbms_checkpoint_interval_time=1s -e NEO4J_dbms_checkpoint_interval_tx=10000".to_string() 38 | } 39 | }, 40 | post_args: "".to_string(), 41 | } 42 | } 43 | 44 | pub(crate) struct Neo4jClientProvider { 45 | graph: Graph, 46 | } 47 | 48 | impl BenchmarkEngine for Neo4jClientProvider { 49 | /// Initiates a new datastore benchmarking engine 50 | async fn setup(_kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 51 | // Get the custom endpoint if specified 52 | let url = options.endpoint.as_deref().unwrap_or(DEFAULT).to_owned(); 53 | // Create a new client with a connection pool. 54 | // The Neo4j client supports connection pooling 55 | // and the recommended advice is to use a single 56 | // graph connection and share that with all async 57 | // tasks. Therefore we create a single connection 58 | // pool and share it with all of the crud-bench 59 | // clients. The Neo4j driver correctly limits the 60 | // number of connections to the number specified 61 | // in the `max_connections` option. 62 | let config = ConfigBuilder::default() 63 | .uri(url) 64 | .db("neo4j") 65 | .user("neo4j") 66 | .password("neo4j") 67 | .fetch_size(500) 68 | .max_connections(options.clients as usize) 69 | .build()?; 70 | // Create the client 71 | Ok(Self { 72 | graph: Graph::connect(config).await?, 73 | }) 74 | } 75 | /// Creates a new client for this benchmarking engine 76 | async fn create_client(&self) -> Result { 77 | Ok(Neo4jClient { 78 | graph: self.graph.clone(), 79 | }) 80 | } 81 | } 82 | 83 | pub(crate) struct Neo4jClient { 84 | graph: Graph, 85 | } 86 | 87 | impl BenchmarkClient for Neo4jClient { 88 | async fn startup(&self) -> Result<()> { 89 | let stm = "CREATE INDEX FOR (r:Record) ON (r.id);"; 90 | self.graph.execute(query(stm)).await?.next().await.ok(); 91 | Ok(()) 92 | } 93 | 94 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 95 | self.create(key, val).await 96 | } 97 | 98 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 99 | self.create(key, val).await 100 | } 101 | 102 | async fn read_u32(&self, key: u32) -> Result<()> { 103 | self.read(key).await 104 | } 105 | 106 | async fn read_string(&self, key: String) -> Result<()> { 107 | self.read(key).await 108 | } 109 | 110 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 111 | self.update(key, val).await 112 | } 113 | 114 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 115 | self.update(key, val).await 116 | } 117 | 118 | async fn delete_u32(&self, key: u32) -> Result<()> { 119 | self.delete(key).await 120 | } 121 | 122 | async fn delete_string(&self, key: String) -> Result<()> { 123 | self.delete(key).await 124 | } 125 | 126 | async fn build_index(&self, spec: &Index, name: &str) -> Result<()> { 127 | // Get the fields 128 | let fields = spec.fields.iter().map(|f| format!("r.{f}")).collect::>().join(", "); 129 | // Check if an index type is specified 130 | let stmt = match &spec.index_type { 131 | Some(kind) if kind == "fulltext" => { 132 | format!("CREATE FULLTEXT INDEX {name} FOR (r:Record) ON EACH [{fields}]") 133 | } 134 | _ => { 135 | format!("CREATE INDEX {name} FOR (r:Record) ON ({fields})") 136 | } 137 | }; 138 | // Create the index 139 | self.graph.execute(query(&stmt)).await?.next().await?; 140 | // Wait for the index to finish building in the background. 141 | // Neo4j indexes build asynchronously, so we need to wait 142 | // for the index to be fully online before proceeding. 143 | self.graph.execute(query("CALL db.awaitIndexes()")).await?.next().await?; 144 | // All ok 145 | Ok(()) 146 | } 147 | 148 | async fn drop_index(&self, name: &str) -> Result<()> { 149 | let stmt = format!("DROP INDEX {name} IF EXISTS"); 150 | self.graph.execute(query(&stmt)).await?.next().await?; 151 | Ok(()) 152 | } 153 | 154 | async fn scan_u32(&self, scan: &Scan, ctx: ScanContext) -> Result { 155 | self.scan(scan, ctx).await 156 | } 157 | 158 | async fn scan_string(&self, scan: &Scan, ctx: ScanContext) -> Result { 159 | self.scan(scan, ctx).await 160 | } 161 | } 162 | 163 | impl Neo4jClient { 164 | async fn create(&self, key: T, val: Value) -> Result<()> 165 | where 166 | T: Into + Sync, 167 | { 168 | let fields = Neo4jDialect::create_clause(val)?; 169 | let stm = format!("CREATE (r:Record {{ id: $id, {fields} }}) RETURN r.id"); 170 | let stm = query(&stm).param("id", key); 171 | let mut res = self.graph.execute(stm).await.unwrap(); 172 | assert!(matches!(res.next().await, Ok(Some(_)))); 173 | assert!(matches!(res.next().await, Ok(None))); 174 | Ok(()) 175 | } 176 | 177 | async fn read(&self, key: T) -> Result<()> 178 | where 179 | T: Into + Sync, 180 | { 181 | let stm = "MATCH (r:Record { id: $id }) RETURN r"; 182 | let stm = query(stm).param("id", key); 183 | let mut res = self.graph.execute(stm).await.unwrap(); 184 | assert!(matches!(black_box(res.next().await), Ok(Some(_)))); 185 | assert!(matches!(res.next().await, Ok(None))); 186 | Ok(()) 187 | } 188 | 189 | async fn update(&self, key: T, val: Value) -> Result<()> 190 | where 191 | T: Into + Sync, 192 | { 193 | let fields = Neo4jDialect::update_clause(val)?; 194 | let stm = format!("MATCH (r:Record {{ id: $id }}) SET {fields} RETURN r.id"); 195 | let stm = query(&stm).param("id", key); 196 | let mut res = self.graph.execute(stm).await.unwrap(); 197 | assert!(matches!(res.next().await, Ok(Some(_)))); 198 | assert!(matches!(res.next().await, Ok(None))); 199 | Ok(()) 200 | } 201 | 202 | async fn delete(&self, key: T) -> Result<()> 203 | where 204 | T: Into + Sync, 205 | { 206 | let stm = "MATCH (r:Record { id: $id }) WITH r, r.id AS id DETACH DELETE r RETURN id"; 207 | let stm = query(stm).param("id", key); 208 | let mut res = self.graph.execute(stm).await.unwrap(); 209 | assert!(matches!(res.next().await, Ok(Some(_)))); 210 | assert!(matches!(res.next().await, Ok(None))); 211 | Ok(()) 212 | } 213 | 214 | async fn scan(&self, scan: &Scan, ctx: ScanContext) -> Result { 215 | // Neo4j requires a full-text index to exist 216 | if ctx == ScanContext::WithoutIndex 217 | && let Some(index) = &scan.index 218 | && let Some(kind) = &index.index_type 219 | && kind == "fulltext" 220 | { 221 | bail!(NOT_SUPPORTED_ERROR); 222 | } 223 | // Extract parameters 224 | let s = scan.start.map(|s| format!("SKIP {s}")).unwrap_or_default(); 225 | let l = scan.limit.map(|s| format!("LIMIT {s}")).unwrap_or_default(); 226 | let c = Neo4jDialect::filter_clause(scan)?; 227 | let p = scan.projection()?; 228 | let n = &scan.name; 229 | // Check if this is a fulltext scan 230 | let fts = scan 231 | .index 232 | .as_ref() 233 | .and_then(|idx| idx.index_type.as_ref()) 234 | .map(|t| t == "fulltext") 235 | .unwrap_or(false); 236 | // Perform the relevant projection scan type 237 | match p { 238 | Projection::Id => { 239 | let stm = match fts { 240 | true => format!( 241 | "CALL db.index.fulltext.queryNodes('{n}', '{c}') YIELD node as r WITH r {s} {l} RETURN r.id" 242 | ), 243 | false => format!("MATCH (r) {c} {s} {l} RETURN r.id"), 244 | }; 245 | let mut res = self.graph.execute(query(&stm)).await.unwrap(); 246 | let mut count = 0; 247 | while let Ok(Some(v)) = res.next().await { 248 | black_box(v); 249 | count += 1; 250 | } 251 | Ok(count) 252 | } 253 | Projection::Full => { 254 | let stm = match fts { 255 | true => format!( 256 | "CALL db.index.fulltext.queryNodes('{n}', '{c}') YIELD node as r WITH r {s} {l} RETURN r" 257 | ), 258 | false => format!("MATCH (r) {c} {s} {l} RETURN r"), 259 | }; 260 | let mut res = self.graph.execute(query(&stm)).await.unwrap(); 261 | let mut count = 0; 262 | while let Ok(Some(v)) = res.next().await { 263 | black_box(v); 264 | count += 1; 265 | } 266 | Ok(count) 267 | } 268 | Projection::Count => { 269 | let stm = match fts { 270 | true => format!( 271 | "CALL db.index.fulltext.queryNodes('{n}', '{c}') YIELD node as r WITH r {s} {l} RETURN count(r) as count" 272 | ), 273 | false => format!("MATCH (r) {c} {s} {l} RETURN count(r) as count"), 274 | }; 275 | let mut res = self.graph.execute(query(&stm)).await.unwrap(); 276 | let count = res.next().await.unwrap().unwrap().get("count").unwrap(); 277 | Ok(count) 278 | } 279 | } 280 | } 281 | } 282 | -------------------------------------------------------------------------------- /src/surrealmx.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "surrealmx")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 5 | use crate::valueprovider::Columns; 6 | use crate::{Benchmark, KeyType, Projection, Scan}; 7 | use anyhow::{Result, bail}; 8 | use serde_json::Value; 9 | use std::hint::black_box; 10 | use std::sync::Arc; 11 | use std::time::Duration; 12 | use surrealmx::Database; 13 | use surrealmx::{AolMode, FsyncMode, SnapshotMode}; 14 | use surrealmx::{DatabaseOptions, PersistenceOptions}; 15 | 16 | const DATABASE_DIR: &str = "surrealmx"; 17 | 18 | pub(crate) struct SurrealMXClientProvider(Arc); 19 | 20 | impl BenchmarkEngine for SurrealMXClientProvider { 21 | /// The number of seconds to wait before connecting 22 | fn wait_timeout(&self) -> Option { 23 | None 24 | } 25 | /// Initiates a new datastore benchmarking engine 26 | async fn setup(_: KeyType, _columns: Columns, options: &Benchmark) -> Result { 27 | // Check if persistence is enabled 28 | if options.persisted { 29 | // Specify the database options 30 | let opts = DatabaseOptions::default(); 31 | // Specify the persistence options 32 | let persistence = match options.sync { 33 | // Write to AOL immediately and fsync when sync is true 34 | true => PersistenceOptions::new(DATABASE_DIR) 35 | .with_snapshot_mode(SnapshotMode::Never) 36 | .with_aol_mode(AolMode::SynchronousOnCommit) 37 | .with_fsync_mode(FsyncMode::EveryAppend), 38 | // Write to AOL in the background and don't fsync when sync is false 39 | false => PersistenceOptions::new(DATABASE_DIR) 40 | .with_snapshot_mode(SnapshotMode::Never) 41 | .with_aol_mode(AolMode::AsynchronousAfterCommit) 42 | .with_fsync_mode(FsyncMode::Never), 43 | }; 44 | // Create the store 45 | return Ok(Self(Arc::new(Database::new_with_persistence(opts, persistence).unwrap()))); 46 | } 47 | // Create the store 48 | Ok(Self(Arc::new(Database::new()))) 49 | } 50 | /// Creates a new client for this benchmarking engine 51 | async fn create_client(&self) -> Result { 52 | Ok(SurrealMXClient { 53 | db: self.0.clone(), 54 | }) 55 | } 56 | } 57 | 58 | pub(crate) struct SurrealMXClient { 59 | db: Arc, 60 | } 61 | 62 | impl BenchmarkClient for SurrealMXClient { 63 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 64 | self.create_bytes(&key.to_ne_bytes(), val).await 65 | } 66 | 67 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 68 | self.create_bytes(&key.into_bytes(), val).await 69 | } 70 | 71 | async fn read_u32(&self, key: u32) -> Result<()> { 72 | self.read_bytes(&key.to_ne_bytes()).await 73 | } 74 | 75 | async fn read_string(&self, key: String) -> Result<()> { 76 | self.read_bytes(&key.into_bytes()).await 77 | } 78 | 79 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 80 | self.update_bytes(&key.to_ne_bytes(), val).await 81 | } 82 | 83 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 84 | self.update_bytes(&key.into_bytes(), val).await 85 | } 86 | 87 | async fn delete_u32(&self, key: u32) -> Result<()> { 88 | self.delete_bytes(&key.to_ne_bytes()).await 89 | } 90 | 91 | async fn delete_string(&self, key: String) -> Result<()> { 92 | self.delete_bytes(&key.into_bytes()).await 93 | } 94 | 95 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 96 | self.scan_bytes(scan).await 97 | } 98 | 99 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 100 | self.scan_bytes(scan).await 101 | } 102 | 103 | async fn batch_create_u32( 104 | &self, 105 | key_vals: impl Iterator + Send, 106 | ) -> Result<()> { 107 | let pairs_iter = key_vals.map(|(key, val)| { 108 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 109 | Ok((key.to_ne_bytes().to_vec(), val)) 110 | }); 111 | self.batch_create_bytes(pairs_iter).await 112 | } 113 | 114 | async fn batch_create_string( 115 | &self, 116 | key_vals: impl Iterator + Send, 117 | ) -> Result<()> { 118 | let pairs_iter = key_vals.map(|(key, val)| { 119 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 120 | Ok((key.into_bytes(), val)) 121 | }); 122 | self.batch_create_bytes(pairs_iter).await 123 | } 124 | 125 | async fn batch_read_u32(&self, keys: impl Iterator + Send) -> Result<()> { 126 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 127 | self.batch_read_bytes(keys_iter).await 128 | } 129 | 130 | async fn batch_read_string(&self, keys: impl Iterator + Send) -> Result<()> { 131 | let keys_iter = keys.map(|key| key.into_bytes()); 132 | self.batch_read_bytes(keys_iter).await 133 | } 134 | 135 | async fn batch_update_u32( 136 | &self, 137 | key_vals: impl Iterator + Send, 138 | ) -> Result<()> { 139 | let pairs_iter = key_vals.map(|(key, val)| { 140 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 141 | Ok((key.to_ne_bytes().to_vec(), val)) 142 | }); 143 | self.batch_update_bytes(pairs_iter).await 144 | } 145 | 146 | async fn batch_update_string( 147 | &self, 148 | key_vals: impl Iterator + Send, 149 | ) -> Result<()> { 150 | let pairs_iter = key_vals.map(|(key, val)| { 151 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 152 | Ok((key.into_bytes(), val)) 153 | }); 154 | self.batch_update_bytes(pairs_iter).await 155 | } 156 | 157 | async fn batch_delete_u32(&self, keys: impl Iterator + Send) -> Result<()> { 158 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 159 | self.batch_delete_bytes(keys_iter).await 160 | } 161 | 162 | async fn batch_delete_string(&self, keys: impl Iterator + Send) -> Result<()> { 163 | let keys_iter = keys.map(|key| key.into_bytes()); 164 | self.batch_delete_bytes(keys_iter).await 165 | } 166 | } 167 | 168 | impl SurrealMXClient { 169 | async fn create_bytes(&self, key: &[u8], val: Value) -> Result<()> { 170 | // Serialise the value 171 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 172 | // Create a new transaction 173 | let mut txn = self.db.transaction(true); 174 | // Process the data 175 | txn.set(key, val)?; 176 | txn.commit()?; 177 | Ok(()) 178 | } 179 | 180 | async fn read_bytes(&self, key: &[u8]) -> Result<()> { 181 | // Create a new transaction 182 | let txn = self.db.transaction(false); 183 | // Process the data 184 | let res = txn.get(key.to_vec())?; 185 | // Check the value exists 186 | assert!(res.is_some()); 187 | // Deserialise the value 188 | black_box(res.unwrap()); 189 | // All ok 190 | Ok(()) 191 | } 192 | 193 | async fn update_bytes(&self, key: &[u8], val: Value) -> Result<()> { 194 | // Serialise the value 195 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 196 | // Create a new transaction 197 | let mut txn = self.db.transaction(true); 198 | // Process the data 199 | txn.set(key, val)?; 200 | txn.commit()?; 201 | Ok(()) 202 | } 203 | 204 | async fn delete_bytes(&self, key: &[u8]) -> Result<()> { 205 | // Create a new transaction 206 | let mut txn = self.db.transaction(true); 207 | // Process the data 208 | txn.del(key)?; 209 | txn.commit()?; 210 | Ok(()) 211 | } 212 | 213 | async fn batch_create_bytes( 214 | &self, 215 | key_vals: impl Iterator, Vec)>>, 216 | ) -> Result<()> { 217 | // Create a new transaction 218 | let mut txn = self.db.transaction(true); 219 | // Process the data 220 | for result in key_vals { 221 | let (key, val) = result?; 222 | txn.set(key, val)?; 223 | } 224 | // Commit the batch 225 | txn.commit()?; 226 | Ok(()) 227 | } 228 | 229 | async fn batch_read_bytes(&self, keys: impl Iterator>) -> Result<()> { 230 | // Create a new transaction 231 | let txn = self.db.transaction(false); 232 | // Process the data 233 | for key in keys { 234 | // Get the current value 235 | let res = txn.get(key)?; 236 | // Check the value exists 237 | assert!(res.is_some()); 238 | // Deserialise the value 239 | black_box(res.unwrap()); 240 | } 241 | // All ok 242 | Ok(()) 243 | } 244 | 245 | async fn batch_update_bytes( 246 | &self, 247 | key_vals: impl Iterator, Vec)>>, 248 | ) -> Result<()> { 249 | // Create a new transaction 250 | let mut txn = self.db.transaction(true); 251 | // Process the data 252 | for result in key_vals { 253 | let (key, val) = result?; 254 | txn.set(key, val)?; 255 | } 256 | // Commit the batch 257 | txn.commit()?; 258 | Ok(()) 259 | } 260 | 261 | async fn batch_delete_bytes(&self, keys: impl Iterator>) -> Result<()> { 262 | // Create a new transaction 263 | let mut txn = self.db.transaction(true); 264 | // Process the data 265 | for key in keys { 266 | txn.del(key)?; 267 | } 268 | // Commit the batch 269 | txn.commit()?; 270 | Ok(()) 271 | } 272 | 273 | async fn scan_bytes(&self, scan: &Scan) -> Result { 274 | // Contional scans are not supported 275 | if scan.condition.is_some() { 276 | bail!(NOT_SUPPORTED_ERROR); 277 | } 278 | // Extract parameters 279 | let p = scan.projection()?; 280 | // Create a new transaction 281 | let txn = self.db.transaction(false); 282 | let beg = [0u8].to_vec(); 283 | let end = [255u8].to_vec(); 284 | // Perform the relevant projection scan type 285 | match p { 286 | Projection::Id => { 287 | // Scan the desired range of keys 288 | let iter = txn.keys(beg..end, scan.start, scan.limit)?; 289 | // Create an iterator starting at the beginning 290 | let iter = iter.into_iter(); 291 | // We use a for loop to iterate over the results, while 292 | // calling black_box internally. This is necessary as 293 | // an iterator with `filter_map` or `map` is optimised 294 | // out by the compiler when calling `count` at the end. 295 | let mut count = 0; 296 | for v in iter { 297 | black_box(v); 298 | count += 1; 299 | } 300 | Ok(count) 301 | } 302 | Projection::Full => { 303 | // Scan the desired range of keys 304 | let iter = txn.scan(beg..end, scan.start, scan.limit)?; 305 | // Create an iterator starting at the beginning 306 | let iter = iter.into_iter(); 307 | // We use a for loop to iterate over the results, while 308 | // calling black_box internally. This is necessary as 309 | // an iterator with `filter_map` or `map` is optimised 310 | // out by the compiler when calling `count` at the end. 311 | let mut count = 0; 312 | for v in iter { 313 | black_box(v.1); 314 | count += 1; 315 | } 316 | Ok(count) 317 | } 318 | Projection::Count => Ok(txn.total(beg..end, scan.start, scan.limit)?), 319 | } 320 | } 321 | } 322 | -------------------------------------------------------------------------------- /src/lmdb.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "lmdb")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 5 | use crate::valueprovider::Columns; 6 | use crate::{Benchmark, KeyType, Projection, Scan}; 7 | use anyhow::{Result, bail}; 8 | use heed::types::Bytes; 9 | use heed::{Database, EnvOpenOptions}; 10 | use heed::{Env, EnvFlags, WithoutTls}; 11 | use serde_json::Value; 12 | use std::hint::black_box; 13 | use std::sync::Arc; 14 | use std::sync::LazyLock; 15 | use std::time::Duration; 16 | 17 | const DATABASE_DIR: &str = "lmdb"; 18 | 19 | const DEFAULT_SIZE: usize = 4_294_967_296; // 4GiB 20 | 21 | static DATABASE_SIZE: LazyLock = LazyLock::new(|| { 22 | std::env::var("CRUD_BENCH_LMDB_DATABASE_SIZE") 23 | .map(|s| s.parse::().unwrap_or(DEFAULT_SIZE)) 24 | .unwrap_or(DEFAULT_SIZE) 25 | }); 26 | 27 | pub(crate) struct LmDBClientProvider(Arc<(Env, Database)>); 28 | 29 | impl BenchmarkEngine for LmDBClientProvider { 30 | /// The number of seconds to wait before connecting 31 | fn wait_timeout(&self) -> Option { 32 | None 33 | } 34 | /// Initiates a new datastore benchmarking engine 35 | async fn setup(_kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 36 | // Cleanup the data directory 37 | std::fs::remove_dir_all(DATABASE_DIR).ok(); 38 | // Recreate the database directory 39 | std::fs::create_dir(DATABASE_DIR)?; 40 | // Configure flags based on options 41 | let mut flags = EnvFlags::NO_READ_AHEAD | EnvFlags::NO_MEM_INIT; 42 | // Configure flags for filesystem sync 43 | if !options.sync { 44 | flags |= EnvFlags::NO_SYNC | EnvFlags::NO_META_SYNC; 45 | } 46 | // Create a new environment 47 | let env = unsafe { 48 | EnvOpenOptions::new() 49 | // Allow more transactions than threads 50 | .read_txn_without_tls() 51 | // Configure database flags 52 | .flags(flags) 53 | // We only use one db for benchmarks 54 | .max_dbs(1) 55 | // Optimize for expected concurrent readers 56 | .max_readers(126) 57 | // Set the database size 58 | .map_size(*DATABASE_SIZE) 59 | // Open the database 60 | .open(DATABASE_DIR) 61 | }?; 62 | // Create the database 63 | let db = { 64 | // Open a new transaction 65 | let mut txn = env.write_txn()?; 66 | // Initiate the database 67 | env.create_database::(&mut txn, None)? 68 | }; 69 | // Create the store 70 | Ok(Self(Arc::new((env, db)))) 71 | } 72 | /// Creates a new client for this benchmarking engine 73 | async fn create_client(&self) -> Result { 74 | Ok(LmDBClient { 75 | db: self.0.clone(), 76 | }) 77 | } 78 | } 79 | 80 | pub(crate) struct LmDBClient { 81 | db: Arc<(Env, Database)>, 82 | } 83 | 84 | impl BenchmarkClient for LmDBClient { 85 | async fn shutdown(&self) -> Result<()> { 86 | // Cleanup the data directory 87 | std::fs::remove_dir_all(DATABASE_DIR).ok(); 88 | // Ok 89 | Ok(()) 90 | } 91 | 92 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 93 | self.create_bytes(&key.to_ne_bytes(), val).await 94 | } 95 | 96 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 97 | self.create_bytes(&key.into_bytes(), val).await 98 | } 99 | 100 | async fn read_u32(&self, key: u32) -> Result<()> { 101 | self.read_bytes(&key.to_ne_bytes()).await 102 | } 103 | 104 | async fn read_string(&self, key: String) -> Result<()> { 105 | self.read_bytes(&key.into_bytes()).await 106 | } 107 | 108 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 109 | self.update_bytes(&key.to_ne_bytes(), val).await 110 | } 111 | 112 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 113 | self.update_bytes(&key.into_bytes(), val).await 114 | } 115 | 116 | async fn delete_u32(&self, key: u32) -> Result<()> { 117 | self.delete_bytes(&key.to_ne_bytes()).await 118 | } 119 | 120 | async fn delete_string(&self, key: String) -> Result<()> { 121 | self.delete_bytes(&key.into_bytes()).await 122 | } 123 | 124 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 125 | self.scan_bytes(scan).await 126 | } 127 | 128 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 129 | self.scan_bytes(scan).await 130 | } 131 | 132 | async fn batch_create_u32( 133 | &self, 134 | key_vals: impl Iterator + Send, 135 | ) -> Result<()> { 136 | let pairs_iter = key_vals.map(|(key, val)| { 137 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 138 | Ok((key.to_ne_bytes().to_vec(), val)) 139 | }); 140 | self.batch_create_bytes(pairs_iter).await 141 | } 142 | 143 | async fn batch_create_string( 144 | &self, 145 | key_vals: impl Iterator + Send, 146 | ) -> Result<()> { 147 | let pairs_iter = key_vals.map(|(key, val)| { 148 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 149 | Ok((key.into_bytes(), val)) 150 | }); 151 | self.batch_create_bytes(pairs_iter).await 152 | } 153 | 154 | async fn batch_read_u32(&self, keys: impl Iterator + Send) -> Result<()> { 155 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 156 | self.batch_read_bytes(keys_iter).await 157 | } 158 | 159 | async fn batch_read_string(&self, keys: impl Iterator + Send) -> Result<()> { 160 | let keys_iter = keys.map(|key| key.into_bytes()); 161 | self.batch_read_bytes(keys_iter).await 162 | } 163 | 164 | async fn batch_update_u32( 165 | &self, 166 | key_vals: impl Iterator + Send, 167 | ) -> Result<()> { 168 | let pairs_iter = key_vals.map(|(key, val)| { 169 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 170 | Ok((key.to_ne_bytes().to_vec(), val)) 171 | }); 172 | self.batch_update_bytes(pairs_iter).await 173 | } 174 | 175 | async fn batch_update_string( 176 | &self, 177 | key_vals: impl Iterator + Send, 178 | ) -> Result<()> { 179 | let pairs_iter = key_vals.map(|(key, val)| { 180 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 181 | Ok((key.into_bytes(), val)) 182 | }); 183 | self.batch_update_bytes(pairs_iter).await 184 | } 185 | 186 | async fn batch_delete_u32(&self, keys: impl Iterator + Send) -> Result<()> { 187 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 188 | self.batch_delete_bytes(keys_iter).await 189 | } 190 | 191 | async fn batch_delete_string(&self, keys: impl Iterator + Send) -> Result<()> { 192 | let keys_iter = keys.map(|key| key.into_bytes()); 193 | self.batch_delete_bytes(keys_iter).await 194 | } 195 | } 196 | 197 | impl LmDBClient { 198 | async fn create_bytes(&self, key: &[u8], val: Value) -> Result<()> { 199 | // Serialise the value 200 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 201 | // Create a new transaction 202 | let mut txn = self.db.0.write_txn()?; 203 | // Process the data 204 | self.db.1.put(&mut txn, key, val.as_ref())?; 205 | txn.commit()?; 206 | Ok(()) 207 | } 208 | 209 | async fn read_bytes(&self, key: &[u8]) -> Result<()> { 210 | // Create a new transaction 211 | let txn = self.db.0.read_txn()?; 212 | // Process the data 213 | let res: Option<_> = self.db.1.get(&txn, key)?; 214 | // Check the value exists 215 | assert!(res.is_some()); 216 | // Deserialise the value 217 | black_box(res.unwrap()); 218 | // All ok 219 | Ok(()) 220 | } 221 | 222 | async fn update_bytes(&self, key: &[u8], val: Value) -> Result<()> { 223 | // Serialise the value 224 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 225 | // Create a new transaction 226 | let mut txn = self.db.0.write_txn()?; 227 | // Process the data 228 | self.db.1.put(&mut txn, key, &val)?; 229 | txn.commit()?; 230 | Ok(()) 231 | } 232 | 233 | async fn delete_bytes(&self, key: &[u8]) -> Result<()> { 234 | // Create a new transaction 235 | let mut txn = self.db.0.write_txn()?; 236 | // Process the data 237 | self.db.1.delete(&mut txn, key)?; 238 | txn.commit()?; 239 | Ok(()) 240 | } 241 | 242 | async fn batch_create_bytes( 243 | &self, 244 | key_vals: impl Iterator, Vec)>>, 245 | ) -> Result<()> { 246 | // Create a new transaction 247 | let mut txn = self.db.0.write_txn()?; 248 | // Process the data 249 | for result in key_vals { 250 | let (key, val) = result?; 251 | self.db.1.put(&mut txn, &key, &val)?; 252 | } 253 | // Commit the batch 254 | txn.commit()?; 255 | Ok(()) 256 | } 257 | 258 | async fn batch_read_bytes(&self, keys: impl Iterator>) -> Result<()> { 259 | // Create a new transaction 260 | let txn = self.db.0.read_txn()?; 261 | // Process the data 262 | for key in keys { 263 | // Get the current value 264 | let res: Option<_> = self.db.1.get(&txn, &key)?; 265 | // Check the value exists 266 | assert!(res.is_some()); 267 | // Deserialise the value 268 | black_box(res.unwrap()); 269 | } 270 | // All ok 271 | Ok(()) 272 | } 273 | 274 | async fn batch_update_bytes( 275 | &self, 276 | key_vals: impl Iterator, Vec)>>, 277 | ) -> Result<()> { 278 | // Create a new transaction 279 | let mut txn = self.db.0.write_txn()?; 280 | // Process the data 281 | for result in key_vals { 282 | let (key, val) = result?; 283 | self.db.1.put(&mut txn, &key, &val)?; 284 | } 285 | // Commit the batch 286 | txn.commit()?; 287 | Ok(()) 288 | } 289 | 290 | async fn batch_delete_bytes(&self, keys: impl Iterator>) -> Result<()> { 291 | // Create a new transaction 292 | let mut txn = self.db.0.write_txn()?; 293 | // Process the data 294 | for key in keys { 295 | self.db.1.delete(&mut txn, &key)?; 296 | } 297 | // Commit the batch 298 | txn.commit()?; 299 | Ok(()) 300 | } 301 | 302 | async fn scan_bytes(&self, scan: &Scan) -> Result { 303 | // Contional scans are not supported 304 | if scan.condition.is_some() { 305 | bail!(NOT_SUPPORTED_ERROR); 306 | } 307 | // Extract parameters 308 | let s = scan.start.unwrap_or(0); 309 | let l = scan.limit.unwrap_or(usize::MAX); 310 | let p = scan.projection()?; 311 | // Create a new transaction 312 | let txn = self.db.0.read_txn()?; 313 | // Create an iterator starting at the beginning 314 | let iter = self.db.1.iter(&txn)?; 315 | // Perform the relevant projection scan type 316 | match p { 317 | Projection::Id => { 318 | // We use a for loop to iterate over the results, while 319 | // calling black_box internally. This is necessary as 320 | // an iterator with `filter_map` or `map` is optimised 321 | // out by the compiler when calling `count` at the end. 322 | let mut count = 0; 323 | for v in iter.skip(s).take(l) { 324 | black_box(v.unwrap().0); 325 | count += 1; 326 | } 327 | Ok(count) 328 | } 329 | Projection::Full => { 330 | // We use a for loop to iterate over the results, while 331 | // calling black_box internally. This is necessary as 332 | // an iterator with `filter_map` or `map` is optimised 333 | // out by the compiler when calling `count` at the end. 334 | let mut count = 0; 335 | for v in iter.skip(s).take(l) { 336 | black_box(v.unwrap().1); 337 | count += 1; 338 | } 339 | Ok(count) 340 | } 341 | Projection::Count => { 342 | Ok(iter 343 | .skip(s) // Skip the first `offset` entries 344 | .take(l) // Take the next `limit` entries 345 | .count()) 346 | } 347 | } 348 | } 349 | } 350 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright © 2021-2025 SurrealDB Ltd. 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /src/mdbx.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "mdbx")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 5 | use crate::valueprovider::Columns; 6 | use crate::{Benchmark, KeyType, Projection, Scan}; 7 | use anyhow::{Result, bail}; 8 | use libmdbx::{ 9 | Database, DatabaseOptions, Mode, NoWriteMap, PageSize, ReadWriteOptions, SyncMode, WriteFlags, 10 | }; 11 | use serde_json::Value; 12 | use std::hint::black_box; 13 | use std::sync::Arc; 14 | use std::time::Duration; 15 | 16 | const DATABASE_DIR: &str = "mdbx"; 17 | 18 | pub(crate) struct MDBXClientProvider(Arc>); 19 | 20 | impl BenchmarkEngine for MDBXClientProvider { 21 | /// The number of seconds to wait before connecting 22 | fn wait_timeout(&self) -> Option { 23 | None 24 | } 25 | /// Initiates a new datastore benchmarking engine 26 | async fn setup(_kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 27 | // Cleanup the data directory 28 | std::fs::remove_dir_all(DATABASE_DIR).ok(); 29 | // Configure database options 30 | let options = DatabaseOptions { 31 | // Configure the read-write options 32 | mode: Mode::ReadWrite(ReadWriteOptions { 33 | sync_mode: if options.sync { 34 | SyncMode::Durable 35 | } else { 36 | SyncMode::UtterlyNoSync 37 | }, 38 | // No maximum database size 39 | max_size: None, 40 | // 64MB minimum database size 41 | min_size: Some(64 * 1024 * 1024), 42 | // Grow in 256MB steps 43 | growth_step: Some(256 * 1024 * 1024), 44 | // Disable shrinking in benchmarks 45 | shrink_threshold: None, 46 | }), 47 | // 16KB pages for better sequential performance 48 | page_size: Some(PageSize::Set(16384)), 49 | // Exclusive mode - no inter-process locking overhead 50 | exclusive: true, 51 | // LIFO garbage collection for better cache performance 52 | liforeclaim: true, 53 | // Disable readahead for better random access 54 | no_rdahead: true, 55 | // Skip memory initialization for performance 56 | no_meminit: true, 57 | // Coalesce transactions for better write performance 58 | coalesce: true, 59 | // Optimize for expected concurrent readers 60 | max_readers: Some(126), 61 | // We only use one table for benchmarks 62 | max_tables: Some(1), 63 | // Use defaults for transaction limits 64 | ..Default::default() 65 | }; 66 | // Create the database 67 | let db = Database::open_with_options(DATABASE_DIR, options)?; 68 | // Begin a new transaction 69 | let tx = db.begin_rw_txn()?; 70 | // Open the default table 71 | let tb = tx.open_table(None)?; 72 | // Prime the table for permaopen 73 | tx.prime_for_permaopen(tb); 74 | // Commit the transaction 75 | tx.commit()?; 76 | // Create the store 77 | Ok(Self(Arc::new(db))) 78 | } 79 | /// Creates a new client for this benchmarking engine 80 | async fn create_client(&self) -> Result { 81 | Ok(MDBXClient { 82 | db: self.0.clone(), 83 | }) 84 | } 85 | } 86 | 87 | pub(crate) struct MDBXClient { 88 | db: Arc>, 89 | } 90 | 91 | impl BenchmarkClient for MDBXClient { 92 | async fn shutdown(&self) -> Result<()> { 93 | // Cleanup the data directory 94 | std::fs::remove_dir_all(DATABASE_DIR).ok(); 95 | // Ok 96 | Ok(()) 97 | } 98 | 99 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 100 | self.create_bytes(&key.to_ne_bytes(), val).await 101 | } 102 | 103 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 104 | self.create_bytes(&key.into_bytes(), val).await 105 | } 106 | 107 | async fn read_u32(&self, key: u32) -> Result<()> { 108 | self.read_bytes(&key.to_ne_bytes()).await 109 | } 110 | 111 | async fn read_string(&self, key: String) -> Result<()> { 112 | self.read_bytes(&key.into_bytes()).await 113 | } 114 | 115 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 116 | self.update_bytes(&key.to_ne_bytes(), val).await 117 | } 118 | 119 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 120 | self.update_bytes(&key.into_bytes(), val).await 121 | } 122 | 123 | async fn delete_u32(&self, key: u32) -> Result<()> { 124 | self.delete_bytes(&key.to_ne_bytes()).await 125 | } 126 | 127 | async fn delete_string(&self, key: String) -> Result<()> { 128 | self.delete_bytes(&key.into_bytes()).await 129 | } 130 | 131 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 132 | self.scan_bytes(scan).await 133 | } 134 | 135 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 136 | self.scan_bytes(scan).await 137 | } 138 | 139 | async fn batch_create_u32( 140 | &self, 141 | key_vals: impl Iterator + Send, 142 | ) -> Result<()> { 143 | let pairs_iter = key_vals.map(|(key, val)| { 144 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 145 | Ok((key.to_ne_bytes().to_vec(), val)) 146 | }); 147 | self.batch_create_bytes(pairs_iter).await 148 | } 149 | 150 | async fn batch_create_string( 151 | &self, 152 | key_vals: impl Iterator + Send, 153 | ) -> Result<()> { 154 | let pairs_iter = key_vals.map(|(key, val)| { 155 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 156 | Ok((key.into_bytes(), val)) 157 | }); 158 | self.batch_create_bytes(pairs_iter).await 159 | } 160 | 161 | async fn batch_read_u32(&self, keys: impl Iterator + Send) -> Result<()> { 162 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 163 | self.batch_read_bytes(keys_iter).await 164 | } 165 | 166 | async fn batch_read_string(&self, keys: impl Iterator + Send) -> Result<()> { 167 | let keys_iter = keys.map(|key| key.into_bytes()); 168 | self.batch_read_bytes(keys_iter).await 169 | } 170 | 171 | async fn batch_update_u32( 172 | &self, 173 | key_vals: impl Iterator + Send, 174 | ) -> Result<()> { 175 | let pairs_iter = key_vals.map(|(key, val)| { 176 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 177 | Ok((key.to_ne_bytes().to_vec(), val)) 178 | }); 179 | self.batch_update_bytes(pairs_iter).await 180 | } 181 | 182 | async fn batch_update_string( 183 | &self, 184 | key_vals: impl Iterator + Send, 185 | ) -> Result<()> { 186 | let pairs_iter = key_vals.map(|(key, val)| { 187 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 188 | Ok((key.into_bytes(), val)) 189 | }); 190 | self.batch_update_bytes(pairs_iter).await 191 | } 192 | 193 | async fn batch_delete_u32(&self, keys: impl Iterator + Send) -> Result<()> { 194 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 195 | self.batch_delete_bytes(keys_iter).await 196 | } 197 | 198 | async fn batch_delete_string(&self, keys: impl Iterator + Send) -> Result<()> { 199 | let keys_iter = keys.map(|key| key.into_bytes()); 200 | self.batch_delete_bytes(keys_iter).await 201 | } 202 | } 203 | 204 | impl MDBXClient { 205 | async fn create_bytes(&self, key: &[u8], val: Value) -> Result<()> { 206 | // Serialise the value 207 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 208 | // Create a new transaction 209 | let txn = self.db.begin_rw_txn()?; 210 | // Open the default table 211 | let table = txn.open_table(None)?; 212 | // Process the data 213 | txn.put(&table, key, &val, WriteFlags::empty())?; 214 | txn.commit()?; 215 | Ok(()) 216 | } 217 | 218 | async fn read_bytes(&self, key: &[u8]) -> Result<()> { 219 | // Create a new transaction 220 | let txn = self.db.begin_ro_txn()?; 221 | // Open the default table 222 | let table = txn.open_table(None)?; 223 | // Process the data 224 | let res: Option> = txn.get(&table, key)?; 225 | // Check the value exists 226 | assert!(res.is_some()); 227 | // Deserialise the value 228 | black_box(res.unwrap()); 229 | // All ok 230 | Ok(()) 231 | } 232 | 233 | async fn update_bytes(&self, key: &[u8], val: Value) -> Result<()> { 234 | // Serialise the value 235 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 236 | // Create a new transaction 237 | let txn = self.db.begin_rw_txn()?; 238 | // Open the default table 239 | let table = txn.open_table(None)?; 240 | // Process the data 241 | txn.put(&table, key, &val, WriteFlags::empty())?; 242 | txn.commit()?; 243 | Ok(()) 244 | } 245 | 246 | async fn delete_bytes(&self, key: &[u8]) -> Result<()> { 247 | // Create a new transaction 248 | let txn = self.db.begin_rw_txn()?; 249 | // Open the default table 250 | let table = txn.open_table(None)?; 251 | // Process the data 252 | txn.del(&table, key, None)?; 253 | txn.commit()?; 254 | Ok(()) 255 | } 256 | 257 | async fn batch_create_bytes( 258 | &self, 259 | key_vals: impl Iterator, Vec)>>, 260 | ) -> Result<()> { 261 | // Create a new transaction 262 | let txn = self.db.begin_rw_txn()?; 263 | // Open the default table 264 | let table = txn.open_table(None)?; 265 | // Process the data 266 | for result in key_vals { 267 | let (key, val) = result?; 268 | txn.put(&table, &key, &val, WriteFlags::empty())?; 269 | } 270 | // Commit the batch 271 | txn.commit()?; 272 | Ok(()) 273 | } 274 | 275 | async fn batch_read_bytes(&self, keys: impl Iterator>) -> Result<()> { 276 | // Create a new transaction 277 | let txn = self.db.begin_ro_txn()?; 278 | // Open the default table 279 | let table = txn.open_table(None)?; 280 | // Process the data 281 | for key in keys { 282 | // Get the current value 283 | let res: Option> = txn.get(&table, &key)?; 284 | // Check the value exists 285 | assert!(res.is_some()); 286 | // Deserialise the value 287 | black_box(res.unwrap()); 288 | } 289 | // All ok 290 | Ok(()) 291 | } 292 | 293 | async fn batch_update_bytes( 294 | &self, 295 | key_vals: impl Iterator, Vec)>>, 296 | ) -> Result<()> { 297 | // Create a new transaction 298 | let txn = self.db.begin_rw_txn()?; 299 | // Open the default table 300 | let table = txn.open_table(None)?; 301 | // Process the data 302 | for result in key_vals { 303 | let (key, val) = result?; 304 | txn.put(&table, &key, &val, WriteFlags::empty())?; 305 | } 306 | // Commit the batch 307 | txn.commit()?; 308 | Ok(()) 309 | } 310 | 311 | async fn batch_delete_bytes(&self, keys: impl Iterator>) -> Result<()> { 312 | // Create a new transaction 313 | let txn = self.db.begin_rw_txn()?; 314 | // Open the default table 315 | let table = txn.open_table(None)?; 316 | // Process the data 317 | for key in keys { 318 | txn.del(&table, &key, None)?; 319 | } 320 | // Commit the batch 321 | txn.commit()?; 322 | Ok(()) 323 | } 324 | 325 | async fn scan_bytes(&self, scan: &Scan) -> Result { 326 | // Conditional scans are not supported 327 | if scan.condition.is_some() { 328 | bail!(NOT_SUPPORTED_ERROR); 329 | } 330 | // Extract parameters 331 | let s = scan.start.unwrap_or(0); 332 | let l = scan.limit.unwrap_or(usize::MAX); 333 | let p = scan.projection()?; 334 | // Create a new transaction 335 | let txn = self.db.begin_ro_txn()?; 336 | // Open the default table 337 | let table = txn.open_table(None)?; 338 | // Create a cursor for iteration 339 | let iter = txn.cursor(&table)?.into_iter(); 340 | // Perform the relevant projection scan type 341 | match p { 342 | Projection::Id => { 343 | // We use a for loop to iterate over the results, while 344 | // calling black_box internally. This is necessary as 345 | // an iterator with `filter_map` or `map` is optimised 346 | // out by the compiler when calling `count` at the end. 347 | let mut count = 0; 348 | for v in iter.skip(s).take(l) { 349 | black_box(v.unwrap().0); 350 | count += 1; 351 | } 352 | Ok(count) 353 | } 354 | Projection::Full => { 355 | // We use a for loop to iterate over the results, while 356 | // calling black_box internally. This is necessary as 357 | // an iterator with `filter_map` or `map` is optimised 358 | // out by the compiler when calling `count` at the end. 359 | let mut count = 0; 360 | for v in iter.skip(s).take(l) { 361 | black_box(v.unwrap().1); 362 | count += 1; 363 | } 364 | Ok(count) 365 | } 366 | Projection::Count => { 367 | Ok(iter 368 | .skip(s) // Skip the first `offset` entries 369 | .take(l) // Take the next `limit` entries 370 | .count()) 371 | } 372 | } 373 | } 374 | } 375 | -------------------------------------------------------------------------------- /src/surrealkv.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "surrealkv")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 5 | use crate::memory::Config; 6 | use crate::valueprovider::Columns; 7 | use crate::{Benchmark, KeyType, Projection, Scan}; 8 | use anyhow::{Result, bail}; 9 | use serde_json::Value; 10 | use std::hint::black_box; 11 | use std::iter::Iterator; 12 | use std::path::PathBuf; 13 | use std::sync::Arc; 14 | use std::time::Duration; 15 | use surrealkv::Durability; 16 | use surrealkv::Mode::{ReadOnly, ReadWrite}; 17 | use surrealkv::Tree; 18 | use surrealkv::TreeBuilder; 19 | 20 | const DATABASE_DIR: &str = "surrealkv"; 21 | 22 | const BLOCK_SIZE: usize = 64 * 1024; 23 | 24 | /// Calculate SurrealKV specific memory allocation 25 | fn calculate_surrealkv_memory() -> u64 { 26 | // Load the system memory 27 | let memory = Config::new(); 28 | // Return configuration 29 | memory.cache_gb * 1024 * 1024 * 1024 30 | } 31 | 32 | pub(crate) struct SurrealKVClientProvider { 33 | store: Arc, 34 | sync: bool, 35 | } 36 | 37 | impl BenchmarkEngine for SurrealKVClientProvider { 38 | /// The number of seconds to wait before connecting 39 | fn wait_timeout(&self) -> Option { 40 | None 41 | } 42 | /// Initiates a new datastore benchmarking engine 43 | async fn setup(_: KeyType, _columns: Columns, options: &Benchmark) -> Result { 44 | // Cleanup the data directory 45 | std::fs::remove_dir_all(DATABASE_DIR).ok(); 46 | // Calculate memory allocation 47 | let block_cache_bytes = calculate_surrealkv_memory(); 48 | // Configure custom options 49 | let builder = TreeBuilder::new(); 50 | // Enable max memtable size 51 | let builder = builder.with_max_memtable_size(256 * 1024 * 1024); 52 | // Enable the block cache capacity 53 | let builder = builder.with_block_cache_capacity(block_cache_bytes); 54 | // Disable versioned queries 55 | let builder = builder.with_versioning(false, 0); 56 | // Enable separated keys and values 57 | let builder = builder.with_enable_vlog(true); 58 | // Set the block size to 64 KiB 59 | let builder = builder.with_block_size(BLOCK_SIZE); 60 | // Set the directory location 61 | let builder = builder.with_path(PathBuf::from(DATABASE_DIR)); 62 | // Create the datastore 63 | let store = builder.build()?; 64 | // Create the store 65 | Ok(Self { 66 | store: Arc::new(store), 67 | sync: options.sync, 68 | }) 69 | } 70 | /// Creates a new client for this benchmarking engine 71 | async fn create_client(&self) -> Result { 72 | Ok(SurrealKVClient { 73 | db: self.store.clone(), 74 | sync: self.sync, 75 | }) 76 | } 77 | } 78 | 79 | pub(crate) struct SurrealKVClient { 80 | db: Arc, 81 | sync: bool, 82 | } 83 | 84 | impl BenchmarkClient for SurrealKVClient { 85 | async fn shutdown(&self) -> Result<()> { 86 | // Close the database 87 | self.db.close().await?; 88 | // Cleanup the data directory 89 | std::fs::remove_dir_all(DATABASE_DIR).ok(); 90 | // Ok 91 | Ok(()) 92 | } 93 | 94 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 95 | self.create_bytes(&key.to_ne_bytes(), val).await 96 | } 97 | 98 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 99 | self.create_bytes(&key.into_bytes(), val).await 100 | } 101 | 102 | async fn read_u32(&self, key: u32) -> Result<()> { 103 | self.read_bytes(&key.to_ne_bytes()).await 104 | } 105 | 106 | async fn read_string(&self, key: String) -> Result<()> { 107 | self.read_bytes(&key.into_bytes()).await 108 | } 109 | 110 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 111 | self.update_bytes(&key.to_ne_bytes(), val).await 112 | } 113 | 114 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 115 | self.update_bytes(&key.into_bytes(), val).await 116 | } 117 | 118 | async fn delete_u32(&self, key: u32) -> Result<()> { 119 | self.delete_bytes(&key.to_ne_bytes()).await 120 | } 121 | 122 | async fn delete_string(&self, key: String) -> Result<()> { 123 | self.delete_bytes(&key.into_bytes()).await 124 | } 125 | 126 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 127 | self.scan_bytes(scan).await 128 | } 129 | 130 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 131 | self.scan_bytes(scan).await 132 | } 133 | 134 | async fn batch_create_u32( 135 | &self, 136 | key_vals: impl Iterator + Send, 137 | ) -> Result<()> { 138 | let pairs_iter = key_vals.map(|(key, val)| { 139 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 140 | Ok((key.to_ne_bytes().to_vec(), val)) 141 | }); 142 | self.batch_create_bytes(pairs_iter).await 143 | } 144 | 145 | async fn batch_create_string( 146 | &self, 147 | key_vals: impl Iterator + Send, 148 | ) -> Result<()> { 149 | let pairs_iter = key_vals.map(|(key, val)| { 150 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 151 | Ok((key.into_bytes(), val)) 152 | }); 153 | self.batch_create_bytes(pairs_iter).await 154 | } 155 | 156 | async fn batch_read_u32(&self, keys: impl Iterator + Send) -> Result<()> { 157 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 158 | self.batch_read_bytes(keys_iter).await 159 | } 160 | 161 | async fn batch_read_string(&self, keys: impl Iterator + Send) -> Result<()> { 162 | let keys_iter = keys.map(|key| key.into_bytes()); 163 | self.batch_read_bytes(keys_iter).await 164 | } 165 | 166 | async fn batch_update_u32( 167 | &self, 168 | key_vals: impl Iterator + Send, 169 | ) -> Result<()> { 170 | let pairs_iter = key_vals.map(|(key, val)| { 171 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 172 | Ok((key.to_ne_bytes().to_vec(), val)) 173 | }); 174 | self.batch_update_bytes(pairs_iter).await 175 | } 176 | 177 | async fn batch_update_string( 178 | &self, 179 | key_vals: impl Iterator + Send, 180 | ) -> Result<()> { 181 | let pairs_iter = key_vals.map(|(key, val)| { 182 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 183 | Ok((key.into_bytes(), val)) 184 | }); 185 | self.batch_update_bytes(pairs_iter).await 186 | } 187 | 188 | async fn batch_delete_u32(&self, keys: impl Iterator + Send) -> Result<()> { 189 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 190 | self.batch_delete_bytes(keys_iter).await 191 | } 192 | 193 | async fn batch_delete_string(&self, keys: impl Iterator + Send) -> Result<()> { 194 | let keys_iter = keys.map(|key| key.into_bytes()); 195 | self.batch_delete_bytes(keys_iter).await 196 | } 197 | } 198 | 199 | impl SurrealKVClient { 200 | async fn create_bytes(&self, key: &[u8], val: Value) -> Result<()> { 201 | // Serialise the value 202 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 203 | // Create a new transaction 204 | let mut txn = self.db.begin_with_mode(ReadWrite)?; 205 | // Set the transaction durability 206 | txn.set_durability(if self.sync { 207 | Durability::Immediate 208 | } else { 209 | Durability::Eventual 210 | }); 211 | // Process the data 212 | txn.set(key, &val)?; 213 | txn.commit().await?; 214 | Ok(()) 215 | } 216 | 217 | async fn read_bytes(&self, key: &[u8]) -> Result<()> { 218 | // Create a new transaction 219 | let txn = self.db.begin_with_mode(ReadOnly)?; 220 | // Process the data 221 | let res = txn.get(key)?; 222 | // Check the value exists 223 | assert!(res.is_some()); 224 | // Deserialise the value 225 | black_box(res.unwrap()); 226 | // All ok 227 | Ok(()) 228 | } 229 | 230 | async fn update_bytes(&self, key: &[u8], val: Value) -> Result<()> { 231 | // Serialise the value 232 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 233 | // Create a new transaction 234 | let mut txn = self.db.begin_with_mode(ReadWrite)?; 235 | // Set the transaction durability 236 | txn.set_durability(if self.sync { 237 | Durability::Immediate 238 | } else { 239 | Durability::Eventual 240 | }); 241 | // Process the data 242 | txn.set(key, &val)?; 243 | txn.commit().await?; 244 | Ok(()) 245 | } 246 | 247 | async fn delete_bytes(&self, key: &[u8]) -> Result<()> { 248 | // Create a new transaction 249 | let mut txn = self.db.begin_with_mode(ReadWrite)?; 250 | // Set the transaction durability 251 | txn.set_durability(if self.sync { 252 | Durability::Immediate 253 | } else { 254 | Durability::Eventual 255 | }); 256 | // Process the data 257 | txn.delete(key)?; 258 | txn.commit().await?; 259 | Ok(()) 260 | } 261 | 262 | async fn batch_create_bytes( 263 | &self, 264 | key_vals: impl Iterator, Vec)>>, 265 | ) -> Result<()> { 266 | // Create a new transaction 267 | let mut txn = self.db.begin_with_mode(ReadWrite)?; 268 | // Set the transaction durability 269 | txn.set_durability(if self.sync { 270 | Durability::Immediate 271 | } else { 272 | Durability::Eventual 273 | }); 274 | // Process the data 275 | for result in key_vals { 276 | let (key, val) = result?; 277 | txn.set(&key, &val)?; 278 | } 279 | // Commit the batch 280 | txn.commit().await?; 281 | Ok(()) 282 | } 283 | 284 | async fn batch_read_bytes(&self, keys: impl Iterator>) -> Result<()> { 285 | // Create a new transaction 286 | let txn = self.db.begin_with_mode(ReadOnly)?; 287 | // Process the data 288 | for key in keys { 289 | // Get the current value 290 | let res = txn.get(&key)?; 291 | // Check the value exists 292 | assert!(res.is_some()); 293 | // Deserialise the value 294 | black_box(res.unwrap()); 295 | } 296 | // All ok 297 | Ok(()) 298 | } 299 | 300 | async fn batch_update_bytes( 301 | &self, 302 | key_vals: impl Iterator, Vec)>>, 303 | ) -> Result<()> { 304 | // Create a new transaction 305 | let mut txn = self.db.begin_with_mode(ReadWrite)?; 306 | // Set the transaction durability 307 | txn.set_durability(if self.sync { 308 | Durability::Immediate 309 | } else { 310 | Durability::Eventual 311 | }); 312 | // Process the data 313 | for result in key_vals { 314 | let (key, val) = result?; 315 | txn.set(&key, &val)?; 316 | } 317 | // Commit the batch 318 | txn.commit().await?; 319 | Ok(()) 320 | } 321 | 322 | async fn batch_delete_bytes(&self, keys: impl Iterator>) -> Result<()> { 323 | // Create a new transaction 324 | let mut txn = self.db.begin_with_mode(ReadWrite)?; 325 | // Set the transaction durability 326 | txn.set_durability(if self.sync { 327 | Durability::Immediate 328 | } else { 329 | Durability::Eventual 330 | }); 331 | // Process the data 332 | for key in keys { 333 | txn.delete(&key)?; 334 | } 335 | // Commit the batch 336 | txn.commit().await?; 337 | Ok(()) 338 | } 339 | 340 | async fn scan_bytes(&self, scan: &Scan) -> Result { 341 | // Contional scans are not supported 342 | if scan.condition.is_some() { 343 | bail!(NOT_SUPPORTED_ERROR); 344 | } 345 | // Extract parameters 346 | let s = scan.start.unwrap_or(0); 347 | let l = scan.limit.unwrap_or(usize::MAX); 348 | let p = scan.projection()?; 349 | // Create a new transaction 350 | let txn = self.db.begin_with_mode(ReadOnly)?; 351 | let beg = [0u8].as_slice(); 352 | let end = [255u8].as_slice(); 353 | // Perform the relevant projection scan type 354 | match p { 355 | Projection::Id => { 356 | // Create an iterator starting at the beginning 357 | let iter = txn.keys(beg, end)?; 358 | // We use a for loop to iterate over the results, while 359 | // calling black_box internally. This is necessary as 360 | // an iterator with `filter_map` or `map` is optimised 361 | // out by the compiler when calling `count` at the end. 362 | let mut count = 0; 363 | for v in iter.skip(s).take(l) { 364 | assert!(v.is_ok()); 365 | black_box(v.unwrap()); 366 | count += 1; 367 | } 368 | Ok(count) 369 | } 370 | Projection::Full => { 371 | // Create an iterator starting at the beginning 372 | let iter = txn.range(beg, end)?; 373 | // We use a for loop to iterate over the results, while 374 | // calling black_box internally. This is necessary as 375 | // an iterator with `filter_map` or `map` is optimised 376 | // out by the compiler when calling `count` at the end. 377 | let mut count = 0; 378 | for v in iter.skip(s).take(l) { 379 | assert!(v.is_ok()); 380 | black_box(v.unwrap().1); 381 | count += 1; 382 | } 383 | Ok(count) 384 | } 385 | Projection::Count => match scan.limit { 386 | Some(_) => bail!(NOT_SUPPORTED_ERROR), 387 | None => Ok(txn.count(beg, end)?), 388 | }, 389 | } 390 | } 391 | } 392 | -------------------------------------------------------------------------------- /src/valueprovider.rs: -------------------------------------------------------------------------------- 1 | use crate::dialect::Dialect; 2 | use anyhow::{Result, anyhow, bail}; 3 | use log::debug; 4 | use rand::prelude::SmallRng; 5 | use rand::{Rng as RandGen, SeedableRng}; 6 | use serde_json::{Map, Number, Value}; 7 | use std::collections::BTreeMap; 8 | use std::fmt::Display; 9 | use std::ops::Range; 10 | use std::str::FromStr; 11 | use uuid::Uuid; 12 | 13 | pub(crate) struct ValueProvider { 14 | generator: ValueGenerator, 15 | rng: SmallRng, 16 | columns: Columns, 17 | } 18 | 19 | impl ValueProvider { 20 | pub(crate) fn new(json: &str) -> Result { 21 | // Decode the JSON string 22 | let val = serde_json::from_str(json)?; 23 | debug!("Value template: {val:#}"); 24 | // Compile a value generator 25 | let generator = ValueGenerator::new(val)?; 26 | // Identifies the field in the top level (used for column oriented DB like Postgresql) 27 | let columns = Columns::new(&generator)?; 28 | Ok(Self { 29 | generator, 30 | columns, 31 | rng: SmallRng::from_os_rng(), 32 | }) 33 | } 34 | 35 | pub(crate) fn columns(&self) -> Columns { 36 | self.columns.clone() 37 | } 38 | 39 | pub(crate) fn generate_value(&mut self) -> Value 40 | where 41 | D: Dialect, 42 | { 43 | self.generator.generate::(&mut self.rng) 44 | } 45 | } 46 | 47 | impl Clone for ValueProvider { 48 | fn clone(&self) -> Self { 49 | Self { 50 | generator: self.generator.clone(), 51 | rng: SmallRng::from_os_rng(), 52 | columns: self.columns.clone(), 53 | } 54 | } 55 | } 56 | 57 | #[derive(Clone, Debug)] 58 | enum ValueGenerator { 59 | Bool, 60 | String(Length), 61 | Text(Length), 62 | Words(Length, Vec), 63 | Integer, 64 | Float, 65 | DateTime, 66 | Uuid, 67 | // We use i32 for better compatibility across DBs 68 | IntegerRange(Range), 69 | // We use f32 by default for better compatibility across DBs 70 | FloatRange(Range), 71 | StringEnum(Vec), 72 | IntegerEnum(Vec), 73 | FloatEnum(Vec), 74 | Array(Vec), 75 | Object(BTreeMap), 76 | } 77 | 78 | const CHARSET: &[u8; 62] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"; 79 | 80 | fn string(rng: &mut SmallRng, size: usize) -> String { 81 | (0..size) 82 | .map(|_| { 83 | let idx = RandGen::random_range(&mut *rng, 0..CHARSET.len()); 84 | CHARSET[idx] as char 85 | }) 86 | .collect() 87 | } 88 | 89 | fn string_range(rng: &mut SmallRng, range: Range) -> String { 90 | let l = RandGen::random_range(rng, range); 91 | string(rng, l) 92 | } 93 | 94 | fn text(rng: &mut SmallRng, size: usize) -> String { 95 | let mut l = 0; 96 | let mut words = Vec::with_capacity(size / 5); 97 | let mut i = 0; 98 | while l < size { 99 | let w = string_range(rng, 2..10); 100 | l += w.len(); 101 | words.push(w); 102 | l += i; 103 | // We ignore the first whitespace, but not the following ones 104 | i = 1; 105 | } 106 | words.join(" ") 107 | } 108 | 109 | fn text_range(rng: &mut SmallRng, range: Range) -> String { 110 | let l = RandGen::random_range(rng, range); 111 | text(rng, l) 112 | } 113 | 114 | fn words(rng: &mut SmallRng, size: usize, dictionary: &[String]) -> String { 115 | let mut l = 0; 116 | let mut words = Vec::with_capacity(size / 5); 117 | let mut i = 0; 118 | while l < size { 119 | let w = dictionary[rng.random_range(0..dictionary.len())].as_str(); 120 | l += w.len(); 121 | words.push(w); 122 | l += i; 123 | // We ignore the first whitespace, but not the following ones 124 | i = 1; 125 | } 126 | words.join(" ") 127 | } 128 | 129 | fn words_range(rng: &mut SmallRng, range: Range, dictionary: &[String]) -> String { 130 | let l = rng.random_range(range); 131 | words(rng, l, dictionary) 132 | } 133 | 134 | impl ValueGenerator { 135 | fn new(value: Value) -> Result { 136 | match value { 137 | Value::Null => bail!("Unsupported type: Null"), 138 | Value::Bool(_) => bail!("Unsupported type: Bool"), 139 | Value::Number(_) => bail!("Unsupported type: Number"), 140 | Value::String(s) => Self::new_string(s), 141 | Value::Array(a) => Self::new_array(a), 142 | Value::Object(o) => Self::new_object(o), 143 | } 144 | } 145 | 146 | fn new_string(s: String) -> Result { 147 | let s = s.to_lowercase(); 148 | let r = if let Some(i) = s.strip_prefix("string:") { 149 | Self::String(Length::new(i)?) 150 | } else if let Some(i) = s.strip_prefix("text:") { 151 | Self::Text(Length::new(i)?) 152 | } else if let Some(i) = s.strip_prefix("words:") { 153 | // Parse format: "words:50;word1,word2,word3" 154 | let parts: Vec<&str> = i.splitn(2, ';').collect(); 155 | if parts.len() != 2 { 156 | bail!( 157 | "Words format requires length and dictionary separated by semicolon: words:50;word1,word2" 158 | ); 159 | } 160 | let length = Length::new(parts[0])?; 161 | let dictionary: Vec = parts[1].split(',').map(|s| s.to_string()).collect(); 162 | if dictionary.is_empty() { 163 | bail!("Words dictionary cannot be empty"); 164 | } 165 | Self::Words(length, dictionary) 166 | } else if let Some(i) = s.strip_prefix("int:") { 167 | if let Length::Range(r) = Length::new(i)? { 168 | Self::IntegerRange(r) 169 | } else { 170 | bail!("Expected a range but got: {i}"); 171 | } 172 | } else if let Some(i) = s.strip_prefix("float:") { 173 | if let Length::Range(r) = Length::new(i)? { 174 | Self::FloatRange(r) 175 | } else { 176 | bail!("Expected a range but got: {i}"); 177 | } 178 | } else if let Some(s) = s.strip_prefix("string_enum:") { 179 | let labels = s.split(",").map(|s| s.to_string()).collect(); 180 | Self::StringEnum(labels) 181 | } else if let Some(s) = s.strip_prefix("int_enum:") { 182 | let split: Vec<&str> = s.split(",").collect(); 183 | let mut numbers = Vec::with_capacity(split.len()); 184 | for s in split { 185 | numbers.push(s.parse::()?.into()); 186 | } 187 | Self::IntegerEnum(numbers) 188 | } else if let Some(s) = s.strip_prefix("float_enum:") { 189 | let split: Vec<&str> = s.split(",").collect(); 190 | let mut numbers = Vec::with_capacity(split.len()); 191 | for s in split { 192 | numbers.push(Number::from_f64(s.parse::()? as f64).unwrap()); 193 | } 194 | Self::FloatEnum(numbers) 195 | } else if s.eq("bool") { 196 | Self::Bool 197 | } else if s.eq("int") { 198 | Self::Integer 199 | } else if s.eq("float") { 200 | Self::Float 201 | } else if s.eq("datetime") { 202 | Self::DateTime 203 | } else if s.eq("uuid") { 204 | Self::Uuid 205 | } else { 206 | bail!("Unsupported type: {s}"); 207 | }; 208 | Ok(r) 209 | } 210 | 211 | fn new_array(a: Vec) -> Result { 212 | let mut array = Vec::with_capacity(a.len()); 213 | for v in a { 214 | array.push(ValueGenerator::new(v)?); 215 | } 216 | Ok(Self::Array(array)) 217 | } 218 | 219 | fn new_object(o: Map) -> Result { 220 | let mut map = BTreeMap::new(); 221 | for (k, v) in o { 222 | map.insert(k, Self::new(v)?); 223 | } 224 | Ok(Self::Object(map)) 225 | } 226 | 227 | fn generate(&self, rng: &mut SmallRng) -> Value 228 | where 229 | D: Dialect, 230 | { 231 | match self { 232 | ValueGenerator::Bool => { 233 | let v = RandGen::random_bool(&mut *rng, 0.5); 234 | Value::Bool(v) 235 | } 236 | ValueGenerator::String(l) => { 237 | let val = match l { 238 | Length::Range(r) => string_range(rng, r.clone()), 239 | Length::Fixed(l) => string(rng, *l), 240 | }; 241 | Value::String(val) 242 | } 243 | ValueGenerator::Text(l) => { 244 | let val = match l { 245 | Length::Range(r) => text_range(rng, r.clone()), 246 | Length::Fixed(l) => text(rng, *l), 247 | }; 248 | Value::String(val) 249 | } 250 | ValueGenerator::Words(l, dictionary) => { 251 | let val = match l { 252 | Length::Range(r) => words_range(rng, r.clone(), dictionary), 253 | Length::Fixed(l) => words(rng, *l, dictionary), 254 | }; 255 | Value::String(val) 256 | } 257 | ValueGenerator::Integer => { 258 | let v = RandGen::random_range(&mut *rng, i32::MIN..i32::MAX); 259 | Value::Number(Number::from(v)) 260 | } 261 | ValueGenerator::Float => { 262 | let v = RandGen::random_range(&mut *rng, f32::MIN..f32::MAX); 263 | Value::Number(Number::from_f64(v as f64).unwrap()) 264 | } 265 | ValueGenerator::DateTime => { 266 | // Number of seconds from Epoch to 31/12/2030 267 | let s = RandGen::random_range(&mut *rng, 0..1_924_991_999); 268 | D::date_time(s) 269 | } 270 | ValueGenerator::Uuid => { 271 | let uuid = Uuid::new_v4(); 272 | D::uuid(uuid) 273 | } 274 | ValueGenerator::IntegerRange(r) => { 275 | let v = rng.random_range(r.start..r.end); 276 | Value::Number(v.into()) 277 | } 278 | ValueGenerator::FloatRange(r) => { 279 | let v = rng.random_range(r.start..r.end); 280 | Value::Number(Number::from_f64(v as f64).unwrap()) 281 | } 282 | ValueGenerator::StringEnum(a) => { 283 | let i = rng.random_range(0..a.len()); 284 | Value::String(a[i].to_string()) 285 | } 286 | ValueGenerator::IntegerEnum(a) => { 287 | let i = rng.random_range(0..a.len()); 288 | Value::Number(a[i].clone()) 289 | } 290 | ValueGenerator::FloatEnum(a) => { 291 | let i = rng.random_range(0..a.len()); 292 | Value::Number(a[i].clone()) 293 | } 294 | ValueGenerator::Array(a) => { 295 | // Generate any array structure values 296 | let mut vec = Vec::with_capacity(a.len()); 297 | for v in a { 298 | vec.push(v.generate::(rng)); 299 | } 300 | Value::Array(vec) 301 | } 302 | ValueGenerator::Object(o) => { 303 | // Generate any object structure values 304 | let mut map = Map::::new(); 305 | for (k, v) in o { 306 | map.insert(k.clone(), v.generate::(rng)); 307 | } 308 | Value::Object(map) 309 | } 310 | } 311 | } 312 | } 313 | 314 | #[derive(Clone, Debug)] 315 | enum Length 316 | where 317 | Idx: FromStr, 318 | { 319 | Range(Range), 320 | Fixed(Idx), 321 | } 322 | 323 | impl Length 324 | where 325 | Idx: FromStr, 326 | { 327 | fn new(s: &str) -> Result 328 | where 329 | ::Err: Display, 330 | { 331 | // Get the length config setting 332 | let parts: Vec<&str> = s.split("..").collect(); 333 | // Check the length parameter 334 | let r = match parts.len() { 335 | 2 => { 336 | let min = Idx::from_str(parts[0]).map_err(|e| anyhow!("{e}"))?; 337 | let max = Idx::from_str(parts[1]).map_err(|e| anyhow!("{e}"))?; 338 | Self::Range(min..max) 339 | } 340 | 1 => Self::Fixed(Idx::from_str(parts[0]).map_err(|e| anyhow!("{e}"))?), 341 | v => { 342 | bail!("Invalid length generation value: {v}"); 343 | } 344 | }; 345 | Ok(r) 346 | } 347 | } 348 | 349 | #[derive(Clone, Debug)] 350 | /// This structures defines the main columns use for create the schema 351 | /// and insert generated data into a column-oriented database (PostreSQL). 352 | pub(crate) struct Columns(pub(crate) Vec<(String, ColumnType)>); 353 | 354 | impl Columns { 355 | fn new(value: &ValueGenerator) -> Result { 356 | if let ValueGenerator::Object(o) = value { 357 | let mut columns = Vec::with_capacity(o.len()); 358 | for (f, g) in o { 359 | columns.push((f.to_string(), ColumnType::new(g)?)); 360 | } 361 | Ok(Columns(columns)) 362 | } else { 363 | bail!("An object was expected, but got: {value:?}"); 364 | } 365 | } 366 | } 367 | 368 | #[derive(Clone, Debug)] 369 | pub(crate) enum ColumnType { 370 | String, 371 | Integer, 372 | Float, 373 | DateTime, 374 | Uuid, 375 | Object, 376 | Bool, 377 | } 378 | 379 | impl ColumnType { 380 | fn new(v: &ValueGenerator) -> Result { 381 | let r = match v { 382 | ValueGenerator::Object(_) => ColumnType::Object, 383 | ValueGenerator::StringEnum(_) 384 | | ValueGenerator::String(_) 385 | | ValueGenerator::Text(_) 386 | | ValueGenerator::Words(_, _) => ColumnType::String, 387 | ValueGenerator::Integer 388 | | ValueGenerator::IntegerRange(_) 389 | | ValueGenerator::IntegerEnum(_) => ColumnType::Integer, 390 | ValueGenerator::Float 391 | | ValueGenerator::FloatRange(_) 392 | | ValueGenerator::FloatEnum(_) => ColumnType::Float, 393 | ValueGenerator::DateTime => ColumnType::DateTime, 394 | ValueGenerator::Bool => ColumnType::Bool, 395 | ValueGenerator::Uuid => ColumnType::Uuid, 396 | t => { 397 | bail!("Invalid data type: {t:?}"); 398 | } 399 | }; 400 | Ok(r) 401 | } 402 | } 403 | 404 | #[cfg(test)] 405 | mod test { 406 | use super::*; 407 | use crate::dialect::AnsiSqlDialect; 408 | use tokio::task; 409 | 410 | #[tokio::test] 411 | async fn check_all_values_are_unique() { 412 | let vp = ValueProvider::new(r#"{ "int": "int", "int_range": "int:1..99"}"#).unwrap(); 413 | let mut v = vp.clone(); 414 | let f1 = task::spawn(async move { 415 | (v.generate_value::(), v.generate_value::()) 416 | }); 417 | let mut v = vp.clone(); 418 | let f2 = task::spawn(async move { 419 | (v.generate_value::(), v.generate_value::()) 420 | }); 421 | let (v1a, v1b) = f1.await.unwrap(); 422 | let (v2a, v2b) = f2.await.unwrap(); 423 | assert_ne!(v1a, v1b); 424 | assert_ne!(v2a, v2b); 425 | assert_ne!(v1a, v2a); 426 | assert_ne!(v1b, v2b); 427 | } 428 | } 429 | -------------------------------------------------------------------------------- /src/fjall.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "fjall")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 5 | use crate::memory::Config as MemoryConfig; 6 | use crate::valueprovider::Columns; 7 | use crate::{Benchmark, KeyType, Projection, Scan}; 8 | use anyhow::{Result, bail}; 9 | use fjall::{ 10 | Config, KvSeparationOptions, PartitionCreateOptions, PersistMode, TransactionalKeyspace, 11 | TxPartitionHandle, 12 | }; 13 | use serde_json::Value; 14 | use std::hint::black_box; 15 | use std::sync::Arc; 16 | use std::time::Duration; 17 | 18 | const DATABASE_DIR: &str = "fjall"; 19 | 20 | /// Calculate Fjall specific memory allocation 21 | fn calculate_fjall_memory() -> u64 { 22 | // Load the system memory 23 | let memory = MemoryConfig::new(); 24 | // Return configuration 25 | memory.cache_gb * 1024 * 1024 * 1024 26 | } 27 | 28 | // Durability will be set dynamically based on sync flag 29 | 30 | pub(crate) struct FjallClientProvider { 31 | keyspace: Arc, 32 | partition: Arc, 33 | sync: bool, 34 | } 35 | 36 | impl BenchmarkEngine for FjallClientProvider { 37 | /// The number of seconds to wait before connecting 38 | fn wait_timeout(&self) -> Option { 39 | None 40 | } 41 | /// Initiates a new datastore benchmarking engine 42 | async fn setup(_kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 43 | // Cleanup the data directory 44 | std::fs::remove_dir_all(DATABASE_DIR).ok(); 45 | // Calculate memory allocation 46 | let memory = calculate_fjall_memory(); 47 | // Configure the key-value separation 48 | let blobopts = KvSeparationOptions::default() 49 | // Separate values if larger than 1 KiB 50 | .separation_threshold(1024); 51 | // Configure and create the keyspace 52 | let keyspace = Config::new(DATABASE_DIR) 53 | // Fsync data every 100 milliseconds 54 | .fsync_ms(if options.sync { 55 | Some(100) 56 | } else { 57 | None 58 | }) 59 | // Handle transaction flushed automatically 60 | .manual_journal_persist(!options.sync) 61 | // Set the amount of data to build up in memory 62 | .max_write_buffer_size(u64::MAX) 63 | // Set the cache size 64 | .cache_size(memory) 65 | // Open a transactional keyspace 66 | .open_transactional()?; 67 | // Configure and create the partition 68 | let partition = PartitionCreateOptions::default() 69 | // Set the data block size to 32 KiB 70 | .block_size(16 * 1_024) 71 | // Set the max memtable size to 256 MiB 72 | .max_memtable_size(256 * 1_024 * 1_024) 73 | // Separate values if larger than 4 KiB 74 | .with_kv_separation(blobopts); 75 | // Create a default data partition 76 | let partition = keyspace.open_partition("default", partition)?; 77 | // Create the store 78 | Ok(Self { 79 | keyspace: Arc::new(keyspace), 80 | partition: Arc::new(partition), 81 | sync: options.sync, 82 | }) 83 | } 84 | /// Creates a new client for this benchmarking engine 85 | async fn create_client(&self) -> Result { 86 | Ok(FjallClient { 87 | keyspace: self.keyspace.clone(), 88 | partition: self.partition.clone(), 89 | sync: self.sync, 90 | }) 91 | } 92 | } 93 | 94 | pub(crate) struct FjallClient { 95 | keyspace: Arc, 96 | partition: Arc, 97 | sync: bool, 98 | } 99 | 100 | impl BenchmarkClient for FjallClient { 101 | async fn shutdown(&self) -> Result<()> { 102 | // Cleanup the data directory 103 | std::fs::remove_dir_all(DATABASE_DIR).ok(); 104 | // Ok 105 | Ok(()) 106 | } 107 | 108 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 109 | self.create_bytes(&key.to_ne_bytes(), val).await 110 | } 111 | 112 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 113 | self.create_bytes(&key.into_bytes(), val).await 114 | } 115 | 116 | async fn read_u32(&self, key: u32) -> Result<()> { 117 | self.read_bytes(&key.to_ne_bytes()).await 118 | } 119 | 120 | async fn read_string(&self, key: String) -> Result<()> { 121 | self.read_bytes(&key.into_bytes()).await 122 | } 123 | 124 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 125 | self.update_bytes(&key.to_ne_bytes(), val).await 126 | } 127 | 128 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 129 | self.update_bytes(&key.into_bytes(), val).await 130 | } 131 | 132 | async fn delete_u32(&self, key: u32) -> Result<()> { 133 | self.delete_bytes(&key.to_ne_bytes()).await 134 | } 135 | 136 | async fn delete_string(&self, key: String) -> Result<()> { 137 | self.delete_bytes(&key.into_bytes()).await 138 | } 139 | 140 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 141 | self.scan_bytes(scan).await 142 | } 143 | 144 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 145 | self.scan_bytes(scan).await 146 | } 147 | 148 | async fn batch_create_u32( 149 | &self, 150 | key_vals: impl Iterator + Send, 151 | ) -> Result<()> { 152 | let pairs_iter = key_vals.map(|(key, val)| { 153 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 154 | Ok((key.to_ne_bytes().to_vec(), val)) 155 | }); 156 | self.batch_create_bytes(pairs_iter).await 157 | } 158 | 159 | async fn batch_create_string( 160 | &self, 161 | key_vals: impl Iterator + Send, 162 | ) -> Result<()> { 163 | let pairs_iter = key_vals.map(|(key, val)| { 164 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 165 | Ok((key.into_bytes(), val)) 166 | }); 167 | self.batch_create_bytes(pairs_iter).await 168 | } 169 | 170 | async fn batch_read_u32(&self, keys: impl Iterator + Send) -> Result<()> { 171 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 172 | self.batch_read_bytes(keys_iter).await 173 | } 174 | 175 | async fn batch_read_string(&self, keys: impl Iterator + Send) -> Result<()> { 176 | let keys_iter = keys.map(|key| key.into_bytes()); 177 | self.batch_read_bytes(keys_iter).await 178 | } 179 | 180 | async fn batch_update_u32( 181 | &self, 182 | key_vals: impl Iterator + Send, 183 | ) -> Result<()> { 184 | let pairs_iter = key_vals.map(|(key, val)| { 185 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 186 | Ok((key.to_ne_bytes().to_vec(), val)) 187 | }); 188 | self.batch_update_bytes(pairs_iter).await 189 | } 190 | 191 | async fn batch_update_string( 192 | &self, 193 | key_vals: impl Iterator + Send, 194 | ) -> Result<()> { 195 | let pairs_iter = key_vals.map(|(key, val)| { 196 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 197 | Ok((key.into_bytes(), val)) 198 | }); 199 | self.batch_update_bytes(pairs_iter).await 200 | } 201 | 202 | async fn batch_delete_u32(&self, keys: impl Iterator + Send) -> Result<()> { 203 | let keys_iter = keys.map(|key| key.to_ne_bytes().to_vec()); 204 | self.batch_delete_bytes(keys_iter).await 205 | } 206 | 207 | async fn batch_delete_string(&self, keys: impl Iterator + Send) -> Result<()> { 208 | let keys_iter = keys.map(|key| key.into_bytes()); 209 | self.batch_delete_bytes(keys_iter).await 210 | } 211 | } 212 | 213 | impl FjallClient { 214 | async fn create_bytes(&self, key: &[u8], val: Value) -> Result<()> { 215 | // Serialise the value 216 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 217 | // Set the transaction durability 218 | let durability = if self.sync { 219 | None 220 | } else { 221 | Some(PersistMode::Buffer) 222 | }; 223 | // Create a new transaction 224 | let mut txn = self.keyspace.write_tx().durability(durability); 225 | // Process the data 226 | txn.insert(&self.partition, key, val); 227 | txn.commit()?; 228 | Ok(()) 229 | } 230 | 231 | async fn read_bytes(&self, key: &[u8]) -> Result<()> { 232 | // Create a new transaction 233 | let txn = self.keyspace.read_tx(); 234 | // Process the data 235 | let res = txn.get(&self.partition, key)?; 236 | // Check the value exists 237 | assert!(res.is_some()); 238 | // Deserialise the value 239 | black_box(res.unwrap()); 240 | // All ok 241 | Ok(()) 242 | } 243 | 244 | async fn update_bytes(&self, key: &[u8], val: Value) -> Result<()> { 245 | // Serialise the value 246 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 247 | // Set the transaction durability 248 | let durability = if self.sync { 249 | None 250 | } else { 251 | Some(PersistMode::Buffer) 252 | }; 253 | // Create a new transaction 254 | let mut txn = self.keyspace.write_tx().durability(durability); 255 | // Process the data 256 | txn.insert(&self.partition, key, val); 257 | txn.commit()?; 258 | Ok(()) 259 | } 260 | 261 | async fn delete_bytes(&self, key: &[u8]) -> Result<()> { 262 | // Set the transaction durability 263 | let durability = if self.sync { 264 | None 265 | } else { 266 | Some(PersistMode::Buffer) 267 | }; 268 | // Create a new transaction 269 | let mut txn = self.keyspace.write_tx().durability(durability); 270 | // Process the data 271 | txn.remove(&self.partition, key); 272 | txn.commit()?; 273 | Ok(()) 274 | } 275 | 276 | async fn batch_create_bytes( 277 | &self, 278 | key_vals: impl Iterator, Vec)>>, 279 | ) -> Result<()> { 280 | // Set the transaction durability 281 | let durability = if self.sync { 282 | None 283 | } else { 284 | Some(PersistMode::Buffer) 285 | }; 286 | // Create a new transaction 287 | let mut txn = self.keyspace.write_tx().durability(durability); 288 | // Process the data 289 | for result in key_vals { 290 | let (key, val) = result?; 291 | txn.insert(&self.partition, &key, val); 292 | } 293 | // Commit the batch 294 | txn.commit()?; 295 | Ok(()) 296 | } 297 | 298 | async fn batch_read_bytes(&self, keys: impl Iterator>) -> Result<()> { 299 | // Create a new transaction 300 | let txn = self.keyspace.read_tx(); 301 | // Process the data 302 | for key in keys { 303 | // Get the current value 304 | let res = txn.get(&self.partition, &key)?; 305 | // Check the value exists 306 | assert!(res.is_some()); 307 | // Deserialise the value 308 | black_box(res.unwrap()); 309 | } 310 | // All ok 311 | Ok(()) 312 | } 313 | 314 | async fn batch_update_bytes( 315 | &self, 316 | key_vals: impl Iterator, Vec)>>, 317 | ) -> Result<()> { 318 | // Set the transaction durability 319 | let durability = if self.sync { 320 | None 321 | } else { 322 | Some(PersistMode::Buffer) 323 | }; 324 | // Create a new transaction 325 | let mut txn = self.keyspace.write_tx().durability(durability); 326 | // Process the data 327 | for result in key_vals { 328 | let (key, val) = result?; 329 | txn.insert(&self.partition, &key, val); 330 | } 331 | // Commit the batch 332 | txn.commit()?; 333 | Ok(()) 334 | } 335 | 336 | async fn batch_delete_bytes(&self, keys: impl Iterator>) -> Result<()> { 337 | // Set the transaction durability 338 | let durability = if self.sync { 339 | None 340 | } else { 341 | Some(PersistMode::Buffer) 342 | }; 343 | // Create a new transaction 344 | let mut txn = self.keyspace.write_tx().durability(durability); 345 | // Process the data 346 | for key in keys { 347 | txn.remove(&self.partition, &key); 348 | } 349 | // Commit the batch 350 | txn.commit()?; 351 | Ok(()) 352 | } 353 | 354 | async fn scan_bytes(&self, scan: &Scan) -> Result { 355 | // Contional scans are not supported 356 | if scan.condition.is_some() { 357 | bail!(NOT_SUPPORTED_ERROR); 358 | } 359 | // Extract parameters 360 | let s = scan.start.unwrap_or(0); 361 | let l = scan.limit.unwrap_or(usize::MAX); 362 | let p = scan.projection()?; 363 | // Create a new transaction 364 | let txn = self.keyspace.read_tx(); 365 | // Perform the relevant projection scan type 366 | match p { 367 | Projection::Id => { 368 | // Create an iterator starting at the beginning 369 | let iter = txn.keys(&self.partition); 370 | // We use a for loop to iterate over the results, while 371 | // calling black_box internally. This is necessary as 372 | // an iterator with `filter_map` or `map` is optimised 373 | // out by the compiler when calling `count` at the end. 374 | let mut count = 0; 375 | for v in iter.skip(s).take(l) { 376 | black_box(v.unwrap()); 377 | count += 1; 378 | } 379 | Ok(count) 380 | } 381 | Projection::Full => { 382 | // Create an iterator starting at the beginning 383 | let iter = txn.iter(&self.partition); 384 | // calling black_box internally. This is necessary as 385 | // an iterator with `filter_map` or `map` is optimised 386 | // out by the compiler when calling `count` at the end. 387 | let mut count = 0; 388 | for v in iter.skip(s).take(l) { 389 | black_box(v.unwrap().1); 390 | count += 1; 391 | } 392 | Ok(count) 393 | } 394 | Projection::Count => { 395 | Ok(txn 396 | .keys(&self.partition) 397 | .skip(s) // Skip the first `offset` entries 398 | .take(l) // Take the next `limit` entries 399 | .count()) 400 | } 401 | } 402 | } 403 | } 404 | -------------------------------------------------------------------------------- /src/redb.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "redb")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 5 | use crate::memory::Config; 6 | use crate::valueprovider::Columns; 7 | use crate::{Benchmark, KeyType, Projection, Scan}; 8 | use anyhow::{Result, bail}; 9 | use redb::{Database, Durability, ReadableDatabase, ReadableTable, TableDefinition}; 10 | use serde_json::Value; 11 | use std::hint::black_box; 12 | use std::sync::Arc; 13 | use std::time::Duration; 14 | 15 | const DATABASE_DIR: &str = "redb"; 16 | 17 | /// Calculate ReDB specific memory allocation 18 | fn calculate_redb_memory() -> u64 { 19 | // Load the system memory 20 | let memory = Config::new(); 21 | // Return configuration 22 | memory.cache_gb * 1024 * 1024 * 1024 23 | } 24 | 25 | const TABLE: TableDefinition<&[u8], Vec> = TableDefinition::new("test"); 26 | 27 | pub(crate) struct ReDBClientProvider { 28 | db: Arc, 29 | sync: bool, 30 | } 31 | 32 | impl BenchmarkEngine for ReDBClientProvider { 33 | /// The number of seconds to wait before connecting 34 | fn wait_timeout(&self) -> Option { 35 | None 36 | } 37 | /// Initiates a new datastore benchmarking engine 38 | async fn setup(_kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 39 | // Cleanup the data directory 40 | std::fs::remove_file(DATABASE_DIR).ok(); 41 | // Calculate memory allocation 42 | let memory = calculate_redb_memory(); 43 | // Configure and create the database 44 | let db = Database::builder() 45 | // Set the cache size 46 | .set_cache_size(memory as usize) 47 | // Create the database directory 48 | .create(DATABASE_DIR)?; 49 | // Create the store 50 | Ok(Self { 51 | db: Arc::new(db), 52 | sync: options.sync, 53 | }) 54 | } 55 | /// Creates a new client for this benchmarking engine 56 | async fn create_client(&self) -> Result { 57 | Ok(ReDBClient { 58 | db: self.db.clone(), 59 | sync: self.sync, 60 | }) 61 | } 62 | } 63 | 64 | pub(crate) struct ReDBClient { 65 | db: Arc, 66 | sync: bool, 67 | } 68 | 69 | impl BenchmarkClient for ReDBClient { 70 | async fn shutdown(&self) -> Result<()> { 71 | // Cleanup the data directory 72 | std::fs::remove_file(DATABASE_DIR).ok(); 73 | // Ok 74 | Ok(()) 75 | } 76 | 77 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 78 | self.create_bytes(&key.to_ne_bytes(), val).await 79 | } 80 | 81 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 82 | self.create_bytes(&key.into_bytes(), val).await 83 | } 84 | 85 | async fn read_u32(&self, key: u32) -> Result<()> { 86 | self.read_bytes(&key.to_ne_bytes()).await 87 | } 88 | 89 | async fn read_string(&self, key: String) -> Result<()> { 90 | self.read_bytes(&key.into_bytes()).await 91 | } 92 | 93 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 94 | self.update_bytes(&key.to_ne_bytes(), val).await 95 | } 96 | 97 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 98 | self.update_bytes(&key.into_bytes(), val).await 99 | } 100 | 101 | async fn delete_u32(&self, key: u32) -> Result<()> { 102 | self.delete_bytes(&key.to_ne_bytes()).await 103 | } 104 | 105 | async fn delete_string(&self, key: String) -> Result<()> { 106 | self.delete_bytes(&key.into_bytes()).await 107 | } 108 | 109 | async fn scan_u32(&self, scan: &Scan, _ctx: ScanContext) -> Result { 110 | self.scan_bytes(scan).await 111 | } 112 | 113 | async fn scan_string(&self, scan: &Scan, _ctx: ScanContext) -> Result { 114 | self.scan_bytes(scan).await 115 | } 116 | 117 | async fn batch_create_u32( 118 | &self, 119 | key_vals: impl Iterator + Send, 120 | ) -> Result<()> { 121 | let pairs: Result> = key_vals 122 | .map(|(key, val)| { 123 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 124 | Ok((key.to_ne_bytes().to_vec(), val)) 125 | }) 126 | .collect(); 127 | self.batch_create_bytes(pairs?.into_iter().map(Ok)).await 128 | } 129 | 130 | async fn batch_create_string( 131 | &self, 132 | key_vals: impl Iterator + Send, 133 | ) -> Result<()> { 134 | let pairs: Result> = key_vals 135 | .map(|(key, val)| { 136 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 137 | Ok((key.into_bytes(), val)) 138 | }) 139 | .collect(); 140 | self.batch_create_bytes(pairs?.into_iter().map(Ok)).await 141 | } 142 | 143 | async fn batch_read_u32(&self, keys: impl Iterator + Send) -> Result<()> { 144 | let keys_vec: Vec<_> = keys.map(|key| key.to_ne_bytes().to_vec()).collect(); 145 | self.batch_read_bytes(keys_vec.into_iter()).await 146 | } 147 | 148 | async fn batch_read_string(&self, keys: impl Iterator + Send) -> Result<()> { 149 | let keys_vec: Vec<_> = keys.map(|key| key.into_bytes()).collect(); 150 | self.batch_read_bytes(keys_vec.into_iter()).await 151 | } 152 | 153 | async fn batch_update_u32( 154 | &self, 155 | key_vals: impl Iterator + Send, 156 | ) -> Result<()> { 157 | let pairs: Result> = key_vals 158 | .map(|(key, val)| { 159 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 160 | Ok((key.to_ne_bytes().to_vec(), val)) 161 | }) 162 | .collect(); 163 | self.batch_update_bytes(pairs?.into_iter().map(Ok)).await 164 | } 165 | 166 | async fn batch_update_string( 167 | &self, 168 | key_vals: impl Iterator + Send, 169 | ) -> Result<()> { 170 | let pairs: Result> = key_vals 171 | .map(|(key, val)| { 172 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 173 | Ok((key.into_bytes(), val)) 174 | }) 175 | .collect(); 176 | self.batch_update_bytes(pairs?.into_iter().map(Ok)).await 177 | } 178 | 179 | async fn batch_delete_u32(&self, keys: impl Iterator + Send) -> Result<()> { 180 | let keys_vec: Vec<_> = keys.map(|key| key.to_ne_bytes().to_vec()).collect(); 181 | self.batch_delete_bytes(keys_vec.into_iter()).await 182 | } 183 | 184 | async fn batch_delete_string(&self, keys: impl Iterator + Send) -> Result<()> { 185 | let keys_vec: Vec<_> = keys.map(|key| key.into_bytes()).collect(); 186 | self.batch_delete_bytes(keys_vec.into_iter()).await 187 | } 188 | } 189 | 190 | impl ReDBClient { 191 | async fn create_bytes(&self, key: &[u8], val: Value) -> Result<()> { 192 | // Clone the datastore and sync flag 193 | let db = self.db.clone(); 194 | let sync = self.sync; 195 | // Execute on the blocking threadpool 196 | affinitypool::spawn_local(move || -> Result<_> { 197 | // Serialise the value 198 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 199 | // Create a new transaction 200 | let mut txn = db.begin_write()?; 201 | // Set the transaction durability 202 | let _ = txn.set_durability(if sync { 203 | Durability::Immediate 204 | } else { 205 | Durability::None 206 | }); 207 | // Open the database table 208 | let mut tab = txn.open_table(TABLE)?; 209 | // Process the data 210 | tab.insert(key, val)?; 211 | drop(tab); 212 | txn.commit()?; 213 | Ok(()) 214 | }) 215 | .await 216 | } 217 | 218 | async fn read_bytes(&self, key: &[u8]) -> Result<()> { 219 | // Clone the datastore and key 220 | let db = self.db.clone(); 221 | let key = key.to_vec(); 222 | // Execute on the blocking threadpool 223 | affinitypool::spawn_local(move || -> Result<_> { 224 | // Create a new transaction 225 | let txn = db.begin_read()?; 226 | // Open the database table 227 | let tab = txn.open_table(TABLE)?; 228 | // Process the data 229 | let res: Option<_> = tab.get(&key[..])?; 230 | // Check the value exists 231 | assert!(res.is_some()); 232 | // Deserialise the value 233 | black_box(res.unwrap().value()); 234 | // All ok 235 | Ok(()) 236 | }) 237 | .await 238 | } 239 | 240 | async fn update_bytes(&self, key: &[u8], val: Value) -> Result<()> { 241 | // Clone the datastore and sync flag 242 | let db = self.db.clone(); 243 | let sync = self.sync; 244 | // Execute on the blocking threadpool 245 | affinitypool::spawn_local(move || -> Result<_> { 246 | // Serialise the value 247 | let val = bincode::serde::encode_to_vec(&val, bincode::config::standard())?; 248 | // Create a new transaction 249 | let mut txn = db.begin_write()?; 250 | // Set the transaction durability 251 | let _ = txn.set_durability(if sync { 252 | Durability::Immediate 253 | } else { 254 | Durability::None 255 | }); 256 | // Open the database table 257 | let mut tab = txn.open_table(TABLE)?; 258 | // Process the data 259 | tab.insert(key, val)?; 260 | drop(tab); 261 | txn.commit()?; 262 | Ok(()) 263 | }) 264 | .await 265 | } 266 | 267 | async fn delete_bytes(&self, key: &[u8]) -> Result<()> { 268 | // Clone the datastore and sync flag 269 | let db = self.db.clone(); 270 | let sync = self.sync; 271 | // Execute on the blocking threadpool 272 | affinitypool::spawn_local(move || -> Result<_> { 273 | // Create a new transaction 274 | let mut txn = db.begin_write()?; 275 | // Set the transaction durability 276 | let _ = txn.set_durability(if sync { 277 | Durability::Immediate 278 | } else { 279 | Durability::None 280 | }); 281 | // Open the database table 282 | let mut tab = txn.open_table(TABLE)?; 283 | // Process the data 284 | tab.remove(key)?; 285 | drop(tab); 286 | txn.commit()?; 287 | Ok(()) 288 | }) 289 | .await 290 | } 291 | 292 | async fn batch_create_bytes( 293 | &self, 294 | key_vals: impl Iterator, Vec)>> + Send + 'static, 295 | ) -> Result<()> { 296 | // Clone the datastore and sync flag 297 | let db = self.db.clone(); 298 | let sync = self.sync; 299 | // Execute on the blocking threadpool 300 | affinitypool::spawn_local(move || -> Result<_> { 301 | // Create a new transaction 302 | let mut txn = db.begin_write()?; 303 | // Set the transaction durability 304 | let _ = txn.set_durability(if sync { 305 | Durability::Immediate 306 | } else { 307 | Durability::None 308 | }); 309 | // Open the database table 310 | let mut tab = txn.open_table(TABLE)?; 311 | // Process all the data in batch 312 | for result in key_vals { 313 | let (key, val) = result?; 314 | tab.insert(&key[..], val)?; 315 | } 316 | drop(tab); 317 | txn.commit()?; 318 | Ok(()) 319 | }) 320 | .await 321 | } 322 | 323 | async fn batch_read_bytes( 324 | &self, 325 | keys: impl Iterator> + Send + 'static, 326 | ) -> Result<()> { 327 | // Clone the datastore 328 | let db = self.db.clone(); 329 | // Execute on the blocking threadpool 330 | affinitypool::spawn_local(move || -> Result<_> { 331 | // Create a new transaction 332 | let txn = db.begin_read()?; 333 | // Open the database table 334 | let tab = txn.open_table(TABLE)?; 335 | // Process all the data in batch 336 | for key in keys { 337 | // Process the data 338 | let res: Option<_> = tab.get(&key[..])?; 339 | // Check the value exists 340 | assert!(res.is_some()); 341 | // Deserialise the value 342 | black_box(res.unwrap().value()); 343 | } 344 | // All ok 345 | Ok(()) 346 | }) 347 | .await 348 | } 349 | 350 | async fn batch_update_bytes( 351 | &self, 352 | key_vals: impl Iterator, Vec)>> + Send + 'static, 353 | ) -> Result<()> { 354 | // Clone the datastore and sync flag 355 | let db = self.db.clone(); 356 | let sync = self.sync; 357 | // Execute on the blocking threadpool 358 | affinitypool::spawn_local(move || -> Result<_> { 359 | // Create a new transaction 360 | let mut txn = db.begin_write()?; 361 | // Set the transaction durability 362 | let _ = txn.set_durability(if sync { 363 | Durability::Immediate 364 | } else { 365 | Durability::None 366 | }); 367 | // Open the database table 368 | let mut tab = txn.open_table(TABLE)?; 369 | // Process all the data in batch 370 | for result in key_vals { 371 | let (key, val) = result?; 372 | tab.insert(&key[..], val)?; 373 | } 374 | drop(tab); 375 | txn.commit()?; 376 | Ok(()) 377 | }) 378 | .await 379 | } 380 | 381 | async fn batch_delete_bytes( 382 | &self, 383 | keys: impl Iterator> + Send + 'static, 384 | ) -> Result<()> { 385 | // Clone the datastore and sync flag 386 | let db = self.db.clone(); 387 | let sync = self.sync; 388 | // Execute on the blocking threadpool 389 | affinitypool::spawn_local(move || -> Result<_> { 390 | // Create a new transaction 391 | let mut txn = db.begin_write()?; 392 | // Set the transaction durability 393 | let _ = txn.set_durability(if sync { 394 | Durability::Immediate 395 | } else { 396 | Durability::None 397 | }); 398 | // Open the database table 399 | let mut tab = txn.open_table(TABLE)?; 400 | // Process all the data in batch 401 | for key in keys { 402 | tab.remove(&key[..])?; 403 | } 404 | drop(tab); 405 | txn.commit()?; 406 | Ok(()) 407 | }) 408 | .await 409 | } 410 | 411 | async fn scan_bytes(&self, scan: &Scan) -> Result { 412 | // Contional scans are not supported 413 | if scan.condition.is_some() { 414 | bail!(NOT_SUPPORTED_ERROR); 415 | } 416 | // Extract parameters 417 | let s = scan.start.unwrap_or(0); 418 | let l = scan.limit.unwrap_or(usize::MAX); 419 | let p = scan.projection()?; 420 | // Clone the datastore 421 | let db = self.db.clone(); 422 | // Execute on the blocking threadpool 423 | affinitypool::spawn_local(|| -> Result<_> { 424 | // Create a new transaction 425 | let txn = db.begin_read()?; 426 | // Open the database table 427 | let tab = txn.open_table(TABLE)?; 428 | // Create an iterator starting at the beginning 429 | let iter = tab.iter()?; 430 | // Perform the relevant projection scan type 431 | match p { 432 | Projection::Id => { 433 | // We use a for loop to iterate over the results, while 434 | // calling black_box internally. This is necessary as 435 | // an iterator with `filter_map` or `map` is optimised 436 | // out by the compiler when calling `count` at the end. 437 | let mut count = 0; 438 | for v in iter.skip(s).take(l) { 439 | black_box(v.unwrap().1.value()); 440 | count += 1; 441 | } 442 | Ok(count) 443 | } 444 | Projection::Full => { 445 | // We use a for loop to iterate over the results, while 446 | // calling black_box internally. This is necessary as 447 | // an iterator with `filter_map` or `map` is optimised 448 | // out by the compiler when calling `count` at the end. 449 | let mut count = 0; 450 | for v in iter.skip(s).take(l) { 451 | black_box(v.unwrap().1.value()); 452 | count += 1; 453 | } 454 | Ok(count) 455 | } 456 | Projection::Count => { 457 | Ok(iter 458 | .skip(s) // Skip the first `offset` entries 459 | .take(l) // Take the next `limit` entries 460 | .count()) 461 | } 462 | } 463 | }) 464 | .await 465 | } 466 | } 467 | -------------------------------------------------------------------------------- /src/mongodb.rs: -------------------------------------------------------------------------------- 1 | #![cfg(feature = "mongodb")] 2 | 3 | use crate::benchmark::NOT_SUPPORTED_ERROR; 4 | use crate::dialect::MongoDBDialect; 5 | use crate::docker::DockerParams; 6 | use crate::engine::{BenchmarkClient, BenchmarkEngine, ScanContext}; 7 | use crate::memory::Config; 8 | use crate::valueprovider::Columns; 9 | use crate::{Benchmark, Index, KeyType, Projection, Scan}; 10 | use anyhow::{Result, bail}; 11 | use futures::{StreamExt, TryStreamExt}; 12 | use mongodb::IndexModel; 13 | use mongodb::Namespace; 14 | use mongodb::bson::{Bson, Document, doc}; 15 | use mongodb::options::ClientOptions; 16 | use mongodb::options::DatabaseOptions; 17 | use mongodb::options::IndexOptions; 18 | use mongodb::options::ReadConcern; 19 | use mongodb::options::{Acknowledgment, ReplaceOneModel, WriteConcern, WriteModel}; 20 | use mongodb::{Client, Collection, Cursor, Database, bson}; 21 | use serde_json::Value; 22 | use std::hint::black_box; 23 | use std::time::Duration; 24 | 25 | pub const DEFAULT: &str = "mongodb://root:root@127.0.0.1:27017"; 26 | 27 | /// Calculate MongoDB specific memory allocation 28 | fn calculate_mongodb_memory() -> u64 { 29 | // Load the system memory 30 | let memory = Config::new(); 31 | // Use ~80% of recommended cache allocation 32 | (memory.cache_gb * 4 / 5).max(1) 33 | } 34 | 35 | pub(crate) fn docker(options: &Benchmark) -> DockerParams { 36 | // Calculate memory allocation 37 | let cache_gb = calculate_mongodb_memory(); 38 | // Return Docker parameters 39 | DockerParams { 40 | image: "mongo", 41 | pre_args: "--ulimit nofile=65536:65536 -p 127.0.0.1:27017:27017 -e MONGO_INITDB_ROOT_USERNAME=root -e MONGO_INITDB_ROOT_PASSWORD=root".to_string(), 42 | post_args: match options.optimised { 43 | // Optimised configuration 44 | true => format!("mongod --wiredTigerCacheSizeGB {cache_gb}"), 45 | // Default configuration 46 | false => "".to_string(), 47 | }, 48 | } 49 | } 50 | 51 | pub(crate) struct MongoDBClientProvider { 52 | sync: bool, 53 | client: Client, 54 | } 55 | 56 | impl BenchmarkEngine for MongoDBClientProvider { 57 | /// Initiates a new datastore benchmarking engine 58 | async fn setup(_kt: KeyType, _columns: Columns, options: &Benchmark) -> Result { 59 | // Get the custom endpoint if specified 60 | let url = options.endpoint.as_deref().unwrap_or(DEFAULT).to_owned(); 61 | // Create a new client with a connection pool. 62 | // The MongoDB client does not correctly limit 63 | // the number of connections in the connection 64 | // pool. Therefore we create a single connection 65 | // pool and share it with all of the crud-bench 66 | // clients. This follows the recommended advice 67 | // for using the MongoDB driver. Note that this 68 | // still creates 2 more connections than has 69 | // been specified in the `max_pool_size` option. 70 | let mut opts = ClientOptions::parse(url).await?; 71 | opts.max_pool_size = Some(options.clients); 72 | opts.min_pool_size = None; 73 | // Set server selection timeout to 60 seconds (default 30s) 74 | opts.server_selection_timeout = Some(Duration::from_secs(60)); 75 | // Set server connect timeout to 30 seconds (default 10s) 76 | opts.connect_timeout = Some(Duration::from_secs(30)); 77 | // Reduce monitoring heartbeats for batch operations (default 10s) 78 | opts.heartbeat_freq = Some(Duration::from_secs(30)); 79 | // Create the client provider 80 | Ok(Self { 81 | sync: options.sync, 82 | client: Client::with_options(opts)?, 83 | }) 84 | } 85 | /// Creates a new client for this benchmarking engine 86 | async fn create_client(&self) -> Result { 87 | let db = self.client.database_with_options( 88 | "crud-bench", 89 | DatabaseOptions::builder() 90 | // Configure the write concern options 91 | .write_concern( 92 | // Configure the write options 93 | WriteConcern::builder() 94 | // Ensure that all writes are written, 95 | // replicated, and acknowledged by the 96 | // majority of nodes in the cluster. 97 | .w(Acknowledgment::Majority) 98 | // Configure journal durability based on sync setting. 99 | // When `true`: writes are acknowledged only after 100 | // being written to the on-disk journal (full durability). 101 | // When `false`: writes are acknowledged after being 102 | // written to memory (faster, less durable). 103 | .journal(self.sync) 104 | // Finalise the write options 105 | .build(), 106 | ) 107 | // Configure the read concern options 108 | .read_concern(ReadConcern::majority()) 109 | // Finalise the database configuration 110 | .build(), 111 | ); 112 | Ok(MongoDBClient { 113 | db, 114 | sync: self.sync, 115 | }) 116 | } 117 | } 118 | 119 | pub(crate) struct MongoDBClient { 120 | db: Database, 121 | sync: bool, 122 | } 123 | 124 | impl BenchmarkClient for MongoDBClient { 125 | async fn compact(&self) -> Result<()> { 126 | // For a database compaction 127 | self.db 128 | .run_command(doc! { 129 | "compact": "record", 130 | "dryRun": false, 131 | "force": true, 132 | }) 133 | .await?; 134 | // Ok 135 | Ok(()) 136 | } 137 | 138 | async fn create_u32(&self, key: u32, val: Value) -> Result<()> { 139 | self.create(key, val).await 140 | } 141 | 142 | async fn create_string(&self, key: String, val: Value) -> Result<()> { 143 | self.create(key, val).await 144 | } 145 | 146 | async fn read_u32(&self, key: u32) -> Result<()> { 147 | let doc = self.read(key).await?; 148 | assert_eq!(doc.unwrap().get("_id").unwrap().as_i64().unwrap() as u32, key); 149 | Ok(()) 150 | } 151 | 152 | async fn read_string(&self, key: String) -> Result<()> { 153 | let doc = self.read(&key).await?; 154 | assert_eq!(doc.unwrap().get_str("_id")?, key); 155 | Ok(()) 156 | } 157 | 158 | async fn update_u32(&self, key: u32, val: Value) -> Result<()> { 159 | self.update(key, val).await 160 | } 161 | 162 | async fn update_string(&self, key: String, val: Value) -> Result<()> { 163 | self.update(key, val).await 164 | } 165 | 166 | async fn delete_u32(&self, key: u32) -> Result<()> { 167 | self.delete(key).await 168 | } 169 | 170 | async fn delete_string(&self, key: String) -> Result<()> { 171 | self.delete(key).await 172 | } 173 | 174 | async fn build_index(&self, spec: &Index, name: &str) -> Result<()> { 175 | // Define the index document 176 | let mut doc = Document::new(); 177 | // Check if an index type is specified 178 | match &spec.index_type { 179 | Some(kind) if kind == "fulltext" => { 180 | // Create a text index 181 | for field in &spec.fields { 182 | doc.insert(field, "text"); 183 | } 184 | } 185 | Some(kind) => { 186 | // Other index types (e.g., "2d", "2dsphere", "hashed") 187 | for field in &spec.fields { 188 | doc.insert(field, kind.as_str()); 189 | } 190 | } 191 | None => { 192 | // Standard ascending index 193 | for field in &spec.fields { 194 | doc.insert(field, 1); 195 | } 196 | } 197 | }; 198 | // Define the index options 199 | let mut options = IndexOptions::default(); 200 | options.name = Some(name.to_string()); 201 | if let Some(unique) = spec.unique { 202 | options.unique = Some(unique); 203 | } 204 | // Create the index model 205 | let index_model = IndexModel::builder().keys(doc).options(options).build(); 206 | // Create the index 207 | self.collection().create_index(index_model).await?; 208 | // All ok 209 | Ok(()) 210 | } 211 | 212 | async fn drop_index(&self, name: &str) -> Result<()> { 213 | self.collection().drop_index(name).await?; 214 | Ok(()) 215 | } 216 | 217 | async fn scan_u32(&self, scan: &Scan, ctx: ScanContext) -> Result { 218 | self.scan(scan, ctx).await 219 | } 220 | 221 | async fn scan_string(&self, scan: &Scan, ctx: ScanContext) -> Result { 222 | self.scan(scan, ctx).await 223 | } 224 | 225 | async fn batch_create_u32( 226 | &self, 227 | key_vals: impl Iterator + Send, 228 | ) -> Result<()> { 229 | self.batch_create(key_vals.collect()).await 230 | } 231 | 232 | async fn batch_create_string( 233 | &self, 234 | key_vals: impl Iterator + Send, 235 | ) -> Result<()> { 236 | self.batch_create(key_vals.collect()).await 237 | } 238 | 239 | async fn batch_read_u32(&self, keys: impl Iterator + Send) -> Result<()> { 240 | self.batch_read(keys.collect()).await 241 | } 242 | 243 | async fn batch_read_string(&self, keys: impl Iterator + Send) -> Result<()> { 244 | self.batch_read(keys.collect()).await 245 | } 246 | 247 | async fn batch_update_u32( 248 | &self, 249 | key_vals: impl Iterator + Send, 250 | ) -> Result<()> { 251 | self.batch_update(key_vals.collect()).await 252 | } 253 | 254 | async fn batch_update_string( 255 | &self, 256 | key_vals: impl Iterator + Send, 257 | ) -> Result<()> { 258 | self.batch_update(key_vals.collect()).await 259 | } 260 | 261 | async fn batch_delete_u32(&self, keys: impl Iterator + Send) -> Result<()> { 262 | self.batch_delete(keys.collect()).await 263 | } 264 | 265 | async fn batch_delete_string(&self, keys: impl Iterator + Send) -> Result<()> { 266 | self.batch_delete(keys.collect()).await 267 | } 268 | } 269 | 270 | impl MongoDBClient { 271 | fn collection(&self) -> Collection { 272 | self.db.collection("record") 273 | } 274 | 275 | fn to_doc(key: K, mut val: Value) -> Result 276 | where 277 | K: Into + Into, 278 | { 279 | let obj = val.as_object_mut().unwrap(); 280 | obj.insert("_id".to_string(), key.into()); 281 | Ok(bson::to_bson(&val)?) 282 | } 283 | 284 | async fn create(&self, key: K, val: Value) -> Result<()> 285 | where 286 | K: Into + Into, 287 | { 288 | let bson = Self::to_doc(key, val)?; 289 | let doc = bson.as_document().unwrap(); 290 | let res = self.collection().insert_one(doc).await?; 291 | assert_ne!(res.inserted_id, Bson::Null); 292 | Ok(()) 293 | } 294 | 295 | async fn read(&self, key: K) -> Result> 296 | where 297 | K: Into, 298 | { 299 | let filter = doc! { "_id": key }; 300 | let doc = self.collection().find_one(filter).await?; 301 | assert!(doc.is_some()); 302 | Ok(doc) 303 | } 304 | 305 | async fn update(&self, key: K, val: Value) -> Result<()> 306 | where 307 | K: Into + Into + Clone, 308 | { 309 | let filter = doc! { "_id": key.clone() }; 310 | let bson = Self::to_doc(key, val)?; 311 | let doc = bson.as_document().unwrap(); 312 | let res = self.collection().replace_one(filter, doc).await?; 313 | assert_eq!(res.modified_count, 1); 314 | Ok(()) 315 | } 316 | 317 | async fn delete(&self, key: K) -> Result<()> 318 | where 319 | K: Into, 320 | { 321 | let filter = doc! { "_id": key }; 322 | let res = self.collection().delete_one(filter).await?; 323 | assert_eq!(res.deleted_count, 1); 324 | Ok(()) 325 | } 326 | 327 | async fn batch_create(&self, key_vals: Vec<(K, Value)>) -> Result<()> 328 | where 329 | K: Into + Into, 330 | { 331 | let mut docs = Vec::with_capacity(key_vals.len()); 332 | for (key, val) in key_vals { 333 | let bson = Self::to_doc(key, val)?; 334 | docs.push(bson.as_document().unwrap().clone()); 335 | } 336 | let docs_len = docs.len(); 337 | let res = self.collection().insert_many(docs).await?; 338 | assert_eq!(res.inserted_ids.len(), docs_len); 339 | Ok(()) 340 | } 341 | 342 | async fn batch_read(&self, keys: Vec) -> Result<()> 343 | where 344 | K: Into, 345 | { 346 | let keys_len = keys.len(); 347 | let ids: Vec = keys.into_iter().map(|k| k.into()).collect(); 348 | let filter = doc! { "_id": { "$in": ids } }; 349 | let cursor = self.collection().find(filter).await?; 350 | let docs: Vec = cursor.try_collect().await?; 351 | assert_eq!(docs.len(), keys_len); 352 | for doc in docs { 353 | black_box(doc); 354 | } 355 | Ok(()) 356 | } 357 | 358 | async fn batch_update(&self, key_vals: Vec<(K, Value)>) -> Result<()> 359 | where 360 | K: Into + Into + Clone, 361 | { 362 | let namespace = Namespace { 363 | db: self.db.name().to_string(), 364 | coll: "record".to_string(), 365 | }; 366 | let mut docs = Vec::with_capacity(key_vals.len()); 367 | for (key, val) in key_vals { 368 | let filter = doc! { "_id": Into::::into(key.clone()) }; 369 | let bson = Self::to_doc(key, val)?; 370 | let replacement = bson.as_document().unwrap().clone(); 371 | let model = ReplaceOneModel::builder() 372 | .namespace(namespace.clone()) 373 | .filter(filter) 374 | .replacement(replacement) 375 | .build(); 376 | docs.push(WriteModel::ReplaceOne(model)); 377 | } 378 | let docs_len = docs.len(); 379 | let res = self 380 | .db 381 | .client() 382 | .bulk_write(docs) 383 | .write_concern( 384 | // Configure the write options 385 | WriteConcern::builder() 386 | // Ensure that all writes are written, 387 | // replicated, and acknowledged by the 388 | // majority of nodes in the cluster. 389 | .w(Acknowledgment::Majority) 390 | // Configure journal durability based on sync setting. 391 | // When `true`: writes are acknowledged only after 392 | // being written to the on-disk journal (full durability). 393 | // When `false`: writes are acknowledged after being 394 | // written to memory (faster, less durable). 395 | .journal(self.sync) 396 | // Finalise the write options 397 | .build(), 398 | ) 399 | .await?; 400 | assert_eq!(res.modified_count, docs_len as i64); 401 | Ok(()) 402 | } 403 | 404 | async fn batch_delete(&self, keys: Vec) -> Result<()> 405 | where 406 | K: Into, 407 | { 408 | let keys_len = keys.len(); 409 | let ids: Vec = keys.into_iter().map(|k| k.into()).collect(); 410 | let filter = doc! { "_id": { "$in": ids } }; 411 | let res = self.collection().delete_many(filter).await?; 412 | assert_eq!(res.deleted_count, keys_len as u64); 413 | Ok(()) 414 | } 415 | 416 | async fn scan(&self, scan: &Scan, ctx: ScanContext) -> Result { 417 | // MongoDB requires a full-text index to use a $text query 418 | if ctx == ScanContext::WithoutIndex 419 | && let Some(index) = &scan.index 420 | && let Some(kind) = &index.index_type 421 | && kind == "fulltext" 422 | { 423 | bail!(NOT_SUPPORTED_ERROR); 424 | } 425 | // Extract parameters 426 | let s = scan.start.unwrap_or(0); 427 | let l = scan.limit.unwrap_or(i64::MAX as usize); 428 | let c = MongoDBDialect::filter_clause(scan)?; 429 | let p = scan.projection()?; 430 | // Consume documents function 431 | let consume = |mut cursor: Cursor| async move { 432 | let mut count = 0; 433 | while let Some(doc) = cursor.try_next().await? { 434 | black_box(doc); 435 | count += 1; 436 | } 437 | Ok(count) 438 | }; 439 | // Perform the relevant projection scan type 440 | match p { 441 | Projection::Id => { 442 | let cursor = self 443 | .collection() 444 | .find(c) 445 | .skip(s as u64) 446 | .limit(l as i64) 447 | .projection(doc! { "_id": 1 }) 448 | .await?; 449 | consume(cursor).await 450 | } 451 | Projection::Full => { 452 | let cursor = self.collection().find(c).skip(s as u64).limit(l as i64).await?; 453 | consume(cursor).await 454 | } 455 | Projection::Count => { 456 | let pipeline = vec![ 457 | doc! { "$skip": s as i64 }, 458 | doc! { "$limit": l as i64 }, 459 | doc! { "$count": "count" }, 460 | ]; 461 | let mut cursor = self.collection().aggregate(pipeline).await?; 462 | if let Some(result) = cursor.next().await { 463 | let doc: Document = result?; 464 | let count = doc.get_i32("count").unwrap_or(0); 465 | Ok(count as usize) 466 | } else { 467 | bail!("No row returned"); 468 | } 469 | } 470 | } 471 | } 472 | } 473 | --------------------------------------------------------------------------------