diff --git a/crates/ecstore/src/store/object.rs b/crates/ecstore/src/store/object.rs index e58b0b2d1..52fdad5ce 100644 --- a/crates/ecstore/src/store/object.rs +++ b/crates/ecstore/src/store/object.rs @@ -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 { diff --git a/crates/rio/src/http_reader.rs b/crates/rio/src/http_reader.rs index 8d454b99e..f5a7e14ef 100644 --- a/crates/rio/src/http_reader.rs +++ b/crates/rio/src/http_reader.rs @@ -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> { @@ -1770,6 +1772,9 @@ impl AsyncWrite for HttpWriter { } fn poll_write_vectored(mut self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>]) -> Poll> { + 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::(); - 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>) { + 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(); diff --git a/rustfs/src/app/multipart_usecase.rs b/rustfs/src/app/multipart_usecase.rs index e92a3be6e..66174374f 100644 --- a/rustfs/src/app/multipart_usecase.rs +++ b/rustfs/src/app/multipart_usecase.rs @@ -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::{ diff --git a/rustfs/src/app/multipart_usecase/tests/copy_admission_tests.rs b/rustfs/src/app/multipart_usecase/tests/copy_admission_tests.rs new file mode 100644 index 000000000..3955fded3 --- /dev/null +++ b/rustfs/src/app/multipart_usecase/tests/copy_admission_tests.rs @@ -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::>(); + 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]); + } + } + }); +} diff --git a/rustfs/src/storage/concurrency/manager.rs b/rustfs/src/storage/concurrency/manager.rs index 183f6d525..70afe0328 100644 --- a/rustfs/src/storage/concurrency/manager.rs +++ b/rustfs/src/storage/concurrency/manager.rs @@ -233,11 +233,7 @@ impl ForegroundWriteAdmissionGate { async fn admit(&self, permits: u32) -> Result { 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() {