diff --git a/Cargo.lock b/Cargo.lock index 52f2e505..35eeef86 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2201,6 +2201,7 @@ dependencies = [ "argh", "axum", "axum-extra", + "bytes", "console", "elegant-departure", "figment", diff --git a/objectstore-server/Cargo.toml b/objectstore-server/Cargo.toml index 7165cbe0..ad18fd41 100644 --- a/objectstore-server/Cargo.toml +++ b/objectstore-server/Cargo.toml @@ -15,6 +15,7 @@ anyhow = { workspace = true } argh = "0.1.13" axum = "0.8.4" axum-extra = "0.10.1" +bytes = { workspace = true } console = "0.16.1" elegant-departure = { version = "0.3.2", features = ["tokio"] } figment = { version = "0.10.19", features = ["env", "test", "yaml"] } diff --git a/objectstore-server/src/endpoints/objects.rs b/objectstore-server/src/endpoints/objects.rs index 66803d74..783551f2 100644 --- a/objectstore-server/src/endpoints/objects.rs +++ b/objectstore-server/src/endpoints/objects.rs @@ -1,8 +1,9 @@ +use anyhow::Context as _; use std::io; use std::time::SystemTime; -use anyhow::Context; use axum::body::Body; +use axum::extract::State; use axum::http::{HeaderMap, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::routing; @@ -15,6 +16,7 @@ use serde::Serialize; use crate::auth::AuthAwareService; use crate::endpoints::common::ApiResult; use crate::extractors::Xt; +use crate::rate_limits::MeteredPayloadStream; use crate::state::ServiceState; pub fn router() -> Router { @@ -39,6 +41,7 @@ pub struct InsertObjectResponse { async fn objects_post( service: AuthAwareService, + State(state): State, Xt(context): Xt, headers: HeaderMap, body: Body, @@ -48,6 +51,8 @@ async fn objects_post( metadata.time_created = Some(SystemTime::now()); let stream = body.into_data_stream().map_err(io::Error::other).boxed(); + let stream = MeteredPayloadStream::from(stream, state.rate_limiter.bytes_accumulator()).boxed(); + let response_id = service .insert_object(context, None, &metadata, stream) .await?; @@ -58,10 +63,15 @@ async fn objects_post( Ok((StatusCode::CREATED, response).into_response()) } -async fn object_get(service: AuthAwareService, Xt(id): Xt) -> ApiResult { +async fn object_get( + service: AuthAwareService, + State(state): State, + Xt(id): Xt, +) -> ApiResult { let Some((metadata, stream)) = service.get_object(&id).await? else { return Ok(StatusCode::NOT_FOUND.into_response()); }; + let stream = MeteredPayloadStream::from(stream, state.rate_limiter.bytes_accumulator()).boxed(); let headers = metadata .to_headers("", false) @@ -83,6 +93,7 @@ async fn object_head(service: AuthAwareService, Xt(id): Xt) -> ApiResu async fn object_put( service: AuthAwareService, + State(state): State, Xt(id): Xt, headers: HeaderMap, body: Body, @@ -93,6 +104,7 @@ async fn object_put( let ObjectId { context, key } = id; let stream = body.into_data_stream().map_err(io::Error::other).boxed(); + let stream = MeteredPayloadStream::from(stream, state.rate_limiter.bytes_accumulator()).boxed(); let response_id = service .insert_object(context, Some(key), &metadata, stream) .await?; diff --git a/objectstore-server/src/rate_limits.rs b/objectstore-server/src/rate_limits.rs index 9972883b..d405a4fd 100644 --- a/objectstore-server/src/rate_limits.rs +++ b/objectstore-server/src/rate_limits.rs @@ -1,6 +1,13 @@ +use std::pin::Pin; +use std::sync::Arc; use std::sync::Mutex; -use std::time::Instant; +use std::sync::atomic::AtomicU64; +use std::task::{Context, Poll}; +use std::time::{Duration, Instant}; +use bytes::Bytes; +use futures_util::Stream; +use objectstore_service::PayloadStream; use objectstore_service::id::ObjectContext; use objectstore_types::scope::Scopes; use serde::{Deserialize, Serialize}; @@ -102,6 +109,7 @@ pub struct BandwidthLimits { #[derive(Debug)] pub struct RateLimiter { config: RateLimits, + bandwidth: BandwidthRateLimiter, global: Option>, // NB: These maps grow unbounded but we accept this as we expect an overall limited // number of usecases and scopes. We emit gauge metrics to monitor their size. @@ -110,6 +118,64 @@ pub struct RateLimiter { rules: papaya::HashMap>, } +#[derive(Debug)] +struct BandwidthRateLimiter { + config: BandwidthLimits, + /// Accumulator that's incremented every time an operation that uses bandwidth is executed. + accumulator: Arc, + /// An estimate of the bandwidth that's currently being utilized in bytes per second. + estimate: Arc, +} + +impl BandwidthRateLimiter { + fn new(config: BandwidthLimits) -> Self { + let accumulator = Arc::new(AtomicU64::new(0)); + let estimate = Arc::new(AtomicU64::new(0)); + + let accumulator_clone = Arc::clone(&accumulator); + let estimate_clone = Arc::clone(&estimate); + tokio::task::spawn(async move { + Self::estimator(accumulator_clone, estimate_clone).await; + }); + + Self { + config, + accumulator, + estimate, + } + } + + /// Estimates the current bandwidth utilization using an exponentially weighted moving average. + /// + /// The calculation is based on the increments of `self.accumulator` happened in the last `TICK`. + /// The estimate is stored in `self.estimate`, which can be queried for bandwidth-based rate-limiting. + async fn estimator(accumulator: Arc, estimate: Arc) { + const TICK: Duration = Duration::from_millis(50); // Recompute EWMA on every TICK + const ALPHA: f64 = 0.2; // EWMA alpha parameter: 20% weight to new sample, 80% to previous average + + let mut interval = tokio::time::interval(TICK); + let to_bps = 1.0 / TICK.as_secs_f64(); // Conversion factor from bytes to bps + let mut ewma: f64 = 0.0; + loop { + interval.tick().await; + let current = accumulator.swap(0, std::sync::atomic::Ordering::Relaxed); + let bps = (current as f64) * to_bps; + ewma = ALPHA * bps + (1.0 - ALPHA) * ewma; + + let ewma_int = ewma.floor() as u64; + estimate.store(ewma_int, std::sync::atomic::Ordering::Relaxed); + merni::gauge!("server.bandwidth.ewma"@b: ewma_int); + } + } + + fn check(&self) -> bool { + let Some(bps) = self.config.global_bps else { + return true; + }; + self.estimate.load(std::sync::atomic::Ordering::Relaxed) <= bps + } +} + impl RateLimiter { pub fn new(config: RateLimits) -> Self { let global = config @@ -118,7 +184,8 @@ impl RateLimiter { .map(|rps| Mutex::new(TokenBucket::new(rps, config.throughput.burst))); Self { - config, + config: config.clone(), + bandwidth: BandwidthRateLimiter::new(config.bandwidth), global, usecases: papaya::HashMap::new(), scopes: papaya::HashMap::new(), @@ -126,11 +193,16 @@ impl RateLimiter { } } + /// Returns a reference to the shared bytes accumulator, used for bandwidth-based rate-limiting. + pub fn bytes_accumulator(&self) -> Arc { + Arc::clone(&self.bandwidth.accumulator) + } + /// Checks if the given context is within the rate limits. /// /// Returns `true` if the context is within the rate limits, `false` otherwise. pub fn check(&self, context: &ObjectContext) -> bool { - self.check_throughput(context) && self.check_bandwidth(context) + self.check_throughput(context) && self.bandwidth.check() } fn check_throughput(&self, context: &ObjectContext) -> bool { @@ -181,11 +253,6 @@ impl RateLimiter { true } - fn check_bandwidth(&self, _context: &ObjectContext) -> bool { - // Not implemented - true - } - fn create_bucket(&self, rps: u32) -> Mutex { Mutex::new(TokenBucket::new(rps, self.config.throughput.burst)) } @@ -273,3 +340,40 @@ impl TokenBucket { } } } + +/// A wrapper around a `PayloadStream` that measures bandwidth usage. +/// +/// This behaves exactly as a `PayloadStream`, except that every time an item is polled, +/// the accumulator is incremented by the size of the returned `Bytes` chunk. +pub(crate) struct MeteredPayloadStream { + inner: PayloadStream, + accumulator: Arc, +} + +impl MeteredPayloadStream { + pub fn from(inner: PayloadStream, accumulator: Arc) -> Self { + Self { inner, accumulator } + } +} + +impl std::fmt::Debug for MeteredPayloadStream { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("MeteredPayloadStream") + .field("accumulator", &self.accumulator) + .finish() + } +} + +impl Stream for MeteredPayloadStream { + type Item = std::io::Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let res = this.inner.as_mut().poll_next(cx); + if let Poll::Ready(Some(Ok(ref bytes))) = res { + this.accumulator + .fetch_add(bytes.len() as u64, std::sync::atomic::Ordering::Relaxed); + } + res + } +}