fix: harden concurrency limits and high-RPM runtime paths

Bound request, stream, queue, and shutdown resource lifetimes. Reduce scheduler and Redis hot-path work and isolate database maintenance. Include regression coverage, load probes, and concurrency audit results.
2026-10-06 01:17:46 +08:00 · 2026-09-10 08:14:58 +08:00
parent 361952ada9
commit ecc16673eb
149 changed files with 27963 additions and 1926 deletions
@@ -0,0 +1,373 @@
+use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+use std::sync::{Arc, LazyLock};
+
+use aether_data_contracts::DataLayerError;
+use aether_runtime_state::RuntimeQueueEntry;
+use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError};
+
+const DEFAULT_READ_PAYLOAD_BUDGET_BYTES: usize = 128 * 1024 * 1024;
+const DEFAULT_READ_BATCH_PAYLOAD_BYTES: usize = 8 * 1024 * 1024;
+
+static READ_BUDGET: LazyLock<Arc<QueueReadBudget>> = LazyLock::new(|| {
+    let limit = std::env::var("AETHER_USAGE_QUEUE_READ_PAYLOAD_BUDGET_BYTES").ok();
+    let batch = std::env::var("AETHER_USAGE_QUEUE_READ_BATCH_PAYLOAD_BYTES").ok();
+    Arc::new(QueueReadBudget::new(
+        configured_bytes(limit.as_deref(), DEFAULT_READ_PAYLOAD_BUDGET_BYTES),
+        configured_bytes(batch.as_deref(), DEFAULT_READ_BATCH_PAYLOAD_BYTES),
+    ))
+});
+
+pub(crate) fn shared_queue_read_budget() -> Arc<QueueReadBudget> {
+    Arc::clone(&READ_BUDGET)
+}
+
+pub(crate) fn queue_read_budget_metrics() -> QueueReadBudgetSnapshot {
+    READ_BUDGET.snapshot()
+}
+
+fn maximum_budget_bytes() -> usize {
+    Semaphore::MAX_PERMITS.min(u32::MAX as usize)
+}
+
+fn configured_bytes(raw: Option<&str>, fallback: usize) -> usize {
+    raw.and_then(|raw| raw.trim().parse::<u128>().ok())
+        .filter(|value| *value > 0)
+        .map(|value| value.min(maximum_budget_bytes() as u128) as usize)
+        .unwrap_or(fallback.min(maximum_budget_bytes()))
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct QueueReadBudgetSnapshot {
+    pub(crate) limit_bytes: usize,
+    pub(crate) batch_limit_bytes: usize,
+    pub(crate) reserved_bytes: usize,
+    pub(crate) waiters: usize,
+    pub(crate) wait_total: u64,
+    /// Observed field names plus values, without cloning the returned strings.
+    pub(crate) actual_field_bytes_total: u64,
+    pub(crate) oversized_entries_total: u64,
+    pub(crate) oversized_batches_total: u64,
+}
+
+/// A process-wide reservation based on the current producer payload limit.
+/// Historical or externally written messages can exceed the estimate. Field names,
+/// allocation capacity, RESP decoding, and decoded JSON are not an RSS bound here.
+pub(crate) struct QueueReadBudget {
+    limit_bytes: usize,
+    batch_limit_bytes: usize,
+    permits: Arc<Semaphore>,
+    reserved_bytes: AtomicUsize,
+    waiters: AtomicUsize,
+    wait_total: AtomicU64,
+    actual_field_bytes_total: AtomicU64,
+    oversized_entries_total: AtomicU64,
+    oversized_batches_total: AtomicU64,
+}
+
+impl QueueReadBudget {
+    pub(crate) fn new(limit_bytes: usize, batch_limit_bytes: usize) -> Self {
+        let limit_bytes = limit_bytes.clamp(1, maximum_budget_bytes());
+        let batch_limit_bytes = batch_limit_bytes.clamp(1, limit_bytes);
+        Self {
+            limit_bytes,
+            batch_limit_bytes,
+            permits: Arc::new(Semaphore::new(limit_bytes)),
+            reserved_bytes: AtomicUsize::new(0),
+            waiters: AtomicUsize::new(0),
+            wait_total: AtomicU64::new(0),
+            actual_field_bytes_total: AtomicU64::new(0),
+            oversized_entries_total: AtomicU64::new(0),
+            oversized_batches_total: AtomicU64::new(0),
+        }
+    }
+
+    pub(crate) fn snapshot(&self) -> QueueReadBudgetSnapshot {
+        QueueReadBudgetSnapshot {
+            limit_bytes: self.limit_bytes,
+            batch_limit_bytes: self.batch_limit_bytes,
+            reserved_bytes: self.reserved_bytes.load(Ordering::Relaxed),
+            waiters: self.waiters.load(Ordering::Relaxed),
+            wait_total: self.wait_total.load(Ordering::Relaxed),
+            actual_field_bytes_total: self.actual_field_bytes_total.load(Ordering::Relaxed),
+            oversized_entries_total: self.oversized_entries_total.load(Ordering::Relaxed),
+            oversized_batches_total: self.oversized_batches_total.load(Ordering::Relaxed),
+        }
+    }
+
+    pub(crate) async fn reserve(
+        self: &Arc<Self>,
+        requested_count: usize,
+        payload_limit: usize,
+    ) -> Result<(usize, QueueReadReservation), DataLayerError> {
+        if payload_limit == 0 || payload_limit > self.limit_bytes {
+            return Err(DataLayerError::InvalidConfiguration(format!(
+                "usage queue payload limit {payload_limit} must be positive and not exceed the {}-byte read payload budget",
+                self.limit_bytes
+            )));
+        }
+        // A single valid payload may exceed the preferred batch target, but never
+        // the total budget. Clamp before multiplying or converting to u32 permits.
+        let count = requested_count
+            .max(1)
+            .min((self.batch_limit_bytes / payload_limit).max(1));
+        let reserved_bytes = count * payload_limit;
+        let permits = reserved_bytes as u32;
+        let permit = match Arc::clone(&self.permits).try_acquire_many_owned(permits) {
+            Ok(permit) => permit,
+            Err(TryAcquireError::NoPermits) => {
+                self.wait_total.fetch_add(1, Ordering::Relaxed);
+                self.waiters.fetch_add(1, Ordering::Relaxed);
+                let _waiting = WaitingReservation { budget: self };
+                Arc::clone(&self.permits)
+                    .acquire_many_owned(permits)
+                    .await
+                    .map_err(|_| closed_budget_error())?
+            }
+            Err(TryAcquireError::Closed) => return Err(closed_budget_error()),
+        };
+        self.reserved_bytes
+            .fetch_add(reserved_bytes, Ordering::Relaxed);
+        Ok((
+            count,
+            QueueReadReservation {
+                budget: Arc::clone(self),
+                reserved_bytes,
+                permit: Some(permit),
+            },
+        ))
+    }
+}
+
+fn closed_budget_error() -> DataLayerError {
+    DataLayerError::InvalidConfiguration("usage queue read payload budget is closed".to_string())
+}
+
+struct WaitingReservation<'a> {
+    budget: &'a QueueReadBudget,
+}
+
+impl Drop for WaitingReservation<'_> {
+    fn drop(&mut self) {
+        self.budget.waiters.fetch_sub(1, Ordering::Relaxed);
+    }
+}
+
+// Deliberately not Clone: every concurrently retained batch needs its own lease.
+pub(crate) struct QueueReadReservation {
+    budget: Arc<QueueReadBudget>,
+    reserved_bytes: usize,
+    permit: Option<OwnedSemaphorePermit>,
+}
+
+impl QueueReadReservation {
+    pub(crate) fn observe_entries(&mut self, entries: &[RuntimeQueueEntry], payload_limit: usize) {
+        let mut value_bytes = 0usize;
+        let mut field_bytes = 0usize;
+        let mut oversized_entries = 0u64;
+        for entry in entries {
+            let mut entry_value_bytes = 0usize;
+            for (key, value) in &entry.fields {
+                entry_value_bytes = entry_value_bytes.saturating_add(value.len());
+                field_bytes = field_bytes
+                    .saturating_add(key.len())
+                    .saturating_add(value.len());
+            }
+            value_bytes = value_bytes.saturating_add(entry_value_bytes);
+            oversized_entries += u64::from(entry_value_bytes > payload_limit);
+        }
+        self.budget.actual_field_bytes_total.fetch_add(
+            u64::try_from(field_bytes).unwrap_or(u64::MAX),
+            Ordering::Relaxed,
+        );
+        self.budget
+            .oversized_entries_total
+            .fetch_add(oversized_entries, Ordering::Relaxed);
+        if value_bytes > self.reserved_bytes {
+            self.budget
+                .oversized_batches_total
+                .fetch_add(1, Ordering::Relaxed);
+        }
+        // Shrink unused payload estimates. Never wait for an upgrade after reading
+        // an oversized historical batch: other batches may hold all remaining bytes.
+        self.shrink_to(value_bytes.min(self.reserved_bytes));
+    }
+
+    fn shrink_to(&mut self, retained_bytes: usize) {
+        let released = self.reserved_bytes.saturating_sub(retained_bytes);
+        if released == 0 {
+            return;
+        }
+        let permit = self
+            .permit
+            .as_mut()
+            .expect("positive reservation must hold a permit")
+            .split(released)
+            .expect("released bytes must belong to this reservation");
+        self.reserved_bytes -= released;
+        self.budget
+            .reserved_bytes
+            .fetch_sub(released, Ordering::Relaxed);
+        drop(permit);
+    }
+}
+
+impl Drop for QueueReadReservation {
+    fn drop(&mut self) {
+        self.budget
+            .reserved_bytes
+            .fetch_sub(self.reserved_bytes, Ordering::Relaxed);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeMap;
+    use std::future::Future;
+    use std::task::Poll;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn queue_read_budget_cancelled_wait_preserves_current_reservations() {
+        let budget = Arc::new(QueueReadBudget::new(32, 16));
+        let (_, first) = budget.reserve(8, 8).await.unwrap();
+        let (_, second) = budget.reserve(8, 8).await.unwrap();
+        let mut pending = Box::pin(budget.reserve(1, 8));
+        std::future::poll_fn(|cx| {
+            assert!(pending.as_mut().poll(cx).is_pending());
+            Poll::Ready(())
+        })
+        .await;
+        assert_eq!(budget.snapshot().reserved_bytes, 32);
+        assert_eq!(budget.snapshot().waiters, 1);
+        drop(pending);
+        assert_eq!(budget.snapshot().waiters, 0);
+        assert_eq!(budget.snapshot().wait_total, 1);
+        drop(first);
+        let (_, replacement) = budget.reserve(2, 8).await.unwrap();
+        assert_eq!(budget.snapshot().reserved_bytes, 32);
+        drop((second, replacement));
+        assert_eq!(budget.snapshot().reserved_bytes, 0);
+        assert_eq!(budget.permits.available_permits(), 32);
+    }
+
+    #[tokio::test]
+    async fn queue_read_budget_counts_payload_values_and_observes_legacy_excess_without_waiting() {
+        let budget = Arc::new(QueueReadBudget::new(16, 16));
+        let (count, mut reservation) = budget.reserve(2, 8).await.unwrap();
+        assert_eq!(count, 2);
+        let entries = [RuntimeQueueEntry {
+            id: "1-0".to_string(),
+            fields: BTreeMap::from([("payload".to_string(), "x".repeat(8))]),
+        }];
+        reservation.observe_entries(&entries, 8);
+        assert_eq!(budget.snapshot().reserved_bytes, 8);
+        assert_eq!(budget.snapshot().actual_field_bytes_total, 15);
+        assert_eq!(budget.snapshot().oversized_entries_total, 0);
+        let (_, mut legacy) = budget.reserve(1, 8).await.unwrap();
+        let legacy_entries = [RuntimeQueueEntry {
+            id: "2-0".to_string(),
+            fields: BTreeMap::from([
+                ("payload".to_string(), "x".repeat(8)),
+                ("extra".to_string(), "y".repeat(24)),
+            ]),
+        }];
+        legacy.observe_entries(&legacy_entries, 8);
+        assert_eq!(budget.snapshot().reserved_bytes, 16);
+        assert_eq!(budget.snapshot().actual_field_bytes_total, 59);
+        assert_eq!(budget.snapshot().oversized_entries_total, 1);
+        assert_eq!(budget.snapshot().oversized_batches_total, 1);
+        assert_eq!(budget.snapshot().wait_total, 0);
+        drop((reservation, legacy));
+        assert_eq!(budget.snapshot().reserved_bytes, 0);
+    }
+
+    #[tokio::test]
+    async fn queue_read_budget_empty_response_releases_all_reserved_bytes() {
+        let budget = Arc::new(QueueReadBudget::new(16, 16));
+        let (_, mut reservation) = budget.reserve(2, 8).await.unwrap();
+        reservation.observe_entries(&[], 8);
+        assert_eq!(budget.snapshot().reserved_bytes, 0);
+        assert_eq!(budget.permits.available_permits(), 16);
+        drop(reservation);
+        assert_eq!(budget.snapshot().reserved_bytes, 0);
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+    async fn queue_read_budget_concurrent_shrink_and_drop_preserve_shared_capacity() {
+        let budget = Arc::new(QueueReadBudget::new(64, 16));
+        let barrier = Arc::new(tokio::sync::Barrier::new(16));
+        let mut tasks = tokio::task::JoinSet::new();
+        for _ in 0..16 {
+            let budget = Arc::clone(&budget);
+            let barrier = Arc::clone(&barrier);
+            tasks.spawn(async move {
+                barrier.wait().await;
+                for _ in 0..16 {
+                    let (count, mut reservation) = budget.reserve(128, 8).await.unwrap();
+                    assert_eq!(count, 2);
+                    assert!(budget.snapshot().reserved_bytes <= 64);
+                    reservation.observe_entries(
+                        &[RuntimeQueueEntry {
+                            id: "1-0".to_string(),
+                            fields: BTreeMap::from([(
+                                "payload".to_string(),
+                                "12345678".to_string(),
+                            )]),
+                        }],
+                        8,
+                    );
+                    tokio::task::yield_now().await;
+                    drop(reservation);
+                }
+            });
+        }
+        tokio::time::timeout(std::time::Duration::from_secs(5), async {
+            while let Some(result) = tasks.join_next().await {
+                result.expect("reservation task");
+            }
+        })
+        .await
+        .expect("all shared reservations complete");
+        assert_eq!(budget.snapshot().reserved_bytes, 0);
+        assert_eq!(budget.snapshot().waiters, 0);
+        assert_eq!(budget.snapshot().actual_field_bytes_total, 16 * 16 * 15);
+        assert_eq!(budget.permits.available_permits(), 64);
+    }
+
+    #[tokio::test]
+    async fn queue_read_budget_large_values_cannot_overflow_or_wait_for_impossible_permits() {
+        let budget = Arc::new(QueueReadBudget::new(usize::MAX, usize::MAX));
+        let maximum = maximum_budget_bytes();
+        assert_eq!(budget.snapshot().limit_bytes, maximum);
+        let (count, reservation) = budget.reserve(usize::MAX, 1).await.unwrap();
+        assert_eq!(count, maximum);
+        assert_eq!(budget.snapshot().reserved_bytes, maximum);
+        drop(reservation);
+        assert!(matches!(
+            budget.reserve(1, maximum + 1).await,
+            Err(DataLayerError::InvalidConfiguration(_))
+        ));
+        assert!(matches!(
+            budget.reserve(1, 0).await,
+            Err(DataLayerError::InvalidConfiguration(_))
+        ));
+        let small_batch = Arc::new(QueueReadBudget::new(32, 4));
+        let (count, reservation) = small_batch.reserve(usize::MAX, 16).await.unwrap();
+        assert_eq!(count, 1);
+        assert_eq!(small_batch.snapshot().reserved_bytes, 16);
+        drop(reservation);
+    }
+
+    #[test]
+    fn queue_read_budget_env_uses_positive_defaults_and_caps_extreme_values() {
+        for raw in [None, Some(""), Some("0"), Some("-1"), Some("bad")] {
+            assert_eq!(configured_bytes(raw, 128), 128);
+        }
+        assert_eq!(configured_bytes(Some(" 42 "), 128), 42);
+        assert_eq!(
+            configured_bytes(Some(&u128::MAX.to_string()), 128),
+            maximum_budget_bytes()
+        );
+    }
+}