diff --git a/Cargo.lock b/Cargo.lock index 86e22ffa5..ca4270e4f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1529,6 +1529,17 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a822ea5bc7590f9d40f1ba12c0dc3c2760f3482c6984db1573ad11031420831" +[[package]] +name = "clocksource" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "129026dd5a8a9592d96916258f3a5379589e513ea5e86aeb0bd2530286e44e9e" +dependencies = [ + "libc", + "time", + "winapi", +] + [[package]] name = "cmake" version = "0.1.57" @@ -6374,6 +6385,17 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" +[[package]] +name = "ratelimit" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36ea961700fd7260e7fa3701c8287d901b2172c51f9c1421fa0f21d7f7e184b7" +dependencies = [ + "clocksource", + "parking_lot", + "thiserror 1.0.69", +] + [[package]] name = "rayon" version = "1.11.0" @@ -7050,6 +7072,7 @@ dependencies = [ "pin-project-lite", "quick-xml 0.39.2", "rand 0.10.0", + "ratelimit", "reed-solomon-simd", "regex", "reqwest 0.13.2", diff --git a/Cargo.toml b/Cargo.toml index c73b47fa0..b32c8296b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -227,6 +227,7 @@ path-clean = "1.0.1" pin-project-lite = "0.2.16" pretty_assertions = "1.4.1" rand = { version = "0.10.0", features = ["serde"] } +ratelimit = "0.10.0" rayon = "1.11.0" reed-solomon-simd = { version = "3.1.0" } regex = { version = "1.12.3" } diff --git a/crates/ecstore/Cargo.toml b/crates/ecstore/Cargo.toml index 0325005cf..8d2d147ec 100644 --- a/crates/ecstore/Cargo.toml +++ b/crates/ecstore/Cargo.toml @@ -111,6 +111,7 @@ google-cloud-storage = { workspace = true } google-cloud-auth = { workspace = true } aws-config = { workspace = true } faster-hex = { workspace = true } +ratelimit = { workspace = true } [dev-dependencies] tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } diff --git a/crates/ecstore/src/bucket/bandwidth/mod.rs b/crates/ecstore/src/bucket/bandwidth/mod.rs new file mode 100644 index 000000000..570229939 --- /dev/null +++ b/crates/ecstore/src/bucket/bandwidth/mod.rs @@ -0,0 +1,16 @@ +// Copyright 2024 RustFS Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +pub mod monitor; +pub mod reader; diff --git a/crates/ecstore/src/bucket/bandwidth/monitor.rs b/crates/ecstore/src/bucket/bandwidth/monitor.rs new file mode 100644 index 000000000..203576789 --- /dev/null +++ b/crates/ecstore/src/bucket/bandwidth/monitor.rs @@ -0,0 +1,321 @@ +// Copyright 2024 RustFS Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::bucket::bandwidth::reader::BucketOptions; +use ratelimit::{Error as RatelimitError, Ratelimiter}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, RwLock}; +use std::time::Duration; +use tracing::warn; + +#[derive(Clone)] +pub struct BucketThrottle { + limiter: Arc>, + pub node_bandwidth_per_sec: i64, +} + +impl BucketThrottle { + fn new(node_bandwidth_per_sec: i64) -> Result { + let node_bandwidth_per_sec = node_bandwidth_per_sec.max(1); + let amount = node_bandwidth_per_sec as u64; + let limiter_inner = Ratelimiter::builder(amount, Duration::from_secs(1)) + .max_tokens(amount) + .build()?; + Ok(Self { + limiter: Arc::new(Mutex::new(limiter_inner)), + node_bandwidth_per_sec, + }) + } + + pub fn burst(&self) -> u64 { + self.limiter.lock().unwrap_or_else(|e| e.into_inner()).max_tokens() + } + + /// The ratelimit crate (0.10.0) does not provide a bulk token consumption API. + /// try_wait() first to consume 1 token AND trigger the internal refill + /// mechanism (tokens are only refilled during try_wait/wait calls). + /// directly adjust available tokens via set_available() to consume the remaining amount. + pub(crate) fn consume(&self, n: u64) -> (u64, f64, u64) { + let guard = self.limiter.lock().unwrap_or_else(|e| { + warn!("bucket throttle mutex poisoned, recovering"); + e.into_inner() + }); + if n == 0 { + return (0, guard.rate(), 0); + } + let mut consumed = 0u64; + if guard.try_wait().is_ok() { + consumed = 1; + } + let available = guard.available(); + let to_consume = n - consumed; + let batch = to_consume.min(available); + if batch > 0 { + let _ = guard.set_available(available - batch); + consumed += batch; + } + let deficit = n.saturating_sub(consumed); + let rate = guard.rate(); + (deficit, rate, consumed) + } +} + +pub struct Monitor { + t_lock: RwLock>, + pub node_count: u64, +} + +impl Monitor { + pub fn new(num_nodes: u64) -> Arc { + let node_cnt = num_nodes.max(1); + Arc::new(Monitor { + t_lock: RwLock::new(HashMap::new()), + node_count: node_cnt, + }) + } + + pub fn delete_bucket(&self, bucket: &str) { + self.t_lock + .write() + .unwrap_or_else(|e| { + warn!("bucket monitor rwlock write poisoned, recovering"); + e.into_inner() + }) + .retain(|opts, _| opts.name != bucket); + } + + pub fn delete_bucket_throttle(&self, bucket: &str, arn: &str) { + let opts = BucketOptions { + name: bucket.to_string(), + replication_arn: arn.to_string(), + }; + self.t_lock + .write() + .unwrap_or_else(|e| { + warn!("bucket monitor rwlock write poisoned, recovering"); + e.into_inner() + }) + .remove(&opts); + } + + pub fn throttle(&self, opts: &BucketOptions) -> Option { + self.t_lock + .read() + .unwrap_or_else(|e| { + warn!("bucket monitor rwlock read poisoned, recovering"); + e.into_inner() + }) + .get(opts) + .cloned() + } + + pub fn set_bandwidth_limit(&self, bucket: &str, arn: &str, limit: i64) { + if limit <= 0 { + warn!( + bucket = bucket, + arn = arn, + limit = limit, + "invalid bandwidth limit, must be positive; ignoring" + ); + return; + } + let limit_bytes = limit / self.node_count as i64; + if limit_bytes == 0 && limit > 0 { + warn!( + bucket = bucket, + arn = arn, + limit = limit, + node_count = self.node_count, + "bandwidth limit too small for cluster size, per-node limit will clamp to 1 byte/s" + ); + } + let opts = BucketOptions { + name: bucket.to_string(), + replication_arn: arn.to_string(), + }; + let throttle = match BucketThrottle::new(limit_bytes) { + Ok(t) => t, + Err(e) => { + warn!( + bucket = bucket, + arn = arn, + limit_bytes = limit_bytes, + err = %e, + "failed to build bandwidth throttle, throttling disabled for this target" + ); + return; + } + }; + self.t_lock + .write() + .unwrap_or_else(|e| { + warn!("bucket monitor rwlock write poisoned, recovering"); + e.into_inner() + }) + .insert(opts, throttle); + } + + pub fn is_throttled(&self, bucket: &str, arn: &str) -> bool { + let opt = BucketOptions { + name: bucket.to_string(), + replication_arn: arn.to_string(), + }; + self.t_lock + .read() + .unwrap_or_else(|e| { + warn!("bucket monitor rwlock read poisoned, recovering"); + e.into_inner() + }) + .contains_key(&opt) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_set_and_get_throttle_with_node_split() { + let monitor = Monitor::new(4); + monitor.set_bandwidth_limit("b1", "arn1", 400); + + let throttle = monitor + .throttle(&BucketOptions { + name: "b1".to_string(), + replication_arn: "arn1".to_string(), + }) + .expect("throttle should exist"); + + assert_eq!(throttle.node_bandwidth_per_sec, 100); + assert!(monitor.is_throttled("b1", "arn1")); + } + + #[test] + fn test_delete_bucket_throttle() { + let monitor = Monitor::new(2); + monitor.set_bandwidth_limit("b1", "arn1", 200); + assert!(monitor.is_throttled("b1", "arn1")); + + monitor.delete_bucket_throttle("b1", "arn1"); + assert!(!monitor.is_throttled("b1", "arn1")); + } + + #[test] + fn test_delete_bucket_removes_all_arns() { + let monitor = Monitor::new(1); + monitor.set_bandwidth_limit("b1", "arn1", 100); + monitor.set_bandwidth_limit("b1", "arn2", 100); + monitor.set_bandwidth_limit("b2", "arn3", 100); + + monitor.delete_bucket("b1"); + + assert!(!monitor.is_throttled("b1", "arn1")); + assert!(!monitor.is_throttled("b1", "arn2")); + assert!(monitor.is_throttled("b2", "arn3")); + } + + #[test] + fn test_set_bandwidth_limit_ignores_non_positive() { + let monitor = Monitor::new(2); + + monitor.set_bandwidth_limit("b1", "arn1", 0); + assert!(!monitor.is_throttled("b1", "arn1")); + + monitor.set_bandwidth_limit("b1", "arn1", -10); + assert!(!monitor.is_throttled("b1", "arn1")); + } + + #[test] + fn test_consume_returns_deficit_when_tokens_exhausted() { + let throttle = BucketThrottle::new(100).expect("test"); + + let (deficit, rate, _consumed) = throttle.consume(200); + + assert!(deficit > 0); + assert!(rate > 0.0); + } + + #[test] + fn test_consume_no_deficit_when_tokens_sufficient() { + let throttle = BucketThrottle::new(10000).expect("test"); + + std::thread::sleep(std::time::Duration::from_millis(1100)); + + let (deficit, _rate, _consumed) = throttle.consume(5000); + assert_eq!(deficit, 0); + } + + #[test] + fn test_burst_equals_bandwidth() { + let throttle = BucketThrottle::new(500).expect("test"); + assert_eq!(throttle.burst(), 500); + } + + #[test] + fn test_concurrent_consume() { + let throttle = BucketThrottle::new(10000).expect("test"); + std::thread::sleep(std::time::Duration::from_millis(1100)); + + let mut handles = vec![]; + for _ in 0..10 { + let t = throttle.clone(); + handles.push(std::thread::spawn(move || t.consume(100))); + } + + let mut total_deficit = 0u64; + let mut total_consumed = 0u64; + for h in handles { + let (deficit, _rate, consumed) = h.join().unwrap(); + total_consumed += consumed; + total_deficit += deficit; + } + + assert_eq!(total_consumed + total_deficit, 1000); + assert!(total_consumed <= 10000); + } + + #[test] + fn test_zero_bandwidth_clamped_to_one() { + let throttle = BucketThrottle::new(0).expect("test"); + assert_eq!(throttle.burst(), 1); + assert_eq!(throttle.node_bandwidth_per_sec, 1); + } + + #[test] + fn test_negative_bandwidth_clamped_to_one() { + let throttle = BucketThrottle::new(-100).expect("test"); + assert_eq!(throttle.burst(), 1); + assert_eq!(throttle.node_bandwidth_per_sec, 1); + } + + #[test] + fn test_update_bandwidth_limit_overrides_previous() { + let monitor = Monitor::new(1); + monitor.set_bandwidth_limit("b1", "arn1", 1000); + + let opts = BucketOptions { + name: "b1".to_string(), + replication_arn: "arn1".to_string(), + }; + let t1 = monitor.throttle(&opts).expect("throttle should exist"); + assert_eq!(t1.burst(), 1000); + assert_eq!(t1.node_bandwidth_per_sec, 1000); + + monitor.set_bandwidth_limit("b1", "arn1", 500); + + let t2 = monitor.throttle(&opts).expect("throttle should exist after update"); + assert_eq!(t2.burst(), 500); + assert_eq!(t2.node_bandwidth_per_sec, 500); + } +} diff --git a/crates/ecstore/src/bucket/bandwidth/reader.rs b/crates/ecstore/src/bucket/bandwidth/reader.rs new file mode 100644 index 000000000..af8b81af4 --- /dev/null +++ b/crates/ecstore/src/bucket/bandwidth/reader.rs @@ -0,0 +1,336 @@ +// Copyright 2024 RustFS Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::bucket::bandwidth::monitor::Monitor; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, ReadBuf}; +use tokio::time::Sleep; +use tracing::{debug, warn}; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct BucketOptions { + pub name: String, + pub replication_arn: String, +} + +pub struct MonitorReaderOptions { + pub bucket_options: BucketOptions, + pub header_size: usize, +} + +struct WaitState { + sleep: Pin>, +} + +pub struct MonitoredReader { + r: R, + m: Arc, + opts: MonitorReaderOptions, + wait_state: std::sync::Mutex>, + temp_buf: Vec, +} + +impl MonitoredReader { + pub fn new(m: Arc, r: R, opts: MonitorReaderOptions) -> Self { + let throttle = m.throttle(&opts.bucket_options); + debug!( + bucket = opts.bucket_options.name, + arn = opts.bucket_options.replication_arn, + throttle_active = throttle.is_some(), + limit_bps = throttle.as_ref().map(|t| t.node_bandwidth_per_sec).unwrap_or(0), + header_size = opts.header_size, + "MonitoredReader created" + ); + MonitoredReader { + r, + m, + opts, + wait_state: std::sync::Mutex::new(None), + temp_buf: Vec::new(), + } + } +} + +impl AsyncRead for MonitoredReader { + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + let this = self.get_mut(); + + { + let mut guard = this.wait_state.lock().unwrap_or_else(|e| { + warn!("MonitoredReader wait_state mutex poisoned, recovering"); + e.into_inner() + }); + if let Some(ref mut ws) = *guard { + match ws.sleep.as_mut().poll(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(()) => { + *guard = None; + drop(guard); + } + } + } + } + + let throttle = match this.m.throttle(&this.opts.bucket_options) { + Some(t) => t, + None => return Pin::new(&mut this.r).poll_read(cx, buf), + }; + + let b = throttle.burst(); + debug_assert!(b >= 1, "burst must be at least 1"); + let (need, tokens) = calc_need_and_tokens(b, buf.remaining(), &mut this.opts.header_size); + let (deficit, rate, consumed) = throttle.consume(tokens); + let need = need.min(consumed as usize); + + if deficit > 0 && rate > 0.0 { + let duration = std::time::Duration::from_secs_f64(deficit as f64 / rate); + debug!( + tokens = tokens, + deficit = deficit, + rate = rate, + sleep_ms = duration.as_millis() as u64, + "bandwidth throttle sleep" + ); + let mut sleep = Box::pin(tokio::time::sleep(duration)); + match sleep.as_mut().poll(cx) { + Poll::Pending => { + *this.wait_state.lock().unwrap_or_else(|e| { + warn!("MonitoredReader wait_state mutex poisoned, recovering"); + e.into_inner() + }) = Some(WaitState { sleep }); + return Poll::Pending; + } + Poll::Ready(()) => {} + } + } + + poll_limited_read(&mut this.r, cx, buf, need, &mut this.temp_buf) + } +} + +fn calc_need_and_tokens(burst: u64, need_upper: usize, header_size: &mut usize) -> (usize, u64) { + let hdr = *header_size; + let mut need = need_upper; + let tokens: u64 = if hdr > 0 { + if (hdr as u64) < burst { + *header_size = 0; + need = ((burst - hdr as u64) as usize).min(need); + need as u64 + hdr as u64 + } else { + *header_size -= burst as usize; + need = 0; + burst + } + } else { + need = need.min(burst as usize); + need as u64 + }; + (need, tokens) +} + +fn poll_limited_read( + r: &mut R, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + limit: usize, + reusable_buf: &mut Vec, +) -> Poll> { + let remaining = buf.remaining(); + if limit == 0 || remaining == 0 { + return Poll::Ready(Ok(())); + } + + if remaining <= limit { + return Pin::new(r).poll_read(cx, buf); + } + + reusable_buf.resize(limit, 0); + let mut temp_buf = ReadBuf::new(&mut reusable_buf[..limit]); + match Pin::new(r).poll_read(cx, &mut temp_buf) { + Poll::Ready(Ok(())) => { + buf.put_slice(temp_buf.filled()); + Poll::Ready(Ok(())) + } + other => other, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures_util::task::noop_waker_ref; + use std::io; + use std::pin::Pin; + use std::task::Context; + use tokio::io::AsyncReadExt; + + #[derive(Default)] + struct TestAsyncReader { + data: Vec, + pos: usize, + } + + impl TestAsyncReader { + fn new(data: &[u8]) -> Self { + Self { + data: data.to_vec(), + pos: 0, + } + } + } + + impl AsyncRead for TestAsyncReader { + fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + if self.pos >= self.data.len() { + return Poll::Ready(Ok(())); + } + + let remaining = self.data.len() - self.pos; + let n = remaining.min(buf.remaining()); + buf.put_slice(&self.data[self.pos..self.pos + n]); + self.pos += n; + Poll::Ready(Ok(())) + } + } + + #[tokio::test] + async fn test_monitored_reader_passthrough_when_throttle_absent() { + let monitor = Monitor::new(1); + let opts = MonitorReaderOptions { + bucket_options: BucketOptions { + name: "b1".to_string(), + replication_arn: "arn1".to_string(), + }, + header_size: 0, + }; + let inner = TestAsyncReader::new(b"hello-world"); + let mut reader = MonitoredReader::new(monitor, inner, opts); + + let mut out = Vec::new(); + reader.read_to_end(&mut out).await.unwrap(); + assert_eq!(out, b"hello-world"); + } + + #[test] + fn test_poll_limited_read_honors_limit() { + let mut inner = TestAsyncReader::new(b"abcdef"); + let mut out = [0u8; 8]; + let mut reusable = Vec::new(); + let waker = noop_waker_ref(); + let mut cx = Context::from_waker(waker); + + let mut read_buf = ReadBuf::new(&mut out); + let first = poll_limited_read(&mut inner, &mut cx, &mut read_buf, 3, &mut reusable); + assert!(matches!(first, Poll::Ready(Ok(())))); + assert_eq!(read_buf.filled(), b"abc"); + + let second = poll_limited_read(&mut inner, &mut cx, &mut read_buf, 3, &mut reusable); + assert!(matches!(second, Poll::Ready(Ok(())))); + assert_eq!(read_buf.filled(), b"abcdef"); + } + + #[tokio::test] + async fn test_header_only_consumption_skips_io() { + let monitor = Monitor::new(1); + monitor.set_bandwidth_limit("b1", "arn1", 10); + + let data = vec![0xAAu8; 20]; + let inner = TestAsyncReader::new(&data); + let opts = MonitorReaderOptions { + bucket_options: BucketOptions { + name: "b1".to_string(), + replication_arn: "arn1".to_string(), + }, + header_size: 25, + }; + let mut reader = MonitoredReader::new(monitor, inner, opts); + + let waker = noop_waker_ref(); + let mut cx = Context::from_waker(waker); + let mut out = [0u8; 16]; + let mut read_buf = ReadBuf::new(&mut out); + + let poll = Pin::new(&mut reader).poll_read(&mut cx, &mut read_buf); + assert!(matches!(poll, Poll::Ready(Ok(())) | Poll::Pending)); + assert!(read_buf.filled().is_empty()); + assert_eq!(reader.opts.header_size, 15); + } + + #[tokio::test] + async fn test_monitored_reader_with_throttle_reads_all_data() { + let monitor = Monitor::new(1); + monitor.set_bandwidth_limit("b1", "arn1", 1024); + + let data = vec![0xABu8; 4096]; + let inner = TestAsyncReader::new(&data); + let opts = MonitorReaderOptions { + bucket_options: BucketOptions { + name: "b1".to_string(), + replication_arn: "arn1".to_string(), + }, + header_size: 0, + }; + let mut reader = MonitoredReader::new(monitor, inner, opts); + + let mut out = Vec::new(); + reader.read_to_end(&mut out).await.unwrap(); + assert_eq!(out.len(), 4096); + assert!(out.iter().all(|&b| b == 0xAB)); + } + + #[tokio::test] + async fn test_monitored_reader_header_size_accounting() { + let monitor = Monitor::new(1); + monitor.set_bandwidth_limit("b1", "arn1", 100); + + let data = vec![0u8; 200]; + let inner = TestAsyncReader::new(&data); + let opts = MonitorReaderOptions { + bucket_options: BucketOptions { + name: "b1".to_string(), + replication_arn: "arn1".to_string(), + }, + header_size: 50, + }; + let mut reader = MonitoredReader::new(monitor, inner, opts); + + let mut out = Vec::new(); + reader.read_to_end(&mut out).await.unwrap(); + assert_eq!(out.len(), 200); + assert_eq!(reader.opts.header_size, 0); + } + + #[tokio::test] + async fn test_monitored_reader_very_small_limit() { + let monitor = Monitor::new(1); + monitor.set_bandwidth_limit("b1", "arn1", 1); + + let data = vec![0xFFu8; 10]; + let inner = TestAsyncReader::new(&data); + let opts = MonitorReaderOptions { + bucket_options: BucketOptions { + name: "b1".to_string(), + replication_arn: "arn1".to_string(), + }, + header_size: 0, + }; + let mut reader = MonitoredReader::new(monitor, inner, opts); + + let mut out = Vec::new(); + reader.read_to_end(&mut out).await.unwrap(); + assert_eq!(out.len(), 10); + } +} diff --git a/crates/ecstore/src/bucket/bucket_target_sys.rs b/crates/ecstore/src/bucket/bucket_target_sys.rs index 761ec5ea2..4f286e35e 100644 --- a/crates/ecstore/src/bucket/bucket_target_sys.rs +++ b/crates/ecstore/src/bucket/bucket_target_sys.rs @@ -21,6 +21,7 @@ use crate::bucket::target::ARN; use crate::bucket::target::BucketTargetType; use crate::bucket::target::{self, BucketTarget, BucketTargets, Credentials}; use crate::bucket::versioning_sys::BucketVersioningSys; +use crate::global::get_global_bucket_monitor; use aws_credential_types::Credentials as SdkCredentials; use aws_sdk_s3::config::Region as SdkRegion; use aws_sdk_s3::error::ProvideErrorMetadata; @@ -667,9 +668,19 @@ impl BucketTargetSys { Ok(true) } - fn update_bandwidth_limit(&self, _bucket: &str, _arn: &str, _limit: i64) { - // Implementation for bandwidth limit update - // This would interact with the global bucket monitor + fn update_bandwidth_limit(&self, bucket: &str, arn: &str, limit: i64) { + if let Some(bucket_monitor) = get_global_bucket_monitor() { + if limit == 0 { + bucket_monitor.delete_bucket_throttle(bucket, arn); + return; + } + bucket_monitor.set_bandwidth_limit(bucket, arn, limit); + } else { + error!( + "Global bucket monitor uninitialized; skipping bandwidth limit update for bucket '{}' and ARN '{}'", + bucket, arn + ); + } } pub async fn get_remote_target_client_by_arn(&self, _bucket: &str, arn: &str) -> Option> { @@ -691,6 +702,7 @@ impl BucketTargetSys { if let Some(existing_targets) = targets_map.remove(bucket) { for target in existing_targets { arn_remotes_map.remove(&target.arn); + self.update_bandwidth_limit(bucket, &target.arn, 0); } } diff --git a/crates/ecstore/src/bucket/mod.rs b/crates/ecstore/src/bucket/mod.rs index 18fe80509..ac04822fc 100644 --- a/crates/ecstore/src/bucket/mod.rs +++ b/crates/ecstore/src/bucket/mod.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod bandwidth; pub mod bucket_target_sys; pub mod error; pub mod lifecycle; diff --git a/crates/ecstore/src/bucket/replication/replication_resyncer.rs b/crates/ecstore/src/bucket/replication/replication_resyncer.rs index ec434de26..6f3c83ae6 100644 --- a/crates/ecstore/src/bucket/replication/replication_resyncer.rs +++ b/crates/ecstore/src/bucket/replication/replication_resyncer.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::bucket::bandwidth::reader::{BucketOptions, MonitorReaderOptions, MonitoredReader}; use crate::bucket::bucket_target_sys::{ AdvancedPutOptions, BucketTargetSys, PutObjectOptions, PutObjectPartOptions, RemoveObjectOptions, TargetClient, }; @@ -28,16 +29,21 @@ use crate::error::{Error, Result, is_err_object_not_found, is_err_version_not_fo use crate::event::name::EventName; use crate::event_notification::{EventArgs, send_event}; use crate::global::GLOBAL_LocalNodeName; +use crate::global::get_global_bucket_monitor; use crate::set_disk::get_lock_acquire_timeout; -use crate::store_api::{DeletedObject, ObjectInfo, ObjectOptions, ObjectToDelete, WalkOptions}; +use crate::store_api::{DeletedObject, HTTPRangeSpec, ObjectInfo, ObjectOptions, ObjectToDelete, WalkOptions}; use crate::{StorageAPI, new_object_layer_fn}; use aws_sdk_s3::error::SdkError; use aws_sdk_s3::operation::head_object::HeadObjectOutput; use aws_sdk_s3::primitives::ByteStream; use aws_sdk_s3::types::{CompletedPart, ObjectLockLegalHoldStatus}; +use aws_smithy_types::body::SdkBody; use byteorder::ByteOrder; use futures::future::join_all; +use futures::stream::StreamExt; use http::HeaderMap; +use http_body::Frame; +use http_body_util::StreamBody; use regex::Regex; use rustfs_filemeta::{ MrfReplicateEntry, REPLICATE_EXISTING, REPLICATE_EXISTING_DELETE, REPLICATION_RESET, ReplicateDecision, ReplicateObjectInfo, @@ -61,10 +67,11 @@ use std::collections::HashMap; use std::sync::Arc; use time::OffsetDateTime; use time::format_description::well_known::Rfc3339; -use tokio::io::{AsyncRead, AsyncReadExt}; +use tokio::io::AsyncRead; use tokio::sync::RwLock; use tokio::task::JoinSet; use tokio::time::Duration as TokioDuration; +use tokio_util::io::ReaderStream; use tokio_util::sync::CancellationToken; use tracing::{error, info, warn}; @@ -74,6 +81,8 @@ const RESYNC_META_FORMAT: u16 = 1; const RESYNC_META_VERSION: u16 = 1; const RESYNC_TIME_INTERVAL: TokioDuration = TokioDuration::from_secs(60); +static WARNED_MONITOR_UNINIT: std::sync::Once = std::sync::Once::new(); + #[derive(Debug, Clone, Default)] pub struct ResyncOpts { pub bucket: String, @@ -1956,20 +1965,16 @@ impl ReplicateObjectInfoExt for ReplicateObjectInfo { let versioned = BucketVersioningSys::prefix_enabled(&bucket, &object).await; let version_suspended = BucketVersioningSys::prefix_suspended(&bucket, &object).await; + let obj_opts = ObjectOptions { + version_id: self.version_id.map(|v| v.to_string()), + version_suspended, + versioned, + replication_request: true, + ..Default::default() + }; + let mut gr = match storage - .get_object_reader( - &bucket, - &object, - None, - HeaderMap::new(), - &ObjectOptions { - version_id: self.version_id.map(|v| v.to_string()), - version_suspended, - versioned, - replication_request: true, - ..Default::default() - }, - ) + .get_object_reader(&bucket, &object, None, HeaderMap::new(), &obj_opts) .await { Ok(gr) => gr, @@ -2078,34 +2083,26 @@ impl ReplicateObjectInfoExt for ReplicateObjectInfo { } }; - // TODO:bandwidth - if let Some(err) = if is_multipart { - replicate_object_with_multipart(tgt_client.clone(), &tgt_client.bucket, &object, gr.stream, &object_info, put_opts) - .await - .err() + drop(gr); + replicate_object_with_multipart(MultipartReplicationContext { + storage: storage.clone(), + cli: tgt_client.clone(), + src_bucket: &bucket, + dst_bucket: &tgt_client.bucket, + object: &object, + object_info: &object_info, + obj_opts: &obj_opts, + arn: &rinfo.arn, + put_opts, + }) + .await + .err() } else { - // TODO: use stream - let body = match gr.read_all().await { - Ok(body) => body, - Err(e) => { - rinfo.replication_status = ReplicationStatusType::Failed; - rinfo.error = Some(e.to_string()); - warn!("failed to read object for bucket:{} arn:{} error:{}", bucket, tgt_client.arn, e); - send_event(EventArgs { - event_name: EventName::ObjectReplicationNotTracked.as_ref().to_string(), - bucket_name: bucket.clone(), - object: object_info.clone(), - host: GLOBAL_LocalNodeName.to_string(), - user_agent: "Internal: [Replication]".to_string(), - ..Default::default() - }); - return rinfo; - } - }; - let reader = ByteStream::from(body); + gr.stream = wrap_with_bandwidth_monitor(gr.stream, &put_opts, &bucket, &rinfo.arn); + let byte_stream = async_read_to_bytestream(gr.stream); tgt_client - .put_object(&tgt_client.bucket, &object, size, reader, &put_opts) + .put_object(&tgt_client.bucket, &object, size, byte_stream, &put_opts) .await .map_err(|e| std::io::Error::other(e.to_string())) .err() @@ -2161,20 +2158,16 @@ impl ReplicateObjectInfoExt for ReplicateObjectInfo { let versioned = BucketVersioningSys::prefix_enabled(&bucket, &object).await; let version_suspended = BucketVersioningSys::prefix_suspended(&bucket, &object).await; + let obj_opts = ObjectOptions { + version_id: self.version_id.map(|v| v.to_string()), + version_suspended, + versioned, + replication_request: true, + ..Default::default() + }; + let mut gr = match storage - .get_object_reader( - &bucket, - &object, - None, - HeaderMap::new(), - &ObjectOptions { - version_id: self.version_id.map(|v| v.to_string()), - version_suspended, - versioned, - replication_request: true, - ..Default::default() - }, - ) + .get_object_reader(&bucket, &object, None, HeaderMap::new(), &obj_opts) .await { Ok(gr) => gr, @@ -2370,39 +2363,27 @@ impl ReplicateObjectInfoExt for ReplicateObjectInfo { return rinfo; } }; + if let Some(err) = if is_multipart { - replicate_object_with_multipart( - tgt_client.clone(), - &tgt_client.bucket, - &object, - gr.stream, - &object_info, + drop(gr); + replicate_object_with_multipart(MultipartReplicationContext { + storage: storage.clone(), + cli: tgt_client.clone(), + src_bucket: &bucket, + dst_bucket: &tgt_client.bucket, + object: &object, + object_info: &object_info, + obj_opts: &obj_opts, + arn: &rinfo.arn, put_opts, - ) + }) .await .err() } else { - let body = match gr.read_all().await { - Ok(body) => body, - Err(e) => { - rinfo.replication_status = ReplicationStatusType::Failed; - rinfo.error = Some(e.to_string()); - warn!("failed to read object for bucket:{} arn:{} error:{}", bucket, tgt_client.arn, e); - send_event(EventArgs { - event_name: EventName::ObjectReplicationNotTracked.as_ref().to_string(), - bucket_name: bucket.clone(), - object: object_info, - host: GLOBAL_LocalNodeName.to_string(), - user_agent: "Internal: [Replication]".to_string(), - ..Default::default() - }); - rinfo.duration = (OffsetDateTime::now_utc() - start_time).unsigned_abs(); - return rinfo; - } - }; - let reader = ByteStream::from(body); + gr.stream = wrap_with_bandwidth_monitor(gr.stream, &put_opts, &bucket, &rinfo.arn); + let byte_stream = async_read_to_bytestream(gr.stream); tgt_client - .put_object(&tgt_client.bucket, &object, size, reader, &put_opts) + .put_object(&tgt_client.bucket, &object, size, byte_stream, &put_opts) .await .map_err(|e| std::io::Error::other(e.to_string())) .err() @@ -2456,6 +2437,63 @@ static STANDARD_HEADERS: &[&str] = &[ headers::AMZ_SERVER_SIDE_ENCRYPTION, ]; +fn calc_put_object_header_size(put_opts: &PutObjectOptions) -> usize { + let mut header_size: usize = 0; + for (key, value) in put_opts.header().iter() { + header_size += key.as_str().len(); + header_size += value.as_bytes().len(); + // Account for HTTP header formatting: ": " (2 bytes) and "\r\n" (2 bytes) + header_size += 4; + } + header_size +} + +fn wrap_with_bandwidth_monitor_with_header( + stream: Box, + bucket: &str, + arn: &str, + header_size: usize, +) -> Box { + if let Some(monitor) = get_global_bucket_monitor() { + Box::new(MonitoredReader::new( + monitor, + stream, + MonitorReaderOptions { + bucket_options: BucketOptions { + name: bucket.to_string(), + replication_arn: arn.to_string(), + }, + header_size, + }, + )) + } else { + WARNED_MONITOR_UNINIT.call_once(|| { + warn!( + "Global bucket monitor uninitialized; proceeding with unthrottled replication (bandwidth limits will be ignored)" + ) + }); + stream + } +} + +fn wrap_with_bandwidth_monitor( + stream: Box, + put_opts: &PutObjectOptions, + bucket: &str, + arn: &str, +) -> Box { + let header_size = calc_put_object_header_size(put_opts); + wrap_with_bandwidth_monitor_with_header(stream, bucket, arn, header_size) +} + +fn async_read_to_bytestream(reader: impl AsyncRead + Send + Sync + Unpin + 'static) -> ByteStream { + // Non-retryable: SDK-level retries are not supported for streaming bodies. + // Replication-level retry handles failures at a higher layer. + let stream = ReaderStream::new(reader); + let body = StreamBody::new(stream.map(|r| r.map(Frame::data))); + ByteStream::new(SdkBody::from_body_1_x(body)) +} + fn is_standard_header(k: &str) -> bool { STANDARD_HEADERS.iter().any(|h| h.eq_ignore_ascii_case(k)) } @@ -2679,17 +2717,56 @@ fn put_replication_opts(sc: &str, object_info: &ObjectInfo) -> Result<(PutObject Ok((put_op, is_multipart)) } -async fn replicate_object_with_multipart( +fn part_range_spec_from_actual_size(offset: i64, part_size: i64) -> std::io::Result<(HTTPRangeSpec, i64)> { + if offset < 0 { + return Err(std::io::Error::other("invalid part offset")); + } + if part_size <= 0 { + return Err(std::io::Error::other(format!("invalid part size {part_size}"))); + } + let end = offset + .checked_add(part_size - 1) + .ok_or_else(|| std::io::Error::other("part range overflow"))?; + let next_offset = end + .checked_add(1) + .ok_or_else(|| std::io::Error::other("part offset overflow"))?; + Ok(( + HTTPRangeSpec { + is_suffix_length: false, + start: offset, + end, + }, + next_offset, + )) +} + +struct MultipartReplicationContext<'a, S: StorageAPI> { + storage: Arc, cli: Arc, - bucket: &str, - object: &str, - reader: Box, - object_info: &ObjectInfo, - opts: PutObjectOptions, -) -> std::io::Result<()> { + src_bucket: &'a str, + dst_bucket: &'a str, + object: &'a str, + object_info: &'a ObjectInfo, + obj_opts: &'a ObjectOptions, + arn: &'a str, + put_opts: PutObjectOptions, +} + +async fn replicate_object_with_multipart(ctx: MultipartReplicationContext<'_, S>) -> std::io::Result<()> { + let MultipartReplicationContext { + storage, + cli, + src_bucket, + dst_bucket, + object, + object_info, + obj_opts, + arn, + put_opts, + } = ctx; let mut attempts = 1; let upload_id = loop { - match cli.create_multipart_upload(bucket, object, &opts).await { + match cli.create_multipart_upload(dst_bucket, object, &put_opts).await { Ok(id) => { break id; } @@ -2708,19 +2785,30 @@ async fn replicate_object_with_multipart( let mut uploaded_parts: Vec = Vec::new(); - let mut reader = reader; + let mut header_size = calc_put_object_header_size(&put_opts); + let mut offset: i64 = 0; for part_info in object_info.parts.iter() { - let mut chunk = vec![0u8; part_info.actual_size as usize]; - AsyncReadExt::read_exact(&mut *reader, &mut chunk).await?; + let part_size = part_info.actual_size; + let (range_spec, next_offset) = part_range_spec_from_actual_size(offset, part_size)?; + offset = next_offset; + + let part_reader = storage + .get_object_reader(src_bucket, object, Some(range_spec), HeaderMap::new(), obj_opts) + .await + .map_err(|e| std::io::Error::other(e.to_string()))?; + + let part_stream = wrap_with_bandwidth_monitor_with_header(part_reader.stream, src_bucket, arn, header_size); + header_size = 0; + let byte_stream = async_read_to_bytestream(part_stream); let object_part = cli .put_object_part( - bucket, + dst_bucket, object, &upload_id, part_info.number as i32, - part_info.actual_size, - ByteStream::from(chunk), + part_size, + byte_stream, &PutObjectPartOptions { ..Default::default() }, ) .await @@ -2748,7 +2836,7 @@ async fn replicate_object_with_multipart( ); cli.complete_multipart_upload( - bucket, + dst_bucket, object, &upload_id, uploaded_parts, @@ -2869,3 +2957,22 @@ fn get_replication_action(oi1: &ObjectInfo, oi2: &HeadObjectOutput, op_type: Rep ReplicationAction::None } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_part_range_spec_from_actual_size() { + let (rs, next) = part_range_spec_from_actual_size(0, 10).unwrap(); + assert_eq!(rs.start, 0); + assert_eq!(rs.end, 9); + assert_eq!(next, 10); + } + + #[test] + fn test_part_range_spec_rejects_non_positive() { + assert!(part_range_spec_from_actual_size(0, 0).is_err()); + assert!(part_range_spec_from_actual_size(0, -1).is_err()); + } +} diff --git a/crates/ecstore/src/global.rs b/crates/ecstore/src/global.rs index 724fe7fda..f73363e19 100644 --- a/crates/ecstore/src/global.rs +++ b/crates/ecstore/src/global.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use crate::bucket::bandwidth::monitor::Monitor; use crate::{ bucket::lifecycle::bucket_lifecycle_ops::LifecycleSys, disk::DiskStore, @@ -29,6 +30,7 @@ use std::{ }; use tokio::sync::{OnceCell, RwLock}; use tokio_util::sync::CancellationToken; +use tracing::warn; use uuid::Uuid; pub const DISK_ASSUME_UNKNOWN_SIZE: u64 = 1 << 30; @@ -58,6 +60,20 @@ lazy_static! { pub static ref GLOBAL_REGION: OnceLock = OnceLock::new(); pub static ref GLOBAL_LOCAL_LOCK_CLIENT: OnceLock> = OnceLock::new(); pub static ref GLOBAL_LOCK_CLIENTS: OnceLock>> = OnceLock::new(); + pub static ref GLOBAL_BUCKET_MONITOR: OnceLock> = OnceLock::new(); +} + +pub fn init_global_bucket_monitor(num_nodes: u64) { + if GLOBAL_BUCKET_MONITOR.set(Monitor::new(num_nodes)).is_err() { + warn!( + "global bucket monitor already initialized, ignoring re-initialization with num_nodes={}", + num_nodes + ); + } +} + +pub fn get_global_bucket_monitor() -> Option> { + GLOBAL_BUCKET_MONITOR.get().cloned() } /// Global cancellation token for background services (data scanner and auto heal) diff --git a/crates/ecstore/src/store.rs b/crates/ecstore/src/store.rs index 9a419c62c..e480b229a 100644 --- a/crates/ecstore/src/store.rs +++ b/crates/ecstore/src/store.rs @@ -39,8 +39,9 @@ use crate::error::{ }; use crate::global::{ DISK_ASSUME_UNKNOWN_SIZE, DISK_FILL_FRACTION, DISK_MIN_INODES, DISK_RESERVE_FRACTION, GLOBAL_BOOT_TIME, - GLOBAL_LOCAL_DISK_MAP, GLOBAL_LOCAL_DISK_SET_DRIVES, GLOBAL_TierConfigMgr, get_global_deployment_id, get_global_endpoints, - is_dist_erasure, is_erasure_sd, set_global_deployment_id, set_object_layer, + GLOBAL_LOCAL_DISK_MAP, GLOBAL_LOCAL_DISK_SET_DRIVES, GLOBAL_TierConfigMgr, get_global_bucket_monitor, + get_global_deployment_id, get_global_endpoints, init_global_bucket_monitor, is_dist_erasure, is_erasure_sd, + set_global_deployment_id, set_object_layer, }; use crate::notification_sys::get_global_notification_sys; use crate::pools::PoolMeta; @@ -405,6 +406,9 @@ impl ECStore { } } + let num_nodes = get_global_endpoints().get_nodes().len() as u64; + init_global_bucket_monitor(num_nodes); + init_background_expiry(self.clone()).await; TransitionState::init(self.clone()).await; @@ -1523,6 +1527,9 @@ impl StorageAPI for ECStore { // Delete the metadata self.delete_all(RUSTFS_META_BUCKET, format!("{BUCKET_META_PREFIX}/{bucket}").as_str()) .await?; + if let Some(monitor) = get_global_bucket_monitor() { + monitor.delete_bucket(bucket); + } Ok(()) }