mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-19 19:16:17 +00:00
Merge branch 'main' into feat/kms-vault-transit2
This commit is contained in:
@@ -0,0 +1,441 @@
|
||||
// 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.
|
||||
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use aws_sdk_s3::config::{Credentials, Region};
|
||||
use aws_sdk_s3::primitives::ByteStream;
|
||||
use aws_sdk_s3::types::{Delete, ObjectIdentifier};
|
||||
use aws_sdk_s3::{Client, Config};
|
||||
use aws_smithy_http_client::Builder as SmithyHttpClientBuilder;
|
||||
use bytes::Bytes;
|
||||
use clap::Parser;
|
||||
use serde::Serialize;
|
||||
use std::path::PathBuf;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "small_put_bench")]
|
||||
#[command(about = "Rust-native small PUT benchmark for RustFS-compatible S3 endpoints")]
|
||||
struct Args {
|
||||
#[arg(long, env = "RUSTFS_BENCH_ENDPOINT")]
|
||||
endpoint: String,
|
||||
|
||||
#[arg(long, env = "RUSTFS_BENCH_ACCESS_KEY", default_value = "rustfsadmin")]
|
||||
access_key: String,
|
||||
|
||||
#[arg(long, env = "RUSTFS_BENCH_SECRET_KEY", default_value = "rustfsadmin")]
|
||||
secret_key: String,
|
||||
|
||||
#[arg(long, env = "RUSTFS_BENCH_REGION", default_value = "us-east-1")]
|
||||
region: String,
|
||||
|
||||
#[arg(long, env = "RUSTFS_BENCH_BUCKET", default_value = "small-put-benchmark")]
|
||||
bucket: String,
|
||||
|
||||
#[arg(long, env = "RUSTFS_BENCH_SIZES", default_value = "4KiB,16KiB,64KiB,256KiB,1MiB")]
|
||||
sizes: String,
|
||||
|
||||
#[arg(long, env = "RUSTFS_BENCH_CONCURRENCY", default_value_t = 8)]
|
||||
concurrency: usize,
|
||||
|
||||
#[arg(long, env = "RUSTFS_BENCH_DURATION_SECS", default_value_t = 10)]
|
||||
duration_secs: u64,
|
||||
|
||||
#[arg(long, env = "RUSTFS_BENCH_TIMEOUT_SECS", default_value_t = 15)]
|
||||
timeout_secs: u64,
|
||||
|
||||
#[arg(long, env = "RUSTFS_BENCH_PREFIX")]
|
||||
prefix: Option<String>,
|
||||
|
||||
#[arg(long)]
|
||||
output_json: Option<PathBuf>,
|
||||
|
||||
#[arg(long, default_value_t = false)]
|
||||
cleanup: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct SizeSpec {
|
||||
label: String,
|
||||
slug: String,
|
||||
bytes: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Sample {
|
||||
ok: bool,
|
||||
duration_ms: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct SizeSummary {
|
||||
label: String,
|
||||
bytes: usize,
|
||||
total: usize,
|
||||
succeeded: usize,
|
||||
failed: usize,
|
||||
wall_secs: f64,
|
||||
object_rate: f64,
|
||||
throughput_mib_per_sec: f64,
|
||||
avg_ms: Option<f64>,
|
||||
p50_ms: Option<f64>,
|
||||
p90_ms: Option<f64>,
|
||||
p99_ms: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct RunSummary {
|
||||
run_id: String,
|
||||
endpoint: String,
|
||||
bucket: String,
|
||||
concurrency: usize,
|
||||
duration_secs: u64,
|
||||
timeout_secs: u64,
|
||||
sizes: Vec<SizeSummary>,
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.context("failed to build tokio runtime")?;
|
||||
runtime.block_on(async_main())
|
||||
}
|
||||
|
||||
async fn async_main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
validate_args(&args)?;
|
||||
|
||||
let sizes = parse_size_list(&args.sizes)?;
|
||||
let run_id = args.prefix.clone().unwrap_or_else(default_run_id);
|
||||
let client = build_s3_client(&args.endpoint, &args.access_key, &args.secret_key, &args.region);
|
||||
|
||||
ensure_bucket(&client, &args.bucket).await?;
|
||||
|
||||
let mut size_summaries = Vec::with_capacity(sizes.len());
|
||||
for size in &sizes {
|
||||
let summary = run_size_benchmark(
|
||||
client.clone(),
|
||||
args.bucket.clone(),
|
||||
run_id.clone(),
|
||||
size.clone(),
|
||||
args.concurrency,
|
||||
Duration::from_secs(args.duration_secs),
|
||||
Duration::from_secs(args.timeout_secs),
|
||||
)
|
||||
.await?;
|
||||
print_size_summary(&summary);
|
||||
size_summaries.push(summary);
|
||||
}
|
||||
|
||||
if args.cleanup {
|
||||
cleanup_prefix(&client, &args.bucket, &run_id).await?;
|
||||
}
|
||||
|
||||
let summary = RunSummary {
|
||||
run_id,
|
||||
endpoint: args.endpoint,
|
||||
bucket: args.bucket,
|
||||
concurrency: args.concurrency,
|
||||
duration_secs: args.duration_secs,
|
||||
timeout_secs: args.timeout_secs,
|
||||
sizes: size_summaries,
|
||||
};
|
||||
|
||||
if let Some(path) = args.output_json {
|
||||
let json = serde_json::to_vec_pretty(&summary).context("failed to serialize benchmark summary")?;
|
||||
std::fs::write(&path, json).with_context(|| format!("failed to write benchmark summary to {}", path.display()))?;
|
||||
println!("Wrote summary to {}", path.display());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_args(args: &Args) -> Result<()> {
|
||||
if args.concurrency == 0 {
|
||||
bail!("--concurrency must be greater than zero");
|
||||
}
|
||||
if args.duration_secs == 0 {
|
||||
bail!("--duration-secs must be greater than zero");
|
||||
}
|
||||
if args.timeout_secs == 0 {
|
||||
bail!("--timeout-secs must be greater than zero");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn build_s3_client(endpoint: &str, access_key: &str, secret_key: &str, region: &str) -> Client {
|
||||
let credentials = Credentials::new(access_key, secret_key, None, None, "small-put-bench");
|
||||
let mut config = Config::builder()
|
||||
.credentials_provider(credentials)
|
||||
.region(Region::new(region.to_string()))
|
||||
.endpoint_url(endpoint)
|
||||
.force_path_style(true)
|
||||
.behavior_version_latest();
|
||||
|
||||
if endpoint.starts_with("http://") {
|
||||
config = config.http_client(SmithyHttpClientBuilder::new().build_http());
|
||||
}
|
||||
|
||||
Client::from_conf(config.build())
|
||||
}
|
||||
|
||||
async fn ensure_bucket(client: &Client, bucket: &str) -> Result<()> {
|
||||
if client.head_bucket().bucket(bucket).send().await.is_ok() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
match client.create_bucket().bucket(bucket).send().await {
|
||||
Ok(_) => Ok(()),
|
||||
Err(err) => {
|
||||
let rendered = err.to_string();
|
||||
if rendered.contains("BucketAlreadyOwnedByYou") || rendered.contains("BucketAlreadyExists") {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(err).with_context(|| format!("failed to create benchmark bucket {bucket}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_size_benchmark(
|
||||
client: Client,
|
||||
bucket: String,
|
||||
run_id: String,
|
||||
size: SizeSpec,
|
||||
concurrency: usize,
|
||||
duration: Duration,
|
||||
timeout: Duration,
|
||||
) -> Result<SizeSummary> {
|
||||
let payload = Bytes::from(vec![0_u8; size.bytes]);
|
||||
let deadline = Instant::now() + duration;
|
||||
let wall_start = Instant::now();
|
||||
|
||||
let mut handles = Vec::with_capacity(concurrency);
|
||||
for worker in 0..concurrency {
|
||||
let client = client.clone();
|
||||
let bucket = bucket.clone();
|
||||
let payload = payload.clone();
|
||||
let prefix = format!("{run_id}/{}/worker-{worker}", size.slug);
|
||||
handles.push(tokio::spawn(async move {
|
||||
let mut samples = Vec::new();
|
||||
let mut idx = 0usize;
|
||||
|
||||
while Instant::now() < deadline {
|
||||
let key = format!("{prefix}/obj-{idx}.bin");
|
||||
let started_at = Instant::now();
|
||||
let request = client
|
||||
.put_object()
|
||||
.bucket(&bucket)
|
||||
.key(key)
|
||||
.body(ByteStream::from(payload.clone()))
|
||||
.content_type("application/octet-stream");
|
||||
|
||||
let ok = matches!(tokio::time::timeout(timeout, request.send()).await, Ok(Ok(_)));
|
||||
samples.push(Sample {
|
||||
ok,
|
||||
duration_ms: started_at.elapsed().as_secs_f64() * 1000.0,
|
||||
});
|
||||
idx += 1;
|
||||
}
|
||||
|
||||
samples
|
||||
}));
|
||||
}
|
||||
|
||||
let mut samples = Vec::new();
|
||||
for handle in handles {
|
||||
samples.extend(handle.await.map_err(|err| anyhow!("benchmark worker join error: {err}"))?);
|
||||
}
|
||||
|
||||
Ok(build_size_summary(&size, samples, wall_start.elapsed()))
|
||||
}
|
||||
|
||||
fn build_size_summary(size: &SizeSpec, mut samples: Vec<Sample>, wall_elapsed: Duration) -> SizeSummary {
|
||||
let total = samples.len();
|
||||
let succeeded = samples.iter().filter(|sample| sample.ok).count();
|
||||
let failed = total.saturating_sub(succeeded);
|
||||
let wall_secs = wall_elapsed.as_secs_f64();
|
||||
let object_rate = if wall_secs > 0.0 { succeeded as f64 / wall_secs } else { 0.0 };
|
||||
let throughput_mib_per_sec = if wall_secs > 0.0 {
|
||||
((size.bytes * succeeded) as f64 / (1024.0 * 1024.0)) / wall_secs
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
let avg_ms = if total > 0 {
|
||||
Some(samples.iter().map(|sample| sample.duration_ms).sum::<f64>() / total as f64)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
samples.sort_by(|lhs, rhs| lhs.duration_ms.total_cmp(&rhs.duration_ms));
|
||||
let durations: Vec<f64> = samples.into_iter().map(|sample| sample.duration_ms).collect();
|
||||
|
||||
SizeSummary {
|
||||
label: size.label.clone(),
|
||||
bytes: size.bytes,
|
||||
total,
|
||||
succeeded,
|
||||
failed,
|
||||
wall_secs,
|
||||
object_rate,
|
||||
throughput_mib_per_sec,
|
||||
avg_ms,
|
||||
p50_ms: percentile(&durations, 0.50),
|
||||
p90_ms: percentile(&durations, 0.90),
|
||||
p99_ms: percentile(&durations, 0.99),
|
||||
}
|
||||
}
|
||||
|
||||
async fn cleanup_prefix(client: &Client, bucket: &str, prefix: &str) -> Result<()> {
|
||||
let mut continuation_token = None;
|
||||
loop {
|
||||
let response = client
|
||||
.list_objects_v2()
|
||||
.bucket(bucket)
|
||||
.prefix(prefix)
|
||||
.set_continuation_token(continuation_token.clone())
|
||||
.send()
|
||||
.await
|
||||
.with_context(|| format!("failed to list objects for cleanup under {bucket}/{prefix}"))?;
|
||||
|
||||
let objects: Vec<ObjectIdentifier> = response
|
||||
.contents
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.filter_map(|object| object.key.map(|key| ObjectIdentifier::builder().key(key).build().ok()))
|
||||
.flatten()
|
||||
.collect();
|
||||
|
||||
for chunk in objects.chunks(1_000) {
|
||||
if chunk.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
client
|
||||
.delete_objects()
|
||||
.bucket(bucket)
|
||||
.delete(
|
||||
Delete::builder()
|
||||
.set_objects(Some(chunk.to_vec()))
|
||||
.quiet(true)
|
||||
.build()
|
||||
.context("failed to build delete request")?,
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.with_context(|| format!("failed to delete cleanup batch under {bucket}/{prefix}"))?;
|
||||
}
|
||||
|
||||
if response.is_truncated.unwrap_or(false) {
|
||||
continuation_token = response.next_continuation_token;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_size_list(input: &str) -> Result<Vec<SizeSpec>> {
|
||||
input
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.filter(|item| !item.is_empty())
|
||||
.map(parse_size_spec)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn parse_size_spec(input: &str) -> Result<SizeSpec> {
|
||||
let normalized = input.trim();
|
||||
let lower = normalized.to_ascii_lowercase();
|
||||
|
||||
let (number_part, multiplier) = if let Some(value) = lower.strip_suffix("kib") {
|
||||
(value, 1024usize)
|
||||
} else if let Some(value) = lower.strip_suffix("mib") {
|
||||
(value, 1024usize * 1024usize)
|
||||
} else if let Some(value) = lower.strip_suffix('b') {
|
||||
(value, 1usize)
|
||||
} else {
|
||||
(lower.as_str(), 1usize)
|
||||
};
|
||||
|
||||
let value = number_part
|
||||
.trim()
|
||||
.parse::<usize>()
|
||||
.with_context(|| format!("invalid size component: {input}"))?;
|
||||
let bytes = value
|
||||
.checked_mul(multiplier)
|
||||
.ok_or_else(|| anyhow!("size overflow for {input}"))?;
|
||||
|
||||
Ok(SizeSpec {
|
||||
label: normalized.to_string(),
|
||||
slug: normalized
|
||||
.chars()
|
||||
.filter(|ch| ch.is_ascii_alphanumeric())
|
||||
.collect::<String>()
|
||||
.to_ascii_lowercase(),
|
||||
bytes,
|
||||
})
|
||||
}
|
||||
|
||||
fn percentile(values: &[f64], percentile: f64) -> Option<f64> {
|
||||
if values.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let index = ((values.len() - 1) as f64 * percentile).floor() as usize;
|
||||
values.get(index).copied()
|
||||
}
|
||||
|
||||
fn default_run_id() -> String {
|
||||
format!("small-put-bench-{}", chrono::Utc::now().format("%Y%m%d-%H%M%S"))
|
||||
}
|
||||
|
||||
fn print_size_summary(summary: &SizeSummary) {
|
||||
println!(
|
||||
"{}: success={} failed={} obj/s={:.3} MiB/s={:.3} avg={:.3?} p50={:.3?} p90={:.3?} p99={:.3?}",
|
||||
summary.label,
|
||||
summary.succeeded,
|
||||
summary.failed,
|
||||
summary.object_rate,
|
||||
summary.throughput_mib_per_sec,
|
||||
summary.avg_ms,
|
||||
summary.p50_ms,
|
||||
summary.p90_ms,
|
||||
summary.p99_ms,
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_size_spec_supports_binary_units() {
|
||||
let four_kib = parse_size_spec("4KiB").expect("4KiB should parse");
|
||||
assert_eq!(four_kib.bytes, 4 * 1024);
|
||||
|
||||
let one_mib = parse_size_spec("1MiB").expect("1MiB should parse");
|
||||
assert_eq!(one_mib.bytes, 1024 * 1024);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn percentile_returns_expected_bucket() {
|
||||
let values = vec![10.0, 20.0, 30.0, 40.0, 50.0];
|
||||
assert_eq!(percentile(&values, 0.50), Some(30.0));
|
||||
assert_eq!(percentile(&values, 0.90), Some(40.0));
|
||||
}
|
||||
}
|
||||
@@ -19,9 +19,13 @@
|
||||
mod tests {
|
||||
use crate::common::{RustFSTestEnvironment, init_logging};
|
||||
use aws_sdk_s3::Client;
|
||||
use aws_sdk_s3::primitives::ByteStream;
|
||||
use aws_sdk_s3::primitives::{ByteStream, SdkBody};
|
||||
use aws_sdk_s3::types::{ChecksumAlgorithm, ChecksumMode, CompletedMultipartUpload, CompletedPart};
|
||||
use base64::Engine;
|
||||
use bytes::Bytes;
|
||||
use futures::StreamExt;
|
||||
use http_body::Frame;
|
||||
use http_body_util::StreamBody;
|
||||
use rustfs_rio::{Checksum, ChecksumType as RioChecksumType};
|
||||
use serial_test::serial;
|
||||
use sha2::{Digest, Sha256};
|
||||
@@ -64,6 +68,53 @@ mod tests {
|
||||
.encoded
|
||||
}
|
||||
|
||||
fn streamed_body_70kib_of_a() -> ByteStream {
|
||||
let bytes = Bytes::from_static(&[b'a'; 1024]);
|
||||
let stream = futures::stream::repeat_with(move || {
|
||||
let frame = Frame::data(bytes.clone());
|
||||
Ok::<_, std::io::Error>(frame)
|
||||
});
|
||||
let body = WithSizeHint::new(StreamBody::new(stream.take(70)), 70 * 1024);
|
||||
ByteStream::new(SdkBody::from_body_1_x(body))
|
||||
}
|
||||
|
||||
struct WithSizeHint<T> {
|
||||
inner: T,
|
||||
size_hint: usize,
|
||||
}
|
||||
|
||||
impl<T> WithSizeHint<T> {
|
||||
fn new(inner: T, size_hint: usize) -> Self {
|
||||
Self { inner, size_hint }
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> http_body::Body for WithSizeHint<T>
|
||||
where
|
||||
T: http_body::Body + Unpin,
|
||||
{
|
||||
type Data = T::Data;
|
||||
type Error = T::Error;
|
||||
|
||||
fn poll_frame(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
|
||||
let this = self.get_mut();
|
||||
std::pin::Pin::new(&mut this.inner).poll_frame(cx)
|
||||
}
|
||||
|
||||
fn is_end_stream(&self) -> bool {
|
||||
self.inner.is_end_stream()
|
||||
}
|
||||
|
||||
fn size_hint(&self) -> http_body::SizeHint {
|
||||
let mut hint = self.inner.size_hint();
|
||||
hint.set_exact(self.size_hint as u64);
|
||||
hint
|
||||
}
|
||||
}
|
||||
|
||||
/// PutObject with Content-MD5: upload succeeds and GetObject returns same content.
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
@@ -136,6 +187,121 @@ mod tests {
|
||||
info!("PASSED: PutObject with checksum_sha256 and GetObject content match");
|
||||
}
|
||||
|
||||
/// Mirrors `s3s-e2e` behavior: only request `checksum_algorithm`, then expect
|
||||
/// both PutObject and GetObject(checksum_mode=enabled) to expose the same checksum.
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn test_put_object_with_checksum_algorithm_only() {
|
||||
init_logging();
|
||||
info!("TEST: PutObject with checksum_algorithm only");
|
||||
|
||||
let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
|
||||
env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");
|
||||
|
||||
let client = create_s3_client(&env);
|
||||
let bucket = "test-checksum-algorithm-only";
|
||||
create_bucket(&client, bucket).await.expect("Failed to create bucket");
|
||||
|
||||
let key = "obj-with-checksum-algorithm-only.txt";
|
||||
let content = vec![b'a'; 70 * 1024];
|
||||
|
||||
let put_resp = client
|
||||
.put_object()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.checksum_algorithm(ChecksumAlgorithm::Crc32)
|
||||
.body(ByteStream::from(content.clone()))
|
||||
.send()
|
||||
.await
|
||||
.expect("PutObject with checksum_algorithm should succeed");
|
||||
|
||||
let put_checksum = put_resp
|
||||
.checksum_crc32()
|
||||
.expect("PutObject should return checksum_crc32 when checksum_algorithm is used")
|
||||
.to_string();
|
||||
|
||||
let mut get_resp = client
|
||||
.get_object()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.checksum_mode(ChecksumMode::Enabled)
|
||||
.send()
|
||||
.await
|
||||
.expect("GetObject should succeed");
|
||||
|
||||
let body_bytes = std::mem::replace(&mut get_resp.body, ByteStream::new(aws_sdk_s3::primitives::SdkBody::empty()))
|
||||
.collect()
|
||||
.await
|
||||
.expect("collect body")
|
||||
.into_bytes();
|
||||
|
||||
assert_eq!(body_bytes.as_ref(), content.as_slice(), "GetObject body must match uploaded content");
|
||||
assert_eq!(
|
||||
get_resp.checksum_crc32().map(str::to_string),
|
||||
Some(put_checksum),
|
||||
"GetObject(checksum_mode=enabled) should expose the stored CRC32 checksum"
|
||||
);
|
||||
}
|
||||
|
||||
/// Matches the `s3s-e2e` streaming upload shape more closely than `ByteStream::from(Vec<u8>)`.
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn test_put_object_with_checksum_algorithm_only_streaming_body() {
|
||||
init_logging();
|
||||
info!("TEST: PutObject with checksum_algorithm only using streaming body");
|
||||
|
||||
let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
|
||||
env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");
|
||||
|
||||
let client = create_s3_client(&env);
|
||||
let bucket = "test-checksum-algorithm-streaming";
|
||||
create_bucket(&client, bucket).await.expect("Failed to create bucket");
|
||||
|
||||
let key = "obj-with-checksum-algorithm-streaming.txt";
|
||||
let expected_content = vec![b'a'; 70 * 1024];
|
||||
|
||||
let put_resp = client
|
||||
.put_object()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.checksum_algorithm(ChecksumAlgorithm::Crc32)
|
||||
.body(streamed_body_70kib_of_a())
|
||||
.send()
|
||||
.await
|
||||
.expect("PutObject with streaming checksum_algorithm should succeed");
|
||||
|
||||
let put_checksum = put_resp
|
||||
.checksum_crc32()
|
||||
.expect("PutObject should return checksum_crc32 for streaming checksum_algorithm uploads")
|
||||
.to_string();
|
||||
|
||||
let mut get_resp = client
|
||||
.get_object()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.checksum_mode(ChecksumMode::Enabled)
|
||||
.send()
|
||||
.await
|
||||
.expect("GetObject should succeed");
|
||||
|
||||
let body_bytes = std::mem::replace(&mut get_resp.body, ByteStream::new(SdkBody::empty()))
|
||||
.collect()
|
||||
.await
|
||||
.expect("collect body")
|
||||
.into_bytes();
|
||||
|
||||
assert_eq!(
|
||||
body_bytes.as_ref(),
|
||||
expected_content.as_slice(),
|
||||
"GetObject body must match uploaded content"
|
||||
);
|
||||
assert_eq!(
|
||||
get_resp.checksum_crc32().map(str::to_string),
|
||||
Some(put_checksum),
|
||||
"GetObject(checksum_mode=enabled) should expose the stored CRC32 checksum for streaming uploads"
|
||||
);
|
||||
}
|
||||
|
||||
/// Multipart upload with checksum: CreateMultipartUpload, UploadPart(s) with checksum_sha256, CompleteMultipartUpload; then GetObject verifies content.
|
||||
/// Uses part size >= 5MB (server minimum) for two parts.
|
||||
#[tokio::test]
|
||||
@@ -234,6 +400,114 @@ mod tests {
|
||||
info!("PASSED: MultipartUpload with checksum and GetObject content match");
|
||||
}
|
||||
|
||||
/// Mirrors `s3s-e2e` multipart behavior: request checksum algorithm at MPU creation,
|
||||
/// rely on auto checksum handling during UploadPart, and expect CompleteMultipartUpload to succeed.
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn test_multipart_upload_with_crc32_algorithm_only() {
|
||||
init_logging();
|
||||
info!("TEST: MultipartUpload with checksum_algorithm only (CRC32)");
|
||||
|
||||
let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
|
||||
env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");
|
||||
|
||||
let client = create_s3_client(&env);
|
||||
let bucket = "test-multipart-checksum-crc32-auto";
|
||||
create_bucket(&client, bucket).await.expect("Failed to create bucket");
|
||||
|
||||
let key = "multipart-with-crc32-auto.bin";
|
||||
let part1_content = "a".repeat(5 * 1024 * 1024 + 1);
|
||||
let part2_content = "b".repeat(1024);
|
||||
|
||||
let create_resp = client
|
||||
.create_multipart_upload()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.checksum_algorithm(ChecksumAlgorithm::Crc32)
|
||||
.send()
|
||||
.await
|
||||
.expect("CreateMultipartUpload should succeed");
|
||||
|
||||
let upload_id = create_resp.upload_id().expect("upload_id should be present");
|
||||
|
||||
let part1_resp = client
|
||||
.upload_part()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.upload_id(upload_id)
|
||||
.part_number(1)
|
||||
.body(ByteStream::from(part1_content.clone().into_bytes()))
|
||||
.send()
|
||||
.await
|
||||
.expect("UploadPart 1 should succeed");
|
||||
let part1_checksum = part1_resp
|
||||
.checksum_crc32()
|
||||
.expect("UploadPart 1 should return checksum_crc32")
|
||||
.to_string();
|
||||
|
||||
let part2_resp = client
|
||||
.upload_part()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.upload_id(upload_id)
|
||||
.part_number(2)
|
||||
.body(ByteStream::from(part2_content.clone().into_bytes()))
|
||||
.send()
|
||||
.await
|
||||
.expect("UploadPart 2 should succeed");
|
||||
let part2_checksum = part2_resp
|
||||
.checksum_crc32()
|
||||
.expect("UploadPart 2 should return checksum_crc32")
|
||||
.to_string();
|
||||
|
||||
let completed_upload = CompletedMultipartUpload::builder()
|
||||
.parts(
|
||||
CompletedPart::builder()
|
||||
.part_number(1)
|
||||
.e_tag(part1_resp.e_tag().expect("etag part 1"))
|
||||
.checksum_crc32(part1_checksum)
|
||||
.build(),
|
||||
)
|
||||
.parts(
|
||||
CompletedPart::builder()
|
||||
.part_number(2)
|
||||
.e_tag(part2_resp.e_tag().expect("etag part 2"))
|
||||
.checksum_crc32(part2_checksum)
|
||||
.build(),
|
||||
)
|
||||
.build();
|
||||
|
||||
client
|
||||
.complete_multipart_upload()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.upload_id(upload_id)
|
||||
.multipart_upload(completed_upload)
|
||||
.send()
|
||||
.await
|
||||
.expect("CompleteMultipartUpload should succeed");
|
||||
|
||||
let body_bytes = client
|
||||
.get_object()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.send()
|
||||
.await
|
||||
.expect("GetObject should succeed")
|
||||
.body
|
||||
.collect()
|
||||
.await
|
||||
.expect("collect body")
|
||||
.into_bytes();
|
||||
|
||||
let expected_content = format!("{part1_content}{part2_content}");
|
||||
assert_eq!(
|
||||
body_bytes.as_ref(),
|
||||
expected_content.as_bytes(),
|
||||
"completed multipart object must match concatenated parts"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression test for issue #2282:
|
||||
/// CRC64NVME full-object checksum should match between direct PutObject and multipart upload.
|
||||
#[tokio::test]
|
||||
|
||||
@@ -344,6 +344,9 @@ async fn test_local_kms_multipart_upload() {
|
||||
test_multipart_upload_with_sse_c(&s3_client, TEST_BUCKET)
|
||||
.await
|
||||
.expect("SSE-C multipart upload test failed");
|
||||
test_multipart_download_with_wrong_sse_c_key_fails(&s3_client, TEST_BUCKET)
|
||||
.await
|
||||
.expect("SSE-C multipart wrong-key download test failed");
|
||||
|
||||
// Test 4: Large multipart upload (test streaming encryption with multiple blocks)
|
||||
// TODO: Re-enable after fixing streaming encryption issues with large files
|
||||
@@ -648,6 +651,99 @@ async fn test_multipart_upload_with_sse_c(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_multipart_download_with_wrong_sse_c_key_fails(
|
||||
s3_client: &aws_sdk_s3::Client,
|
||||
bucket: &str,
|
||||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
let object_key = "multipart-sse-c-bad-download-test";
|
||||
let part_size = 5 * 1024 * 1024;
|
||||
let total_parts = 2;
|
||||
let total_size = part_size * total_parts;
|
||||
|
||||
let encryption_key = "01234567890123456789012345678901";
|
||||
let key_b64 = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, encryption_key);
|
||||
let key_md5 = sse_customer_key_md5_base64(encryption_key);
|
||||
|
||||
let wrong_key = "abcdefghijklmnopqrstuvwxyz012345";
|
||||
let wrong_key_b64 = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, wrong_key);
|
||||
let wrong_key_md5 = sse_customer_key_md5_base64(wrong_key);
|
||||
|
||||
let test_data: Vec<u8> = (0..total_size).map(|i| ((i * 5) % 256) as u8).collect();
|
||||
|
||||
let create_multipart_output = s3_client
|
||||
.create_multipart_upload()
|
||||
.bucket(bucket)
|
||||
.key(object_key)
|
||||
.sse_customer_algorithm("AES256")
|
||||
.sse_customer_key(&key_b64)
|
||||
.sse_customer_key_md5(&key_md5)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let upload_id = create_multipart_output.upload_id().unwrap();
|
||||
let mut completed_parts = Vec::new();
|
||||
|
||||
for part_number in 1..=total_parts {
|
||||
let start = (part_number - 1) * part_size;
|
||||
let end = std::cmp::min(start + part_size, total_size);
|
||||
let part_data = &test_data[start..end];
|
||||
|
||||
let upload_part_output = s3_client
|
||||
.upload_part()
|
||||
.bucket(bucket)
|
||||
.key(object_key)
|
||||
.upload_id(upload_id)
|
||||
.part_number(part_number as i32)
|
||||
.body(aws_sdk_s3::primitives::ByteStream::from(part_data.to_vec()))
|
||||
.sse_customer_algorithm("AES256")
|
||||
.sse_customer_key(&key_b64)
|
||||
.sse_customer_key_md5(&key_md5)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
completed_parts.push(
|
||||
aws_sdk_s3::types::CompletedPart::builder()
|
||||
.part_number(part_number as i32)
|
||||
.e_tag(upload_part_output.e_tag().unwrap())
|
||||
.build(),
|
||||
);
|
||||
}
|
||||
|
||||
let completed_multipart_upload = aws_sdk_s3::types::CompletedMultipartUpload::builder()
|
||||
.set_parts(Some(completed_parts))
|
||||
.build();
|
||||
|
||||
s3_client
|
||||
.complete_multipart_upload()
|
||||
.bucket(bucket)
|
||||
.key(object_key)
|
||||
.upload_id(upload_id)
|
||||
.multipart_upload(completed_multipart_upload)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let err = s3_client
|
||||
.get_object()
|
||||
.bucket(bucket)
|
||||
.key(object_key)
|
||||
.sse_customer_algorithm("AES256")
|
||||
.sse_customer_key(&wrong_key_b64)
|
||||
.sse_customer_key_md5(&wrong_key_md5)
|
||||
.send()
|
||||
.await
|
||||
.expect_err("multipart SSE-C download with the wrong key should fail");
|
||||
|
||||
let service_err = err.into_service_error();
|
||||
assert_eq!(
|
||||
service_err.meta().code(),
|
||||
Some("InvalidRequest"),
|
||||
"wrong-key multipart SSE-C download should return InvalidRequest, got {:?}",
|
||||
service_err.meta().code()
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test large multipart upload to verify streaming encryption works correctly
|
||||
#[allow(dead_code)]
|
||||
async fn test_large_multipart_upload(
|
||||
|
||||
@@ -97,6 +97,10 @@ mod cluster_concurrency_test;
|
||||
#[cfg(test)]
|
||||
mod checksum_upload_test;
|
||||
|
||||
// Range request regression tests
|
||||
#[cfg(test)]
|
||||
mod range_request_test;
|
||||
|
||||
// Group deletion tests
|
||||
#[cfg(test)]
|
||||
mod group_delete_test;
|
||||
|
||||
@@ -313,6 +313,31 @@ async fn list_target_arns(env: &RustFSTestEnvironment) -> Result<Vec<String>, Bo
|
||||
Ok(serde_json::from_slice(&body)?)
|
||||
}
|
||||
|
||||
async fn delete_webhook_target(env: &RustFSTestEnvironment, target_name: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
let url = format!("{}/rustfs/admin/v3/target/notify_webhook/{target_name}/reset", env.url);
|
||||
let response = signed_request(http::Method::DELETE, &url, &env.access_key, &env.secret_key, None, None).await?;
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
if status != StatusCode::OK {
|
||||
return Err(format!("failed to delete webhook target {target_name}: {status} {body}").into());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn notification_target_is_listed(targets: &serde_json::Value, target_name: &str) -> bool {
|
||||
targets["notification_endpoints"]
|
||||
.as_array()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.any(|entry| {
|
||||
entry["account_id"].as_str() == Some(target_name)
|
||||
&& entry["service"]
|
||||
.as_str()
|
||||
.is_some_and(|service| service == "webhook" || service.starts_with("webhook-"))
|
||||
})
|
||||
}
|
||||
|
||||
async fn wait_for_target_visibility(
|
||||
env: &RustFSTestEnvironment,
|
||||
target_name: &str,
|
||||
@@ -324,18 +349,7 @@ async fn wait_for_target_visibility(
|
||||
last_targets = list_notification_targets(env).await?;
|
||||
last_arns = list_target_arns(env).await?;
|
||||
|
||||
let listed = last_targets["notification_endpoints"]
|
||||
.as_array()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.any(|entry| {
|
||||
entry["account_id"].as_str() == Some(target_name)
|
||||
&& entry["service"]
|
||||
.as_str()
|
||||
.is_some_and(|service| service == "webhook" || service.starts_with("webhook-"))
|
||||
});
|
||||
|
||||
if listed {
|
||||
if notification_target_is_listed(&last_targets, target_name) {
|
||||
return Ok((last_targets, last_arns));
|
||||
}
|
||||
|
||||
@@ -345,10 +359,51 @@ async fn wait_for_target_visibility(
|
||||
Err(format!("target {target_name} did not become visible in admin APIs; targets={last_targets}, arns={last_arns:?}").into())
|
||||
}
|
||||
|
||||
async fn wait_for_target_absence(
|
||||
env: &RustFSTestEnvironment,
|
||||
target_name: &str,
|
||||
) -> Result<(serde_json::Value, Vec<String>), Box<dyn Error + Send + Sync>> {
|
||||
let mut last_targets = serde_json::Value::Null;
|
||||
let mut last_arns = Vec::new();
|
||||
|
||||
for _ in 0..20 {
|
||||
last_targets = list_notification_targets(env).await?;
|
||||
last_arns = list_target_arns(env).await?;
|
||||
|
||||
let listed = notification_target_is_listed(&last_targets, target_name);
|
||||
let arn_listed = last_arns.iter().any(|arn| arn.ends_with(&format!(":{target_name}:webhook")));
|
||||
if !listed && !arn_listed {
|
||||
return Ok((last_targets, last_arns));
|
||||
}
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(250)).await;
|
||||
}
|
||||
|
||||
Err(format!("target {target_name} remained visible in admin APIs; targets={last_targets}, arns={last_arns:?}").into())
|
||||
}
|
||||
|
||||
async fn restart_rustfs_server(env: &mut RustFSTestEnvironment) -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
env.stop_server();
|
||||
env.start_rustfs_server_without_cleanup(vec![]).await
|
||||
}
|
||||
|
||||
async fn read_persisted_server_config(env: &RustFSTestEnvironment) -> String {
|
||||
let path = format!("{}/.rustfs.sys/config/config.json", env.temp_dir);
|
||||
match tokio::fs::read_to_string(&path).await {
|
||||
Ok(content) => content,
|
||||
Err(err) if err.kind() == std::io::ErrorKind::IsADirectory => {
|
||||
let mut entries = Vec::new();
|
||||
match tokio::fs::read_dir(&path).await {
|
||||
Ok(mut dir) => {
|
||||
while let Ok(Some(entry)) = dir.next_entry().await {
|
||||
entries.push(entry.file_name().to_string_lossy().to_string());
|
||||
}
|
||||
entries.sort();
|
||||
format!("persisted config stored as object directory at {path}; entries={entries:?}")
|
||||
}
|
||||
Err(dir_err) => format!("persisted config directory exists at {path} but could not be listed: {dir_err}"),
|
||||
}
|
||||
}
|
||||
Err(err) => format!("failed to read persisted config at {path}: {err}"),
|
||||
}
|
||||
}
|
||||
@@ -400,6 +455,66 @@ async fn read_listen_notification_event(
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn test_notification_target_persists_across_restart_and_delete() -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
init_logging();
|
||||
|
||||
let (webhook_url, _request_rx, webhook_handle) = spawn_object_lambda_webhook_server().await?;
|
||||
|
||||
let mut env = RustFSTestEnvironment::new().await?;
|
||||
env.start_rustfs_server(vec![]).await?;
|
||||
|
||||
let target_name = "restart-target";
|
||||
configure_webhook_target(&env, target_name, &webhook_url, "secret-token").await?;
|
||||
|
||||
let (visible_targets, visible_arns) = wait_for_target_visibility(&env, target_name).await?;
|
||||
assert!(notification_target_is_listed(&visible_targets, target_name));
|
||||
assert!(
|
||||
visible_arns
|
||||
.iter()
|
||||
.any(|arn| arn.ends_with(&format!(":{target_name}:webhook"))),
|
||||
"target ARN missing after initial configure: {visible_arns:?}"
|
||||
);
|
||||
|
||||
restart_rustfs_server(&mut env).await?;
|
||||
|
||||
let (targets_after_restart, arns_after_restart) = wait_for_target_visibility(&env, target_name).await?;
|
||||
assert!(notification_target_is_listed(&targets_after_restart, target_name));
|
||||
assert!(
|
||||
arns_after_restart
|
||||
.iter()
|
||||
.any(|arn| arn.ends_with(&format!(":{target_name}:webhook"))),
|
||||
"target ARN missing after restart: {arns_after_restart:?}"
|
||||
);
|
||||
|
||||
delete_webhook_target(&env, target_name).await?;
|
||||
let (targets_after_delete, arns_after_delete) = wait_for_target_absence(&env, target_name).await?;
|
||||
assert!(!notification_target_is_listed(&targets_after_delete, target_name));
|
||||
assert!(
|
||||
!arns_after_delete
|
||||
.iter()
|
||||
.any(|arn| arn.ends_with(&format!(":{target_name}:webhook"))),
|
||||
"target ARN still visible after delete: {arns_after_delete:?}"
|
||||
);
|
||||
|
||||
restart_rustfs_server(&mut env).await?;
|
||||
|
||||
let (targets_after_delete_restart, arns_after_delete_restart) = wait_for_target_absence(&env, target_name).await?;
|
||||
assert!(!notification_target_is_listed(&targets_after_delete_restart, target_name));
|
||||
assert!(
|
||||
!arns_after_delete_restart
|
||||
.iter()
|
||||
.any(|arn| arn.ends_with(&format!(":{target_name}:webhook"))),
|
||||
"target ARN still visible after delete + restart: {arns_after_delete_restart:?}"
|
||||
);
|
||||
|
||||
webhook_handle.abort();
|
||||
let _ = webhook_handle.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn test_get_object_lambda_accepts_presigned_requests() -> Result<(), Box<dyn Error + Send + Sync>> {
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
// 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.
|
||||
|
||||
//! End-to-end regression test for invalid GET object ranges.
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::common::{RustFSTestEnvironment, init_logging};
|
||||
use aws_sdk_s3::Client;
|
||||
use aws_sdk_s3::error::SdkError;
|
||||
use aws_sdk_s3::primitives::ByteStream;
|
||||
use serial_test::serial;
|
||||
use tracing::info;
|
||||
|
||||
fn create_s3_client(env: &RustFSTestEnvironment) -> Client {
|
||||
env.create_s3_client()
|
||||
}
|
||||
|
||||
async fn create_bucket(client: &Client, bucket: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||||
match client.create_bucket().bucket(bucket).send().await {
|
||||
Ok(_) => Ok(()),
|
||||
Err(err) => {
|
||||
if err.to_string().contains("BucketAlreadyOwnedByYou") || err.to_string().contains("BucketAlreadyExists") {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(Box::new(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial]
|
||||
async fn test_get_object_invalid_range_returns_416_issue_s3_implemented_tests() {
|
||||
init_logging();
|
||||
info!("TEST: GetObject invalid range should return InvalidRange/416");
|
||||
|
||||
let mut env = RustFSTestEnvironment::new().await.expect("Failed to create test environment");
|
||||
env.start_rustfs_server(vec![]).await.expect("Failed to start RustFS");
|
||||
|
||||
let client = create_s3_client(&env);
|
||||
let bucket = "test-invalid-range";
|
||||
let key = "range.txt";
|
||||
let content = b"testcontent";
|
||||
|
||||
create_bucket(&client, bucket).await.expect("Failed to create bucket");
|
||||
client
|
||||
.put_object()
|
||||
.bucket(bucket)
|
||||
.key(key)
|
||||
.body(ByteStream::from_static(content))
|
||||
.send()
|
||||
.await
|
||||
.expect("PutObject should succeed");
|
||||
|
||||
let result = client.get_object().bucket(bucket).key(key).range("bytes=40-50").send().await;
|
||||
|
||||
let err = result.expect_err("GetObject with an unsatisfiable range should fail");
|
||||
match err {
|
||||
SdkError::ServiceError(service_err) => {
|
||||
assert_eq!(service_err.raw().status().as_u16(), 416, "invalid range should return HTTP 416");
|
||||
|
||||
let s3_err = service_err.into_err();
|
||||
assert_eq!(s3_err.meta().code(), Some("InvalidRange"), "invalid range should map to InvalidRange");
|
||||
}
|
||||
other_err => panic!("Expected S3 service error, got: {other_err:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -21,7 +21,7 @@ use rustfs_lock::{
|
||||
LockClient, LockError, LockId, LockInfo, LockRequest, LockResponse, LockStats, LockStatus, LockType, Result,
|
||||
types::{LockMetadata, LockPriority},
|
||||
};
|
||||
use rustfs_protos::proto_gen::node_service::{GenerallyLockRequest, PingRequest};
|
||||
use rustfs_protos::proto_gen::node_service::{BatchGenerallyLockRequest, GenerallyLockRequest, PingRequest};
|
||||
use tonic::Request;
|
||||
use tracing::{info, warn};
|
||||
|
||||
@@ -64,6 +64,44 @@ impl GrpcLockClient {
|
||||
suppress_contention_logs: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_lock_info(request: &LockRequest, lock_info_json: Option<String>) -> LockInfo {
|
||||
if let Some(lock_info_json) = lock_info_json {
|
||||
match serde_json::from_str::<LockInfo>(&lock_info_json) {
|
||||
Ok(info) => info,
|
||||
Err(e) => {
|
||||
warn!("Failed to deserialize lock_info from response: {}, using request data", e);
|
||||
LockInfo {
|
||||
id: request.lock_id.clone(),
|
||||
resource: request.resource.clone(),
|
||||
lock_type: request.lock_type,
|
||||
status: LockStatus::Acquired,
|
||||
owner: request.owner.clone(),
|
||||
acquired_at: std::time::SystemTime::now(),
|
||||
expires_at: std::time::SystemTime::now() + request.ttl,
|
||||
last_refreshed: std::time::SystemTime::now(),
|
||||
metadata: request.metadata.clone(),
|
||||
priority: request.priority,
|
||||
wait_start_time: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
LockInfo {
|
||||
id: request.lock_id.clone(),
|
||||
resource: request.resource.clone(),
|
||||
lock_type: request.lock_type,
|
||||
status: LockStatus::Acquired,
|
||||
owner: request.owner.clone(),
|
||||
acquired_at: std::time::SystemTime::now(),
|
||||
expires_at: std::time::SystemTime::now() + request.ttl,
|
||||
last_refreshed: std::time::SystemTime::now(),
|
||||
metadata: request.metadata.clone(),
|
||||
priority: request.priority,
|
||||
wait_start_time: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -89,46 +127,10 @@ impl LockClient for GrpcLockClient {
|
||||
|
||||
// Check if the lock acquisition was successful
|
||||
if resp.success {
|
||||
// Try to deserialize lock_info from response
|
||||
let lock_info = if let Some(lock_info_json) = resp.lock_info {
|
||||
match serde_json::from_str::<LockInfo>(&lock_info_json) {
|
||||
Ok(info) => info,
|
||||
Err(e) => {
|
||||
// If deserialization fails, fall back to constructing from request
|
||||
warn!("Failed to deserialize lock_info from response: {}, using request data", e);
|
||||
LockInfo {
|
||||
id: request.lock_id.clone(),
|
||||
resource: request.resource.clone(),
|
||||
lock_type: request.lock_type,
|
||||
status: LockStatus::Acquired,
|
||||
owner: request.owner.clone(),
|
||||
acquired_at: std::time::SystemTime::now(),
|
||||
expires_at: std::time::SystemTime::now() + request.ttl,
|
||||
last_refreshed: std::time::SystemTime::now(),
|
||||
metadata: request.metadata.clone(),
|
||||
priority: request.priority,
|
||||
wait_start_time: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// If lock_info is not provided, construct from request
|
||||
LockInfo {
|
||||
id: request.lock_id.clone(),
|
||||
resource: request.resource.clone(),
|
||||
lock_type: request.lock_type,
|
||||
status: LockStatus::Acquired,
|
||||
owner: request.owner.clone(),
|
||||
acquired_at: std::time::SystemTime::now(),
|
||||
expires_at: std::time::SystemTime::now() + request.ttl,
|
||||
last_refreshed: std::time::SystemTime::now(),
|
||||
metadata: request.metadata.clone(),
|
||||
priority: request.priority,
|
||||
wait_start_time: None,
|
||||
}
|
||||
};
|
||||
|
||||
Ok(LockResponse::success(lock_info, std::time::Duration::ZERO))
|
||||
Ok(LockResponse::success(
|
||||
Self::build_lock_info(request, resp.lock_info),
|
||||
std::time::Duration::ZERO,
|
||||
))
|
||||
} else {
|
||||
// Lock acquisition failed
|
||||
Ok(LockResponse::failure(
|
||||
@@ -138,6 +140,45 @@ impl LockClient for GrpcLockClient {
|
||||
}
|
||||
}
|
||||
|
||||
async fn acquire_locks_batch(&self, requests: &[LockRequest]) -> Result<Vec<LockResponse>> {
|
||||
let mut client = self.get_client().await?;
|
||||
let req = Request::new(BatchGenerallyLockRequest {
|
||||
args: requests
|
||||
.iter()
|
||||
.map(|request| {
|
||||
serde_json::to_string(request).map_err(|e| LockError::internal(format!("Failed to serialize request: {e}")))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?,
|
||||
});
|
||||
|
||||
let resp = client
|
||||
.lock_batch(req)
|
||||
.await
|
||||
.map_err(|e| LockError::internal(e.to_string()))?
|
||||
.into_inner();
|
||||
|
||||
Ok(requests
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, request)| match resp.results.get(idx) {
|
||||
Some(result) if result.success => {
|
||||
LockResponse::success(Self::build_lock_info(request, result.lock_info.clone()), std::time::Duration::ZERO)
|
||||
}
|
||||
Some(result) => LockResponse::failure(
|
||||
result
|
||||
.error_info
|
||||
.clone()
|
||||
.unwrap_or_else(|| "Lock acquisition failed on remote server".to_string()),
|
||||
std::time::Duration::ZERO,
|
||||
),
|
||||
None => LockResponse::failure(
|
||||
format!("Lock batch response missing entry for request index {idx}"),
|
||||
std::time::Duration::ZERO,
|
||||
),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn release(&self, lock_id: &LockId) -> Result<bool> {
|
||||
info!("grpc release for {}", lock_id);
|
||||
|
||||
@@ -161,6 +202,31 @@ impl LockClient for GrpcLockClient {
|
||||
Ok(resp.success)
|
||||
}
|
||||
|
||||
async fn release_locks_batch(&self, lock_ids: &[LockId]) -> Result<Vec<bool>> {
|
||||
let mut client = self.get_client().await?;
|
||||
let req = Request::new(BatchGenerallyLockRequest {
|
||||
args: lock_ids
|
||||
.iter()
|
||||
.map(|lock_id| {
|
||||
serde_json::to_string(&Self::create_unlock_request(lock_id))
|
||||
.map_err(|e| LockError::internal(format!("Failed to serialize request: {e}")))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?,
|
||||
});
|
||||
|
||||
let resp = client
|
||||
.un_lock_batch(req)
|
||||
.await
|
||||
.map_err(|e| LockError::internal(e.to_string()))?
|
||||
.into_inner();
|
||||
|
||||
Ok(lock_ids
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, _)| resp.results.get(idx).map(|result| result.success).unwrap_or(false))
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn refresh(&self, lock_id: &LockId) -> Result<bool> {
|
||||
info!("grpc refresh for {}", lock_id);
|
||||
let refresh_request = Self::create_unlock_request(lock_id);
|
||||
|
||||
@@ -21,7 +21,8 @@ use rustfs_lock::{LockClient, LockRequest};
|
||||
use rustfs_protos::{
|
||||
models::PingBodyBuilder,
|
||||
proto_gen::node_service::{
|
||||
GenerallyLockRequest, GenerallyLockResponse, PingRequest, PingResponse, node_service_server::NodeService,
|
||||
BatchGenerallyLockRequest, BatchGenerallyLockResponse, GenerallyLockRequest, GenerallyLockResponse, GenerallyLockResult,
|
||||
PingRequest, PingResponse, node_service_server::NodeService,
|
||||
},
|
||||
};
|
||||
use std::pin::Pin;
|
||||
@@ -33,6 +34,22 @@ use tracing::debug;
|
||||
|
||||
type ResponseStream<T> = Pin<Box<dyn Stream<Item = Result<T, Status>> + Send>>;
|
||||
|
||||
fn lock_result_from_response(response: rustfs_lock::LockResponse) -> GenerallyLockResult {
|
||||
GenerallyLockResult {
|
||||
success: response.success,
|
||||
error_info: response.error,
|
||||
lock_info: response.lock_info.and_then(|info| serde_json::to_string(&info).ok()),
|
||||
}
|
||||
}
|
||||
|
||||
fn lock_result_from_error(error: impl Into<String>) -> GenerallyLockResult {
|
||||
GenerallyLockResult {
|
||||
success: false,
|
||||
error_info: Some(error.into()),
|
||||
lock_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimal NodeService implementation that only supports Lock RPCs
|
||||
/// Used for testing distributed lock scenarios with real gRPC
|
||||
#[derive(Debug)]
|
||||
@@ -187,6 +204,92 @@ impl NodeService for MinimalLockNodeService {
|
||||
}
|
||||
}
|
||||
|
||||
async fn lock_batch(
|
||||
&self,
|
||||
request: Request<BatchGenerallyLockRequest>,
|
||||
) -> Result<Response<BatchGenerallyLockResponse>, Status> {
|
||||
let request = request.into_inner();
|
||||
let mut results = vec![lock_result_from_error("request was not processed"); request.args.len()];
|
||||
let mut valid_requests = Vec::with_capacity(request.args.len());
|
||||
let mut valid_indices = Vec::with_capacity(request.args.len());
|
||||
|
||||
for (idx, arg) in request.args.iter().enumerate() {
|
||||
match serde_json::from_str::<LockRequest>(arg) {
|
||||
Ok(args) => {
|
||||
valid_requests.push(args);
|
||||
valid_indices.push(idx);
|
||||
}
|
||||
Err(err) => {
|
||||
results[idx] = lock_result_from_error(format!("can not decode args, err: {err}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !valid_requests.is_empty() {
|
||||
match self.lock_client.acquire_locks_batch(&valid_requests).await {
|
||||
Ok(batch_results) => {
|
||||
for (result_idx, response) in batch_results.into_iter().enumerate() {
|
||||
if let Some(request_idx) = valid_indices.get(result_idx) {
|
||||
results[*request_idx] = lock_result_from_response(response);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
for request_idx in valid_indices {
|
||||
results[request_idx] = lock_result_from_error(format!("can not batch lock, err: {err}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Response::new(BatchGenerallyLockResponse { results }))
|
||||
}
|
||||
|
||||
async fn un_lock_batch(
|
||||
&self,
|
||||
request: Request<BatchGenerallyLockRequest>,
|
||||
) -> Result<Response<BatchGenerallyLockResponse>, Status> {
|
||||
let request = request.into_inner();
|
||||
let mut results = vec![lock_result_from_error("request was not processed"); request.args.len()];
|
||||
let mut lock_ids = Vec::with_capacity(request.args.len());
|
||||
let mut valid_indices = Vec::with_capacity(request.args.len());
|
||||
|
||||
for (idx, arg) in request.args.iter().enumerate() {
|
||||
match serde_json::from_str::<LockRequest>(arg) {
|
||||
Ok(args) => {
|
||||
lock_ids.push(args.lock_id);
|
||||
valid_indices.push(idx);
|
||||
}
|
||||
Err(err) => {
|
||||
results[idx] = lock_result_from_error(format!("can not decode args, err: {err}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !lock_ids.is_empty() {
|
||||
match self.lock_client.release_locks_batch(&lock_ids).await {
|
||||
Ok(batch_results) => {
|
||||
for (result_idx, success) in batch_results.into_iter().enumerate() {
|
||||
if let Some(request_idx) = valid_indices.get(result_idx) {
|
||||
results[*request_idx] = GenerallyLockResult {
|
||||
success,
|
||||
error_info: None,
|
||||
lock_info: None,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
for request_idx in valid_indices {
|
||||
results[request_idx] = lock_result_from_error(format!("can not batch unlock, err: {err}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Response::new(BatchGenerallyLockResponse { results }))
|
||||
}
|
||||
|
||||
// All other methods return unimplemented
|
||||
async fn heal_bucket(
|
||||
&self,
|
||||
|
||||
@@ -14,8 +14,10 @@
|
||||
// limitations under the License.
|
||||
|
||||
use super::{grpc_lock_client::GrpcLockClient, grpc_lock_server::spawn_lock_server};
|
||||
use rustfs_lock::client::local::LocalClient;
|
||||
use rustfs_lock::{GlobalLockManager, LockError, LockInfo, LockResponse, LockStats, NamespaceLock, ObjectKey};
|
||||
use rustfs_lock::client::{LockClient, local::LocalClient};
|
||||
use rustfs_lock::{
|
||||
GlobalLockManager, LockError, LockInfo, LockRequest, LockResponse, LockStats, LockType, NamespaceLock, ObjectKey,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -223,6 +225,56 @@ async fn test_distributed_lock_2_nodes_grpc_read_survives_failed_node() {
|
||||
failing_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_grpc_lock_client_batch_acquire_and_release() {
|
||||
let manager = Arc::new(GlobalLockManager::new());
|
||||
let local_client: Arc<dyn rustfs_lock::LockClient> = Arc::new(LocalClient::with_manager(manager));
|
||||
|
||||
let (addr, handle) = spawn_lock_server(local_client).await.expect("Failed to spawn server");
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
|
||||
let grpc_client = GrpcLockClient::new(addr);
|
||||
let requests = vec![
|
||||
LockRequest::new(test_resource(), LockType::Exclusive, "owner-a").with_acquire_timeout(Duration::from_secs(2)),
|
||||
LockRequest::new(
|
||||
ObjectKey {
|
||||
bucket: Arc::from("test-bucket"),
|
||||
object: Arc::from("test-object-2"),
|
||||
version: None,
|
||||
},
|
||||
LockType::Exclusive,
|
||||
"owner-a",
|
||||
)
|
||||
.with_acquire_timeout(Duration::from_secs(2)),
|
||||
];
|
||||
|
||||
let responses = grpc_client
|
||||
.acquire_locks_batch(&requests)
|
||||
.await
|
||||
.expect("batch acquire should succeed");
|
||||
assert_eq!(responses.len(), requests.len());
|
||||
assert!(responses.iter().all(|response| response.success));
|
||||
|
||||
let lock_ids = responses
|
||||
.iter()
|
||||
.map(|response| {
|
||||
response
|
||||
.lock_info
|
||||
.as_ref()
|
||||
.expect("batch response should include lock info")
|
||||
.id
|
||||
.clone()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let released = grpc_client
|
||||
.release_locks_batch(&lock_ids)
|
||||
.await
|
||||
.expect("batch release should succeed");
|
||||
assert_eq!(released, vec![true, true]);
|
||||
|
||||
handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_distributed_lock_4_nodes_grpc_read_write_quorum_split_with_two_failed_nodes() {
|
||||
let manager1 = Arc::new(GlobalLockManager::new());
|
||||
|
||||
Reference in New Issue
Block a user