test(rustfs): move manual bench tools to examples, gate Swift suites (#6712)

This commit is contained in:
Zhengchao An
2026-08-27 10:23:14 +08:00
committed by GitHub
parent 5f3620a00d
commit a6c5d80069
7 changed files with 46 additions and 66 deletions
@@ -0,0 +1,65 @@
# Concurrent Download Tool (example)
This tool downloads multiple URLs concurrently and saves files to a target directory.
Saved filename format:
`<nanoseconds>_<index>_<original_filename>`
All downloaded files are written into one output directory.
## Environment variables
- `DOWNLOAD_URLS` (required): comma-separated URLs.
- `DOWNLOAD_OUTPUT_DIR` (optional): output directory, default `target/tmp/concurrent_downloads`.
- `DOWNLOAD_CONCURRENCY` (optional): max concurrent downloads, default `8`.
- `DOWNLOAD_REPEAT` (optional): repeat count per URL, default `1`.
- `DOWNLOAD_MAX_RETRIES` (optional): retry count per task after first failure, default `0`.
- `DOWNLOAD_RETRY_BACKOFF_MS` (optional): fixed backoff between retries, default `200`.
## Statistics output
After run, the tool prints:
- total tasks
- succeeded
- failed
- total bytes
- elapsed ms
- throughput bps
- total attempts
- retried tasks
- retry attempts
- latency p50 ms
- latency p95 ms
- failure details (`[index] url => error`) when failures exist
If any task fails, the tool exits with an error after printing the summary.
Retry is triggered only for recoverable cases:
- network/request timeout/connect errors
- HTTP `429`
- HTTP `5xx`
## Compile check
```bash
cargo build -p rustfs --example concurrent_download_tool
```
## Manual run example
The commands below are for manual execution only.
The tool is an example binary, so it is built on demand and never runs as part of the test suites.
```bash
DOWNLOAD_URLS="http://127.0.0.1:9001/demo/google-cloud-aiplugin-1.46.1-253.zip?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Content-Sha256=UNSIGNED-PAYLOAD&X-Amz-Credential=HAXVOTZK9MLBJT8KWI4E%2F20260329%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20260329T105159Z&X-Amz-Expires=86400&X-Amz-Security-Token=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzUxMiJ9.eyJwYXJlbnQiOiJydXN0ZnNhZG1pbiIsImV4cCI6MTc3NDgyMDgyMX0.tYhQoPRcg0Ysx4KVw9ez7ZpYxsqGgqomtsP_iaeTsKzoii8EVNt74BZm2wbUjXW-FbGXc1pqEYX6wZ5Ncpk9Iw&X-Amz-Signature=15f47b19832f53b34f9e0fe1862d53d71660bbf8f1a512669bb2d041ac8d0697&X-Amz-SignedHeaders=host&x-amz-checksum-mode=ENABLED&x-id=GetObject" \
DOWNLOAD_OUTPUT_DIR="/Users/zhi/Documents/code/rust/rustfs/rustfs/target/tmp/concurrent_downloads" \
DOWNLOAD_CONCURRENCY="40" \
DOWNLOAD_REPEAT="40" \
DOWNLOAD_MAX_RETRIES="2" \
DOWNLOAD_RETRY_BACKOFF_MS="300" \
cargo run -p rustfs --example concurrent_download_tool
```
+410
View File
@@ -0,0 +1,410 @@
// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! Manual concurrent-download performance tool: requires a running RustFS server
//! configured via `DOWNLOAD_*` env vars. Run with
//! `cargo run -p rustfs --example concurrent_download_tool`.
use anyhow::{Context, Result, anyhow};
use futures::stream::{self, StreamExt};
use reqwest::{Client, Url};
use std::env;
use std::path::{Path, PathBuf};
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use tokio::time::{Duration, sleep};
#[derive(Debug)]
struct DownloadSettings {
urls: Vec<String>,
output_dir: PathBuf,
concurrency: usize,
repeat: usize,
max_retries: usize,
retry_backoff_ms: u64,
}
#[derive(Debug)]
struct DownloadSuccess {
path: PathBuf,
bytes: usize,
attempts_used: usize,
elapsed_ms: u128,
}
#[derive(Debug)]
struct DownloadAttemptError {
attempts_used: usize,
error: String,
elapsed_ms: u128,
}
#[derive(Debug)]
struct DownloadFailure {
index: usize,
url: String,
attempts_used: usize,
error: String,
}
#[derive(Debug)]
struct DownloadSummary {
saved_files: Vec<PathBuf>,
total_tasks: usize,
succeeded: usize,
failed: usize,
total_bytes: usize,
elapsed_ms: u128,
throughput_bps: f64,
total_attempts: usize,
retried_tasks: usize,
retry_attempts: usize,
latency_p50_ms: u128,
latency_p95_ms: u128,
failures: Vec<DownloadFailure>,
}
fn should_retry_status(status: reqwest::StatusCode) -> bool {
status.as_u16() == 429 || status.is_server_error()
}
fn should_retry_reqwest_error(err: &reqwest::Error) -> bool {
if err.is_timeout() || err.is_connect() || err.is_request() {
return true;
}
match err.status() {
Some(status) => should_retry_status(status),
None => false,
}
}
fn percentile(values: &[u128], p: f64) -> u128 {
if values.is_empty() {
return 0;
}
let mut sorted = values.to_vec();
sorted.sort_unstable();
let rank = ((sorted.len() as f64 - 1.0) * p).round() as usize;
sorted[rank]
}
impl DownloadSettings {
fn from_env() -> Result<Self> {
let urls_raw = env::var("DOWNLOAD_URLS").context("missing DOWNLOAD_URLS, expected comma-separated URLs")?;
let urls: Vec<String> = urls_raw
.split(',')
.map(str::trim)
.filter(|v| !v.is_empty())
.map(ToString::to_string)
.collect();
if urls.is_empty() {
return Err(anyhow!("DOWNLOAD_URLS is empty, expected comma-separated URLs"));
}
let output_dir = env::var("DOWNLOAD_OUTPUT_DIR")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from("target/tmp/concurrent_downloads"));
let concurrency = env::var("DOWNLOAD_CONCURRENCY")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.filter(|v| *v > 0)
.unwrap_or(8);
let repeat = env::var("DOWNLOAD_REPEAT")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.filter(|v| *v > 0)
.unwrap_or(1);
let max_retries = env::var("DOWNLOAD_MAX_RETRIES")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.unwrap_or(0);
let retry_backoff_ms = env::var("DOWNLOAD_RETRY_BACKOFF_MS")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.filter(|v| *v > 0)
.unwrap_or(200);
Ok(Self {
urls,
output_dir,
concurrency,
repeat,
max_retries,
retry_backoff_ms,
})
}
}
fn original_filename(url: &str) -> String {
Url::parse(url)
.ok()
.and_then(|parsed| {
parsed
.path_segments()
.and_then(|mut segments| segments.rfind(|s| !s.is_empty()))
.map(ToString::to_string)
})
.filter(|name| !name.is_empty())
.unwrap_or_else(|| "download.bin".to_string())
}
fn nanos_prefix() -> Result<u128> {
Ok(SystemTime::now()
.duration_since(UNIX_EPOCH)
.context("system clock is before UNIX_EPOCH")?
.as_nanos())
}
async fn download_one(
client: &Client,
output_dir: &Path,
index: usize,
url: String,
max_retries: usize,
retry_backoff_ms: u64,
) -> std::result::Result<DownloadSuccess, DownloadAttemptError> {
let task_started_at = Instant::now();
let mut attempt = 0usize;
let mut last_error = String::new();
let mut retryable = false;
while attempt <= max_retries {
attempt += 1;
let response = match client.get(&url).send().await {
Ok(resp) => resp,
Err(err) => {
retryable = should_retry_reqwest_error(&err);
last_error = format!("failed request: {url}, error: {err}");
if retryable && attempt <= max_retries {
sleep(Duration::from_millis(retry_backoff_ms)).await;
continue;
}
break;
}
};
let status = response.status();
if !status.is_success() {
retryable = should_retry_status(status);
last_error = format!("non-success status for URL: {url}, status: {status}");
if retryable && attempt <= max_retries {
sleep(Duration::from_millis(retry_backoff_ms)).await;
continue;
}
break;
}
let body = match response.bytes().await {
Ok(bytes) => bytes,
Err(err) => {
retryable = should_retry_reqwest_error(&err);
last_error = format!("failed to read response body: {url}, error: {err}");
if retryable && attempt <= max_retries {
sleep(Duration::from_millis(retry_backoff_ms)).await;
continue;
}
break;
}
};
let source_name = original_filename(&url);
let nanos = match nanos_prefix() {
Ok(v) => v,
Err(err) => {
last_error = err.to_string();
retryable = false;
break;
}
};
let target_name = format!("{}_{}_{}", nanos, index, source_name);
let target_path = output_dir.join(target_name);
let result: Result<DownloadSuccess> = async {
tokio::fs::write(&target_path, &body)
.await
.with_context(|| format!("failed to write file: {}", target_path.display()))?;
Ok(DownloadSuccess {
path: target_path,
bytes: body.len(),
attempts_used: attempt,
elapsed_ms: task_started_at.elapsed().as_millis(),
})
}
.await;
match result {
Ok(success) => return Ok(success),
Err(err) => {
last_error = err.to_string();
retryable = false;
if retryable && attempt <= max_retries {
sleep(Duration::from_millis(retry_backoff_ms)).await;
}
break;
}
}
}
Err(DownloadAttemptError {
attempts_used: attempt,
error: if retryable {
last_error
} else {
format!("{} (non-retryable)", last_error)
},
elapsed_ms: task_started_at.elapsed().as_millis(),
})
}
async fn run_concurrent_downloads(settings: DownloadSettings) -> Result<DownloadSummary> {
let started_at = Instant::now();
tokio::fs::create_dir_all(&settings.output_dir)
.await
.with_context(|| format!("failed to create output dir: {}", settings.output_dir.display()))?;
let client = Client::new();
let tasks = settings
.urls
.into_iter()
.flat_map(|url| (0..settings.repeat).map(move |_| url.clone()))
.enumerate();
let results = stream::iter(tasks)
.map(|(index, url)| {
let client = client.clone();
let output_dir = settings.output_dir.clone();
let max_retries = settings.max_retries;
let retry_backoff_ms = settings.retry_backoff_ms;
async move {
let current_url = url.clone();
let result = download_one(&client, &output_dir, index, url, max_retries, retry_backoff_ms).await;
(index, current_url, result)
}
})
.buffer_unordered(settings.concurrency)
.collect::<Vec<(usize, String, std::result::Result<DownloadSuccess, DownloadAttemptError>)>>()
.await;
let mut saved_files = Vec::new();
let mut total_bytes = 0usize;
let mut total_attempts = 0usize;
let mut retried_tasks = 0usize;
let mut latencies_ms = Vec::new();
let mut failures = Vec::new();
for (index, url, item) in results {
match item {
Ok(success) => {
total_bytes += success.bytes;
total_attempts += success.attempts_used;
if success.attempts_used > 1 {
retried_tasks += 1;
}
latencies_ms.push(success.elapsed_ms);
saved_files.push(success.path);
}
Err(err) => {
total_attempts += err.attempts_used;
if err.attempts_used > 1 {
retried_tasks += 1;
}
latencies_ms.push(err.elapsed_ms);
failures.push(DownloadFailure {
index,
url,
attempts_used: err.attempts_used,
error: err.error,
});
}
}
}
let total_tasks = saved_files.len() + failures.len();
let retry_attempts = total_attempts.saturating_sub(total_tasks);
let elapsed_ms = started_at.elapsed().as_millis();
let throughput_bps = if elapsed_ms == 0 {
0.0
} else {
(total_bytes as f64) / ((elapsed_ms as f64) / 1000.0)
};
let latency_p50_ms = percentile(&latencies_ms, 0.50);
let latency_p95_ms = percentile(&latencies_ms, 0.95);
Ok(DownloadSummary {
total_tasks,
succeeded: saved_files.len(),
failed: failures.len(),
total_bytes,
elapsed_ms,
throughput_bps,
total_attempts,
retried_tasks,
retry_attempts,
latency_p50_ms,
latency_p95_ms,
saved_files,
failures,
})
}
#[tokio::main]
async fn main() -> Result<()> {
let settings = DownloadSettings::from_env()?;
let summary = run_concurrent_downloads(settings).await?;
for path in &summary.saved_files {
println!("saved: {}", path.display());
}
println!("download complete");
println!("total tasks: {}", summary.total_tasks);
println!("succeeded: {}", summary.succeeded);
println!("failed: {}", summary.failed);
println!("total bytes: {}", summary.total_bytes);
println!("elapsed ms: {}", summary.elapsed_ms);
println!("throughput bps: {:.2}", summary.throughput_bps);
println!("total attempts: {}", summary.total_attempts);
println!("retried tasks: {}", summary.retried_tasks);
println!("retry attempts: {}", summary.retry_attempts);
println!("latency p50 ms: {}", summary.latency_p50_ms);
println!("latency p95 ms: {}", summary.latency_p95_ms);
if !summary.failures.is_empty() {
println!("failure details:");
for failure in &summary.failures {
println!(
" [{}] attempts={} {} => {}",
failure.index, failure.attempts_used, failure.url, failure.error
);
}
return Err(anyhow!("download finished with {} failures", summary.failures.len()));
}
Ok(())
}
+542
View File
@@ -0,0 +1,542 @@
// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! Manual >1GiB GET benchmark: requires a running RustFS server configured via
//! `GT1G_GET_*` env vars. Run with `cargo run -p rustfs --example gt1g_get_benchmark_tool`.
use anyhow::{Context, Result, anyhow};
use aws_config::BehaviorVersion;
use aws_config::meta::region::RegionProviderChain;
use aws_sdk_s3::Client;
use aws_sdk_s3::config::{Builder as S3ConfigBuilder, Credentials, Region};
use aws_sdk_s3::primitives::ByteStream;
use std::env;
use std::path::PathBuf;
use std::time::Instant;
use tokio::fs::OpenOptions;
use tokio::io::AsyncWriteExt;
use tokio::time::{Duration, sleep};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ToolAction {
Prepare,
Bench,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum BenchMode {
Sequential,
RangedParallel,
}
impl BenchMode {
fn as_str(self) -> &'static str {
match self {
Self::Sequential => "sequential",
Self::RangedParallel => "ranged_parallel",
}
}
}
#[derive(Clone, Debug)]
struct ObjectSpec {
label: String,
key: String,
size_bytes: i64,
}
#[derive(Clone, Debug)]
struct ToolSettings {
action: ToolAction,
endpoint: String,
access_key: String,
secret_key: String,
bucket: String,
region: String,
out_dir: PathBuf,
objects: Vec<ObjectSpec>,
modes: Vec<BenchMode>,
concurrencies: Vec<usize>,
range_workers: Vec<usize>,
rounds: usize,
cooldown_secs: u64,
force: bool,
}
type BenchGroupKey = (String, String, String, usize, usize);
type BenchGroupValue = (bool, u64, f64, f64, f64);
fn parse_action() -> Result<ToolAction> {
match env::var("GT1G_GET_ACTION")
.unwrap_or_else(|_| "bench".to_string())
.trim()
.to_ascii_lowercase()
.as_str()
{
"prepare" => Ok(ToolAction::Prepare),
"bench" => Ok(ToolAction::Bench),
other => Err(anyhow!("unsupported GT1G_GET_ACTION: {other}")),
}
}
fn parse_size_label_to_bytes(label: &str) -> Result<i64> {
let normalized = label.trim();
if let Some(raw) = normalized.strip_suffix("GiB") {
let value = raw.trim().parse::<i64>()?;
return Ok(value * 1024 * 1024 * 1024);
}
if let Some(raw) = normalized.strip_suffix("MiB") {
let value = raw.trim().parse::<i64>()?;
return Ok(value * 1024 * 1024);
}
if let Some(raw) = normalized.strip_suffix("KiB") {
let value = raw.trim().parse::<i64>()?;
return Ok(value * 1024);
}
Err(anyhow!("unsupported size label: {label}"))
}
fn parse_object_specs(raw: &str) -> Result<Vec<ObjectSpec>> {
raw.split(',')
.map(str::trim)
.filter(|entry| !entry.is_empty())
.map(|entry| {
let (label, key) = entry.split_once('=').ok_or_else(|| anyhow!("invalid object spec: {entry}"))?;
Ok(ObjectSpec {
label: label.trim().to_string(),
key: key.trim().to_string(),
size_bytes: parse_size_label_to_bytes(label.trim())?,
})
})
.collect()
}
fn parse_csv_usize(raw: &str, field: &str) -> Result<Vec<usize>> {
let values: Vec<usize> = raw
.split(',')
.map(str::trim)
.filter(|entry| !entry.is_empty())
.map(|entry| {
entry
.parse::<usize>()
.with_context(|| format!("invalid {field} value: {entry}"))
})
.collect::<Result<Vec<_>>>()?;
if values.is_empty() {
return Err(anyhow!("{field} cannot be empty"));
}
Ok(values)
}
fn parse_modes(raw: &str) -> Result<Vec<BenchMode>> {
let modes: Vec<BenchMode> = raw
.split(',')
.map(str::trim)
.filter(|entry| !entry.is_empty())
.map(|entry| match entry {
"sequential" => Ok(BenchMode::Sequential),
"ranged_parallel" => Ok(BenchMode::RangedParallel),
other => Err(anyhow!("unsupported GT1G_GET_MODES value: {other}")),
})
.collect::<Result<Vec<_>>>()?;
if modes.is_empty() {
return Err(anyhow!("GT1G_GET_MODES cannot be empty"));
}
Ok(modes)
}
impl ToolSettings {
fn from_env() -> Result<Self> {
let action = parse_action()?;
let endpoint = env::var("GT1G_GET_ENDPOINT").context("missing GT1G_GET_ENDPOINT")?;
let access_key = env::var("GT1G_GET_ACCESS_KEY").context("missing GT1G_GET_ACCESS_KEY")?;
let secret_key = env::var("GT1G_GET_SECRET_KEY").context("missing GT1G_GET_SECRET_KEY")?;
let bucket = env::var("GT1G_GET_BUCKET").context("missing GT1G_GET_BUCKET")?;
let region = env::var("GT1G_GET_REGION").unwrap_or_else(|_| "us-east-1".to_string());
let out_dir = env::var("GT1G_GET_OUT_DIR")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from(format!("target/bench/gt1g-get-tool-{}", chrono_like_timestamp())));
let objects = parse_object_specs(&env::var("GT1G_GET_OBJECTS").context("missing GT1G_GET_OBJECTS")?)?;
let modes = parse_modes(&env::var("GT1G_GET_MODES").unwrap_or_else(|_| "sequential,ranged_parallel".to_string()))?;
let concurrencies = parse_csv_usize(
&env::var("GT1G_GET_CONCURRENCIES").unwrap_or_else(|_| "1,4,8".to_string()),
"GT1G_GET_CONCURRENCIES",
)?;
let range_workers = parse_csv_usize(
&env::var("GT1G_GET_RANGE_WORKERS").unwrap_or_else(|_| "4".to_string()),
"GT1G_GET_RANGE_WORKERS",
)?;
let rounds = env::var("GT1G_GET_ROUNDS")
.unwrap_or_else(|_| "3".to_string())
.parse::<usize>()
.context("invalid GT1G_GET_ROUNDS")?;
let cooldown_secs = env::var("GT1G_GET_COOLDOWN_SECS")
.unwrap_or_else(|_| "15".to_string())
.parse::<u64>()
.context("invalid GT1G_GET_COOLDOWN_SECS")?;
let force = env::var("GT1G_GET_FORCE")
.ok()
.map(|raw| matches!(raw.as_str(), "1" | "true" | "TRUE" | "yes" | "YES"))
.unwrap_or(false);
Ok(Self {
action,
endpoint,
access_key,
secret_key,
bucket,
region,
out_dir,
objects,
modes,
concurrencies,
range_workers,
rounds,
cooldown_secs,
force,
})
}
}
fn chrono_like_timestamp() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
now.to_string()
}
async fn build_client(settings: &ToolSettings) -> Result<Client> {
let region = Region::new(settings.region.clone());
let region_provider = RegionProviderChain::first_try(Some(region.clone()));
let shared_config = aws_config::defaults(BehaviorVersion::latest())
.region(region_provider)
.credentials_provider(Credentials::new(
settings.access_key.clone(),
settings.secret_key.clone(),
None,
None,
"issue713-gt1g-get-tool",
))
.load()
.await;
let s3_config = S3ConfigBuilder::from(&shared_config)
.endpoint_url(settings.endpoint.clone())
.force_path_style(true)
.region(region)
.build();
Ok(Client::from_conf(s3_config))
}
async fn ensure_bucket(client: &Client, bucket: &str) -> Result<()> {
if client.head_bucket().bucket(bucket).send().await.is_ok() {
return Ok(());
}
client.create_bucket().bucket(bucket).send().await?;
Ok(())
}
async fn object_matches_size(client: &Client, bucket: &str, object: &ObjectSpec) -> Result<bool> {
let out = match client.head_object().bucket(bucket).key(&object.key).send().await {
Ok(out) => out,
Err(_) => return Ok(false),
};
Ok(out.content_length().unwrap_or_default() == object.size_bytes)
}
async fn create_sparse_file(path: &PathBuf, size_bytes: i64) -> Result<()> {
let file = OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(path)
.await
.with_context(|| format!("failed to open sparse file: {}", path.display()))?;
file.set_len(size_bytes as u64)
.await
.with_context(|| format!("failed to set sparse file length: {}", path.display()))?;
Ok(())
}
async fn prepare_objects(settings: &ToolSettings, client: &Client) -> Result<()> {
ensure_bucket(client, &settings.bucket).await?;
tokio::fs::create_dir_all(&settings.out_dir).await?;
let tmp_root = settings.out_dir.join("prepare-tmp");
tokio::fs::create_dir_all(&tmp_root).await?;
for object in &settings.objects {
if !settings.force && object_matches_size(client, &settings.bucket, object).await? {
println!("skip existing: s3://{}/{}", settings.bucket, object.key);
continue;
}
let local_path = tmp_root.join(object.key.replace('/', "_"));
create_sparse_file(&local_path, object.size_bytes).await?;
println!("uploading: s3://{}/{} ({})", settings.bucket, object.key, object.label);
client
.put_object()
.bucket(&settings.bucket)
.key(&object.key)
.body(ByteStream::from_path(&local_path).await?)
.send()
.await
.with_context(|| format!("failed to upload {}", object.key))?;
let _ = tokio::fs::remove_file(&local_path).await;
}
Ok(())
}
async fn drain_body(mut body: ByteStream) -> Result<u64> {
let mut total = 0u64;
while let Some(chunk) = body.try_next().await? {
total += chunk.len() as u64;
}
Ok(total)
}
async fn run_sequential_task(client: Client, bucket: String, key: String, expected_size: i64) -> Result<(u64, u128)> {
let started = Instant::now();
let output = client.get_object().bucket(bucket).key(key).send().await?;
let bytes = drain_body(output.body).await?;
if bytes != expected_size as u64 {
return Err(anyhow!("downloaded bytes mismatch: expected {}, got {}", expected_size, bytes));
}
Ok((bytes, started.elapsed().as_millis()))
}
async fn run_range_worker(client: Client, bucket: String, key: String, start: i64, end: i64) -> Result<u64> {
let range = format!("bytes={start}-{end}");
let output = client.get_object().bucket(bucket).key(key).range(range).send().await?;
let bytes = drain_body(output.body).await?;
Ok(bytes)
}
async fn run_ranged_parallel_task(
client: Client,
bucket: String,
key: String,
expected_size: i64,
range_workers: usize,
) -> Result<(u64, u128)> {
let started = Instant::now();
let chunk_size = (expected_size + range_workers as i64 - 1) / range_workers as i64;
let mut join_set = tokio::task::JoinSet::new();
for worker_index in 0..range_workers {
let start = worker_index as i64 * chunk_size;
if start >= expected_size {
continue;
}
let end = (start + chunk_size - 1).min(expected_size - 1);
join_set.spawn(run_range_worker(client.clone(), bucket.clone(), key.clone(), start, end));
}
let mut total = 0u64;
while let Some(result) = join_set.join_next().await {
total += result??;
}
if total != expected_size as u64 {
return Err(anyhow!("ranged download bytes mismatch: expected {}, got {}", expected_size, total));
}
Ok((total, started.elapsed().as_millis()))
}
async fn run_bench(settings: &ToolSettings, client: &Client) -> Result<()> {
tokio::fs::create_dir_all(&settings.out_dir).await?;
let round_csv = settings.out_dir.join("round_results.csv");
let median_csv = settings.out_dir.join("median_summary.csv");
tokio::fs::write(
&round_csv,
"object_label,object_key,mode,range_workers,concurrency,round,status,total_bytes,elapsed_ms,throughput_bps,avg_client_latency_ms\n",
)
.await?;
tokio::fs::write(
&median_csv,
"object_label,object_key,mode,range_workers,concurrency,successful_rounds,failed_rounds,median_total_bytes,median_elapsed_ms,median_throughput_bps,median_avg_client_latency_ms\n",
)
.await?;
let mut rows = Vec::new();
for object in &settings.objects {
for &mode in &settings.modes {
for &concurrency in &settings.concurrencies {
let worker_values = if mode == BenchMode::Sequential {
vec![1usize]
} else {
settings.range_workers.clone()
};
for range_workers in worker_values {
for round in 1..=settings.rounds {
let started = Instant::now();
let mut join_set = tokio::task::JoinSet::new();
for _ in 0..concurrency {
let task_client = client.clone();
let bucket = settings.bucket.clone();
let key = object.key.clone();
let size = object.size_bytes;
if mode == BenchMode::Sequential {
join_set.spawn(run_sequential_task(task_client, bucket, key, size));
} else {
join_set.spawn(run_ranged_parallel_task(task_client, bucket, key, size, range_workers));
}
}
let mut status = "ok".to_string();
let mut total_bytes = 0u64;
let mut latencies_ms = Vec::new();
while let Some(result) = join_set.join_next().await {
match result {
Ok(Ok((bytes, elapsed_ms))) => {
total_bytes += bytes;
latencies_ms.push(elapsed_ms);
}
Ok(Err(err)) => {
status = format!("failed:{err}");
}
Err(err) => {
status = format!("failed:{err}");
}
}
}
let elapsed_ms = started.elapsed().as_millis();
let avg_client_latency_ms = if latencies_ms.is_empty() {
0.0
} else {
latencies_ms.iter().sum::<u128>() as f64 / latencies_ms.len() as f64
};
let throughput_bps = if elapsed_ms == 0 {
0.0
} else {
total_bytes as f64 / (elapsed_ms as f64 / 1000.0)
};
let row = format!(
"{},{},{},{},{},{},{},{},{},{:.6},{:.3}\n",
object.label,
object.key,
mode.as_str(),
range_workers,
concurrency,
round,
status,
total_bytes,
elapsed_ms,
throughput_bps,
avg_client_latency_ms
);
tokio::fs::OpenOptions::new()
.append(true)
.open(&round_csv)
.await?
.write_all(row.as_bytes())
.await?;
rows.push((
object.label.clone(),
object.key.clone(),
mode.as_str().to_string(),
range_workers,
concurrency,
status,
total_bytes,
elapsed_ms as f64,
throughput_bps,
avg_client_latency_ms,
));
if settings.cooldown_secs > 0 && round < settings.rounds {
sleep(Duration::from_secs(settings.cooldown_secs)).await;
}
}
}
}
}
}
use std::collections::BTreeMap;
let mut grouped: BTreeMap<BenchGroupKey, Vec<BenchGroupValue>> = BTreeMap::new();
for (label, key, mode, range_workers, concurrency, status, total_bytes, elapsed_ms, throughput_bps, avg_latency_ms) in rows {
grouped
.entry((label, key, mode, range_workers, concurrency))
.or_default()
.push((status == "ok", total_bytes, elapsed_ms, throughput_bps, avg_latency_ms));
}
for ((label, key, mode, range_workers, concurrency), values) in grouped {
let successful_rounds = values.iter().filter(|v| v.0).count();
let failed_rounds = values.len().saturating_sub(successful_rounds);
let mut bytes = Vec::new();
let mut elapsed = Vec::new();
let mut throughput = Vec::new();
let mut latency = Vec::new();
for (ok, total_bytes, elapsed_ms, throughput_bps, avg_latency_ms) in values {
if ok {
bytes.push(total_bytes as f64);
elapsed.push(elapsed_ms);
throughput.push(throughput_bps);
latency.push(avg_latency_ms);
}
}
let median = |input: &mut Vec<f64>| -> String {
if input.is_empty() {
return "N/A".to_string();
}
input.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let n = input.len();
if n % 2 == 1 {
format!("{:.6}", input[n / 2])
} else {
format!("{:.6}", (input[n / 2 - 1] + input[n / 2]) / 2.0)
}
};
let row = format!(
"{},{},{},{},{},{},{},{},{},{},{}\n",
label,
key,
mode,
range_workers,
concurrency,
successful_rounds,
failed_rounds,
median(&mut bytes),
median(&mut elapsed),
median(&mut throughput),
median(&mut latency)
);
tokio::fs::OpenOptions::new()
.append(true)
.open(&median_csv)
.await?
.write_all(row.as_bytes())
.await?;
}
Ok(())
}
#[tokio::main]
async fn main() -> Result<()> {
let settings = ToolSettings::from_env()?;
let client = build_client(&settings).await?;
match settings.action {
ToolAction::Prepare => prepare_objects(&settings, &client).await?,
ToolAction::Bench => run_bench(&settings, &client).await?,
}
Ok(())
}