fix: automatically manage gateway runtime and memory (#134)

Automatically size CPU execution and buffering from usable CPUs and container memory. Bound AI, MCP, cache, discovery and audit buffering, retain admission through cancellation and slow consumers, and protect Argon2 workspaces. Verified with 939 unit tests and 492 integration tests.

Author: fylorn

Cargo.lock

@@ -70,10 +70,10 @@ opt-level = 3
 # never re-exported through a local shim. And the reverse: something only
 # this side uses (the at-rest crypto, IMDSv2 credentials, the gateway error)
 # lives here, not in core.
-tw-bedrock = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.86.0" }
-tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.86.0" }
-tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.86.0" }
-tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.86.0" }
+tw-bedrock = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.89.0" }
+tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.89.0" }
+tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.89.0" }
+tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.89.0" }
 
 # Web framework
 axum = { version = "0.8", features = ["macros", "ws"] }

Cargo.toml

@@ -98,7 +98,15 @@ pub async fn exchange_authorization_code(
         .map_err(TokenExchangeError::Unreachable)?;
 
     let status = resp.status();
-    let resp_text = resp.text().await.map_err(TokenExchangeError::Unreachable)?;
+    let bytes = think_watch_common::resources::read_response(resp, 64 * 1024)
+        .await
+        .map_err(|e| match e {
+            think_watch_common::resources::ReadError::Transport(e) => {
+                TokenExchangeError::Unreachable(e)
+            }
+            other => TokenExchangeError::InvalidResponse(other.to_string()),
+        })?;
+    let resp_text = String::from_utf8_lossy(&bytes);
 
     // Two cases land here:
     //   * non-2xx status (always treated as failure)

crates/auth/src/oauth/client.rs

@@ -1,10 +1,7 @@
-use std::sync::{Arc, LazyLock};
-
 use argon2::{
     Argon2,
     password_hash::{PasswordHasher, PasswordVerifier, phc::PasswordHash},
 };
-use tokio::sync::Semaphore;
 
 const RANDOM_PASSWORD_LEN: usize = 16;
 const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$%&*";
@@ -60,24 +57,13 @@ pub fn verify_password_or_dummy(password: &str, hash: Option<&str>) -> bool {
 // every other request scheduled on that thread, so the handlers use
 // these versions, which run it on the blocking pool.
 
-/// How many Argon2 runs go at once: one per CPU, as many as the async
-/// workers allowed when they ran it. Each run holds 19 MiB (`m=19456`)
-/// while it lasts and the blocking pool grows to hundreds of threads, so
-/// without this a burst of logins would hold gigabytes at once.
-static ARGON2_PERMITS: LazyLock<Arc<Semaphore>> = LazyLock::new(|| {
-    Arc::new(Semaphore::new(
-        std::thread::available_parallelism().map_or(1, |n| n.get()),
-    ))
-});
-
-/// `f` on the blocking pool once an [`ARGON2_PERMITS`] place is free. The
-/// task keeps the place until it ends, even when the caller has gone.
+/// Share CPU admission with the data plane and limit Argon2 memory using
+/// the container allowance. The running task owns both places even when
+/// its caller is cancelled.
 async fn on_blocking_pool<T: Send + 'static>(
     f: impl FnOnce() -> T + Send + 'static,
 ) -> Result<T, tokio::task::JoinError> {
-    // Never closed, so the acquire cannot fail; the permit inside the
-    // result is held all the same.
-    let permit = Arc::clone(&ARGON2_PERMITS).acquire_owned().await;
+    let permit = think_watch_common::resources::password_slot().await;
     tokio::task::spawn_blocking(move || {
         let _permit = permit;
         f()
@@ -111,6 +97,48 @@ pub async fn verify_password_or_dummy_off_runtime(password: String, hash: Option
 mod tests {
     use super::*;
 
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn cancelling_callers_does_not_admit_more_running_password_jobs() {
+        use std::time::Duration;
+        let policy = think_watch_common::resources::policy();
+        let jobs = policy.argon_jobs.min(policy.cpu_jobs);
+        let (started, mut beginnings) = tokio::sync::mpsc::unbounded_channel();
+        let mut release = Vec::new();
+        let mut callers = Vec::new();
+        for _ in 0..jobs {
+            let (tx, rx) = std::sync::mpsc::channel::<()>();
+            release.push(tx);
+            let started = started.clone();
+            callers.push(tokio::spawn(on_blocking_pool(move || {
+                started.send(()).unwrap();
+                let _ = rx.recv_timeout(Duration::from_secs(10));
+            })));
+        }
+        for _ in 0..jobs {
+            tokio::time::timeout(Duration::from_secs(5), beginnings.recv())
+                .await
+                .unwrap()
+                .unwrap();
+        }
+        for caller in callers {
+            caller.abort();
+            assert!(caller.await.unwrap_err().is_cancelled());
+        }
+        let next = tokio::spawn(on_blocking_pool(move || started.send(()).unwrap()));
+        assert!(
+            tokio::time::timeout(Duration::from_millis(50), beginnings.recv())
+                .await
+                .is_err()
+        );
+        drop(release);
+        tokio::time::timeout(Duration::from_secs(5), next)
+            .await
+            .unwrap()
+            .unwrap()
+            .unwrap();
+        assert_eq!(beginnings.recv().await, Some(()));
+    }
+
     #[tokio::test]
     async fn the_off_runtime_versions_agree_with_the_blocking_ones() {
         let hash = hash_password_off_runtime("pw".into()).await.unwrap();
@@ -133,7 +161,9 @@ mod tests {
         use std::sync::atomic::{AtomicUsize, Ordering};
         static RUNNING: AtomicUsize = AtomicUsize::new(0);
         static PEAK: AtomicUsize = AtomicUsize::new(0);
-        let permits = std::thread::available_parallelism().map_or(1, |n| n.get());
+        let permits = think_watch_common::resources::policy()
+            .argon_jobs
+            .min(think_watch_common::resources::policy().cpu_jobs);
         let tasks: Vec<_> = (0..permits * 4)
             .map(|_| {
                 tokio::spawn(on_blocking_pool(|| {

crates/auth/src/password.rs

@@ -48,3 +48,11 @@ aws-sigv4 = { workspace = true }
 aws-credential-types = { workspace = true }
 http_1x = { package = "http", version = "1" }
 async-trait = "0.1"
+http-body-util = { workspace = true }
+http-body = "1"
+sysinfo = { version = "0.36", default-features = false, features = ["system"] }
+
+[dev-dependencies]
+async-stream = { workspace = true }
+tempfile = "3"
+tower = { workspace = true, features = ["util"] }

crates/common/Cargo.toml

@@ -202,7 +202,10 @@ pub(super) async fn send_kafka(
     if resp.status().is_success() {
         Ok(())
     } else {
-        let body = resp.text().await.unwrap_or_default();
+        let bytes = crate::resources::read_response(resp, 64 * 1024)
+            .await
+            .map_err(|e| e.to_string())?;
+        let body = String::from_utf8_lossy(&bytes);
         Err(format!("Kafka REST proxy returned error: {body}"))
     }
 }
@@ -300,7 +303,10 @@ pub(super) async fn send_webhook(
         Ok(resp) if resp.status().is_success() => Ok(()),
         Ok(resp) => {
             let status = resp.status();
-            let rtext = resp.text().await.unwrap_or_default();
+            let bytes = crate::resources::read_response(resp, 64 * 1024)
+                .await
+                .map_err(|e| e.to_string())?;
+            let rtext = String::from_utf8_lossy(&bytes);
             Err(format!("Webhook returned {status}: {rtext}"))
         }
         Err(e) => Err(format!("Webhook request failed: {e}")),

crates/common/src/audit/forwarders.rs

@@ -59,17 +59,15 @@ pub struct AuditLogger {
 
 /// Bounded audit channel capacity.
 ///
-/// At ~80 bytes per entry that's about 8 MB of in-memory backlog,
-/// which is fine for any reasonable host. 100k entries means a
-/// 30-second ClickHouse outage at 3k req/s still survives; drops
-/// are surfaced via a metric.
+/// Secondary count limit. Entries include captured bodies, so a separate
+/// process-wide byte budget is the primary memory bound, including retries.
 const AUDIT_CHANNEL_CAPACITY: usize = 100_000;
 
 /// Throttle window for the structured "audit drop" log line. The
 /// metric still increments on every drop, so dashboards see the
 /// real rate; the log is just for human inspection and one line
 /// per second is plenty.
-fn log_audit_drop_throttled(err: &tokio::sync::mpsc::error::TrySendError<AuditEntry>) {
+fn log_audit_drop_throttled(err: &impl std::fmt::Display) {
     use std::sync::atomic::{AtomicU64, Ordering};
     use std::time::{SystemTime, UNIX_EPOCH};
     static LAST_LOG_SECS: AtomicU64 = AtomicU64::new(0);
@@ -272,6 +270,8 @@ impl AuditLogger {
                         // the REMAINING capacity, so depth = total - capacity.
                         let depth = AUDIT_CHANNEL_CAPACITY.saturating_sub(tx_for_gauge.capacity());
                         metrics::gauge!("audit_log_queue_depth").set(depth as f64);
+                        metrics::gauge!("audit_log_buffered_bytes")
+                            .set(crate::resources::audit_budget().used() as f64);
                     }
                 }
             });
@@ -297,7 +297,7 @@ impl AuditLogger {
             .store(bps, std::sync::atomic::Ordering::Relaxed);
     }
 
-    pub fn log(&self, entry: AuditEntry) {
+    pub fn log(&self, mut entry: AuditEntry) {
         // Sampling: consult the atomic rate and skip the entry when a
         // uniform draw falls outside the keep window. A full keep
         // (10000) short-circuits the RNG so the hot path pays nothing
@@ -309,6 +309,12 @@ impl AuditLogger {
             metrics::counter!("audit_log_sampled_out_total").increment(1);
             return;
         }
+        if !entry.reserve_memory(crate::resources::audit_budget()) {
+            metrics::counter!("audit_log_dropped_total").increment(1);
+            metrics::counter!("audit_log_memory_dropped_total").increment(1);
+            log_audit_drop_throttled(&"audit byte budget occupied (queue and retained batches)");
+            return;
+        }
         if let Err(e) = self.tx.try_send(entry) {
             // Compliance signal: dropped audit entries are a real
             // operational incident, not a debug warning. Bump the

crates/common/src/audit/logger.rs

@@ -132,6 +132,10 @@ impl std::fmt::Display for BodyCaptureStatus {
 /// back out and feeds them into `send_webhook` again.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct AuditEntry {
+    /// Account for queued and retained data through its final clone. Never part
+    /// of ClickHouse rows or the durable webhook JSON contract.
+    #[serde(skip)]
+    memory: Option<std::sync::Arc<crate::resources::Reservation>>,
     pub id: String,
     #[serde(skip, default = "default_log_type")]
     pub log_type: LogType,
@@ -407,6 +411,7 @@ impl AuditEntry {
     )]
     pub fn new(action: impl Into<String>) -> Self {
         Self {
+            memory: None,
             id: Uuid::new_v4().to_string(),
             log_type: LogType::Audit,
             user_id: None,
@@ -430,6 +435,61 @@ impl AuditEntry {
         }
     }
 
+    pub(super) fn reserve_memory(
+        &mut self,
+        budget: &std::sync::Arc<crate::resources::ByteBudget>,
+    ) -> bool {
+        fn value_bytes(value: &serde_json::Value) -> usize {
+            use serde_json::Value;
+            match value {
+                Value::String(s) => s.capacity(),
+                Value::Array(values) => values
+                    .capacity()
+                    .saturating_mul(std::mem::size_of::<Value>())
+                    .saturating_add(values.iter().map(value_bytes).sum::<usize>()),
+                Value::Object(values) => values.iter().fold(0usize, |total, (key, value)| {
+                    total
+                        .saturating_add(key.capacity())
+                        .saturating_add(std::mem::size_of::<(String, Value)>() * 2)
+                        .saturating_add(value_bytes(value))
+                }),
+                _ => 0,
+            }
+        }
+        let strings = [
+            Some(&self.id),
+            Some(&self.action),
+            Some(&self.created_at),
+            self.user_id.as_ref(),
+            self.user_email.as_ref(),
+            self.api_key_id.as_ref(),
+            self.api_key_lineage_id.as_ref(),
+            self.resource.as_ref(),
+            self.resource_id.as_ref(),
+            self.ip_address.as_ref(),
+            self.user_agent.as_ref(),
+            self.trace_id.as_ref(),
+            self.session_id.as_ref(),
+            self.request_body.as_ref(),
+            self.response_body.as_ref(),
+            self.body_capture_status.as_ref(),
+        ];
+        let bytes = strings
+            .into_iter()
+            .flatten()
+            .fold(std::mem::size_of::<Self>(), |n, s| {
+                n.saturating_add(s.capacity())
+            })
+            .saturating_add(self.detail.as_ref().map_or(0, value_bytes));
+        // Flush mappers and forwarders temporarily clone/serialize one entry.
+        let mut lease = budget.reserve();
+        if !lease.grow(bytes.saturating_mul(2)) {
+            return false;
+        }
+        self.memory = Some(std::sync::Arc::new(lease));
+        true
+    }
+
     /// Create entry for gateway request logs.
     ///
     /// Bare constructor — same deprecation reasoning as `::new`. Use
@@ -680,3 +740,39 @@ mod tests {
         assert_eq!(decoded.detail, original.detail);
     }
 }
+
+#[cfg(test)]
+#[allow(deprecated)]
+mod memory_tests {
+    use super::*;
+    use crate::resources::ByteBudget;
+
+    #[test]
+    fn queue_and_retry_clones_retain_body_admission_without_changing_wire_json() {
+        let budget = ByteBudget::new(32 * 1024);
+        let mut entry = AuditEntry::new("tools.call").request_body("x".repeat(4096));
+        let before = serde_json::to_value(&entry).unwrap();
+        assert!(entry.reserve_memory(&budget));
+        assert!(budget.used() >= 8192);
+        assert_eq!(before, serde_json::to_value(&entry).unwrap());
+        let retained = entry.clone();
+        drop(entry);
+        assert!(budget.used() >= 8192);
+        drop(retained);
+        assert_eq!(budget.used(), 0);
+    }
+
+    #[test]
+    fn a_few_large_entries_hit_the_byte_limit_before_the_count_limit() {
+        let budget = ByteBudget::new(16 * 1024);
+        let mut first = AuditEntry::new("gateway").response_body("x".repeat(4096));
+        let mut second = first.clone();
+        assert!(first.reserve_memory(&budget));
+        let used = budget.used();
+        assert!(!second.reserve_memory(&budget));
+        assert_eq!(used, budget.used());
+        drop(first);
+        assert_eq!(budget.used(), 0);
+        assert!(second.reserve_memory(&budget));
+    }
+}

crates/common/src/audit/types.rs

@@ -378,8 +378,12 @@ impl S3Store {
             .await
             .map_err(|e| BlobError::Transport(e.to_string()))?;
         let status = resp.status();
-        let bytes = resp
-            .bytes()
+        let max = if status.is_success() {
+            crate::resources::policy().max_body_bytes
+        } else {
+            64 * 1024
+        };
+        let bytes = crate::resources::read_response(resp, max)
             .await
             .map_err(|e| BlobError::Transport(e.to_string()))?;
         if !status.is_success() {

crates/common/src/blob_store.rs

@@ -43,6 +43,7 @@ pub mod clickhouse_client;
 pub mod cost_decimal; // Decimal ↔ raw i64/i128 helpers for CH Decimal(18, 10) columns
 pub mod lifecycle; // Surface-agnostic request pipeline (see lifecycle::mod docs)
 pub mod limits; // rate-limit & budget evaluation
+pub mod resources; // automatic CPU admission and process-wide byte budgets
 
 // --- Utilities ---
 pub mod crypto; // AES-256-GCM envelope for secrets at rest