This is an automated email from the ASF dual-hosted git repository. CritasWang pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/iotdb-client-rust.git
commit bb1a1d0ac6f6b387cee589a1024bed1f60cd02d5 Author: CritasWang <[email protected]> AuthorDate: Mon Jul 13 13:53:34 2026 +0800 V6: write benchmark suite (examples/benchmark.rs) mirroring Node.js/iot-benchmark methodology; tree 9.7M pts/s @ 10M points, 8 clients, 1000-row tablets --- README.md | 25 ++ README_ZH.md | 25 ++ examples/benchmark.rs | 726 ++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 776 insertions(+) diff --git a/README.md b/README.md index 746f805..a6bc084 100644 --- a/README.md +++ b/README.md @@ -208,6 +208,31 @@ docker compose up -d # standalone IoTDB (see docker-compose-1c1d.yml for a 1C1 cargo test # now includes the live-server tests ``` +## Benchmark + +`examples/benchmark.rs` is a write-performance benchmark modeled on the Node.js client's `benchmark/` suite (which follows [thulab/iot-benchmark](https://github.com/thulab/iot-benchmark)); metric definitions match, so results are comparable across the SDKs. Tablets are pre-generated outside the timed section; N worker threads each own a pooled session and insert `insert_tablet` batches round-robin over their devices. Timestamps are sequential per device from a fixed base, so runs are det [...] + +```sh +# tree model, defaults: 100 devices × 10 sensors × 20 batches × 1000 rows = 20M points, 8 clients +cargo run --release --example benchmark -- --mode tree + +# table model at a custom scale, dropping the database afterwards +cargo run --release --example benchmark -- --mode table \ + --devices 20 --sensors 10 --batches 100 --batch-size 100 --clients 8 --cleanup +``` + +Knobs: `--mode tree|table`, `--devices`, `--sensors`, `--batches` (per device), `--batch-size` (rows per tablet), `--clients` (worker threads = pool size), `--host/--port/--user/--password` (also via `IOTDB_HOST/PORT/USER/PASSWORD`), `--base-ts`, `--point-step`, `--cleanup`. Sensor types follow the Node.js default distribution (30% FLOAT, 20% DOUBLE, 20% INT32, 10% INT64, 10% TEXT, 10% BOOLEAN). The report includes total points, wall time, points/sec, per-batch latency p50/p90/p95/p99/ma [...] + +Measured on an Apple M2 Pro (10 cores), IoTDB 2.0.6 standalone in Docker on the same machine, release build, 8 clients: + +| Mode | Devices × Sensors × Batches × Rows | Points | Throughput | p50 / p99 latency | +| --- | --- | --- | --- | --- | +| tree | 20 × 10 × 100 × 100 | 2M | ~1.98M pts/s | 2.46 ms / 8.38 ms | +| table | 20 × 10 × 100 × 100 | 2M | ~1.97M pts/s | 2.13 ms / 9.97 ms | +| tree | 100 × 10 × 10 × 1000 | 10M | ~9.73M pts/s | 4.36 ms / 72.18 ms | + +Throughput scales with tablet size (rows × sensors per RPC); 1000-row tablets give ~5× the throughput of 100-row tablets at the same client count. Numbers are client+server on one machine — treat them as an upper bound on client overhead, not a server capacity measurement. + ## Project layout | Path | Contents | diff --git a/README_ZH.md b/README_ZH.md index d743d31..2a5cb36 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -208,6 +208,31 @@ docker compose up -d # 单机版 IoTDB（1C1D 集群拓扑见 docker-compose-1 cargo test # 此时包含在线测试 ``` +## 性能基准测试 + +`examples/benchmark.rs` 是写入性能基准，参考 Node.js 客户端的 `benchmark/` 套件（后者遵循 [thulab/iot-benchmark](https://github.com/thulab/iot-benchmark)）；指标定义一致，因此各语言 SDK 的结果可直接对比。Tablet 在计时区间之外预先生成；N 个工作线程各持有一个池化会话，按批次轮询各自的设备执行 `insert_tablet`。时间戳从固定基准按设备顺序递增，运行结果可复现。 + +```sh +# 树模型，默认规模：100 设备 × 10 传感器 × 20 批 × 1000 行 = 2000 万点，8 客户端 +cargo run --release --example benchmark -- --mode tree + +# 表模型，自定义规模，结束后删除数据库 +cargo run --release --example benchmark -- --mode table \ + --devices 20 --sensors 10 --batches 100 --batch-size 100 --clients 8 --cleanup +``` + +参数：`--mode tree|table`、`--devices`、`--sensors`、`--batches`（每设备批数）、`--batch-size`（每 tablet 行数）、`--clients`（工作线程数 = 池大小）、`--host/--port/--user/--password`（亦支持 `IOTDB_HOST/PORT/USER/PASSWORD` 环境变量）、`--base-ts`、`--point-step`、`--cleanup`。传感器类型分布沿用 Node.js 默认比例（30% FLOAT、20% DOUBLE、20% INT32、10% INT64、10% TEXT、10% BOOLEAN）。报告包含总点数、耗时、points/sec、单批延迟 p50/p90/p95/p99/max、错误数，以及读回行数校验。 + +实测环境：Apple M2 Pro（10 核），IoTDB 2.0.6 standalone（Docker，与客户端同机），release 构建，8 客户端： + +| 模式 | 设备 × 传感器 × 批数 × 行数 | 总点数 | 吞吐量 | p50 / p99 延迟 | +| --- | --- | --- | --- | --- | +| tree | 20 × 10 × 100 × 100 | 200 万 | ~198 万 pts/s | 2.46 ms / 8.38 ms | +| table | 20 × 10 × 100 × 100 | 200 万 | ~197 万 pts/s | 2.13 ms / 9.97 ms | +| tree | 100 × 10 × 10 × 1000 | 1000 万 | ~973 万 pts/s | 4.36 ms / 72.18 ms | + +吞吐量随 tablet 大小（每 RPC 的行 × 传感器数）增长；同等客户端数下 1000 行 tablet 的吞吐约为 100 行的 5 倍。数据为客户端与服务器同机测得——应视为客户端开销的上界，而非服务器容量。 + ## 项目结构 | 路径 | 内容 | diff --git a/examples/benchmark.rs b/examples/benchmark.rs new file mode 100644 index 0000000..9c1fe3b --- /dev/null +++ b/examples/benchmark.rs @@ -0,0 +1,726 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Write-performance benchmark for the tree and table models, modeled on the +//! Node.js client's `benchmark/` suite (which in turn follows +//! [thulab/iot-benchmark](https://github.com/thulab/iot-benchmark)). Metric +//! definitions match the Node.js `MetricsCollector` so results are directly +//! comparable: per-`insert_tablet` latency percentiles (p = ceil(p% × n) − 1 +//! over the sorted samples), points/sec over the wall time of the timed +//! insert phase only. Schema setup and data generation happen **outside** +//! the timed section ("Pre-generated Test Data"); timestamps are sequential +//! per device from a fixed base, so runs are deterministic. +//! +//! Sensor data types follow the Node.js default distribution: +//! 30% FLOAT, 20% DOUBLE, 20% INT32, 10% INT64, 10% TEXT, 10% BOOLEAN. +//! +//! Run against a live IoTDB (release mode, or the client dominates): +//! +//! ```sh +//! cargo run --release --example benchmark -- --mode tree \ +//! --devices 20 --sensors 10 --batches 100 --batch-size 100 --clients 8 +//! cargo run --release --example benchmark -- --mode table --cleanup +//! ``` +//! +//! Connection defaults honor `IOTDB_HOST` / `IOTDB_PORT` / `IOTDB_USER` / +//! `IOTDB_PASSWORD` like the e2e tests; CLI flags override env vars. + +use std::env; +use std::process::exit; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::thread; +use std::time::{Duration, Instant}; + +use iotdb_client::{ + ColumnCategory, PooledSession, Result, SessionPool, SessionPoolConfig, TSDataType, + TableSessionPool, Tablet, Value, +}; + +const TREE_DB: &str = "root.benchmark"; +const TABLE_DB: &str = "benchmark_db"; +const TABLE_NAME: &str = "benchmark_table"; +/// Node.js `STRING_LENGTH` default. +const TEXT_LEN: usize = 16; +/// Progress report interval (Node.js `REPORT_INTERVAL` default 5000 ms). +const REPORT_INTERVAL: Duration = Duration::from_secs(5); + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Mode { + Tree, + Table, +} + +struct BenchConfig { + mode: Mode, + devices: usize, + sensors: usize, + batches: usize, + batch_size: usize, + clients: usize, + host: String, + port: u16, + user: String, + password: String, + /// Fixed base timestamp (epoch ms) — deterministic, no wall clock in + /// the data path. + base_ts: i64, + /// Time interval between consecutive points per device (ms), Node.js + /// `POINT_STEP`. + point_step: i64, + cleanup: bool, +} + +impl BenchConfig { + fn total_points(&self) -> u64 { + (self.devices * self.sensors * self.batches * self.batch_size) as u64 + } +} + +const USAGE: &str = "IoTDB Rust client write benchmark + +USAGE: + cargo run --release --example benchmark -- [OPTIONS] + +OPTIONS: + --mode <tree|table> Data model to benchmark (default: tree) + --devices <N> Number of devices (default: 100) + --sensors <N> Sensors (FIELD columns) per device (default: 10) + --batches <N> Batches (tablets) per device (default: 20) + --batch-size <N> Rows per tablet (default: 1000) + --clients <N> Worker threads = session pool size (default: 8) + --host <HOST> Server host (default: $IOTDB_HOST or 127.0.0.1) + --port <PORT> Server port (default: $IOTDB_PORT or 6667) + --user <USER> Username (default: $IOTDB_USER or root) + --password <PASS> Password (default: $IOTDB_PASSWORD or root) + --base-ts <MS> Base epoch-ms timestamp (default: 1720000000000) + --point-step <MS> Interval between points per device (default: 1000) + --cleanup Drop the benchmark database after the run + --help Print this help + +Total points = devices x sensors x batches x batch-size +(defaults: 100 x 10 x 20 x 1000 = 20,000,000)."; + +fn env_or(key: &str, default: &str) -> String { + env::var(key).unwrap_or_else(|_| default.to_string()) +} + +fn parse_args() -> BenchConfig { + let mut config = BenchConfig { + mode: Mode::Tree, + devices: 100, + sensors: 10, + batches: 20, + batch_size: 1000, + clients: 8, + host: env_or("IOTDB_HOST", "127.0.0.1"), + port: env_or("IOTDB_PORT", "6667").parse().unwrap_or(6667), + user: env_or("IOTDB_USER", "root"), + password: env_or("IOTDB_PASSWORD", "root"), + base_ts: 1_720_000_000_000, + point_step: 1000, + cleanup: false, + }; + + let args: Vec<String> = env::args().skip(1).collect(); + let mut i = 0; + let value = |i: &mut usize, flag: &str| -> String { + *i += 1; + args.get(*i).cloned().unwrap_or_else(|| { + eprintln!("missing value for {flag}\n\n{USAGE}"); + exit(2); + }) + }; + while i < args.len() { + let flag = args[i].as_str(); + match flag { + "--mode" => { + config.mode = match value(&mut i, flag).as_str() { + "tree" => Mode::Tree, + "table" => Mode::Table, + other => { + eprintln!("invalid --mode `{other}` (expected tree|table)"); + exit(2); + } + } + } + "--devices" => config.devices = parse_num(&value(&mut i, flag), flag), + "--sensors" => config.sensors = parse_num(&value(&mut i, flag), flag), + "--batches" => config.batches = parse_num(&value(&mut i, flag), flag), + "--batch-size" => config.batch_size = parse_num(&value(&mut i, flag), flag), + "--clients" => config.clients = parse_num(&value(&mut i, flag), flag), + "--host" => config.host = value(&mut i, flag), + "--port" => config.port = parse_num(&value(&mut i, flag), flag) as u16, + "--user" => config.user = value(&mut i, flag), + "--password" => config.password = value(&mut i, flag), + "--base-ts" => config.base_ts = parse_num(&value(&mut i, flag), flag) as i64, + "--point-step" => config.point_step = parse_num(&value(&mut i, flag), flag) as i64, + "--cleanup" => config.cleanup = true, + "--help" | "-h" => { + println!("{USAGE}"); + exit(0); + } + other => { + eprintln!("unknown flag `{other}`\n\n{USAGE}"); + exit(2); + } + } + i += 1; + } + + for (name, v) in [ + ("--devices", config.devices), + ("--sensors", config.sensors), + ("--batches", config.batches), + ("--batch-size", config.batch_size), + ("--clients", config.clients), + ] { + if v == 0 { + eprintln!("{name} must be positive"); + exit(2); + } + } + config +} + +fn parse_num(s: &str, flag: &str) -> usize { + s.parse().unwrap_or_else(|_| { + eprintln!("invalid numeric value `{s}` for {flag}"); + exit(2); + }) +} + +// --------------------------------------------------------------------------- +// Deterministic data generation (Node.js data-generator.js distribution) +// --------------------------------------------------------------------------- + +/// Sensor `i` of `n` gets its type from the Node.js default +/// `INSERT_DATATYPE_PROPORTION`: FLOAT 0.3, DOUBLE 0.2, INT32 0.2, +/// INT64 0.1, TEXT 0.1, BOOLEAN 0.1. +fn sensor_type(i: usize, n: usize) -> TSDataType { + let f = i as f64 / n as f64; + if f < 0.3 { + TSDataType::Float + } else if f < 0.5 { + TSDataType::Double + } else if f < 0.7 { + TSDataType::Int32 + } else if f < 0.8 { + TSDataType::Int64 + } else if f < 0.9 { + TSDataType::Text + } else { + TSDataType::Boolean + } +} + +fn type_name(ty: TSDataType) -> &'static str { + match ty { + TSDataType::Float => "FLOAT", + TSDataType::Double => "DOUBLE", + TSDataType::Int32 => "INT32", + TSDataType::Int64 => "INT64", + TSDataType::Text => "TEXT", + TSDataType::Boolean => "BOOLEAN", + other => unreachable!("benchmark does not generate {other:?}"), + } +} + +/// Deterministic cell value from (device, sensor, row) — no RNG dependency. +fn cell_value(ty: TSDataType, device: usize, sensor: usize, row: usize) -> Value { + let seed = (row as u64) + .wrapping_mul(31) + .wrapping_add((sensor as u64).wrapping_mul(7)) + .wrapping_add(device as u64); + match ty { + TSDataType::Float => Value::Float((seed % 1000) as f32 * 0.1), + TSDataType::Double => Value::Double((seed % 10_000) as f64 * 0.01), + TSDataType::Int32 => Value::Int32((seed % 100_000) as i32), + TSDataType::Int64 => Value::Int64(seed as i64), + TSDataType::Text => Value::Text(format!( + "v{:0width$}", + seed % 1_000_000_000, + width = TEXT_LEN - 1 + )), + TSDataType::Boolean => Value::Boolean(seed % 2 == 0), + other => unreachable!("benchmark does not generate {other:?}"), + } +} + +/// One tablet = one batch for one device. Timestamps are sequential per +/// device: `base_ts + (batch*batch_size + row) * point_step` — already +/// sorted, so the client's sort pass is a no-op. +fn build_tablet( + config: &BenchConfig, + device: usize, + batch: usize, + sensor_names: &[String], + sensor_types: &[TSDataType], +) -> Result<Tablet> { + let mut tablet = match config.mode { + Mode::Tree => Tablet::new( + format!("{TREE_DB}.d{device}"), + sensor_names.to_vec(), + sensor_types.to_vec(), + )?, + Mode::Table => { + let mut names = Vec::with_capacity(sensor_names.len() + 1); + names.push("device_id".to_string()); + names.extend_from_slice(sensor_names); + let mut types = Vec::with_capacity(sensor_types.len() + 1); + types.push(TSDataType::String); + types.extend_from_slice(sensor_types); + let mut categories = vec![ColumnCategory::Tag]; + categories.extend(vec![ColumnCategory::Field; sensor_types.len()]); + Tablet::new_table(TABLE_NAME, names, types, categories)? + } + }; + + for r in 0..config.batch_size { + let row_index = batch * config.batch_size + r; + let ts = config.base_ts + row_index as i64 * config.point_step; + let mut row: Vec<Option<Value>> = Vec::with_capacity(sensor_types.len() + 1); + if config.mode == Mode::Table { + row.push(Some(Value::String(format!("d{device}")))); + } + for (s, &ty) in sensor_types.iter().enumerate() { + row.push(Some(cell_value(ty, device, s, row_index))); + } + tablet.add_row(ts, row)?; + } + Ok(tablet) +} + +// --------------------------------------------------------------------------- +// Pool abstraction (tree vs table) — both hand out the same PooledSession +// --------------------------------------------------------------------------- + +enum Pool { + Tree(SessionPool), + Table(TableSessionPool), +} + +impl Pool { + fn acquire(&self) -> Result<PooledSession<'_>> { + match self { + Pool::Tree(p) => p.acquire(), + Pool::Table(p) => p.acquire(), + } + } + + fn execute_non_query(&self, sql: &str) -> Result<()> { + match self { + Pool::Tree(p) => p.execute_non_query(sql), + Pool::Table(p) => p.execute_non_query(sql), + } + } + + fn close(&self) { + match self { + Pool::Tree(p) => p.close(), + Pool::Table(p) => p.close(), + } + } +} + +fn create_pool(config: &BenchConfig) -> Result<Pool> { + let mut pool_config = SessionPoolConfig { + // One dedicated session per worker; open all of them eagerly so + // connection setup stays outside the timed section. + max_size: config.clients, + min_size: config.clients, + ..SessionPoolConfig::default() + } + .with_node_urls(&[format!("{}:{}", config.host, config.port)])?; + pool_config.session.username = config.user.clone(); + pool_config.session.password = config.password.clone(); + Ok(match config.mode { + Mode::Tree => Pool::Tree(SessionPool::new(pool_config)?), + Mode::Table => Pool::Table(TableSessionPool::new(pool_config)?), + }) +} + +// --------------------------------------------------------------------------- +// Schema setup / cleanup (outside the timed section) +// --------------------------------------------------------------------------- + +fn setup_schema( + pool: &Pool, + config: &BenchConfig, + sensor_names: &[String], + sensor_types: &[TSDataType], +) -> Result<()> { + match config.mode { + Mode::Tree => { + // Fresh database each run; ignore "does not exist" on the drop. + let _ = pool.execute_non_query(&format!("DELETE DATABASE {TREE_DB}")); + pool.execute_non_query(&format!("CREATE DATABASE {TREE_DB}"))?; + // Pre-register every timeseries so metadata creation cost stays + // out of the write path (Node.js schema-manager behavior). + let mut session = pool.acquire()?; + for device in 0..config.devices { + for (name, &ty) in sensor_names.iter().zip(sensor_types) { + session.execute_non_query(&format!( + "CREATE TIMESERIES {TREE_DB}.d{device}.{name} WITH DATATYPE={}, ENCODING=PLAIN", + type_name(ty) + ))?; + } + } + } + Mode::Table => { + let _ = pool.execute_non_query(&format!("DROP DATABASE IF EXISTS {TABLE_DB}")); + pool.execute_non_query(&format!("CREATE DATABASE {TABLE_DB}"))?; + // The pool replays the last USE on every acquire, so all worker + // sessions land in the right database. + pool.execute_non_query(&format!("USE {TABLE_DB}"))?; + let columns: Vec<String> = std::iter::once("device_id STRING TAG".to_string()) + .chain( + sensor_names + .iter() + .zip(sensor_types) + .map(|(name, &ty)| format!("{name} {} FIELD", type_name(ty))), + ) + .collect(); + pool.execute_non_query(&format!( + "CREATE TABLE IF NOT EXISTS {TABLE_NAME} ({})", + columns.join(", ") + ))?; + } + } + Ok(()) +} + +fn cleanup_schema(pool: &Pool, mode: Mode) -> Result<()> { + match mode { + Mode::Tree => pool.execute_non_query(&format!("DELETE DATABASE {TREE_DB}")), + Mode::Table => pool.execute_non_query(&format!("DROP DATABASE {TABLE_DB}")), + } +} + +/// Post-run sanity check: read back the row count for the whole run. +fn verify_row_count(pool: &Pool, config: &BenchConfig) -> Result<()> { + let expected_rows = (config.devices * config.batches * config.batch_size) as i64; + let sql = match config.mode { + // COUNT over one representative sensor across all devices. + Mode::Tree => format!("SELECT COUNT(s_0) FROM {TREE_DB}.*"), + Mode::Table => format!("SELECT COUNT(*) FROM {TABLE_NAME}"), + }; + let mut session = pool.acquire()?; + let mut dataset = session.execute_query(&sql)?; + let mut total: i64 = 0; + while let Some(row) = dataset.next_row()? { + for v in &row.values { + if let Value::Int64(n) = v { + total += n; + } + } + } + let status = if total == expected_rows { + "OK" + } else { + "MISMATCH" + }; + println!("[Verify] rows on server: {total} (expected {expected_rows}) — {status}"); + Ok(()) +} + +// --------------------------------------------------------------------------- +// Metrics (definitions match Node.js benchmark-core.js MetricsCollector) +// --------------------------------------------------------------------------- + +#[derive(Default)] +struct WorkerStats { + /// One entry per insert_tablet call (ms), success or failure. + latencies_ms: Vec<f64>, + ops: u64, + failures: u64, + /// Points from successful operations only. + points: u64, + error_samples: Vec<String>, +} + +/// Node.js `getPercentile`: index = ceil(p/100 × n) − 1 over ascending samples. +fn percentile(sorted: &[f64], p: f64) -> f64 { + if sorted.is_empty() { + return 0.0; + } + let index = ((p / 100.0 * sorted.len() as f64).ceil() as usize).saturating_sub(1); + sorted[index.min(sorted.len() - 1)] +} + +fn print_summary(config: &BenchConfig, wall: Duration, workers: &[WorkerStats]) { + let mut latencies: Vec<f64> = workers + .iter() + .flat_map(|w| w.latencies_ms.iter().copied()) + .collect(); + latencies.sort_by(|a, b| a.partial_cmp(b).expect("latency is never NaN")); + let ops: u64 = workers.iter().map(|w| w.ops).sum(); + let failures: u64 = workers.iter().map(|w| w.failures).sum(); + let points: u64 = workers.iter().map(|w| w.points).sum(); + let secs = wall.as_secs_f64(); + let avg = if latencies.is_empty() { + 0.0 + } else { + latencies.iter().sum::<f64>() / latencies.len() as f64 + }; + + let sep = "=".repeat(80); + println!("\n{sep}\nBENCHMARK RESULTS\n{sep}"); + println!("\n[Execution Time]"); + println!( + " Duration: {:.2}s ({:.0}ms)", + secs, + secs * 1000.0 + ); + println!("\n[Operations]"); + println!(" Total Operations: {ops}"); + println!(" Successful: {}", ops - failures); + println!(" Failed: {failures}"); + println!( + " Success Rate: {:.2}%", + if ops == 0 { + 0.0 + } else { + (ops - failures) as f64 / ops as f64 * 100.0 + } + ); + println!("\n[Data Points]"); + println!(" Total Points Written: {points}"); + println!("\n[Throughput]"); + println!(" Operations/sec: {:.2}", ops as f64 / secs); + println!(" Points/sec: {:.0}", points as f64 / secs); + println!("\n[Latency (ms)]"); + println!( + " Min: {:.2}ms", + latencies.first().copied().unwrap_or(0.0) + ); + println!( + " Max: {:.2}ms", + latencies.last().copied().unwrap_or(0.0) + ); + println!(" Average: {avg:.2}ms"); + println!( + " P50 (Median): {:.2}ms", + percentile(&latencies, 50.0) + ); + println!( + " P90: {:.2}ms", + percentile(&latencies, 90.0) + ); + println!( + " P95: {:.2}ms", + percentile(&latencies, 95.0) + ); + println!( + " P99: {:.2}ms", + percentile(&latencies, 99.0) + ); + + let samples: Vec<&String> = workers.iter().flat_map(|w| &w.error_samples).collect(); + if !samples.is_empty() { + println!("\n[Error Samples]"); + for (i, err) in samples.iter().take(5).enumerate() { + println!(" {}. {err}", i + 1); + } + } + println!("\n{sep}"); + println!( + "Config: mode={} devices={} sensors={} batches={} batch-size={} clients={} → {} points", + if config.mode == Mode::Tree { + "tree" + } else { + "table" + }, + config.devices, + config.sensors, + config.batches, + config.batch_size, + config.clients, + config.total_points(), + ); +} + +// --------------------------------------------------------------------------- +// Main +// --------------------------------------------------------------------------- + +fn main() -> Result<()> { + env_logger::init(); + let config = parse_args(); + + let sensor_names: Vec<String> = (0..config.sensors).map(|i| format!("s_{i}")).collect(); + let sensor_types: Vec<TSDataType> = (0..config.sensors) + .map(|i| sensor_type(i, config.sensors)) + .collect(); + + let sep = "=".repeat(80); + println!("{sep}\nIoTDB Rust Client Write Benchmark\n{sep}"); + println!( + " Mode: {}", + if config.mode == Mode::Tree { + "tree" + } else { + "table" + } + ); + println!(" Server: {}:{}", config.host, config.port); + println!(" Devices: {}", config.devices); + println!( + " Sensors: {} ({})", + config.sensors, + sensor_types + .iter() + .map(|&t| type_name(t)) + .collect::<Vec<_>>() + .join(", ") + ); + println!(" Batches: {} per device", config.batches); + println!(" Batch size: {} rows", config.batch_size); + println!(" Clients: {} (pool size)", config.clients); + println!(" Total: {} points", config.total_points()); + println!("{sep}\n"); + + // --- Pool + schema (untimed) ----------------------------------------- + println!( + "[Setup] opening session pool ({} eager sessions)...", + config.clients + ); + let pool = create_pool(&config)?; + println!("[Setup] creating schema..."); + let t0 = Instant::now(); + setup_schema(&pool, &config, &sensor_names, &sensor_types)?; + println!("[Setup] schema ready in {:.2}s", t0.elapsed().as_secs_f64()); + + // --- Pre-generate all tablets (untimed) ------------------------------- + // Worker w owns devices where device % clients == w and walks them + // batch-major, i.e. round-robin over its devices. + println!("[Setup] pre-generating test data..."); + let t0 = Instant::now(); + let mut worker_tablets: Vec<Vec<Tablet>> = Vec::with_capacity(config.clients); + for w in 0..config.clients { + let devices: Vec<usize> = (0..config.devices) + .filter(|d| d % config.clients == w) + .collect(); + let mut tablets = Vec::with_capacity(devices.len() * config.batches); + for batch in 0..config.batches { + for &device in &devices { + tablets.push(build_tablet( + &config, + device, + batch, + &sensor_names, + &sensor_types, + )?); + } + } + worker_tablets.push(tablets); + } + println!( + "[Setup] {} tablets generated in {:.2}s", + worker_tablets.iter().map(Vec::len).sum::<usize>(), + t0.elapsed().as_secs_f64() + ); + + // --- Timed insert phase ------------------------------------------------ + println!( + "\n[Test Phase] running with {} concurrent clients...\n", + config.clients + ); + let ops_done = AtomicU64::new(0); + let points_done = AtomicU64::new(0); + let stop_reporter = AtomicBool::new(false); + + let started = Instant::now(); + let worker_stats: Vec<WorkerStats> = thread::scope(|scope| { + // Progress reporter (Node.js ProgressReporter equivalent). + let reporter = scope.spawn(|| { + let mut last_ops = 0u64; + let mut last_tick = Instant::now(); + while !stop_reporter.load(Ordering::Relaxed) { + thread::sleep(Duration::from_millis(200)); + if last_tick.elapsed() >= REPORT_INTERVAL { + let ops = ops_done.load(Ordering::Relaxed); + let rate = (ops - last_ops) as f64 / last_tick.elapsed().as_secs_f64(); + println!( + "[Progress] Operations: {ops}, Rate: {rate:.2} ops/s, Total Points: {}", + points_done.load(Ordering::Relaxed) + ); + last_ops = ops; + last_tick = Instant::now(); + } + } + }); + + let handles: Vec<_> = worker_tablets + .iter() + .map(|tablets| { + let pool = &pool; + let ops_done = &ops_done; + let points_done = &points_done; + scope.spawn(move || -> Result<WorkerStats> { + let mut session = pool.acquire()?; + let mut stats = WorkerStats { + latencies_ms: Vec::with_capacity(tablets.len()), + ..WorkerStats::default() + }; + for tablet in tablets { + let points = (tablet.row_count() * config.sensors) as u64; + let start = Instant::now(); + let outcome = session.insert_tablet(tablet); + stats + .latencies_ms + .push(start.elapsed().as_secs_f64() * 1000.0); + stats.ops += 1; + ops_done.fetch_add(1, Ordering::Relaxed); + match outcome { + Ok(()) => { + stats.points += points; + points_done.fetch_add(points, Ordering::Relaxed); + } + Err(e) => { + stats.failures += 1; + if stats.error_samples.len() < 5 { + stats.error_samples.push(e.to_string()); + } + } + } + } + Ok(stats) + }) + }) + .collect(); + + let stats = handles + .into_iter() + .map(|h| h.join().expect("worker thread panicked")) + .collect::<Result<Vec<_>>>(); + stop_reporter.store(true, Ordering::Relaxed); + reporter.join().expect("reporter thread panicked"); + stats + })?; + let wall = started.elapsed(); + + // --- Results + verification + cleanup --------------------------------- + print_summary(&config, wall, &worker_stats); + verify_row_count(&pool, &config)?; + + if config.cleanup { + println!("[Cleanup] dropping benchmark database..."); + cleanup_schema(&pool, config.mode)?; + } + pool.close(); + Ok(()) +}
