diff --git a/crates/ecstore/src/api/mod.rs b/crates/ecstore/src/api/mod.rs index 547448a25..617df5bd9 100644 --- a/crates/ecstore/src/api/mod.rs +++ b/crates/ecstore/src/api/mod.rs @@ -445,7 +445,7 @@ pub mod disk { }; pub use crate::disk::{ BackgroundMemoryBudget, PhysicalDiskBinding, PhysicalDiskPlan, PhysicalDiskSample, PhysicalDiskSampler, - PhysicalResourcePlan, RemoteBackgroundIo, with_remote_background_io, + PhysicalResourcePlan, RemoteBackgroundIo, with_background_resource_unit, with_remote_background_io, }; pub use bytes::Bytes; pub use endpoint::Endpoint; diff --git a/crates/ecstore/src/cluster/rpc/internode_data_transport.rs b/crates/ecstore/src/cluster/rpc/internode_data_transport.rs index eabd230bc..859286d9c 100644 --- a/crates/ecstore/src/cluster/rpc/internode_data_transport.rs +++ b/crates/ecstore/src/cluster/rpc/internode_data_transport.rs @@ -33,7 +33,9 @@ use rustfs_config::{ DEFAULT_INTERNODE_DATA_TRANSPORT, ENV_RUSTFS_INTERNODE_DATA_TRANSPORT, INTERNODE_DATA_TRANSPORT_TCP, KNOWN_INTERNODE_DATA_TRANSPORT_BACKENDS, }; -use rustfs_rio::{ChunkReaderBox, HttpChunkReader, HttpReader, HttpWriter}; +use rustfs_rio::{ + ChunkReaderBox, HttpChunkReader, HttpReader, HttpWriter, HttpWriterMemoryReservation, HttpWriterMemoryReserver, +}; use sha2::{Digest, Sha256}; use std::collections::HashMap; use std::future::Future; @@ -66,6 +68,19 @@ const PUT_FILE_CAPABILITY_PROBE_TIMEOUT: Duration = Duration::from_secs(5); const CONTENT_TYPE_JSON: &str = "application/json"; const CONTENT_TYPE_MSGPACK: &str = "application/msgpack"; +fn http_writer_memory_reserver() -> Option { + crate::disk::current_resource_unit() + .and_then(|resources| resources.memory_budget()) + .map(|budget| { + // Payloads can outlive the operation; never retain its node or disk units. + Arc::new(move |capacity| { + budget + .reserve(capacity) + .map(|reservation| Box::new(reservation) as HttpWriterMemoryReservation) + }) as HttpWriterMemoryReserver + }) +} + fn unsupported_transport_message(transport: &str) -> String { format!( "invalid {ENV_RUSTFS_INTERNODE_DATA_TRANSPORT}={transport:?}; supported values: {}", @@ -346,13 +361,17 @@ impl InternodeDataTransport for TcpHttpInternodeDataTransport { } async fn open_write(&self, request: WriteStreamRequest) -> Result { + let memory_reserver = http_writer_memory_reserver(); let server_epoch = self.put_file_auth_capability(&request.endpoint).await?; let auth_scope = server_epoch.map(|server_epoch| (Uuid::new_v4(), server_epoch)); let url = build_put_file_stream_url(&request, auth_scope); let endpoint = request.endpoint; let mut headers = json_headers(); build_auth_headers(&url, &Method::PUT, &mut headers)?; - let writer = HttpWriter::new(url.clone(), Method::PUT, headers).await?; + let writer = match memory_reserver { + Some(reserve) => HttpWriter::new_with_memory_reserver(url.clone(), Method::PUT, headers, reserve).await?, + None => HttpWriter::new(url.clone(), Method::PUT, headers).await?, + }; match auth_scope { Some((nonce, server_epoch)) => Ok(Box::new(PutFileAuthWriter::new(writer, url, nonce, endpoint, server_epoch))), None => Ok(Box::new(writer)), @@ -948,6 +967,104 @@ mod tests { use tokio::io::AsyncWriteExt; use tokio::sync::{Barrier, Notify}; + fn memory_test_write_request(authenticated: bool) -> WriteStreamRequest { + crate::runtime::sources::ensure_test_rpc_secret(); + let endpoint = format!("http://writer-memory-{}.invalid", Uuid::new_v4()); + put_file_capability_cache_entry(&endpoint).write().cached = Some(if authenticated { + PutFileCapabilityState::V1 { + server_epoch: Uuid::new_v4(), + revalidate_after: Instant::now() + PUT_FILE_V1_CAPABILITY_TTL, + } + } else { + PutFileCapabilityState::LegacyUntil(Instant::now() + PUT_FILE_LEGACY_CAPABILITY_TTL) + }); + WriteStreamRequest { + endpoint, + disk: "disk-a".to_owned(), + volume: "bucket".to_owned(), + path: "part.1".to_owned(), + append: false, + size: 4, + } + } + + #[tokio::test] + async fn http_write_memory_reserver_outlives_only_the_memory_scope() { + use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit}; + + assert!(http_writer_memory_reserver().is_none(), "foreground must use the existing writer"); + let memory = BackgroundMemoryBudget::new(); + memory.set_sampled_limit(2 * 1024 * 1024, Duration::from_secs(30)); + let (resources, gates) = test_resource_units(2); + let request = memory_test_write_request(true); + let endpoint = request.endpoint.clone(); + let writer = with_optional_resource_unit( + Some(resources.with_memory_budget(memory.clone())), + TcpHttpInternodeDataTransport.open_write(request), + ) + .await; + PUT_FILE_CAPABILITY_CACHE.write().remove(&endpoint); + let mut writer = writer.expect("admitted transport writer"); + assert!( + gates.iter().all(|gate| gate.snapshot().active == 0), + "the callback must not retain disk units" + ); + assert_eq!(memory.snapshot().used, 1024 * 1024); + writer.write_all(b"data").await.expect("buffered bytes"); + writer.flush().await.expect("transfer charged allocation to queue"); + memory.set_sampled_limit(2 * 1024 * 1024, Duration::ZERO); + let error = writer + .write(b"rejected") + .await + .expect_err("expired observation must reject new storage"); + assert_eq!(error.kind(), io::ErrorKind::WouldBlock); + assert!( + crate::disk::background_memory::is_memory_pressure(&error), + "preserve the typed memory-pressure marker" + ); + assert_eq!(memory.snapshot().used, 1024 * 1024, "queued storage remains charged after expiry"); + drop(writer); + tokio::task::yield_now().await; + assert_eq!(memory.snapshot().used, 0); + } + + #[tokio::test] + async fn http_write_memory_denial_preserves_the_budget_error() { + use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit}; + + for (authenticated, sample) in [false, true].into_iter().flat_map(|authenticated| { + [ + None, + Some((1024 * 1024 - 1, Duration::from_secs(30))), + Some((1024 * 1024, Duration::ZERO)), + ] + .into_iter() + .map(move |sample| (authenticated, sample)) + }) { + let memory = BackgroundMemoryBudget::new(); + if let Some((limit, valid_for)) = sample { + memory.set_sampled_limit(limit, valid_for); + } + let (resources, _) = test_resource_units(1); + let request = memory_test_write_request(authenticated); + let endpoint = request.endpoint.clone(); + let result = with_optional_resource_unit( + Some(resources.with_memory_budget(memory.clone())), + TcpHttpInternodeDataTransport.open_write(request), + ) + .await; + PUT_FILE_CAPABILITY_CACHE.write().remove(&endpoint); + let error = match result { + Err(Error::Io(error)) => error, + Err(error) => panic!("unexpected transport error: {error}"), + Ok(_) => panic!("unavailable memory must reject construction"), + }; + assert!(crate::disk::background_memory::is_memory_pressure(&error)); + assert_eq!(error.kind(), io::ErrorKind::WouldBlock); + assert_eq!(memory.snapshot().used, 0); + } + } + #[tokio::test] async fn admitted_remote_reads_bind_receiver_admission_to_the_http_signature() { use crate::disk::with_background_resource_unit; diff --git a/crates/ecstore/src/disk/io_admission/memory.rs b/crates/ecstore/src/disk/io_admission/memory.rs index eaeecb619..1e3a5f280 100644 --- a/crates/ecstore/src/disk/io_admission/memory.rs +++ b/crates/ecstore/src/disk/io_admission/memory.rs @@ -179,7 +179,9 @@ impl BackgroundBuffer { self.grow(length, 0) } - fn grow(&mut self, length: usize, preserve: usize) -> io::Result<()> { + /// Preserve an initialized prefix when replacing storage. Callers keep + /// `preserve` within both the existing allocation and the requested length. + pub(crate) fn grow(&mut self, length: usize, preserve: usize) -> io::Result<()> { if self.bytes.len() >= length { return Ok(()); } @@ -288,6 +290,37 @@ mod tests { assert_eq!(owner.snapshot().used, 0, "empty results free unused storage before refund"); } + #[test] + fn preserving_growth_keeps_existing_payload_on_denial_and_reuse() { + let owner = BackgroundMemoryBudget::new(); + owner.set_sampled_limit(11, Duration::from_secs(30)); + let mut buffer = BackgroundBuffer::new(Some(owner.clone())); + buffer.ensure_len(4).expect("allocate initial payload"); + buffer.as_mut_slice().copy_from_slice(b"data"); + + let error = buffer + .grow(8, 4) + .expect_err("old and replacement allocations need twelve bytes"); + assert!(is_memory_pressure(&error)); + assert_eq!(buffer.as_slice(), b"data"); + assert_eq!(owner.snapshot().used, 4); + + owner.set_sampled_limit(12, Duration::from_secs(30)); + buffer.grow(8, 4).expect("preserve the payload after admission"); + assert_eq!(buffer.as_slice(), b"data\0\0\0\0"); + assert_eq!(owner.snapshot().used, 8); + owner.set_sampled_limit(0, Duration::ZERO); + buffer.grow(8, 4).expect("reuse retained storage after expiry"); + assert_eq!(buffer.as_slice(), b"data\0\0\0\0"); + + let bytes = buffer.take_bytes(4); + drop(buffer); + assert_eq!(bytes.as_ref(), b"data"); + assert_eq!(owner.snapshot().used, 8); + drop(bytes); + assert_eq!(owner.snapshot().used, 0); + } + #[test] fn whole_file_growth_denial_and_source_failure_refund_storage() { let owner = BackgroundMemoryBudget::new(); diff --git a/crates/ecstore/src/erasure/coding/bitrot.rs b/crates/ecstore/src/erasure/coding/bitrot.rs index 0bddca374..3c5d270e8 100644 --- a/crates/ecstore/src/erasure/coding/bitrot.rs +++ b/crates/ecstore/src/erasure/coding/bitrot.rs @@ -744,15 +744,74 @@ pub async fn bitrot_verify( /// Custom writer enum that supports inline buffer storage pub enum CustomWriter { /// Inline buffer writer - stores data in memory - InlineBuffer(Vec), + InlineBuffer(InlineBuffer), /// Disk-based writer using tokio file Other(Box), } +/// Inline payload storage retains its background allocation through export. +pub struct InlineBuffer { + foreground: Vec, + background: Option, + filled: usize, +} + +impl InlineBuffer { + fn new() -> Self { + let background = crate::disk::current_resource_unit() + .and_then(|unit| unit.memory_budget()) + .map(|budget| crate::disk::background_memory::BackgroundBuffer::new(Some(budget))); + Self { + foreground: Vec::new(), + background, + filled: 0, + } + } + + fn as_slice(&self) -> &[u8] { + match &self.background { + Some(buffer) => &buffer.as_slice()[..self.filled], + None => &self.foreground, + } + } + + fn append(&mut self, bufs: &[IoSlice<'_>]) -> std::io::Result { + let total = bufs + .iter() + .try_fold(0usize, |length, buf| length.checked_add(buf.len()).ok_or(std::io::ErrorKind::OutOfMemory))?; + if let Some(buffer) = &mut self.background { + let length = self.filled.checked_add(total).ok_or(std::io::ErrorKind::OutOfMemory)?; + if length > buffer.as_slice().len() { + let capacity = buffer.as_slice().len().checked_mul(2).unwrap_or(length).max(length); + // Admission includes both allocations until the prefix has been copied. + buffer.grow(capacity, self.filled)?; + } + let mut offset = self.filled; + for buf in bufs { + buffer.as_mut_slice()[offset..offset + buf.len()].copy_from_slice(buf); + offset += buf.len(); + } + self.filled = length; + } else { + for buf in bufs { + self.foreground.extend_from_slice(buf); + } + } + Ok(total) + } + + fn into_bytes(self) -> bytes::Bytes { + match self.background { + Some(mut buffer) => buffer.take_bytes(self.filled), + None => bytes::Bytes::from(self.foreground), + } + } +} + impl CustomWriter { /// Create a new inline buffer writer pub fn new_inline_buffer() -> Self { - Self::InlineBuffer(Vec::new()) + Self::InlineBuffer(InlineBuffer::new()) } /// Create a new disk writer from any AsyncWrite implementation @@ -766,15 +825,15 @@ impl CustomWriter { /// Get the inline buffer data if this is an inline buffer writer pub fn get_inline_data(&self) -> Option<&[u8]> { match self { - Self::InlineBuffer(data) => Some(data), + Self::InlineBuffer(data) => Some(data.as_slice()), Self::Other(_) => None, } } /// Extract the inline buffer data, consuming the writer - pub fn into_inline_data(self) -> Option> { + pub fn into_inline_data(self) -> Option { match self { - Self::InlineBuffer(data) => Some(data), + Self::InlineBuffer(data) => Some(data.into_bytes()), Self::Other(_) => None, } } @@ -787,10 +846,7 @@ impl AsyncWrite for CustomWriter { buf: &[u8], ) -> std::task::Poll> { match self.get_mut() { - Self::InlineBuffer(data) => { - data.extend_from_slice(buf); - std::task::Poll::Ready(Ok(buf.len())) - } + Self::InlineBuffer(data) => std::task::Poll::Ready(data.append(&[IoSlice::new(buf)])), Self::Other(writer) => { let pinned_writer = std::pin::Pin::new(writer.as_mut()); pinned_writer.poll_write(cx, buf) @@ -824,13 +880,7 @@ impl AsyncWrite for CustomWriter { bufs: &[IoSlice<'_>], ) -> std::task::Poll> { match self.get_mut() { - Self::InlineBuffer(data) => { - let total = bufs.iter().map(|buf| buf.len()).sum::(); - for buf in bufs { - data.extend_from_slice(buf); - } - std::task::Poll::Ready(Ok(total)) - } + Self::InlineBuffer(data) => std::task::Poll::Ready(data.append(bufs)), Self::Other(writer) => { let pinned_writer = std::pin::Pin::new(writer.as_mut()); pinned_writer.poll_write_vectored(cx, bufs) @@ -905,7 +955,7 @@ impl BitrotWriterWrapper { } /// Extract the inline buffer data, consuming the wrapper - pub fn into_inline_data(self) -> Option> { + pub fn into_inline_data(self) -> Option { match self.writer_type { WriterType::InlineBuffer => { let writer = self.bitrot_writer.into_inner(); @@ -2627,4 +2677,128 @@ mod tests { assert_eq!(output, vec![9u8]); } } + + mod inline_memory_regressions { + use super::*; + use crate::disk::background_memory::is_memory_pressure; + use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit}; + use std::sync::Arc; + use std::time::Duration; + use tokio::io::AsyncWriteExt; + + async fn writer_for(memory: Arc) -> CustomWriter { + let (resources, _) = test_resource_units(1); + with_optional_resource_unit(Some(resources.with_memory_budget(memory)), async { CustomWriter::new_inline_buffer() }) + .await + } + + #[tokio::test] + async fn inline_writer_rejects_closed_and_expired_owners_before_first_allocation() { + for expired in [false, true] { + let memory = BackgroundMemoryBudget::new(); + if expired { + memory.set_sampled_limit(128, Duration::ZERO); + } + let mut writer = writer_for(memory.clone()).await; + let error = writer.write_all(&[7; 64]).await.expect_err("no live grant"); + assert_eq!(error.kind(), io::ErrorKind::WouldBlock); + assert!(is_memory_pressure(&error)); + assert_eq!(writer.get_inline_data(), Some(&[][..])); + assert_eq!(memory.snapshot().used, 0); + } + } + + #[tokio::test] + async fn inline_writer_growth_charges_old_plus_new_and_is_atomic_for_vectored_writes() { + let memory = BackgroundMemoryBudget::new(); + memory.set_sampled_limit(128, Duration::from_secs(30)); + let mut writer = writer_for(memory.clone()).await; + writer.write_all(&[1; 64]).await.expect("first allocation"); + assert_eq!(memory.snapshot().used, 64); + let second = [2; 32]; + let third = [3; 32]; + let chunks = [IoSlice::new(&second), IoSlice::new(&third)]; + let error = writer.write_vectored(&chunks).await.expect_err("64 old + 128 new exceed 128"); + assert!(is_memory_pressure(&error)); + assert_eq!(writer.get_inline_data(), Some(&[1; 64][..])); + assert_eq!(memory.snapshot().used, 64); + memory.set_sampled_limit(192, Duration::from_secs(30)); + assert_eq!(writer.write_vectored(&chunks).await.expect("old and new both fit"), 64); + assert_eq!(memory.snapshot().used, 128); + let data = writer.get_inline_data().expect("inline payload"); + assert_eq!(&data[..64], &[1; 64]); + assert_eq!(&data[64..96], &second); + assert_eq!(&data[96..], &third); + drop(writer); + assert_eq!(memory.snapshot().used, 0); + } + + #[tokio::test] + async fn inline_writer_expiry_preserves_charged_storage_but_blocks_growth() { + let memory = BackgroundMemoryBudget::new(); + memory.set_sampled_limit(192, Duration::from_secs(30)); + let mut writer = writer_for(memory.clone()).await; + writer.write_all(&[5; 64]).await.expect("existing storage"); + memory.set_sampled_limit(192, Duration::ZERO); + let error = writer + .write_all(&[9]) + .await + .expect_err("expired sample cannot authorize replacement"); + assert!(is_memory_pressure(&error)); + assert_eq!(writer.get_inline_data(), Some(&[5; 64][..])); + assert_eq!(memory.snapshot().used, 64); + writer.write_all(&[]).await.expect("empty write needs no allocation"); + writer.flush().await.expect("flush does not revoke storage"); + writer.shutdown().await.expect("shutdown does not revoke storage"); + let bytes = writer.into_inline_data().expect("export after expiry"); + assert_eq!(bytes.as_ref(), &[5; 64]); + assert_eq!(memory.snapshot().used, 64); + drop(bytes); + assert_eq!(memory.snapshot().used, 0); + } + + #[tokio::test] + async fn inline_writer_export_keeps_receipt_through_clones_and_the_last_slice() { + let memory = BackgroundMemoryBudget::new(); + memory.set_sampled_limit(64, Duration::from_secs(30)); + let mut writer = writer_for(memory.clone()).await; + writer.write_all(&[11; 64]).await.expect("allocate payload"); + let bytes = writer.into_inline_data().expect("export bytes"); + let clone = bytes.clone(); + let slice = bytes.slice(3..7); + memory.set_sampled_limit(0, Duration::ZERO); + drop(bytes); + drop(clone); + assert_eq!(slice.as_ref(), &[11; 4]); + assert_eq!(memory.snapshot().used, 64); + drop(slice); + assert_eq!(memory.snapshot().used, 0); + } + + #[tokio::test] + async fn inline_writer_keeps_the_constructing_owner_after_leaving_its_scope() { + let first = BackgroundMemoryBudget::new(); + first.set_sampled_limit(64, Duration::from_secs(30)); + let mut writer = writer_for(first.clone()).await; + let second = BackgroundMemoryBudget::new(); + let (resources, _) = test_resource_units(1); + with_optional_resource_unit(Some(resources.with_memory_budget(second.clone())), async { + writer.write_all(&[13; 64]).await.expect("writer still owns the first grant"); + }) + .await; + assert_eq!(first.snapshot().used, 64); + assert_eq!(second.snapshot().used, 0); + drop(writer); + assert_eq!(first.snapshot().used, 0); + } + + #[tokio::test] + async fn foreground_inline_writer_preserves_payload_and_empty_writes() { + let mut writer = CustomWriter::new_inline_buffer(); + writer.write_all(&[]).await.expect("empty foreground write"); + writer.write_all(b"inline body").await.expect("foreground write"); + let bytes = writer.into_inline_data().expect("foreground export"); + assert_eq!(bytes.as_ref(), b"inline body"); + } + } } diff --git a/crates/ecstore/src/io_support/bitrot.rs b/crates/ecstore/src/io_support/bitrot.rs index 081dc963b..0920f863c 100644 --- a/crates/ecstore/src/io_support/bitrot.rs +++ b/crates/ecstore/src/io_support/bitrot.rs @@ -1536,7 +1536,7 @@ mod tests { writer.write(chunk).await.expect("write chunk"); } - let inline_data = Bytes::from(writer.into_inline_data().expect("inline buffer")); + let inline_data = writer.into_inline_data().expect("inline buffer"); let mut reader = create_bitrot_reader_from_bytes( Some(inline_data), None, @@ -1584,7 +1584,7 @@ mod tests { let inline_data = writer.into_inline_data().expect("inline buffer"); let mut reader = create_deferred_bitrot_reader( - Some(inline_data.into()), + Some(inline_data), None, "test-bucket", "test-path", @@ -1626,7 +1626,7 @@ mod tests { for stripe in 0..payload.len() / shard_size { let (mut reader, handle) = create_deferred_bitrot_reader_with_stripe_handle( - Some(inline_data.clone().into()), + Some(inline_data.clone()), None, "test-bucket", "test-path", @@ -1671,7 +1671,7 @@ mod tests { for chunk in payload.chunks(shard_size) { writer.write(chunk).await.expect("write chunk"); } - let mut inline_data = writer.into_inline_data().expect("inline buffer"); + let mut inline_data = writer.into_inline_data().expect("inline buffer").to_vec(); // Corrupt one payload byte of block 2 (blocks are hash + data). let block_len = algo.size() + shard_size; @@ -1845,4 +1845,122 @@ mod tests { assert_eq!(output, data); } + + mod inline_memory_regressions { + use super::*; + use crate::disk::background_memory::Reservation; + use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit}; + use std::time::Duration; + + #[tokio::test] + async fn owned_inline_reader_retains_existing_receipt_without_a_new_grant() { + let memory = BackgroundMemoryBudget::new(); + memory.set_sampled_limit(64, Duration::from_secs(30)); + let receipt = memory.reserve(64).expect("existing metadata payload grant"); + let bytes = Reservation::wrap(Bytes::from(vec![7; 64]), Some(receipt)); + memory.set_sampled_limit(0, Duration::ZERO); + let (resources, _) = test_resource_units(1); + let mut reader = with_optional_resource_unit( + Some(resources.with_memory_budget(memory.clone())), + create_bitrot_reader_from_bytes( + Some(bytes.clone()), + None, + "bucket", + "inline", + 0, + 64, + 64, + HashAlgorithm::None, + false, + false, + ), + ) + .await + .expect("owned data requires no allocation") + .expect("inline reader"); + drop(bytes); + assert_eq!(memory.snapshot().used, 64, "the source owner must remain attached to the reader"); + let mut output = [0; 64]; + assert_eq!(reader.read(&mut output).await.expect("read existing payload"), 64); + assert_eq!(output, [7; 64]); + assert_eq!(memory.snapshot().used, 64); + drop(reader); + assert_eq!(memory.snapshot().used, 0); + } + + #[tokio::test] + async fn accounted_inline_writer_round_trips_bitrot_and_preserves_export_lifetime() { + let memory = BackgroundMemoryBudget::new(); + memory.set_sampled_limit(192, Duration::from_secs(30)); + let (resources, _) = test_resource_units(1); + let bytes = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async { + let mut writer = create_bitrot_writer(true, None, "bucket", "inline", 64, 64, HashAlgorithm::HighwayHash256S) + .await + .expect("inline writer"); + assert_eq!(writer.write(&[19; 64]).await.expect("complete shard"), 64); + writer.shutdown().await.expect("shutdown"); + writer.into_inline_data().expect("bitrot bytes") + }) + .await; + assert_eq!(bytes.len(), 96); + assert_eq!(memory.snapshot().used, 96, "exported [hash][shard] remains charged"); + let (resources, _) = test_resource_units(1); + let mut reader = with_optional_resource_unit( + Some(resources.with_memory_budget(memory.clone())), + create_bitrot_reader_from_bytes( + Some(bytes.clone()), + None, + "bucket", + "inline", + 0, + 64, + 64, + HashAlgorithm::HighwayHash256S, + false, + false, + ), + ) + .await + .expect("reader") + .expect("inline reader"); + let mut output = [0; 64]; + assert_eq!(reader.read(&mut output).await.expect("bitrot remains verified"), 64); + assert_eq!(output, [19; 64]); + assert_eq!(memory.snapshot().used, 192, "input owner and verification scratch are both live"); + drop(reader); + assert_eq!(memory.snapshot().used, 96); + let slice = bytes.slice(32..36); + drop(bytes); + assert_eq!(memory.snapshot().used, 96); + drop(slice); + assert_eq!(memory.snapshot().used, 0); + } + + #[tokio::test] + async fn inline_heal_writer_memory_denial_is_retryable_after_read_and_verify_fit() { + use crate::erasure::coding::Erasure; + let erasure = Erasure::new(2, 1, 16); + let encoded = erasure.encode_data(&[37; 16]).expect("foreground test fixture"); + let readers = encoded + .into_iter() + .map(|shard| Some(BitrotReader::new(Cursor::new(shard), 8, HashAlgorithm::None, false))) + .collect(); + let memory = BackgroundMemoryBudget::new(); + // Read shards: 24 bytes; source parity: 8; verify scratch: 8. + // After verification, input shards still use 24. The target needs a + // new 40-byte [hash][shard] allocation and must fail before copying. + memory.set_sampled_limit(40, Duration::from_secs(30)); + let (resources, _) = test_resource_units(1); + let result = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async { + let writer = create_bitrot_writer(true, None, "bucket", "inline", 8, 8, HashAlgorithm::HighwayHash256S) + .await + .expect("writer setup does not allocate payload"); + let mut writers = vec![Some(writer), None, None]; + erasure.heal(&mut writers, readers, 16, &[]).await + }) + .await; + assert!(matches!(result, Err(DiskError::Io(error)) if error.kind() == io::ErrorKind::WouldBlock)); + assert_eq!(memory.snapshot().used, 0); + } + } } diff --git a/crates/ecstore/src/set_disk/core/io_primitives.rs b/crates/ecstore/src/set_disk/core/io_primitives.rs index c16cee586..8614f4b37 100644 --- a/crates/ecstore/src/set_disk/core/io_primitives.rs +++ b/crates/ecstore/src/set_disk/core/io_primitives.rs @@ -8716,11 +8716,11 @@ mod tests { writer.shutdown().await.expect("inline writer should shutdown"); let mut fi = base.clone(); fi.erasure.index = index + 1; - fi.data = Some(Bytes::from( + fi.data = Some( writer .into_inline_data() .expect("inline bitrot writer should retain encoded data"), - )); + ); files.push(fi); } files diff --git a/crates/ecstore/src/set_disk/mod.rs b/crates/ecstore/src/set_disk/mod.rs index 4cd848226..9f78c1236 100644 --- a/crates/ecstore/src/set_disk/mod.rs +++ b/crates/ecstore/src/set_disk/mod.rs @@ -14519,7 +14519,7 @@ mod tests { file.erasure.index = files.len() + 1; file.erasure.block_size = erasure.block_size; file.uses_legacy_checksum = uses_legacy; - file.data = Some(Bytes::from(data)); + file.data = Some(data); files.push(file); } @@ -14980,7 +14980,7 @@ mod tests { } let encoded = bitrot_writer.into_inline_data().expect("bitrot encoded data should exist"); - disk.write_all(bucket, &shard_path, Bytes::from(encoded)) + disk.write_all(bucket, &shard_path, encoded) .await .expect("encoded shard should be stored"); @@ -15064,7 +15064,7 @@ mod tests { } let encoded = bitrot_writer.into_inline_data().expect("bitrot encoded data should exist"); - disk.write_all(bucket, &shard_path, Bytes::from(encoded)) + disk.write_all(bucket, &shard_path, encoded) .await .expect("encoded shard should be stored"); @@ -15177,7 +15177,7 @@ mod tests { } let encoded = bitrot_writer.into_inline_data().expect("bitrot encoded data should exist"); - disk.write_all(bucket, &shard_path, Bytes::from(encoded)) + disk.write_all(bucket, &shard_path, encoded) .await .expect("encoded shard should be stored"); } diff --git a/crates/ecstore/src/set_disk/ops/heal.rs b/crates/ecstore/src/set_disk/ops/heal.rs index a0c0712f1..55fcdcbf0 100644 --- a/crates/ecstore/src/set_disk/ops/heal.rs +++ b/crates/ecstore/src/set_disk/ops/heal.rs @@ -13,14 +13,14 @@ // limitations under the License. use super::super::{ - Bytes, CHECK_PART_FILE_CORRUPT, CHECK_PART_SUCCESS, DeleteOptions, DiskError, DiskStore, DriveState, EVENT_SET_DISK_HEAL, - Error, FileInfo, HashAlgorithm, HashMap, HealDriveInfo, HealItemType, HealOpts, HealResultItem, HealScanMode, Infos, + CHECK_PART_FILE_CORRUPT, CHECK_PART_SUCCESS, DeleteOptions, DiskError, DiskStore, DriveState, EVENT_SET_DISK_HEAL, Error, + FileInfo, HashAlgorithm, HashMap, HealDriveInfo, HealItemType, HealOpts, HealResultItem, HealScanMode, Infos, LOG_SUBSYSTEM_SET_DISK, ObjectInfo, ObjectOptions, ObjectPartInfo, Path, RUSTFS_META_TMP_BUCKET, ReadOptions, Result, - SLASH_SEPARATOR, SetDisks, StorageError, Uuid, coding, count_errs, count_part_not_success, create_bitrot_reader, - create_bitrot_writer, debug, disk, disks_with_all_parts, encode_dir_object, error, file_info_is_valid_for_metadata, - formats_match_reference_slots, get_format_erasure_in_quorum, get_lock_acquire_timeout, has_suffix, - heal_bucket_local_on_disks, is_object_dir_dangling, join_all, load_format_erasure_all, path_join_buf, save_format_file, - should_heal_object_on_disk, stat_all_dirs, to_object_err, warn, + SLASH_SEPARATOR, SetDisks, StorageError, Uuid, coding, count_errs, count_part_not_success, create_bitrot_writer, debug, disk, + disks_with_all_parts, encode_dir_object, error, file_info_is_valid_for_metadata, formats_match_reference_slots, + get_format_erasure_in_quorum, get_lock_acquire_timeout, has_suffix, heal_bucket_local_on_disks, is_object_dir_dangling, + join_all, load_format_erasure_all, path_join_buf, save_format_file, should_heal_object_on_disk, stat_all_dirs, to_object_err, + warn, }; use crate::bucket::retirement::MarkerRetirementContext; use crate::diagnostics::object_lock::ObjectLockAttempt; @@ -44,7 +44,7 @@ struct ExplicitVersionHeal<'a> { use crate::disk::DataDirDeleteStatus; use crate::disk::DiskAPI; use crate::disk::local::{DELETE_DATA_DIR_MARKER_PREFIX, metadata_less_part_file}; -use crate::io_support::bitrot::object_mmap_read_enabled; +use crate::io_support::bitrot::{create_bitrot_reader_from_bytes, object_mmap_read_enabled}; use crate::storage_api_contracts::namespace::NamespaceLocking as _; use rustfs_common::mrf_channel::MrfDeleteMarkerPurge; use rustfs_common::trace_bus::{TraceEvent, TraceFunc, TraceKind, trace_emit}; @@ -1793,8 +1793,8 @@ impl SetDisks { checksum_info.algorithm }; - match create_bitrot_reader( - metadata.data.as_deref(), + match create_bitrot_reader_from_bytes( + metadata.data.clone(), Some(disk), bucket, &path_join_buf(&[object, &src_data_dir, &format!("part.{}", part.number)]), @@ -1891,8 +1891,7 @@ impl SetDisks { // if let Some(w) = writer.as_any().downcast_ref::() { // parts_metadata[index].data = Some(w.inline_data().to_vec()); // } - parts_metadata[index].data = - Some(writer.into_inline_data().map(Bytes::from).unwrap_or_default()); + parts_metadata[index].data = Some(writer.into_inline_data().unwrap_or_default()); } parts_metadata[index].set_inline_data(); } else { diff --git a/crates/ecstore/src/set_disk/ops/object.rs b/crates/ecstore/src/set_disk/ops/object.rs index 0fe4e4643..48a5a636a 100644 --- a/crates/ecstore/src/set_disk/ops/object.rs +++ b/crates/ecstore/src/set_disk/ops/object.rs @@ -4713,7 +4713,7 @@ impl SetDisks { .ok_or_else(|| Error::other(format!("inline encoder omitted disk shard {i}")))?, ); } else if let Some(writer) = writers[i].take() { - pfi.data = Some(writer.into_inline_data().map(Bytes::from).unwrap_or_default()); + pfi.data = Some(writer.into_inline_data().unwrap_or_default()); } pfi.set_inline_data(); diff --git a/crates/rio/src/http_reader.rs b/crates/rio/src/http_reader.rs index 8534603e0..0879320b8 100644 --- a/crates/rio/src/http_reader.rs +++ b/crates/rio/src/http_reader.rs @@ -30,8 +30,8 @@ use std::io::{self, Error}; use std::net::IpAddr; use std::ops::Not as _; use std::pin::Pin; -use std::sync::LazyLock; use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, LazyLock}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; @@ -1409,6 +1409,24 @@ impl Stream for ReceiverStream { } } +/// Owns a caller's memory reservation until the associated allocation is freed. +pub type HttpWriterMemoryReservation = Box; + +/// Reserves a full buffer capacity before the writer allocates or copies data. +pub type HttpWriterMemoryReserver = Arc io::Result + Send + Sync + 'static>; + +struct ReservedHttpBytes { + // Release the allocation before its reservation, including its last slice. + bytes: Bytes, + _reservation: HttpWriterMemoryReservation, +} + +impl AsRef<[u8]> for ReservedHttpBytes { + fn as_ref(&self) -> &[u8] { + &self.bytes + } +} + // Not pin-projected: every field is `Unpin` and the `AsyncWrite` impl accesses // them through `get_mut()`, so a plain struct lets us add a manual `Drop` // (pin-project forbids one) to abort the background HTTP task — see below. @@ -1421,6 +1439,8 @@ pub struct HttpWriter { sender: PollSender>, handle: tokio::task::JoinHandle>, pending_chunk: BytesMut, + pending_reservation: Option, + memory_reserver: Option, finish: bool, track_internode_metrics: bool, internode_operation: Option<&'static str>, @@ -1432,6 +1452,31 @@ const HTTP_WRITER_BUFFER_SIZE: usize = 1024 * 1024; impl HttpWriter { /// Create a new HttpWriter for the given URL. The HTTP request is performed in the background. pub async fn new(url: String, method: Method, headers: HeaderMap) -> io::Result { + Self::with_memory_reserver(url, method, headers, None).await + } + + /// Keep each pending and queued allocation charged through its last byte slice. + pub async fn new_with_memory_reserver( + url: String, + method: Method, + headers: HeaderMap, + memory_reserver: HttpWriterMemoryReserver, + ) -> io::Result { + Self::with_memory_reserver(url, method, headers, Some(memory_reserver)).await + } + + async fn with_memory_reserver( + url: String, + method: Method, + headers: HeaderMap, + memory_reserver: Option, + ) -> io::Result { + // A rejected reservation must not leave an HTTP task or payload allocation. + let pending_reservation = memory_reserver + .as_ref() + .map(|reserve| reserve(HTTP_WRITER_BUFFER_SIZE)) + .transpose()?; + let pending_chunk = BytesMut::with_capacity(HTTP_WRITER_BUFFER_SIZE); // http_log!("[HttpWriter::new] url: {url}, method: {method:?}, headers: {headers:?}"); let url_clone = url.clone(); let method_clone = method.clone(); @@ -1523,7 +1568,9 @@ impl HttpWriter { start_tx: Some(start_tx), sender: PollSender::new(sender), handle, - pending_chunk: BytesMut::with_capacity(HTTP_WRITER_BUFFER_SIZE), + pending_chunk, + pending_reservation, + memory_reserver, finish: false, track_internode_metrics, internode_operation, @@ -1718,6 +1765,34 @@ fn send_error_to_io(err: tokio_util::sync::PollSendError, context: &str) - } impl HttpWriter { + fn ensure_pending_capacity(&mut self) -> io::Result<()> { + if let Some(reserve) = &self.memory_reserver + && self.pending_reservation.is_none() + { + let reservation = reserve(HTTP_WRITER_BUFFER_SIZE)?; + self.pending_chunk = BytesMut::with_capacity(HTTP_WRITER_BUFFER_SIZE); + self.pending_reservation = Some(reservation); + } + Ok(()) + } + + fn take_pending_chunk(&mut self) -> io::Result { + if self.memory_reserver.is_none() { + return Ok(self.pending_chunk.split().freeze()); + } + let reservation = self + .pending_reservation + .take() + .ok_or_else(|| io::Error::other("HTTP writer buffer has no memory reservation"))?; + // split() leaves the unused tail sharing the allocation. Transfer it all + // so a consumed output cannot refund storage still held by this writer. + let bytes = std::mem::take(&mut self.pending_chunk).freeze(); + Ok(Bytes::from_owner(ReservedHttpBytes { + bytes, + _reservation: reservation, + })) + } + fn start_request(&mut self) { if let Some(start_tx) = self.start_tx.take() { let _ = start_tx.send(()); @@ -1738,7 +1813,7 @@ impl HttpWriter { match self.sender.poll_reserve(cx) { Poll::Ready(Ok(())) => { - let chunk = self.pending_chunk.split().freeze(); + let chunk = self.take_pending_chunk()?; self.sender .send_item(Some(chunk)) .map_err(|e| send_error_to_io(e, "HttpWriter send error"))?; @@ -1774,8 +1849,18 @@ impl AsyncWrite for HttpWriter { if buf.len() >= HTTP_WRITER_BUFFER_SIZE && this.pending_chunk.is_empty() { match this.sender.poll_reserve(cx) { Poll::Ready(Ok(())) => { + let chunk = if this.memory_reserver.is_some() { + if let Err(error) = this.ensure_pending_capacity() { + this.sender.abort_send(); + return Poll::Ready(Err(error)); + } + this.pending_chunk.extend_from_slice(&buf[..HTTP_WRITER_BUFFER_SIZE]); + this.take_pending_chunk()? + } else { + Bytes::copy_from_slice(&buf[..HTTP_WRITER_BUFFER_SIZE]) + }; this.sender - .send_item(Some(Bytes::copy_from_slice(&buf[..HTTP_WRITER_BUFFER_SIZE]))) + .send_item(Some(chunk)) .map_err(|e| send_error_to_io(e, "HttpWriter send error"))?; this.start_request(); return Poll::Ready(Ok(HTTP_WRITER_BUFFER_SIZE)); @@ -1787,6 +1872,9 @@ impl AsyncWrite for HttpWriter { // 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()); + if written > 0 { + this.ensure_pending_capacity()?; + } this.pending_chunk.extend_from_slice(&buf[..written]); if written > 0 { this.start_request(); @@ -1825,6 +1913,9 @@ impl AsyncWrite for HttpWriter { let mut written = 0; for buf in bufs { let len = buf.len().min(available - written); + if len > 0 { + this.ensure_pending_capacity()?; + } this.pending_chunk.extend_from_slice(&buf[..len]); written += len; if written == available { @@ -2943,6 +3034,16 @@ mod tests { } fn paused_http_writer() -> (HttpWriter, tokio::sync::mpsc::Receiver>) { + paused_http_writer_with_memory(None).expect("unrestricted writer") + } + + fn paused_http_writer_with_memory( + memory_reserver: Option, + ) -> io::Result<(HttpWriter, tokio::sync::mpsc::Receiver>)> { + let pending_reservation = memory_reserver + .as_ref() + .map(|reserve| reserve(HTTP_WRITER_BUFFER_SIZE)) + .transpose()?; let (sender, receiver) = tokio::sync::mpsc::channel(HTTP_WRITER_CHANNEL_CAPACITY); let (_, err_rx) = tokio::sync::oneshot::channel(); let writer = HttpWriter { @@ -2954,11 +3055,261 @@ mod tests { sender: PollSender::new(sender), handle: tokio::spawn(std::future::pending()), pending_chunk: BytesMut::with_capacity(HTTP_WRITER_BUFFER_SIZE), + pending_reservation, + memory_reserver, finish: false, track_internode_metrics: false, internode_operation: None, }; - (writer, receiver) + Ok((writer, receiver)) + } + + struct HttpWriterTestMemory { + limit: AtomicUsize, + used: AtomicUsize, + calls: AtomicUsize, + } + + struct HttpWriterTestReservation { + memory: Arc, + capacity: usize, + } + + impl Drop for HttpWriterTestReservation { + fn drop(&mut self) { + self.memory.used.fetch_sub(self.capacity, Ordering::SeqCst); + } + } + + impl HttpWriterTestMemory { + fn new(limit: usize) -> Arc { + Arc::new(Self { + limit: AtomicUsize::new(limit), + used: AtomicUsize::new(0), + calls: AtomicUsize::new(0), + }) + } + + fn reserver(self: &Arc) -> HttpWriterMemoryReserver { + let memory = self.clone(); + Arc::new(move |capacity| { + assert_eq!(capacity, HTTP_WRITER_BUFFER_SIZE, "reserve full allocation capacity"); + memory.calls.fetch_add(1, Ordering::SeqCst); + memory + .used + .try_update(Ordering::SeqCst, Ordering::SeqCst, |used| { + used.checked_add(capacity) + .filter(|next| *next <= memory.limit.load(Ordering::SeqCst)) + }) + .map_err(|_| io::Error::from(io::ErrorKind::WouldBlock))?; + Ok(Box::new(HttpWriterTestReservation { + memory: memory.clone(), + capacity, + }) as HttpWriterMemoryReservation) + }) + } + + fn used(&self) -> usize { + self.used.load(Ordering::SeqCst) + } + } + + #[test] + fn http_writer_memory_denial_precedes_runtime_or_allocation() { + use std::future::Future; + + for limit in [0, HTTP_WRITER_BUFFER_SIZE - 1] { + let memory = HttpWriterTestMemory::new(limit); + let mut create = Box::pin(HttpWriter::new_with_memory_reserver( + "http://127.0.0.1/unpolled".to_owned(), + Method::PUT, + HeaderMap::new(), + memory.reserver(), + )); + let mut cx = Context::from_waker(std::task::Waker::noop()); + // Spawning before admission would panic without a Tokio runtime. + assert!( + matches!(create.as_mut().poll(&mut cx), Poll::Ready(Err(error)) if error.kind() == io::ErrorKind::WouldBlock) + ); + assert_eq!(memory.calls.load(Ordering::SeqCst), 1); + assert_eq!(memory.used(), 0); + } + } + + #[tokio::test] + async fn http_writer_memory_tracks_pending_and_last_output_slice() { + let memory = HttpWriterTestMemory::new(2 * HTTP_WRITER_BUFFER_SIZE); + let (mut writer, mut receiver) = paused_http_writer_with_memory(Some(memory.reserver())).expect("reserved writer"); + assert_eq!(writer.pending_chunk.capacity(), HTTP_WRITER_BUFFER_SIZE); + assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE); + memory.limit.store(0, Ordering::SeqCst); + writer.write_all(b"first").await.expect("pending payload"); + writer.flush().await.expect("emit partial allocation"); + assert_eq!(writer.pending_chunk.capacity(), 0, "no shared tail may remain after emission"); + let frame = receiver.recv().await.expect("queued frame").expect("payload"); + let clone = frame.clone(); + let last = frame.slice(4..); + memory.limit.store(2 * HTTP_WRITER_BUFFER_SIZE, Ordering::SeqCst); + writer.write_all(b"next").await.expect("a separately reserved allocation"); + assert_eq!(memory.used(), 2 * HTTP_WRITER_BUFFER_SIZE); + drop(frame); + drop(clone); + drop(writer); + assert_eq!(last.as_ref(), b"t"); + assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE); + drop(last); + assert_eq!(memory.used(), 0); + } + + #[tokio::test] + async fn http_writer_memory_shrink_rejects_new_scalar_and_vectored_storage() { + for mode in ["scalar", "fast", "vectored"] { + let memory = HttpWriterTestMemory::new(2 * HTTP_WRITER_BUFFER_SIZE); + let (mut writer, mut receiver) = paused_http_writer_with_memory(Some(memory.reserver())).expect("reserved writer"); + writer.write_all(b"retained").await.expect("first payload"); + writer.flush().await.expect("first allocation transferred"); + let retained = receiver.recv().await.expect("first frame").expect("payload"); + memory.limit.store(0, Ordering::SeqCst); + let error = match mode { + "vectored" => writer.write_vectored(&[IoSlice::new(b"lost"), IoSlice::new(b"bytes")]).await, + "fast" => writer.write(&vec![0xff; HTTP_WRITER_BUFFER_SIZE]).await, + _ => writer.write(b"lost bytes").await, + } + .expect_err("new storage must be rejected after shrink"); + assert_eq!(error.kind(), io::ErrorKind::WouldBlock); + assert_eq!(writer.pending_chunk.capacity(), 0); + assert!(receiver.try_recv().is_err(), "rejected bytes must not be queued"); + assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE); + let calls = memory.calls.load(Ordering::SeqCst); + assert_eq!(writer.write(&[]).await.expect("empty write"), 0); + assert_eq!(memory.calls.load(Ordering::SeqCst), calls); + memory.limit.store(2 * HTTP_WRITER_BUFFER_SIZE, Ordering::SeqCst); + writer.write_all(b"accepted").await.expect("restored capacity"); + writer.flush().await.expect("emit accepted bytes"); + let accepted = receiver.recv().await.expect("accepted frame").expect("payload"); + assert_eq!(accepted.as_ref(), b"accepted"); + drop(writer); + drop(retained); + drop(accepted); + assert_eq!(memory.used(), 0); + } + } + + #[tokio::test] + async fn http_writer_memory_full_queue_keeps_backpressure_and_capacity() { + for vectored in [false, true] { + let memory = HttpWriterTestMemory::new((HTTP_WRITER_CHANNEL_CAPACITY + 1) * HTTP_WRITER_BUFFER_SIZE); + let (mut writer, mut receiver) = paused_http_writer_with_memory(Some(memory.reserver())).expect("reserved writer"); + let payload = vec![0xa5; HTTP_WRITER_BUFFER_SIZE]; + for _ in 0..HTTP_WRITER_CHANNEL_CAPACITY { + let written = if vectored { + writer + .write_vectored(&[IoSlice::new(&payload[..32]), IoSlice::new(&payload[32..])]) + .await + } else { + writer.write(&payload).await + } + .expect("free queue slot"); + assert_eq!(written, HTTP_WRITER_BUFFER_SIZE); + writer.flush().await.expect("queue the complete allocation"); + } + assert_eq!(memory.used(), HTTP_WRITER_CHANNEL_CAPACITY * HTTP_WRITER_BUFFER_SIZE); + assert_eq!(writer.pending_chunk.capacity(), 0); + if vectored { + let written = writer + .write_vectored(&[IoSlice::new(&payload[..32]), IoSlice::new(&payload[32..])]) + .await + .expect("one bounded pending allocation beside the full queue"); + assert_eq!(written, HTTP_WRITER_BUFFER_SIZE); + assert_eq!(memory.used(), (HTTP_WRITER_CHANNEL_CAPACITY + 1) * HTTP_WRITER_BUFFER_SIZE); + } + let calls = memory.calls.load(Ordering::SeqCst); + { + let write = writer.write(&payload); + tokio::pin!(write); + assert!( + futures::poll!(&mut write).is_pending(), + "a full queue must not accept another large write" + ); + } + assert_eq!( + memory.calls.load(Ordering::SeqCst), + calls, + "backpressure must not allocate a fast-path copy" + ); + if !vectored { + drop(receiver.recv().await.expect("free a queue slot")); + assert_eq!(writer.write(&payload).await.expect("resumed write"), HTTP_WRITER_BUFFER_SIZE); + } + drop(writer); + assert_eq!(memory.used(), HTTP_WRITER_CHANNEL_CAPACITY * HTTP_WRITER_BUFFER_SIZE); + while let Ok(frame) = receiver.try_recv() { + let frame = frame.expect("payload"); + assert_eq!(frame.as_ref(), payload.as_slice()); + } + assert_eq!(memory.used(), 0); + } + } + + #[test] + fn http_writer_memory_cancellation_and_runtime_shutdown_release_queued_storage() { + for cancel_writer in [false, true] { + let memory = HttpWriterTestMemory::new(2 * HTTP_WRITER_BUFFER_SIZE); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("writer runtime"); + let writer = runtime.block_on(async { + let (mut writer, receiver) = paused_http_writer_with_memory(Some(memory.reserver())).expect("reserved writer"); + writer.handle.abort(); + writer.handle = tokio::spawn(async move { + std::future::pending::<()>().await; + drop(receiver); + Ok(()) + }); + writer.write_all(b"queued").await.expect("pending bytes"); + writer.flush().await.expect("queued allocation"); + writer + }); + assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE); + if cancel_writer { + drop(writer); + runtime.block_on(tokio::task::yield_now()); + assert_eq!(memory.used(), 0, "cancelled task must release its queue"); + drop(runtime); + } else { + drop(runtime); + assert_eq!(memory.used(), 0, "runtime shutdown must release queued allocations"); + drop(writer); + } + } + } + + #[tokio::test] + async fn http_writer_memory_round_trip_preserves_vectored_body_and_shutdown() { + let state = TestState::default(); + let (url, handle) = start_test_server(state.clone()).await.expect("memory writer test server"); + let memory = HttpWriterTestMemory::new((HTTP_WRITER_CHANNEL_CAPACITY + 1) * HTTP_WRITER_BUFFER_SIZE); + let header = [0x11; 32]; + let payload = vec![0xa5; 3 * HTTP_WRITER_BUFFER_SIZE + 17]; + let mut writer = HttpWriter::new_with_memory_reserver(url, Method::PUT, HeaderMap::new(), memory.reserver()) + .await + .expect("reserved HTTP writer"); + let mut slices = [IoSlice::new(&[]), IoSlice::new(&header), IoSlice::new(&payload)]; + let mut remaining = slices.as_mut_slice(); + while remaining.iter().any(|slice| !slice.is_empty()) { + let written = writer.write_vectored(remaining).await.expect("reserved vectored write"); + assert!(written > 0 && written <= HTTP_WRITER_BUFFER_SIZE); + IoSlice::advance_slices(&mut remaining, written); + } + writer.flush().await.expect("flush final reserved allocation"); + writer.shutdown().await.expect("complete reserved HTTP request"); + 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()); + assert_eq!(memory.used(), 0, "completed transport retains no payload reservation"); + handle.abort(); } #[tokio::test] diff --git a/docs/operations/scanner-runtime-controls.md b/docs/operations/scanner-runtime-controls.md index 13970b633..e3856fd76 100644 --- a/docs/operations/scanner-runtime-controls.md +++ b/docs/operations/scanner-runtime-controls.md @@ -498,7 +498,9 @@ Direct writers reserve their full aligned staging allocation when opened and ret Admitted bitrot verification and reader scratch, Heal read and reconstruction shard buffers, legacy codec workspaces, and modern parity-verification scratch also reserve their payload storage before allocation. Cached legacy workspaces retain their charge while cached and can only reuse it for the same memory owner. Heal shard output retains its charge through the last `Bytes` clone or slice. HTTP read-file, walk-directory and namespace-scan response streams charge each staging chunk separately from the source reader; a retained output chunk can prevent the next allocation after a cap reduction. -This cap does not yet cover inline Heal input copies and output, other encode/decode paths, subsequent RPC serialization and transport copies, allocator overhead or total process RSS. Already resident format-cache bytes are shared without a new payload allocation. The `rustfs_background_buffer_memory_limit_bytes` and `rustfs_background_buffer_memory_retained_bytes` gauges expose the sampled cap and retained payload. Like the node gauges, they can remain stale when the sampler stalls; the in-process admission deadline still expires after three seconds. +Inline Heal readers share existing metadata bytes. Inline writers charge replacement storage while the old allocation remains live, and their returned `Bytes` retain the charge. Admitted remote Heal HTTP writers reserve each full 1 MiB staging chunk before allocation. Queued chunks retain their reservations through the HTTP body and its final clone or slice, independently of the sender's disk capacity. A denied allocation accepts no new input bytes; existing storage remains usable after a cap reduction. + +This cap does not yet cover Scanner duplex buffers, other encode/decode paths, subsequent RPC serialization and transport copies, allocator overhead or total process RSS. Already resident format-cache bytes are shared without a new payload allocation. The `rustfs_background_buffer_memory_limit_bytes` and `rustfs_background_buffer_memory_retained_bytes` gauges expose the sampled cap and retained payload. Like the node gauges, they can remain stale when the sampler stalls; the in-process admission deadline still expires after three seconds. Each complete sample leases admission for three seconds. Missing resource caps block new work; missing or low-confidence feedback retains at most one maintenance unit. CPU use above 90%, throttling or pressure above 1% of elapsed time, memory headroom at most 10%, foreground errors/cancellations, a foreground queue or a comparable foreground p99 regression reduce concurrency immediately. Recovery requires CPU use at most 70%, throttling/stalls at most 0.1% of elapsed time, at least 20% memory headroom, fresh foreground evidence and actual successful work. The existing controller probes one additional unit only after eight consecutive healthy seconds and cooldown. Limits affect new admission and do not revoke active work. diff --git a/rustfs/src/server/http.rs b/rustfs/src/server/http.rs index fa47296d4..033c2f66c 100644 --- a/rustfs/src/server/http.rs +++ b/rustfs/src/server/http.rs @@ -4926,6 +4926,85 @@ mod tests { assert_eq!(physical.completed(), Some(3)); assert_eq!(memory.snapshot().used, 0); } + let warm_writer = remote + .create_file("bucket", "bucket", "writer-capability", 0) + .await + .expect("negotiate writer authentication before testing payload admission"); + drop(warm_writer); + let source_disk = AdaptiveAdmission::new(NonZeroU16::MIN, 1); + memory.set_sampled_limit(0, Duration::ZERO); + let denied = storage::with_background_resource_unit( + source_disk.try_acquire().expect("sender disk admission"), + storage::with_remote_background_io( + Some(storage::RemoteBackgroundIo::new(sender_owner.clone())), + remote.create_file("bucket", "bucket", "client-memory-closed", i64::try_from(DATA.len()).expect("size")), + ), + ) + .await; + let error = match denied { + Ok(_) => panic!("closed sender memory budget must reject the HTTP writer"), + Err(error) => error, + }; + let error = std::io::Error::from(error); + assert_eq!( + error.kind(), + std::io::ErrorKind::WouldBlock, + "sender memory denial remains retryable: {error}" + ); + assert_eq!(source_disk.snapshot().active, 0); + assert_eq!(memory.snapshot().used, 0); + assert!(!root.path().join("bucket/client-memory-closed").exists()); + assert_eq!(node_owner.completed(), Some(3)); + + memory.set_sampled_limit(4 * 1024 * 1024, Duration::from_secs(30)); + node_owner.update( + &mut controller, + Some(FeedbackObservation { + generation: 6, + max_concurrency: 1, + pressure: ResourcePressure::Unknown, + completed_units: None, + }), + ); + disk.physical_admission() + .expect("physical binding") + .set_test_capacity(vec![physical.clone()]) + .expect("renew upload receiver device plan"); + let mut writer = storage::with_background_resource_unit( + source_disk.try_acquire().expect("sender disk admission"), + storage::with_remote_background_io( + Some(storage::RemoteBackgroundIo::new(sender_owner.clone())), + remote.create_file("bucket", "bucket", "accounted-upload", i64::try_from(DATA.len()).expect("size")), + ), + ) + .await + .expect("admitted production HTTP writer"); + assert_eq!(source_disk.snapshot().active, 0, "escaped HTTP writer retains only memory ownership"); + assert_eq!(memory.snapshot().used, 1024 * 1024, "sender reserves the whole pending chunk"); + writer + .write_all(DATA) + .await + .expect("write authenticated body after leaving sender scope"); + tokio::time::timeout(Duration::from_secs(5), writer.shutdown()) + .await + .expect("bounded authenticated upload") + .expect("authenticated upload completes"); + drop(writer); + tokio::time::timeout(Duration::from_secs(5), async { + while memory.snapshot().used != 0 || physical.snapshot().active != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("sender frames and receiver native writes release their storage"); + assert_eq!( + tokio::fs::read(root.path().join("bucket/accounted-upload")) + .await + .expect("uploaded bytes"), + DATA + ); + assert_eq!(node_owner.completed(), Some(4), "complete authenticated upload earns one completion"); + assert_eq!(physical.completed(), Some(4)); remote.close().await.expect("close remote client"); disk.close().await.expect("stop receiver disk monitors"); server.abort(); diff --git a/rustfs/src/storage/storage_api.rs b/rustfs/src/storage/storage_api.rs index 3a9a567a6..7c44420d6 100644 --- a/rustfs/src/storage/storage_api.rs +++ b/rustfs/src/storage/storage_api.rs @@ -506,7 +506,7 @@ pub(crate) mod ecstore_disk { validate_batch_read_version_item_count, with_remote_background_io, }; #[cfg(test)] - pub(crate) use rustfs_ecstore::api::disk::{CurrentDeleteGeneration, DiskOption, new_disk}; + pub(crate) use rustfs_ecstore::api::disk::{CurrentDeleteGeneration, DiskOption, new_disk, with_background_resource_unit}; pub(crate) use rustfs_ecstore::api::disk::{endpoint, error, error_reduce}; } diff --git a/rustfs/src/storage_api.rs b/rustfs/src/storage_api.rs index 19a22f340..5baf3ab44 100644 --- a/rustfs/src/storage_api.rs +++ b/rustfs/src/storage_api.rs @@ -181,7 +181,9 @@ pub(crate) mod server { pub(crate) mod http { #[cfg(test)] - pub(crate) use crate::storage::storage_api::ecstore_disk::{DiskAPI, DiskOption, new_disk}; + pub(crate) use crate::storage::storage_api::ecstore_disk::{ + DiskAPI, DiskOption, new_disk, with_background_resource_unit, + }; pub(crate) use crate::storage::storage_api::{ RemoteBackgroundIo, ServerContextSlot, TONIC_RPC_PREFIX, normalize_tonic_rpc_audience, rpc_background_admission, tonic_boot_epoch_challenge, tonic_boot_epoch_response_headers, tonic_rpc_auth_failure_reason,