Merge branch 'main' into fix/list-parts-pagination-infinite-loop

This commit is contained in:
chapman
2026-10-02 02:13:18 +08:00
committed by GitHub
5 changed files with 612 additions and 43 deletions
+13
View File
@@ -985,6 +985,19 @@ impl PreparedGetObjectReader {
reader.body_source = crate::object_api::GetObjectBodySource::HookMissed;
Ok(ECStore::attach_read_lock_guard(reader, self.read_lock_guard))
}
/// Open the prepared source with bounded copy prefetch and request-owned
/// cancellation, including cancellation during reader construction.
pub async fn into_reader_for_copy(self) -> Result<(GetObjectReader, tokio_util::sync::DropGuard)> {
let cancellation = tokio_util::sync::CancellationToken::new();
let guard = cancellation.clone().drop_guard();
let reader = crate::set_disk::with_get_object_read_cancellation(
cancellation,
crate::set_disk::with_get_object_read_policy(crate::set_disk::GetObjectReadPolicy::CopySource, self.into_reader()),
)
.await?;
Ok((reader, guard))
}
}
struct LockGuardedReader {
+145 -27
View File
@@ -1742,22 +1742,24 @@ impl AsyncWrite for HttpWriter {
match this.sender.poll_reserve(cx) {
Poll::Ready(Ok(())) => {
this.sender
.send_item(Some(Bytes::copy_from_slice(buf)))
.send_item(Some(Bytes::copy_from_slice(&buf[..HTTP_WRITER_BUFFER_SIZE])))
.map_err(|e| send_error_to_io(e, "HttpWriter send error"))?;
this.start_request();
return Poll::Ready(Ok(buf.len()));
return Poll::Ready(Ok(HTTP_WRITER_BUFFER_SIZE));
}
Poll::Ready(Err(err)) => return Poll::Ready(Err(poll_send_error_to_io(err, "HttpWriter send error"))),
Poll::Pending => return Poll::Pending,
}
}
this.pending_chunk.extend_from_slice(buf);
if !buf.is_empty() {
// A slot bounds memory only when every queued and pending chunk is bounded.
let written = buf.len().min(HTTP_WRITER_BUFFER_SIZE - this.pending_chunk.len());
this.pending_chunk.extend_from_slice(&buf[..written]);
if written > 0 {
this.start_request();
}
Poll::Ready(Ok(buf.len()))
Poll::Ready(Ok(written))
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
@@ -1770,6 +1772,9 @@ impl AsyncWrite for HttpWriter {
}
fn poll_write_vectored(mut self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>]) -> Poll<io::Result<usize>> {
if bufs.len() == 1 {
return self.poll_write(cx, &bufs[0]);
}
let this = self.as_mut().get_mut();
if let Err(err) = this.take_background_error() {
return Poll::Ready(Err(err));
@@ -1783,33 +1788,21 @@ impl AsyncWrite for HttpWriter {
}
}
let total_len = bufs.iter().map(|buf| buf.len()).sum::<usize>();
if total_len == 0 {
return Poll::Ready(Ok(0));
}
if bufs.len() == 1 && this.pending_chunk.is_empty() && total_len >= HTTP_WRITER_BUFFER_SIZE {
match this.sender.poll_reserve(cx) {
Poll::Ready(Ok(())) => {
this.sender
.send_item(Some(Bytes::copy_from_slice(bufs[0].as_ref())))
.map_err(|e| send_error_to_io(e, "HttpWriter send error"))?;
this.start_request();
return Poll::Ready(Ok(total_len));
}
Poll::Ready(Err(err)) => return Poll::Ready(Err(poll_send_error_to_io(err, "HttpWriter send error"))),
Poll::Pending => return Poll::Pending,
let available = HTTP_WRITER_BUFFER_SIZE - this.pending_chunk.len();
let mut written = 0;
for buf in bufs {
let len = buf.len().min(available - written);
this.pending_chunk.extend_from_slice(&buf[..len]);
written += len;
if written == available {
break;
}
}
for buf in bufs {
this.pending_chunk.extend_from_slice(buf);
}
if total_len > 0 {
if written > 0 {
this.start_request();
}
Poll::Ready(Ok(total_len))
Poll::Ready(Ok(written))
}
fn is_write_vectored(&self) -> bool {
@@ -2916,6 +2909,131 @@ mod tests {
handle.abort();
}
fn paused_http_writer() -> (HttpWriter, tokio::sync::mpsc::Receiver<Option<Bytes>>) {
let (sender, receiver) = tokio::sync::mpsc::channel(HTTP_WRITER_CHANNEL_CAPACITY);
let (_, err_rx) = tokio::sync::oneshot::channel();
let writer = HttpWriter {
url: "http://127.0.0.1/paused".to_string(),
method: Method::PUT,
headers: HeaderMap::new(),
err_rx,
start_tx: None,
sender: PollSender::new(sender),
handle: tokio::spawn(std::future::pending()),
pending_chunk: BytesMut::with_capacity(HTTP_WRITER_BUFFER_SIZE),
finish: false,
track_internode_metrics: false,
internode_operation: None,
};
(writer, receiver)
}
#[tokio::test]
async fn http_writer_bounds_large_writes_under_backpressure() {
let payload = vec![0xa5; 3 * HTTP_WRITER_BUFFER_SIZE + 17];
for vectored in [false, true] {
let (mut writer, mut receiver) = paused_http_writer();
for _ in 0..HTTP_WRITER_CHANNEL_CAPACITY {
let written = std::future::poll_fn(|cx| {
if vectored {
Pin::new(&mut writer).poll_write_vectored(cx, &[IoSlice::new(&payload)])
} else {
Pin::new(&mut writer).poll_write(cx, &payload)
}
})
.await
.expect("a free queue slot should accept one bounded chunk");
assert_eq!(written, HTTP_WRITER_BUFFER_SIZE, "large writes must be partial");
}
{
let write = std::future::poll_fn(|cx| {
if vectored {
Pin::new(&mut writer).poll_write_vectored(cx, &[IoSlice::new(&payload)])
} else {
Pin::new(&mut writer).poll_write(cx, &payload)
}
});
tokio::pin!(write);
assert!(futures::poll!(&mut write).is_pending(), "a stalled receiver must exert backpressure");
let chunk = receiver.recv().await.expect("queued body chunk").expect("body, not EOF");
assert_eq!(chunk.as_ref(), &payload[..HTTP_WRITER_BUFFER_SIZE]);
let written = tokio::time::timeout(Duration::from_secs(1), write)
.await
.expect("draining a queue slot should wake the blocked writer")
.expect("resumed write should succeed");
assert_eq!(written, HTTP_WRITER_BUFFER_SIZE);
}
assert!(writer.pending_chunk.len() <= HTTP_WRITER_BUFFER_SIZE);
for _ in 0..HTTP_WRITER_CHANNEL_CAPACITY {
let chunk = receiver.try_recv().expect("queued body chunk").expect("body, not EOF");
assert_eq!(chunk.as_ref(), &payload[..HTTP_WRITER_BUFFER_SIZE]);
}
assert!(receiver.try_recv().is_err());
}
}
#[tokio::test]
async fn http_writer_bounds_mixed_writes_and_preserves_bytes() {
let header = [0x11; 32];
let payload = vec![0xa5; 3 * HTTP_WRITER_BUFFER_SIZE + 17];
for vectored in [false, true] {
let (mut writer, mut receiver) = paused_http_writer();
writer.write_all(&header).await.expect("small prefix should be buffered");
let mut remaining = payload.as_slice();
while !remaining.is_empty() {
let written = if vectored {
writer
.write_vectored(&[IoSlice::new(&[]), IoSlice::new(remaining), IoSlice::new(&[])])
.await
} else {
writer.write(remaining).await
}
.expect("payload write should succeed");
assert!(written > 0 && written <= HTTP_WRITER_BUFFER_SIZE);
assert!(writer.pending_chunk.len() <= HTTP_WRITER_BUFFER_SIZE, "pending bytes must stay bounded");
remaining = &remaining[written..];
}
writer.flush().await.expect("remaining prefix should flush");
let mut actual = Vec::new();
while let Ok(chunk) = receiver.try_recv() {
let chunk = chunk.expect("body, not EOF");
assert!(chunk.len() <= HTTP_WRITER_BUFFER_SIZE, "queued bytes must stay bounded");
actual.extend_from_slice(&chunk);
}
assert_eq!(&actual[..header.len()], &header);
assert_eq!(&actual[header.len()..], payload.as_slice());
}
}
#[tokio::test]
async fn http_writer_large_vectored_body_preserves_hash_and_payload() {
let state = TestState::default();
let (url, handle) = start_test_server(state.clone()).await.expect("test server should bind");
let header = [0x11; 32];
let payload: Vec<_> = (0..3 * HTTP_WRITER_BUFFER_SIZE + 17).map(|i| (i % 251) as u8).collect();
let mut writer = HttpWriter::new(url, Method::PUT, HeaderMap::new())
.await
.expect("HTTP writer should open");
let mut slices = [
IoSlice::new(&[]),
IoSlice::new(&header),
IoSlice::new(&payload),
IoSlice::new(&[]),
];
let mut remaining = slices.as_mut_slice();
while remaining.iter().any(|slice| !slice.is_empty()) {
let written = writer.write_vectored(remaining).await.expect("body write should succeed");
assert!(written > 0 && written <= HTTP_WRITER_BUFFER_SIZE);
IoSlice::advance_slices(&mut remaining, written);
}
writer.shutdown().await.expect("request should finish");
let bodies = state.put_bodies.lock().await;
assert_eq!(bodies.len(), 1);
assert_eq!(&bodies[0][..header.len()], &header);
assert_eq!(&bodies[0][header.len()..], payload.as_slice());
handle.abort();
}
#[tokio::test]
async fn http_writer_does_not_send_empty_preflight_put() {
let state = TestState::default();
+37 -8
View File
@@ -1769,18 +1769,25 @@ impl DefaultMultipartUsecase {
.into());
}
let (src_reader, _source_cancellation) = store
.get_object_reader_for_copy(&src_bucket, &src_key, rs.clone(), h, &get_opts)
let prepared_source = store
.prepare_get_object_reader(&src_bucket, &src_key, rs.clone(), h, &get_opts)
.await
.map_err(map_get_object_reader_error)?;
let src_info = src_reader.object_info;
let src_info = prepared_source.object_info();
if src_info.delete_marker {
let error = if src_opts.version_id.is_none() {
StorageError::ObjectNotFound(src_bucket.clone(), src_key.clone())
} else {
StorageError::MethodNotAllowed
};
return Err(map_get_object_reader_error(error).into());
}
// Same shape as CopyObject: the part copy reads the source plaintext, and the source
// read resolves its material inside the object layer, which carries no request identity.
authorize_sse_kms_object_read(copy_principal.as_ref(), &src_info.user_defined).await?;
let src_stream = src_reader.stream;
let resolved_src_version_id = src_info.version_id.map(|version_id| {
if version_id == Uuid::nil() {
"null".to_string()
@@ -1811,12 +1818,13 @@ impl DefaultMultipartUsecase {
return Err(s3_error!(PreconditionFailed));
}
let source_logical_size = match src_info.get_actual_size() {
Ok(size) if size >= 0 => size,
Ok(_) | Err(_) if destination_size_limit.is_some() => {
let source_actual_size = src_info.get_actual_size().ok().filter(|size| *size >= 0);
let source_logical_size = match source_actual_size {
Some(size) => size,
None if destination_size_limit.is_some() => {
return Err(S3Error::new(S3ErrorCode::UnexpectedContent));
}
Ok(_) | Err(_) => src_info.size,
None => src_info.size,
};
let (_start_offset, length) = if let Some(ref range_spec) = rs {
@@ -1832,6 +1840,26 @@ impl DefaultMultipartUsecase {
(0, source_logical_size)
};
// Resolve the range under the prepared source's read lock, then try the
// shared budget without waiting while holding namespace/lifecycle locks.
// Keep both the budget and source cancellation alive through the write.
let admission_length = source_actual_size.map_or(-1, |_| length);
let _copy_admission = match self.concurrency_manager().try_admit_multipart_part_copy(admission_length) {
ForegroundWriteAdmission::Disabled => None,
ForegroundWriteAdmission::Admitted(permit) => Some(permit),
ForegroundWriteAdmission::Rejected => {
return Err(S3Error::with_message(
S3ErrorCode::SlowDown,
"foreground write concurrency limit reached, please reduce your request rate",
));
}
};
let (src_reader, _source_cancellation) = prepared_source
.into_reader_for_copy()
.await
.map_err(map_get_object_reader_error)?;
let src_stream = src_reader.stream;
let is_disk_compressed =
rustfs_utils::http::contains_key_str(&mp_info.user_defined, rustfs_utils::http::SUFFIX_COMPRESSION);
@@ -2035,6 +2063,7 @@ mod tests {
use super::*;
mod body_read_tests;
mod copy_admission_tests;
use http::{Extensions, HeaderMap, Method, Uri, header::HeaderValue};
use rustfs_filemeta::ObjectPartInfo;
use rustfs_utils::http::{
@@ -0,0 +1,316 @@
// Copyright 2026 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 super::*;
use crate::app::storage_api::test::contract::bucket::{BucketOperations, MakeBucketOptions};
use crate::app::storage_api::test::contract::object::{ObjectIO, ObjectOperations};
#[test]
#[serial_test::serial]
fn upload_part_copy_shares_foreground_admission_with_upload_part() {
crate::app::gating_test_env::run_large_stack_test("copy-part-admission", || async {
struct RestoreGetMetrics(bool);
impl Drop for RestoreGetMetrics {
fn drop(&mut self) {
rustfs_io_metrics::set_get_stage_metrics_enabled(self.0);
}
}
let store = crate::app::gating_test_env::shared_gating_ecstore().await;
let ambient = crate::app::gating_test_env::shared_gating_ambient().await;
let context = Arc::new(AppContext::new(Arc::clone(&store), ambient.iam(), ambient.kms()));
let bucket = format!("copy-part-admission-{}", Uuid::new_v4().simple());
store
.make_bucket(&bucket, &MakeBucketOptions::default())
.await
.expect("create bucket");
let bytes = vec![0x5a; 2 * 1024 * 1024];
store
.put_object(&bucket, "source", &mut PutObjReader::from_vec(bytes.clone()), &ObjectOptions::default())
.await
.expect("write copy source");
let upload = store
.new_multipart_upload(&bucket, "object", &ObjectOptions::default())
.await
.expect("create session");
let manager = Arc::new(ConcurrencyManager::with_large_put_admission_for_test(true, 1, 1, Duration::ZERO));
let held = manager.admit_multipart_part(1024).await.expect("hold UploadPart permit");
assert!(matches!(held, ForegroundWriteAdmission::Admitted(_)));
let usecase = DefaultMultipartUsecase::with_context_and_concurrency_manager(Some(context), Arc::clone(&manager));
let input = UploadPartCopyInput::builder()
.bucket(bucket.clone())
.key("object".to_owned())
.copy_source(CopySource::Bucket {
bucket: bucket.clone().into(),
key: "source".into(),
version_id: None,
})
.part_number(1)
.upload_id(upload.upload_id.clone())
.build()
.expect("copy request");
let recorder = metrics_util::debugging::DebuggingRecorder::new();
let snapshotter = recorder.snapshotter();
let _recorder = metrics::set_default_local_recorder(&recorder);
let _restore = RestoreGetMetrics(rustfs_io_metrics::get_stage_metrics_enabled());
rustfs_io_metrics::set_get_stage_metrics_enabled(true);
let source_size_bucket = rustfs_io_metrics::get_object_size_bucket(i64::try_from(bytes.len()).expect("source size"));
let source_reader_observed = || {
snapshotter.snapshot().into_vec().iter().any(|(key, _, _, _)| {
key.key().name() == "rustfs_io_get_object_reader_path_by_size_total"
&& key
.key()
.labels()
.any(|label| label.key() == "size_bucket" && label.value() == source_size_bucket)
})
};
let error = Box::pin(usecase.execute_upload_part_copy(build_request(input.clone(), Method::PUT)))
.await
.expect_err("copy must not bypass the occupied UploadPart pool");
assert_eq!(error.code(), &S3ErrorCode::SlowDown);
assert!(!source_reader_observed(), "saturated copy must not construct a source body reader");
let parts = store
.list_object_parts(&bucket, "object", &upload.upload_id, None, 1000, &ObjectOptions::default())
.await
.expect("list parts");
assert!(parts.parts.is_empty(), "rejected copy must not create a destination part");
drop(held);
let copied = Box::pin(usecase.execute_upload_part_copy(build_request(input, Method::PUT)))
.await
.expect("retry copy after releasing UploadPart permit");
assert!(source_reader_observed());
assert_eq!(
manager.put_object_admission_snapshot().active,
Some(0),
"success releases the copy permit"
);
let etag = copied
.output
.copy_part_result
.expect("part result")
.e_tag
.expect("copy ETag")
.value()
.trim_matches('"')
.to_owned();
Arc::clone(&store)
.complete_multipart_upload(
&bucket,
"object",
&upload.upload_id,
vec![CompletePart {
part_num: 1,
etag: Some(etag),
..Default::default()
}],
&ObjectOptions::default(),
)
.await
.expect("commit copied object");
let mut reader = store
.get_object_reader(&bucket, "object", None, HeaderMap::new(), &ObjectOptions::default())
.await
.expect("read committed copy");
let mut actual = Vec::new();
reader.stream.read_to_end(&mut actual).await.expect("read all copied bytes");
assert_eq!(actual, bytes);
drop(reader);
let marker = store
.delete_object(
&bucket,
"source",
ObjectOptions {
versioned: true,
..Default::default()
},
)
.await
.expect("create source delete marker");
assert!(marker.delete_marker);
let marker_version = marker.version_id.expect("marker version").to_string();
let deleted_upload = store
.new_multipart_upload(&bucket, "deleted-copy", &ObjectOptions::default())
.await
.expect("deleted-source destination session");
let _held = manager
.admit_multipart_part(1024)
.await
.expect("saturate gate for deleted source");
for (version_id, expected) in [
(None, S3ErrorCode::NoSuchKey),
(Some(marker_version), S3ErrorCode::MethodNotAllowed),
] {
let input = UploadPartCopyInput::builder()
.bucket(bucket.clone())
.key("deleted-copy".to_owned())
.copy_source(CopySource::Bucket {
bucket: bucket.clone().into(),
key: "source".into(),
version_id: version_id.map(Into::into),
})
.copy_source_if_match(Some("\"wrong-etag\"".parse().expect("source condition")))
.part_number(1)
.upload_id(deleted_upload.upload_id.clone())
.build()
.expect("deleted-source copy request");
let error = Box::pin(usecase.execute_upload_part_copy(build_request(input, Method::PUT)))
.await
.expect_err("deleted source must fail before conditions or admission");
assert_eq!(error.code(), &expected);
}
});
}
#[test]
#[serial_test::serial]
fn upload_part_copy_range_error_and_cancellation_release_foreground_permit() {
crate::app::gating_test_env::run_large_stack_test("copy-part-range-admission", || async {
use rustfs_utils::http::SUFFIX_MAX_TOTAL_OBJECT_SIZE;
let store = crate::app::gating_test_env::shared_gating_ecstore().await;
let ambient = crate::app::gating_test_env::shared_gating_ambient().await;
let context = Arc::new(AppContext::new(Arc::clone(&store), ambient.iam(), ambient.kms()));
let bucket = format!("copy-part-range-{}", Uuid::new_v4().simple());
store
.make_bucket(&bucket, &MakeBucketOptions::default())
.await
.expect("create range bucket");
let bytes = (0..2 * 1024 * 1024)
.map(|i| u8::try_from(i % 251).expect("byte pattern"))
.collect::<Vec<_>>();
store
.put_object(&bucket, "source", &mut PutObjReader::from_vec(bytes.clone()), &ObjectOptions::default())
.await
.expect("write non-inline source");
let manager = Arc::new(ConcurrencyManager::with_large_put_admission_for_test(true, 1, 1, Duration::ZERO));
let usecase = DefaultMultipartUsecase::with_context_and_concurrency_manager(Some(context), Arc::clone(&manager));
let cancelled_upload = store
.new_multipart_upload(&bucket, "cancelled", &ObjectOptions::default())
.await
.expect("cancelled destination session");
let cancelled_input = UploadPartCopyInput::builder()
.bucket(bucket.clone())
.key("cancelled".to_owned())
.copy_source(CopySource::Bucket {
bucket: bucket.clone().into(),
key: "source".into(),
version_id: None,
})
.part_number(1)
.upload_id(cancelled_upload.upload_id.clone())
.build()
.expect("cancelled copy request");
let mut cancelled_copy = Box::pin(usecase.execute_upload_part_copy(build_request(cancelled_input, Method::PUT)));
tokio::time::timeout(Duration::from_secs(10), async {
tokio::select! {
biased;
() = async {
while manager.put_object_admission_snapshot().active != Some(1) {
tokio::task::yield_now().await;
}
} => {}
result = &mut cancelled_copy => panic!("copy completed before cancellation: {result:?}"),
}
})
.await
.expect("copy acquires its permit");
drop(cancelled_copy);
assert_eq!(manager.put_object_admission_snapshot().active, Some(0));
tokio::time::timeout(
Duration::from_secs(10),
store.put_object(&bucket, "source", &mut PutObjReader::from_vec(bytes.clone()), &ObjectOptions::default()),
)
.await
.expect("cancelled source reader releases its namespace lock")
.expect("overwrite source after cancellation");
store
.abort_multipart_upload(&bucket, "cancelled", &cancelled_upload.upload_id, &ObjectOptions::default())
.await
.expect("abort cancelled session");
for capped in [true, false] {
let mut opts = ObjectOptions::default();
if capped {
insert_str(&mut opts.user_defined, SUFFIX_MAX_TOTAL_OBJECT_SIZE, "1024".to_owned());
}
let upload = store
.new_multipart_upload(&bucket, "object", &opts)
.await
.expect("range destination session");
let input = UploadPartCopyInput::builder()
.bucket(bucket.clone())
.key("object".to_owned())
.copy_source(CopySource::Bucket {
bucket: bucket.clone().into(),
key: "source".into(),
version_id: None,
})
.copy_source_range(Some("bytes=12345-1060920".to_owned()))
.part_number(1)
.upload_id(upload.upload_id.clone())
.build()
.expect("range request");
let response = Box::pin(usecase.execute_upload_part_copy(build_request(input, Method::PUT))).await;
assert_eq!(
manager.put_object_admission_snapshot().active,
Some(0),
"copy releases its permit on every outcome"
);
if capped {
assert_eq!(
response.expect_err("range exceeds destination size limit").code(),
&S3ErrorCode::EntityTooLarge
);
let parts = store
.list_object_parts(&bucket, "object", &upload.upload_id, None, 1000, &ObjectOptions::default())
.await
.expect("failed copy parts");
assert!(parts.parts.is_empty());
store
.abort_multipart_upload(&bucket, "object", &upload.upload_id, &ObjectOptions::default())
.await
.expect("abort capped session");
} else {
let copied = response.expect("range copy after failed request released its source and permit");
let etag = copied
.output
.copy_part_result
.expect("range result")
.e_tag
.expect("range ETag")
.value()
.trim_matches('"')
.to_owned();
Arc::clone(&store)
.complete_multipart_upload(
&bucket,
"object",
&upload.upload_id,
vec![CompletePart {
part_num: 1,
etag: Some(etag),
..Default::default()
}],
&ObjectOptions::default(),
)
.await
.expect("commit range copy");
let mut reader = store
.get_object_reader(&bucket, "object", None, HeaderMap::new(), &ObjectOptions::default())
.await
.expect("open range copy");
let mut actual = Vec::new();
reader.stream.read_to_end(&mut actual).await.expect("read full range");
assert_eq!(actual, bytes[12345..=1060920]);
}
}
});
}
+101 -8
View File
@@ -233,11 +233,7 @@ impl ForegroundWriteAdmissionGate {
async fn admit(&self, permits: u32) -> Result<ForegroundWriteAdmission, tokio::sync::AcquireError> {
if self.wait_timeout.is_zero() {
return Ok(match self.semaphore.clone().try_acquire_many_owned(permits) {
Ok(permit) => ForegroundWriteAdmission::Admitted(permit),
Err(tokio::sync::TryAcquireError::NoPermits) => ForegroundWriteAdmission::Rejected,
Err(tokio::sync::TryAcquireError::Closed) => ForegroundWriteAdmission::Rejected,
});
return Ok(self.try_admit(permits));
}
match tokio::time::timeout(self.wait_timeout, self.semaphore.clone().acquire_many_owned(permits)).await {
@@ -245,6 +241,13 @@ impl ForegroundWriteAdmissionGate {
Err(_) => Ok(ForegroundWriteAdmission::Rejected),
}
}
fn try_admit(&self, permits: u32) -> ForegroundWriteAdmission {
match self.semaphore.clone().try_acquire_many_owned(permits) {
Ok(permit) => ForegroundWriteAdmission::Admitted(permit),
Err(_) => ForegroundWriteAdmission::Rejected,
}
}
}
#[derive(Clone)]
@@ -787,6 +790,27 @@ impl ConcurrencyManager {
.await
}
/// Admit a copy part through the UploadPart pool using its resolved logical length.
/// Copy holds source metadata and bucket locks here, so it must never queue
/// behind a writer that may need those locks. Saturation returns SlowDown.
pub fn try_admit_multipart_part_copy(&self, size: i64) -> ForegroundWriteAdmission {
match &self.foreground_write_admission_policy {
ForegroundWriteAdmissionPolicy::Disabled | ForegroundWriteAdmissionPolicy::LegacyCounterOnly => {
ForegroundWriteAdmission::Disabled
}
ForegroundWriteAdmissionPolicy::Strict(gate) => gate.try_admit(1),
ForegroundWriteAdmissionPolicy::Large {
gate,
large_request_permits,
multipart_part_min_size_bytes,
..
} if should_gate_foreground_write(size, *multipart_part_min_size_bytes) => {
gate.try_admit(multipart_admission_permits(size, *large_request_permits))
}
_ => ForegroundWriteAdmission::Disabled,
}
}
// ============================================
// Adaptive I/O Strategy Methods
// ============================================
@@ -1271,9 +1295,9 @@ mod integration_tests {
use super::super::io_schedule::{IoLoadLevel, IoPriority};
use super::super::request_guard::GetObjectGuard;
use super::{
ConcurrencyManager, ForegroundWriteAdmission, ForegroundWriteAdmissionPolicy, SNOWBALL_ARCHIVE_DECODER_LIMIT,
SNOWBALL_MEMBER_COMMIT_LIMIT, SNOWBALL_STAGING_BYTES_LIMIT, derive_large_put_admission_limit,
derive_multipart_admission_max_pending,
ConcurrencyManager, ForegroundWriteAdmission, ForegroundWriteAdmissionGate, ForegroundWriteAdmissionPolicy,
SNOWBALL_ARCHIVE_DECODER_LIMIT, SNOWBALL_MEMBER_COMMIT_LIMIT, SNOWBALL_STAGING_BYTES_LIMIT,
derive_large_put_admission_limit, derive_multipart_admission_max_pending,
};
use crate::storage::storage_api::concurrency_consumer::PutObjectGuard;
use rustfs_concurrency::{AdmissionState, WorkloadAdmissionSnapshotProvider, WorkloadClass};
@@ -1715,6 +1739,75 @@ mod integration_tests {
assert_eq!(manager.put_object_admission_snapshot().queued, Some(0));
}
#[tokio::test]
#[serial]
async fn test_concurrency_manager_multipart_copy_never_queues_with_source_locks() {
for (manager, expected_queued) in [
(
ConcurrencyManager::with_multipart_admission_queue_for_test(1, Duration::from_secs(30), 4),
Some(0),
),
(ConcurrencyManager::with_put_admission_for_test(true, 1, Duration::from_secs(30)), None),
] {
let held = manager
.admit_multipart_part(32 * 1024 * 1024)
.await
.expect("hold write budget");
assert!(matches!(
manager.try_admit_multipart_part_copy(32 * 1024 * 1024),
ForegroundWriteAdmission::Rejected
));
assert_eq!(manager.put_object_admission_snapshot().queued, expected_queued);
drop(held);
let copy = manager.try_admit_multipart_part_copy(32 * 1024 * 1024);
assert!(matches!(copy, ForegroundWriteAdmission::Admitted(_)));
assert!(matches!(manager.try_admit_multipart_part_copy(1), ForegroundWriteAdmission::Rejected));
drop(copy);
assert!(matches!(manager.try_admit_multipart_part_copy(1), ForegroundWriteAdmission::Admitted(_)));
}
}
#[test]
#[serial]
fn test_concurrency_manager_multipart_copy_uses_logical_size_threshold_and_weight() {
let mut manager = ConcurrencyManager::new();
manager.foreground_write_admission_policy = ForegroundWriteAdmissionPolicy::Large {
gate: ForegroundWriteAdmissionGate::new(4, Duration::from_secs(30)),
large_request_permits: 4,
put_object_min_size_bytes: 32 * 1024 * 1024,
multipart_part_min_size_bytes: 8 * 1024 * 1024,
multipart_wait_timeout: Duration::from_secs(30),
multipart_max_pending: 4,
};
assert!(matches!(
manager.try_admit_multipart_part_copy(8 * 1024 * 1024 - 1),
ForegroundWriteAdmission::Disabled
));
let small = manager.try_admit_multipart_part_copy(8 * 1024 * 1024);
assert_eq!(manager.put_object_admission_snapshot().active, Some(1));
assert!(matches!(
manager.try_admit_multipart_part_copy(32 * 1024 * 1024),
ForegroundWriteAdmission::Rejected
));
drop(small);
let full = manager.try_admit_multipart_part_copy(32 * 1024 * 1024);
assert!(matches!(full, ForegroundWriteAdmission::Admitted(_)));
assert_eq!(manager.put_object_admission_snapshot().active, Some(4));
drop(full);
assert_eq!(manager.put_object_admission_snapshot().active, Some(0));
let unknown = manager.try_admit_multipart_part_copy(-1);
assert!(matches!(unknown, ForegroundWriteAdmission::Admitted(_)));
assert_eq!(manager.put_object_admission_snapshot().active, Some(4));
drop(unknown);
for enabled in [false, true] {
let disabled = ConcurrencyManager::with_put_admission_for_test(enabled, 0, Duration::ZERO);
assert!(matches!(
disabled.try_admit_multipart_part_copy(32 * 1024 * 1024),
ForegroundWriteAdmission::Disabled
));
}
}
#[tokio::test(start_paused = true)]
#[serial]
async fn test_concurrency_manager_multipart_part_rejects_after_queue_wait_timeout() {