feat(replication): add bandwidth-aware reporting for bucket replication metrics (#2141)

crates/ecstore/src/bucket/bandwidth/monitor.rs

@@ -14,11 +14,16 @@
 
 use crate::bucket::bandwidth::reader::BucketOptions;
 use ratelimit::{Error as RatelimitError, Ratelimiter};
+use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, Mutex, RwLock};
-use std::time::Duration;
+use std::time::{Duration, Instant};
 use tracing::warn;
 
+/// BETA_BUCKET is the weight used to calculate exponential moving average
+const BETA_BUCKET: f64 = 0.1;
+
 #[derive(Clone)]
 pub struct BucketThrottle {
     limiter: Arc<Mutex<Ratelimiter>>,
@@ -71,25 +76,202 @@ impl BucketThrottle {
     }
 }
 
+#[derive(Debug)]
+pub struct BucketMeasurement {
+    bytes_since_last_window: AtomicU64,
+    start_time: Mutex<Option<Instant>>,
+    exp_moving_avg: Mutex<f64>,
+}
+
+impl BucketMeasurement {
+    pub fn new(init_time: Instant) -> Self {
+        Self {
+            bytes_since_last_window: AtomicU64::new(0),
+            start_time: Mutex::new(Some(init_time)),
+            exp_moving_avg: Mutex::new(0.0),
+        }
+    }
+
+    pub fn increment_bytes(&self, bytes: u64) {
+        self.bytes_since_last_window.fetch_add(bytes, Ordering::Relaxed);
+    }
+
+    pub fn update_exponential_moving_average(&self, end_time: Instant) {
+        let mut start_time = self.start_time.lock().unwrap_or_else(|e| {
+            warn!("bucket measurement start_time mutex poisoned, recovering");
+            e.into_inner()
+        });
+        let previous_start = *start_time;
+        *start_time = Some(end_time);
+
+        let Some(prev_start) = previous_start else {
+            return;
+        };
+
+        if prev_start > end_time {
+            return;
+        }
+
+        let duration = end_time.duration_since(prev_start);
+        if duration.is_zero() {
+            return;
+        }
+
+        let bytes_since_last_window = self.bytes_since_last_window.swap(0, Ordering::Relaxed);
+        let increment = bytes_since_last_window as f64 / duration.as_secs_f64();
+
+        let mut exp_moving_avg = self.exp_moving_avg.lock().unwrap_or_else(|e| {
+            warn!("bucket measurement exp_moving_avg mutex poisoned, recovering");
+            e.into_inner()
+        });
+
+        if *exp_moving_avg == 0.0 {
+            *exp_moving_avg = increment;
+            return;
+        }
+
+        *exp_moving_avg = exponential_moving_average(BETA_BUCKET, *exp_moving_avg, increment);
+    }
+
+    pub fn get_exp_moving_avg_bytes_per_second(&self) -> f64 {
+        *self.exp_moving_avg.lock().unwrap_or_else(|e| {
+            warn!("bucket measurement exp_moving_avg mutex poisoned, recovering");
+            e.into_inner()
+        })
+    }
+}
+
+fn exponential_moving_average(beta: f64, previous_avg: f64, increment_avg: f64) -> f64 {
+    (1f64 - beta) * increment_avg + beta * previous_avg
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct BandwidthDetails {
+    pub limit_bytes_per_sec: i64,
+    pub current_bandwidth_bytes_per_sec: f64,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct BucketBandwidthReport {
+    pub bucket_stats: HashMap<BucketOptions, BandwidthDetails>,
+}
+
 pub struct Monitor {
     t_lock: RwLock<HashMap<BucketOptions, BucketThrottle>>,
+    m_lock: RwLock<HashMap<BucketOptions, BucketMeasurement>>,
     pub node_count: u64,
 }
 
 impl Monitor {
     pub fn new(num_nodes: u64) -> Arc<Self> {
         let node_cnt = num_nodes.max(1);
-        Arc::new(Monitor {
+        let m = Arc::new(Monitor {
             t_lock: RwLock::new(HashMap::new()),
+            m_lock: RwLock::new(HashMap::new()),
             node_count: node_cnt,
-        })
+        });
+        if let Ok(handle) = tokio::runtime::Handle::try_current() {
+            let weak = Arc::downgrade(&m);
+            handle.spawn(async move {
+                let mut interval = tokio::time::interval(Duration::from_secs(2));
+                loop {
+                    interval.tick().await;
+                    let Some(monitor) = weak.upgrade() else { break };
+                    monitor.update_moving_avg();
+                }
+            });
+        }
+        m
+    }
+
+    pub fn init_measurement(&self, opts: &BucketOptions) {
+        let mut guard = self.m_lock.write().unwrap_or_else(|e| {
+            warn!("bucket monitor measurement rwlock write poisoned, recovering");
+            e.into_inner()
+        });
+        guard
+            .entry(opts.clone())
+            .or_insert_with(|| BucketMeasurement::new(Instant::now()));
+    }
+
+    pub fn update_measurement(&self, opts: &BucketOptions, bytes: u64) {
+        {
+            let guard = self.m_lock.read().unwrap_or_else(|e| {
+                warn!("bucket monitor measurement rwlock read poisoned, recovering");
+                e.into_inner()
+            });
+            if let Some(measurement) = guard.get(opts) {
+                measurement.increment_bytes(bytes);
+                return;
+            }
+        }
+
+        // Miss path: write lock + insert once, then increment.
+        let mut guard = self.m_lock.write().unwrap_or_else(|e| {
+            warn!("bucket monitor measurement rwlock write poisoned, recovering");
+            e.into_inner()
+        });
+
+        // Double-check after lock upgrade in case another thread inserted it.
+        let measurement = guard
+            .entry(opts.clone())
+            .or_insert_with(|| BucketMeasurement::new(Instant::now()));
+
+        measurement.increment_bytes(bytes);
+    }
+
+    pub fn update_moving_avg(&self) {
+        let now = Instant::now();
+        let guard = self.m_lock.read().unwrap_or_else(|e| {
+            warn!("bucket monitor measurement rwlock read poisoned, recovering");
+            e.into_inner()
+        });
+        for measurement in guard.values() {
+            measurement.update_exponential_moving_average(now);
+        }
+    }
+
+    pub fn get_report(&self, select_bucket: impl Fn(&str) -> bool) -> BucketBandwidthReport {
+        let t_guard = self.t_lock.read().unwrap_or_else(|e| {
+            warn!("bucket monitor throttle rwlock read poisoned, recovering");
+            e.into_inner()
+        });
+        let m_guard = self.m_lock.read().unwrap_or_else(|e| {
+            warn!("bucket monitor measurement rwlock read poisoned, recovering");
+            e.into_inner()
+        });
+        let mut bucket_stats = HashMap::new();
+        for (opts, throttle) in t_guard.iter() {
+            if !select_bucket(&opts.name) {
+                continue;
+            }
+            let mut current_bandwidth_bytes_per_sec = 0.0;
+            if let Some(measurement) = m_guard.get(opts) {
+                current_bandwidth_bytes_per_sec = measurement.get_exp_moving_avg_bytes_per_second();
+            }
+            bucket_stats.insert(
+                opts.clone(),
+                BandwidthDetails {
+                    limit_bytes_per_sec: throttle.node_bandwidth_per_sec * self.node_count as i64,
+                    current_bandwidth_bytes_per_sec,
+                },
+            );
+        }
+        BucketBandwidthReport { bucket_stats }
     }
 
     pub fn delete_bucket(&self, bucket: &str) {
         self.t_lock
             .write()
             .unwrap_or_else(|e| {
-                warn!("bucket monitor rwlock write poisoned, recovering");
+                warn!("bucket monitor throttle rwlock write poisoned, recovering");
+                e.into_inner()
+            })
+            .retain(|opts, _| opts.name != bucket);
+        self.m_lock
+            .write()
+            .unwrap_or_else(|e| {
+                warn!("bucket monitor measurement rwlock write poisoned, recovering");
                 e.into_inner()
             })
             .retain(|opts, _| opts.name != bucket);
@@ -103,7 +285,14 @@ impl Monitor {
         self.t_lock
             .write()
             .unwrap_or_else(|e| {
-                warn!("bucket monitor rwlock write poisoned, recovering");
+                warn!("bucket monitor throttle rwlock write poisoned, recovering");
+                e.into_inner()
+            })
+            .remove(&opts);
+        self.m_lock
+            .write()
+            .unwrap_or_else(|e| {
+                warn!("bucket monitor measurement rwlock write poisoned, recovering");
                 e.into_inner()
             })
             .remove(&opts);
@@ -113,7 +302,7 @@ impl Monitor {
         self.t_lock
             .read()
             .unwrap_or_else(|e| {
-                warn!("bucket monitor rwlock read poisoned, recovering");
+                warn!("bucket monitor throttle rwlock read poisoned, recovering");
                 e.into_inner()
             })
             .get(opts)
@@ -160,7 +349,7 @@ impl Monitor {
         self.t_lock
             .write()
             .unwrap_or_else(|e| {
-                warn!("bucket monitor rwlock write poisoned, recovering");
+                warn!("bucket monitor throttle rwlock write poisoned, recovering");
                 e.into_inner()
             })
             .insert(opts, throttle);
@@ -174,7 +363,7 @@ impl Monitor {
         self.t_lock
             .read()
             .unwrap_or_else(|e| {
-                warn!("bucket monitor rwlock read poisoned, recovering");
+                warn!("bucket monitor throttle rwlock read poisoned, recovering");
                 e.into_inner()
             })
             .contains_key(&opt)
@@ -184,6 +373,7 @@ impl Monitor {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::panic::{AssertUnwindSafe, catch_unwind};
 
     #[test]
     fn test_set_and_get_throttle_with_node_split() {
@@ -318,4 +508,82 @@ mod tests {
         assert_eq!(t2.burst(), 500);
         assert_eq!(t2.node_bandwidth_per_sec, 500);
     }
+
+    #[test]
+    fn test_bucket_measurement_recovers_from_poisoned_mutexes() {
+        let measurement = BucketMeasurement::new(Instant::now());
+
+        let _ = catch_unwind(AssertUnwindSafe(|| {
+            let _guard = measurement.start_time.lock().unwrap();
+            panic!("poison start_time mutex");
+        }));
+        measurement.increment_bytes(64);
+        measurement.update_exponential_moving_average(Instant::now() + Duration::from_secs(1));
+
+        let _ = catch_unwind(AssertUnwindSafe(|| {
+            let _guard = measurement.exp_moving_avg.lock().unwrap();
+            panic!("poison exp_moving_avg mutex");
+        }));
+        measurement.increment_bytes(32);
+        measurement.update_exponential_moving_average(Instant::now() + Duration::from_secs(2));
+
+        let value = measurement.get_exp_moving_avg_bytes_per_second();
+        assert!(value.is_finite());
+        assert!(value >= 0.0);
+    }
+
+    #[test]
+    fn test_get_report_limit_and_current_bandwidth_after_measurement() {
+        let monitor = Monitor::new(4);
+        monitor.set_bandwidth_limit("b1", "arn1", 400);
+        let opts = BucketOptions {
+            name: "b1".to_string(),
+            replication_arn: "arn1".to_string(),
+        };
+        monitor.init_measurement(&opts);
+        monitor.update_measurement(&opts, 500);
+        monitor.update_measurement(&opts, 500);
+        std::thread::sleep(Duration::from_millis(110));
+        monitor.update_moving_avg();
+
+        let report = monitor.get_report(|name| name == "b1");
+        let details = report.bucket_stats.get(&opts).expect("report should contain b1/arn1");
+        assert_eq!(details.limit_bytes_per_sec, 400);
+        assert!(
+            details.current_bandwidth_bytes_per_sec > 0.0,
+            "current_bandwidth should be positive after update_measurement and update_moving_avg"
+        );
+        assert!(
+            details.current_bandwidth_bytes_per_sec < 20000.0,
+            "current_bandwidth should be in reasonable range"
+        );
+    }
+
+    #[test]
+    fn test_get_report_select_bucket_filters() {
+        let monitor = Monitor::new(2);
+        monitor.set_bandwidth_limit("b1", "arn1", 100);
+        monitor.set_bandwidth_limit("b2", "arn2", 200);
+        let opts_b1 = BucketOptions {
+            name: "b1".to_string(),
+            replication_arn: "arn1".to_string(),
+        };
+        let opts_b2 = BucketOptions {
+            name: "b2".to_string(),
+            replication_arn: "arn2".to_string(),
+        };
+        monitor.init_measurement(&opts_b1);
+        monitor.init_measurement(&opts_b2);
+
+        let report_all = monitor.get_report(|_| true);
+        assert_eq!(report_all.bucket_stats.len(), 2);
+
+        let report_b1 = monitor.get_report(|name| name == "b1");
+        assert_eq!(report_b1.bucket_stats.len(), 1);
+        assert_eq!(report_b1.bucket_stats.get(&opts_b1).unwrap().limit_bytes_per_sec, 100);
+
+        let report_b2 = monitor.get_report(|name| name == "b2");
+        assert_eq!(report_b2.bucket_stats.len(), 1);
+        assert_eq!(report_b2.bucket_stats.get(&opts_b2).unwrap().limit_bytes_per_sec, 200);
+    }
 }

crates/ecstore/src/bucket/bandwidth/reader.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use crate::bucket::bandwidth::monitor::Monitor;
+use serde::{Deserialize, Serialize};
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
@@ -20,7 +21,7 @@ use tokio::io::{AsyncRead, ReadBuf};
 use tokio::time::Sleep;
 use tracing::{debug, warn};
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct BucketOptions {
     pub name: String,
     pub replication_arn: String,
@@ -54,6 +55,9 @@ impl<R> MonitoredReader<R> {
             header_size = opts.header_size,
             "MonitoredReader created"
         );
+        if throttle.is_some() {
+            m.init_measurement(&opts.bucket_options);
+        }
         MonitoredReader {
             r,
             m,
@@ -117,7 +121,15 @@ impl<R: AsyncRead + Unpin> AsyncRead for MonitoredReader<R> {
             }
         }
 
-        poll_limited_read(&mut this.r, cx, buf, need, &mut this.temp_buf)
+        let filled_before = buf.filled().len();
+        let result = poll_limited_read(&mut this.r, cx, buf, need, &mut this.temp_buf);
+        if let Poll::Ready(Ok(())) = result {
+            let read_bytes = buf.filled().len().saturating_sub(filled_before) as u64;
+            if read_bytes > 0 {
+                this.m.update_measurement(&this.opts.bucket_options, read_bytes);
+            }
+        }
+        result
     }
 }