diff --git a/Cargo.lock b/Cargo.lock index 9c607902..32c07700 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -443,6 +443,12 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bech32" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d86b93f97252c47b41663388e6d155714a9d0c398b99f1005cbc5f978b29f445" + [[package]] name = "bitflags" version = "1.3.2" @@ -610,7 +616,7 @@ dependencies = [ "num-traits", "serde", "wasm-bindgen", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -1335,6 +1341,12 @@ dependencies = [ "smallvec", ] +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "ff" version = "0.13.1" @@ -1575,8 +1587,8 @@ dependencies = [ "libc", "log", "rustversion", - "windows-link", - "windows-result", + "windows-link 0.1.3", + "windows-result 0.3.4", ] [[package]] @@ -1919,7 +1931,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core", + "windows-core 0.62.2", ] [[package]] @@ -2108,7 +2120,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -2162,7 +2174,7 @@ dependencies = [ "simd_cesu8", "thiserror 2.0.21", "walkdir", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -2274,6 +2286,15 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "libmimalloc-sys" +version = "0.1.49" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" +dependencies = [ + "cc", +] + [[package]] name = "libsqlite3-sys" version = "0.30.1" @@ -2284,6 +2305,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.3" @@ -2409,6 +2436,15 @@ dependencies = [ "sketches-ddsketch", ] +[[package]] +name = "mimalloc" +version = "0.1.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "mime" version = "0.3.17" @@ -2468,6 +2504,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "ntapi" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3b335231dfd352ffb0f8017f3b6027a4917f7df785ea2143d8af2adc66980ae" +dependencies = [ + "winapi", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2568,6 +2613,25 @@ dependencies = [ "url", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.13.2", +] + +[[package]] +name = "objc2-io-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33fafba39597d6dc1fb709123dfa8289d39406734be322956a69f0931c73bb15" +dependencies = [ + "libc", + "objc2-core-foundation", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -2689,7 +2753,7 @@ dependencies = [ "libc", "redox_syscall", "smallvec", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -2944,7 +3008,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2 0.6.5", + "socket2 0.5.10", "thiserror 2.0.21", "tokio", "tracing", @@ -2983,7 +3047,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.5", + "socket2 0.5.10", "tracing", "windows-sys 0.61.2", ] @@ -3362,6 +3426,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5e76742e5294153d4dde1207dbec5bdd81eef6fabdc0d2118babadabfd858aa" +dependencies = [ + "bitflags 2.13.2", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.46" @@ -4166,6 +4243,20 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "sysinfo" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252800745060e7b9ffb7b2badbd8b31cfa4aa2e61af879d0a3bf2a317c20217d" +dependencies = [ + "libc", + "memchr", + "ntapi", + "objc2-core-foundation", + "objc2-io-kit", + "windows", +] + [[package]] name = "system-configuration" version = "0.7.0" @@ -4187,6 +4278,19 @@ dependencies = [ "libc", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.3.4", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "think-watch-auth" version = "3.10.0" @@ -4223,6 +4327,7 @@ version = "3.10.0" dependencies = [ "aes-gcm", "anyhow", + "async-stream", "async-trait", "aws-credential-types", "aws-sigv4", @@ -4236,6 +4341,8 @@ dependencies = [ "hex", "hmac 0.13.0", "http", + "http-body", + "http-body-util", "hyper-util", "metrics", "rand 0.10.3", @@ -4248,8 +4355,11 @@ dependencies = [ "sha1 0.11.0", "sha2 0.11.0", "sqlx", + "sysinfo", + "tempfile", "thiserror 2.0.21", "tokio", + "tower", "tracing", "tw-bedrock", "tw-breaker", @@ -4344,8 +4454,10 @@ dependencies = [ "futures", "hex", "hmac 0.13.0", + "libc", "metrics", "metrics-exporter-prometheus", + "mimalloc", "openidconnect", "p256 0.14.0", "rand 0.10.3", @@ -4838,8 +4950,8 @@ dependencies = [ [[package]] name = "tw-bedrock" -version = "0.86.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.86.0#62aec779f28b8f2f087670cc5096af829bd5ad20" +version = "0.89.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.89.0#c66a3b04ba297f608f2fc21501c73f21a78ab8ed" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -4853,16 +4965,16 @@ dependencies = [ [[package]] name = "tw-breaker" -version = "0.86.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.86.0#62aec779f28b8f2f087670cc5096af829bd5ad20" +version = "0.89.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.89.0#c66a3b04ba297f608f2fc21501c73f21a78ab8ed" dependencies = [ "serde", ] [[package]] name = "tw-dialect" -version = "0.86.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.86.0#62aec779f28b8f2f087670cc5096af829bd5ad20" +version = "0.89.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.89.0#c66a3b04ba297f608f2fc21501c73f21a78ab8ed" dependencies = [ "memchr", "serde", @@ -4871,12 +4983,14 @@ dependencies = [ [[package]] name = "tw-guard" -version = "0.86.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.86.0#62aec779f28b8f2f087670cc5096af829bd5ad20" +version = "0.89.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.89.0#c66a3b04ba297f608f2fc21501c73f21a78ab8ed" dependencies = [ "aho-corasick", "base64 0.22.1", + "bech32", "bytes", + "curve25519-dalek", "memchr", "regex", "regex-syntax", @@ -5238,6 +5352,41 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows" +version = "0.61.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893" +dependencies = [ + "windows-collections", + "windows-core 0.61.2", + "windows-future", + "windows-link 0.1.3", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3beeceb5e5cfd9eb1d76b381630e82c4241ccd0d27f1a39ed41b2760b255c5e8" +dependencies = [ + "windows-core 0.61.2", +] + +[[package]] +name = "windows-core" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link 0.1.3", + "windows-result 0.3.4", + "windows-strings 0.4.2", +] + [[package]] name = "windows-core" version = "0.62.2" @@ -5246,9 +5395,20 @@ checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" dependencies = [ "windows-implement", "windows-interface", - "windows-link", - "windows-result", - "windows-strings", + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-future" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", + "windows-threading", ] [[package]] @@ -5273,21 +5433,46 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "windows-link" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-numerics" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9150af68066c4c5c07ddc0ce30421554771e528bde427614c61038bc2c92c2b1" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", +] + [[package]] name = "windows-registry" version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" dependencies = [ - "windows-link", - "windows-result", - "windows-strings", + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-result" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6" +dependencies = [ + "windows-link 0.1.3", ] [[package]] @@ -5296,7 +5481,16 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" dependencies = [ - "windows-link", + "windows-link 0.2.1", +] + +[[package]] +name = "windows-strings" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57" +dependencies = [ + "windows-link 0.1.3", ] [[package]] @@ -5305,7 +5499,7 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" dependencies = [ - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -5323,7 +5517,7 @@ version = "0.61.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" dependencies = [ - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -5342,6 +5536,15 @@ dependencies = [ "windows_x86_64_msvc", ] +[[package]] +name = "windows-threading" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b66463ad2e0ea3bbf808b7f1d371311c80e115c0b71d60efc142cafbcfb057a6" +dependencies = [ + "windows-link 0.1.3", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" diff --git a/Cargo.toml b/Cargo.toml index 79fa1aa8..8c6fc137 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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"] } diff --git a/crates/auth/src/oauth/client.rs b/crates/auth/src/oauth/client.rs index d4f62426..553914c7 100644 --- a/crates/auth/src/oauth/client.rs +++ b/crates/auth/src/oauth/client.rs @@ -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) diff --git a/crates/auth/src/password.rs b/crates/auth/src/password.rs index 4a2c3dae..4cb9446c 100644 --- a/crates/auth/src/password.rs +++ b/crates/auth/src/password.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> = 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( f: impl FnOnce() -> T + Send + 'static, ) -> Result { - // 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(|| { diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index b196e8bc..ca37e675 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -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"] } diff --git a/crates/common/src/audit/forwarders.rs b/crates/common/src/audit/forwarders.rs index 90023b4e..31a2dd14 100644 --- a/crates/common/src/audit/forwarders.rs +++ b/crates/common/src/audit/forwarders.rs @@ -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}")), diff --git a/crates/common/src/audit/logger.rs b/crates/common/src/audit/logger.rs index a0a194cd..df17df10 100644 --- a/crates/common/src/audit/logger.rs +++ b/crates/common/src/audit/logger.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) { +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 diff --git a/crates/common/src/audit/types.rs b/crates/common/src/audit/types.rs index 61202580..212a6094 100644 --- a/crates/common/src/audit/types.rs +++ b/crates/common/src/audit/types.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>, 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) -> 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, + ) -> 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::()) + .saturating_add(values.iter().map(value_bytes).sum::()), + 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::(), |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)); + } +} diff --git a/crates/common/src/blob_store.rs b/crates/common/src/blob_store.rs index 5211d7a8..8a057190 100644 --- a/crates/common/src/blob_store.rs +++ b/crates/common/src/blob_store.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() { diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 334950bd..495963c5 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.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 diff --git a/crates/common/src/resources.rs b/crates/common/src/resources.rs new file mode 100644 index 00000000..5f8b3987 --- /dev/null +++ b/crates/common/src/resources.rs @@ -0,0 +1,475 @@ +//! Resource admission shared by AI, MCP, authentication and audit dispatch. +//! Limits describe this process (including cgroups), never the host alone. +//! Byte reservations survive cancellation, shared buffers and audit retries. + +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; + +use bytes::Bytes; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +#[cfg(any(target_os = "linux", test))] +mod cgroup; +pub mod http; + +const MIB: usize = 1024 * 1024; +// Compact JSON arrays can use many times their wire size as Value nodes. This +// also covers conversion/capture copies rather than budgeting raw bytes alone. +pub const BODY_HEADROOM: usize = 32; +pub const IDLE_TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Debug)] +pub struct Policy { + pub io_workers: usize, + pub cpu_jobs: usize, + pub argon_jobs: usize, + pub body_bytes: usize, + pub audit_bytes: usize, + pub max_body_bytes: usize, +} + +impl Policy { + fn for_machine(cpus: usize, memory: u64) -> Self { + let memory = usize::try_from(memory).unwrap_or(usize::MAX); + let cpus = cpus.max(1); + let body_bytes = (memory / 4).clamp(8 * MIB, 8 * 1024 * MIB); + Self { + io_workers: cpus.min(8), + cpu_jobs: cpus.min((memory / (128 * MIB)).max(1)), + // Argon2's 19 MiB workspace, allocator headroom, and at most 1/8 + // of the process allowance. CPU admission is shared with scanning. + argon_jobs: cpus.min((memory / 8 / (32 * MIB)).max(1)), + body_bytes, + audit_bytes: (memory / 16).clamp(2 * MIB, 128 * MIB), + max_body_bytes: (body_bytes / BODY_HEADROOM).min(64 * MIB), + } + } +} + +pub fn policy() -> &'static Policy { + static POLICY: OnceLock = OnceLock::new(); + POLICY.get_or_init(|| { + let mut system = sysinfo::System::new(); + system.refresh_memory(); + let physical = match system.total_memory() { + 0 => 1 << 30, + memory => memory, + }; + let limit = system.cgroup_limits().map(|l| l.total_memory); + #[cfg(target_os = "linux")] + let limit = cgroup::memory_limit().or(limit); + let memory = limit + .filter(|&n| n > 0) + .map_or(physical, |n| n.min(physical)); + Policy::for_machine( + std::thread::available_parallelism().map_or(1, usize::from), + memory, + ) + }) +} + +struct Scheduler { + slots: Arc, + inline: Semaphore, +} + +impl Scheduler { + fn new(policy: &Policy) -> Self { + Self { + slots: Arc::new(Semaphore::new(policy.cpu_jobs)), + inline: Semaphore::new(policy.io_workers.div_ceil(2)), + } + } + + async fn run(&self, bytes: usize, f: impl FnOnce() -> T) -> T { + let _slot = self.slots.clone().acquire_owned().await.unwrap(); + let inline = (bytes < MIB) + .then(|| self.inline.try_acquire().ok()) + .flatten(); + let multi = tokio::runtime::Handle::current().runtime_flavor() + == tokio::runtime::RuntimeFlavor::MultiThread; + if inline.is_some() || !multi { + f() + } else { + tokio::task::block_in_place(f) + } + } +} + +fn scheduler() -> &'static Scheduler { + static SCHEDULER: OnceLock = OnceLock::new(); + SCHEDULER.get_or_init(|| Scheduler::new(policy())) +} + +pub async fn cpu(bytes: usize, f: impl FnOnce() -> T) -> T { + scheduler().run(bytes, f).await +} + +/// Acquire before spawning; cancellation of the caller must not release a +/// running password operation's CPU or memory admission. +pub async fn password_slot() -> (OwnedSemaphorePermit, OwnedSemaphorePermit) { + static PASSWORDS: OnceLock> = OnceLock::new(); + let memory = PASSWORDS.get_or_init(|| Arc::new(Semaphore::new(policy().argon_jobs))); + let memory = memory.clone().acquire_owned().await.unwrap(); + let cpu = scheduler().slots.clone().acquire_owned().await.unwrap(); + (memory, cpu) +} + +#[derive(Debug)] +pub struct ByteBudget { + limit: usize, + used: Mutex, +} + +impl ByteBudget { + pub fn new(limit: usize) -> Arc { + Arc::new(Self { + limit, + used: Mutex::new(0), + }) + } + + pub fn reserve(self: &Arc) -> Reservation { + Reservation { + budget: self.clone(), + charged: 0, + } + } + + pub fn used(&self) -> usize { + *self.used.lock().unwrap_or_else(|e| e.into_inner()) + } +} + +#[derive(Debug)] +pub struct Reservation { + budget: Arc, + charged: usize, +} + +impl Reservation { + pub fn grow(&mut self, bytes: usize) -> bool { + let extra = bytes.saturating_sub(self.charged); + let mut used = self.budget.used.lock().unwrap_or_else(|e| e.into_inner()); + if extra > self.budget.limit.saturating_sub(*used) { + return false; + } + *used += extra; + self.charged += extra; + true + } + + pub fn wrap(self, data: Bytes) -> Bytes { + Bytes::from_owner(Buffered { + data, + reservation: Arc::new(self), + }) + } +} + +impl Drop for Reservation { + fn drop(&mut self) { + *self.budget.used.lock().unwrap_or_else(|e| e.into_inner()) -= self.charged; + } +} + +struct Buffered { + data: Bytes, + reservation: Arc, +} + +impl AsRef<[u8]> for Buffered { + fn as_ref(&self) -> &[u8] { + let _ = &self.reservation; + &self.data + } +} + +pub fn body_budget() -> &'static Arc { + static BUDGET: OnceLock> = OnceLock::new(); + BUDGET.get_or_init(|| ByteBudget::new(policy().body_bytes)) +} + +pub fn audit_budget() -> &'static Arc { + static BUDGET: OnceLock> = OnceLock::new(); + BUDGET.get_or_init(|| ByteBudget::new(policy().audit_bytes)) +} + +/// A stream's retained timeline, shared with its detached post-invoke tail. +pub struct StreamBudget { + reservation: Reservation, + received: usize, + max: usize, +} + +impl StreamBudget { + pub fn new(budget: &Arc, max: usize) -> Self { + Self { + reservation: budget.reserve(), + received: 0, + max, + } + } + + pub fn append(&mut self, bytes: usize) -> bool { + let received = self.received.saturating_add(bytes); + if received > self.max + || !self + .reservation + .grow(received.saturating_mul(BODY_HEADROOM)) + { + return false; + } + self.received = received; + true + } +} + +pub type SharedStreamBudget = Arc>; + +pub fn stream_budget() -> SharedStreamBudget { + Arc::new(Mutex::new(StreamBudget::new( + body_budget(), + policy().max_body_bytes, + ))) +} + +#[derive(Debug, thiserror::Error)] +pub enum ReadError { + #[error("The response exceeds the gateway buffering limit")] + TooLarge, + #[error("The gateway buffering budget is occupied; retry shortly")] + Busy, + #[error(transparent)] + Transport(#[from] reqwest::Error), +} + +/// Charge before appending, including chunked and decompressed responses. Keep +/// room for JSON, conversion and allocator capacity; never wait holding a partial +/// body, since concurrent readers could otherwise deadlock each other. +pub async fn read_response(response: reqwest::Response, max: usize) -> Result { + let max = max.min(policy().max_body_bytes); + read_response_with_budget(response, max, body_budget()).await +} + +/// The caller keeps the returned Bytes while it uses the decoded value. +pub async fn read_json( + response: reqwest::Response, + max: usize, +) -> anyhow::Result<(T, Bytes)> { + let bytes = read_response(response, max).await?; + let value = cpu(bytes.len(), || serde_json::from_slice(&bytes)).await?; + Ok((value, bytes)) +} + +async fn read_response_with_budget( + mut response: reqwest::Response, + max: usize, + budget: &Arc, +) -> Result { + if response.content_length().is_some_and(|n| n > max as u64) { + return Err(ReadError::TooLarge); + } + let mut reservation = budget.reserve(); + let mut data = Vec::new(); + while let Some(chunk) = response.chunk().await? { + let size = data + .len() + .checked_add(chunk.len()) + .ok_or(ReadError::TooLarge)?; + if size > max { + return Err(ReadError::TooLarge); + } + if !reservation.grow(size.saturating_mul(BODY_HEADROOM)) { + metrics::counter!("gateway_memory_busy_total").increment(1); + return Err(ReadError::Busy); + } + data.extend_from_slice(&chunk); + } + Ok(reservation.wrap(data.into())) +} + +/// A transformed/captured body still owns the original read reservation. +pub fn retain(data: Bytes, source: Bytes) -> Bytes { + struct Retained { + data: Bytes, + _source: Bytes, + } + impl AsRef<[u8]> for Retained { + fn as_ref(&self) -> &[u8] { + &self.data + } + } + Bytes::from_owner(Retained { + data, + _source: source, + }) +} + +/// Reserve before fetching a Redis cache entry. STRLEN followed by bounded +/// GETRANGE also stays bounded if another writer replaces the key in between. +pub async fn cached_body(redis: &fred::clients::Client, key: &str) -> Option { + use fred::interfaces::KeysInterface; + let length: usize = redis.strlen(key).await.ok()?; + if length == 0 || length > policy().max_body_bytes { + return None; + } + let mut lease = body_budget().reserve(); + if !lease.grow(length.saturating_mul(BODY_HEADROOM)) { + return None; + } + let data: Vec = redis.getrange(key, 0, length).await.ok()?; + if data.len() > length { + return None; + } + Some(lease.wrap(data.into())) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn memory_limits_allocation_without_capping_a_well_provisioned_eighty_core_server() { + let large = Policy::for_machine(80, 128 << 30); + assert_eq!( + (large.io_workers, large.cpu_jobs, large.argon_jobs), + (8, 80, 80) + ); + let small = Policy::for_machine(80, 256 << 20); + assert_eq!((small.cpu_jobs, small.argon_jobs), (2, 1)); + assert!(small.body_bytes + small.audit_bytes <= 256 * MIB / 2); + } + + #[test] + fn reservations_follow_shared_slices_and_failed_growth_is_atomic() { + let budget = ByteBudget::new(10); + let mut lease = budget.reserve(); + assert!(lease.grow(8)); + assert!(!lease.grow(11)); + let data = lease.wrap(Bytes::from_static(b"hello")); + let slice = data.slice(1..3); + drop(data); + assert_eq!(budget.used(), 8); + drop(slice); + assert_eq!(budget.used(), 0); + } + + #[tokio::test] + async fn chunked_responses_obey_size_and_shared_memory_limits() { + async fn response() -> (reqwest::Response, tokio::task::JoinHandle<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/", listener.local_addr().unwrap()); + let router = axum::Router::new().route( + "/", + axum::routing::get(|| async { + axum::body::Body::from_stream(futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"12345678")), + Ok(Bytes::from_static(b"abcdefgh")), + ])) + }), + ); + let server = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + (reqwest::get(url).await.unwrap(), server) + } + for (limit, max, expected) in [ + (10 * BODY_HEADROOM, 32, "busy"), + (32 * BODY_HEADROOM, 12, "large"), + (32 * BODY_HEADROOM, 32, "ok"), + ] { + let budget = ByteBudget::new(limit); + let (response, server) = response().await; + assert!(response.content_length().is_none()); + match read_response_with_budget(response, max, &budget).await { + Err(ReadError::Busy) => assert_eq!(expected, "busy"), + Err(ReadError::TooLarge) => assert_eq!(expected, "large"), + Ok(bytes) => { + assert_eq!(expected, "ok"); + assert_eq!(&bytes[..], b"12345678abcdefgh"); + assert_eq!(budget.used(), 16 * BODY_HEADROOM); + let copy = bytes.clone(); + drop(bytes); + assert_eq!(budget.used(), 16 * BODY_HEADROOM); + drop(copy); + } + Err(error) => panic!("{error}"), + } + assert_eq!(budget.used(), 0); + server.abort(); + } + } + + #[test] + fn stream_timeline_reservations_cover_retained_history_and_reject_growth() { + let budget = ByteBudget::new(25 * BODY_HEADROOM); + let mut stream = StreamBudget::new(&budget, 64); + assert!(stream.append(20)); + assert!(!stream.append(10)); + assert_eq!(budget.used(), 20 * BODY_HEADROOM); + drop(stream); + assert_eq!(budget.used(), 0); + } + + #[test] + fn eighty_cpu_jobs_expand_beyond_io_workers_then_threads_retire() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Instant; + let alive = Arc::new(AtomicUsize::new(0)); + let (started, stopped) = (alive.clone(), alive.clone()); + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_keep_alive(Duration::from_millis(50)) + .on_thread_start(move || { + started.fetch_add(1, Ordering::SeqCst); + }) + .on_thread_stop(move || { + stopped.fetch_sub(1, Ordering::SeqCst); + }) + .enable_all() + .build() + .unwrap(); + let scheduler = Arc::new(Scheduler::new(&Policy { + io_workers: 2, + ..Policy::for_machine(80, 128 << 30) + })); + runtime.block_on(async { + let release = Arc::new((Mutex::new(false), std::sync::Condvar::new())); + let (began, mut beginnings) = tokio::sync::mpsc::unbounded_channel(); + let mut jobs = Vec::new(); + for _ in 0..80 { + let (scheduler, release, began) = + (scheduler.clone(), release.clone(), began.clone()); + jobs.push(tokio::spawn(async move { + scheduler + .run(4096, || { + began.send(()).unwrap(); + let guard = release.0.lock().unwrap(); + let _guard = release + .1 + .wait_timeout_while(guard, Duration::from_secs(10), |done| !*done) + .unwrap(); + }) + .await; + })); + } + for _ in 0..80 { + tokio::time::timeout(Duration::from_secs(5), beginnings.recv()) + .await + .expect("all 80 calls must start") + .unwrap(); + } + assert_eq!(scheduler.slots.available_permits(), 0); + tokio::time::sleep(Duration::from_millis(1)).await; + *release.0.lock().unwrap() = true; + release.1.notify_all(); + for job in jobs { + job.await.unwrap(); + } + }); + let deadline = Instant::now() + Duration::from_secs(5); + while alive.load(Ordering::SeqCst) != 2 && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(10)); + } + assert_eq!(alive.load(Ordering::SeqCst), 2); + } +} diff --git a/crates/common/src/resources/cgroup.rs b/crates/common/src/resources/cgroup.rs new file mode 100644 index 00000000..72ae6aa1 --- /dev/null +++ b/crates/common/src/resources/cgroup.rs @@ -0,0 +1,89 @@ +//! 同时看当前 cgroup 和它的祖先，systemd 的子组也不能越过父组的内存限额。 + +use std::path::{Component, Path}; + +#[cfg(target_os = "linux")] +pub(super) fn memory_limit() -> Option { + let membership = std::fs::read_to_string("/proc/self/cgroup").ok()?; + membership + .lines() + .filter_map(|line| { + let mut fields = line.splitn(3, ':'); + let hierarchy = fields.next()?; + let controllers = fields.next()?; + let path = fields.next()?; + if hierarchy == "0" && controllers.is_empty() { + lowest_limit(Path::new("/sys/fs/cgroup"), Path::new(path), "memory.max") + } else if controllers.split(',').any(|c| c == "memory") { + lowest_limit( + Path::new("/sys/fs/cgroup/memory"), + Path::new(path), + "memory.limit_in_bytes", + ) + } else { + None + } + }) + .min() +} + +fn lowest_limit(root: &Path, membership: &Path, file: &str) -> Option { + // cgroup namespace 可以把当前组挂成根目录。深路径不存在时，仍然检查挂载根的限额。 + if membership + .components() + .any(|c| matches!(c, Component::ParentDir | Component::Prefix(_))) + { + return None; + } + let relative = membership.strip_prefix("/").unwrap_or(membership); + let current = root.join(relative); + current + .ancestors() + .take_while(|p| p.starts_with(root)) + .filter_map(|p| { + std::fs::read_to_string(p.join(file)) + .ok()? + .trim() + .parse::() + .ok() + .filter(|&n| n > 0) + }) + .min() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn nested_groups_obey_the_smallest_ancestor_limit() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + std::fs::create_dir_all(root.join("service/child")).unwrap(); + std::fs::write(root.join("memory.max"), "max\n").unwrap(); + std::fs::write(root.join("service/memory.max"), "536870912\n").unwrap(); + std::fs::write(root.join("service/child/memory.max"), "1073741824\n").unwrap(); + assert_eq!( + lowest_limit(root, Path::new("/service/child"), "memory.max"), + Some(512 << 20) + ); + } + + #[test] + fn a_namespaced_mount_root_and_v1_limits_are_supported() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("memory.limit_in_bytes"), "268435456").unwrap(); + assert_eq!( + lowest_limit( + dir.path(), + Path::new("/unavailable/host/path"), + "memory.limit_in_bytes" + ), + Some(256 << 20) + ); + assert_eq!( + lowest_limit(dir.path(), Path::new("/../escape"), "memory.limit_in_bytes"), + None + ); + } +} diff --git a/crates/common/src/resources/http.rs b/crates/common/src/resources/http.rs new file mode 100644 index 00000000..e07f7eb0 --- /dev/null +++ b/crates/common/src/resources/http.rs @@ -0,0 +1,227 @@ +//! Body admission precedes extraction; checking in the handler is too late. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use axum::body::{Body, Bytes, HttpBody}; +use axum::extract::Request; +use axum::middleware::Next; +use axum::response::{IntoResponse, Response}; +use futures::StreamExt; + +use super::{ByteBudget, policy}; + +/// Preserve Axum's JSON errors while moving large/concurrent decoding off I/O. +pub struct ScheduledJson(pub T); + +impl axum::extract::FromRequest for ScheduledJson +where + S: Send + Sync, + T: serde::de::DeserializeOwned, +{ + type Rejection = Response; + + async fn from_request(request: Request, state: &S) -> Result { + let json = request + .headers() + .get(axum::http::header::CONTENT_TYPE) + .and_then(|h| h.to_str().ok()) + .and_then(|h| h.split(';').next()) + .is_some_and(|h| { + let h = h.trim().to_ascii_lowercase(); + h == "application/json" || (h.starts_with("application/") && h.ends_with("+json")) + }); + if !json { + return Err(( + axum::http::StatusCode::UNSUPPORTED_MEDIA_TYPE, + "Expected request with Content-Type: application/json", + ) + .into_response()); + } + let bytes = Bytes::from_request(request, state) + .await + .map_err(IntoResponse::into_response)?; + super::cpu(bytes.len(), || axum::Json::::from_bytes(&bytes)) + .await + .map(|axum::Json(value)| Self(value)) + .map_err(IntoResponse::into_response) + } +} + +pub async fn admit(request: Request, next: Next) -> Response { + admit_with_budget(request, next, super::body_budget(), policy().max_body_bytes).await +} + +async fn admit_with_budget( + request: Request, + next: Next, + budget: &Arc, + max: usize, +) -> Response { + let (parts, body) = request.into_parts(); + if body.size_hint().upper().is_some_and(|n| n > max as u64) { + return rejection(axum::http::StatusCode::PAYLOAD_TOO_LARGE); + } + let mut data = Vec::new(); + let mut lease = budget.reserve(); + let mut chunks = body.into_data_stream(); + while let Some(chunk) = chunks.next().await { + let chunk = match chunk { + Ok(chunk) => chunk, + Err(error) => { + use std::error::Error; + let status = if error + .source() + .is_some_and(|e| e.is::()) + { + axum::http::StatusCode::PAYLOAD_TOO_LARGE + } else { + axum::http::StatusCode::BAD_REQUEST + }; + return rejection(status); + } + }; + let size = data.len().saturating_add(chunk.len()); + if size > max { + return rejection(axum::http::StatusCode::PAYLOAD_TOO_LARGE); + } + if !lease.grow(size.saturating_mul(super::BODY_HEADROOM)) { + return rejection(axum::http::StatusCode::TOO_MANY_REQUESTS); + } + data.extend_from_slice(&chunk); + } + let data = lease.wrap(data.into()); + let mut response = next + .run(Request::from_parts(parts, Body::from(data.clone()))) + .await; + let body = std::mem::replace(response.body_mut(), Body::empty()); + *response.body_mut() = Body::new(HeldBody { + inner: body, + _request: data, + }); + response +} + +fn rejection(status: axum::http::StatusCode) -> Response { + let busy = status == axum::http::StatusCode::TOO_MANY_REQUESTS; + let message = if busy { + "The gateway is buffering other requests. Retry shortly." + } else if status == axum::http::StatusCode::PAYLOAD_TOO_LARGE { + "The request body exceeds the gateway buffering limit." + } else { + "The request body could not be read." + }; + let mut response = (status, axum::Json(serde_json::json!({ + "error": {"type": if busy { "rate_limit_error" } else { "invalid_request_error" }, "message": message} + }))).into_response(); + if busy { + response + .headers_mut() + .insert(axum::http::header::RETRY_AFTER, "1".parse().unwrap()); + metrics::counter!("gateway_memory_busy_total").increment(1); + } + response +} + +/// A JSON extractor discards the input Bytes while its decoded object still +/// exists. Retain admission through the response, including streaming and drop. +struct HeldBody { + inner: Body, + _request: Bytes, +} + +impl HttpBody for HeldBody { + type Data = Bytes; + type Error = axum::Error; + + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + Pin::new(&mut self.inner).poll_frame(cx) + } + + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + fn size_hint(&self) -> http_body::SizeHint { + self.inner.size_hint() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::{Router, routing::post}; + use tower::ServiceExt; + + fn router(budget: Arc, max: usize) -> Router { + Router::new() + .route("/", post(|body: Bytes| async move { body })) + .layer(axum::middleware::from_fn( + move |request: Request, next: Next| { + let budget = budget.clone(); + async move { admit_with_budget(request, next, &budget, max).await } + }, + )) + } + + #[tokio::test] + async fn chunked_pressure_returns_retryable_429_and_frees_partial_buffers() { + let budget = ByteBudget::new(25 * super::super::BODY_HEADROOM); + let app = router(budget.clone(), 64); + let chunks = futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from(vec![b'a'; 20])), + Ok(Bytes::from(vec![b'b'; 10])), + ]); + let request = Request::post("/").body(Body::from_stream(chunks)).unwrap(); + let response = app.clone().oneshot(request).await.unwrap(); + assert_eq!(response.status(), 429); + assert_eq!(response.headers()["retry-after"], "1"); + assert_eq!(budget.used(), 0); + + let request = Request::post("/").body(Body::from(vec![b'x'; 16])).unwrap(); + let response = app.oneshot(request).await.unwrap(); + assert_eq!(response.status(), 200); + assert_eq!( + budget.used(), + 16 * super::super::BODY_HEADROOM, + "input admission outlives the extractor" + ); + let body = axum::body::to_bytes(response.into_body(), 100) + .await + .unwrap(); + assert_eq!(body.len(), 16); + drop(body); + assert_eq!(budget.used(), 0); + } + + #[tokio::test] + async fn an_oversize_body_is_413_and_cancellation_releases_partial_admission() { + let budget = ByteBudget::new(25 * super::super::BODY_HEADROOM); + let app = router(budget.clone(), 20); + let response = app + .clone() + .oneshot(Request::post("/").body(Body::from(vec![0; 21])).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), 413); + assert_eq!(budget.used(), 0); + + let sent = Arc::new(tokio::sync::Notify::new()); + let started = sent.clone(); + let chunks = async_stream::stream! { + yield Ok::<_, std::io::Error>(Bytes::from_static(b"12345678")); + started.notify_one(); + std::future::pending::<()>().await; + }; + let task = + tokio::spawn(app.oneshot(Request::post("/").body(Body::from_stream(chunks)).unwrap())); + sent.notified().await; + assert_eq!(budget.used(), 8 * super::super::BODY_HEADROOM); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert_eq!(budget.used(), 0); + } +} diff --git a/crates/gateway/src/bedrock/sigv4.rs b/crates/gateway/src/bedrock/sigv4.rs index 632d9940..f6dc4f31 100644 --- a/crates/gateway/src/bedrock/sigv4.rs +++ b/crates/gateway/src/bedrock/sigv4.rs @@ -195,27 +195,29 @@ impl Signer { let fail = |what: &str, e: reqwest::Error| SignError::Credentials(format!("{what}: {e}")); let imds = &self.imds; - let token = client + let token_response = client .put(format!("{imds}/latest/api/token")) .header("X-aws-ec2-metadata-token-ttl-seconds", "300") .send() .await .and_then(reqwest::Response::error_for_status) - .map_err(|e| fail("IMDSv2 token request", e))? - .text() + .map_err(|e| fail("IMDSv2 token request", e))?; + let token_bytes = think_watch_common::resources::read_response(token_response, 64 * 1024) .await - .map_err(|e| fail("IMDSv2 token read", e))?; + .map_err(|e| SignError::Credentials(format!("IMDSv2 token read: {e}")))?; + let token = String::from_utf8_lossy(&token_bytes); - let role = client + let role_response = client .get(format!("{imds}/latest/meta-data/iam/security-credentials/")) - .header("X-aws-ec2-metadata-token", &token) + .header("X-aws-ec2-metadata-token", token.as_ref()) .send() .await .and_then(reqwest::Response::error_for_status) - .map_err(|e| fail("IMDSv2 role lookup", e))? - .text() + .map_err(|e| fail("IMDSv2 role lookup", e))?; + let role_bytes = think_watch_common::resources::read_response(role_response, 64 * 1024) .await - .map_err(|e| fail("IMDSv2 role read", e))?; + .map_err(|e| SignError::Credentials(format!("IMDSv2 role read: {e}")))?; + let role = String::from_utf8_lossy(&role_bytes); // One role per instance profile; the listing is one name per line let role = role.lines().next().unwrap_or_default().trim(); if role.is_empty() { @@ -224,18 +226,21 @@ impl Signer { )); } - let creds: serde_json::Value = client + let credentials_response = client .get(format!( "{imds}/latest/meta-data/iam/security-credentials/{role}" )) - .header("X-aws-ec2-metadata-token", &token) + .header("X-aws-ec2-metadata-token", token.as_ref()) .send() .await .and_then(reqwest::Response::error_for_status) - .map_err(|e| fail("IMDSv2 credentials fetch", e))? - .json() - .await - .map_err(|e| fail("IMDSv2 credentials parse", e))?; + .map_err(|e| fail("IMDSv2 credentials fetch", e))?; + let credentials_bytes = + think_watch_common::resources::read_response(credentials_response, 64 * 1024) + .await + .map_err(|e| SignError::Credentials(format!("IMDSv2 credentials read: {e}")))?; + let creds: serde_json::Value = serde_json::from_slice(&credentials_bytes) + .map_err(|e| SignError::Credentials(format!("IMDSv2 credentials parse: {e}")))?; let field = |k: &str| { creds[k] diff --git a/crates/gateway/src/cache.rs b/crates/gateway/src/cache.rs index e8e3b933..42eef3c9 100644 --- a/crates/gateway/src/cache.rs +++ b/crates/gateway/src/cache.rs @@ -7,6 +7,7 @@ use xxhash_rust::xxh3::xxh3_128; /// A cached answer, in the caller's format with redaction placeholders intact. pub struct Cached { + pub memory: bytes::Bytes, pub body: Vec, /// Kept so a hit counts toward token limits and budgets the way the /// original call did. @@ -135,17 +136,20 @@ impl ResponseCache { /// Look up a cached answer. pub async fn get(&self, fingerprint: &[u8]) -> Option { let key = Self::cache_key_for(fingerprint); - let stored: Option = self.redis.get(&key).await.ok().flatten(); - stored.and_then(|json| { - serde_json::from_str::(&json) + let json = think_watch_common::resources::cached_body(&self.redis, &key).await?; + let memory = json.clone(); + think_watch_common::resources::cpu(json.len(), || { + serde_json::from_slice::(&json) .map_err(|e| tracing::warn!("Failed to read a cached response: {e}")) .ok() .map(|s| Cached { + memory, body: serde_json::to_vec(&s.body).unwrap_or_default(), prompt_tokens: s.prompt_tokens, completion_tokens: s.completion_tokens, }) }) + .await } /// Invalidate all cached responses by deleting keys matching the diff --git a/crates/gateway/src/lifecycle/mod.rs b/crates/gateway/src/lifecycle/mod.rs index 0d574b43..a844499d 100644 --- a/crates/gateway/src/lifecycle/mod.rs +++ b/crates/gateway/src/lifecycle/mod.rs @@ -51,6 +51,8 @@ pub(crate) struct ChatCompletionSurface; /// A whole answer, in the caller's format. pub struct Completed { + /// Retain the upstream byte reservation through conversion and audit. + pub memory: Option, /// As the upstream answered, in the caller's format, with redaction /// placeholders still in place and the answer's text not yet through /// the content rules — the form the cache stores, so a later caller @@ -96,6 +98,7 @@ pub struct ChatStreamCaptured { /// attempt and by the streaming tail task. Owned rather than borrowed so /// the detached tail needs no `'static` borrow. pub(crate) struct ChatRequestSnapshot { + pub memory: think_watch_common::resources::SharedStreamBudget, pub identity: GatewayRequestIdentity, /// The one id every audit row and gateway log carries for this request. pub trace_id: String, @@ -108,7 +111,7 @@ pub(crate) struct ChatRequestSnapshot { /// stripped anything, before outbound redaction: the audit row is the /// record of what the user wrote. Body capture applies its own /// redaction toggle on top. - pub request_for_audit: Vec, + pub request_for_audit: Bytes, /// Where the cache keeps this request's answer. `None` when the /// request must not be cached. pub cache_fingerprint: Option>, @@ -260,6 +263,7 @@ pub(crate) fn build_chat_pump( provider: &str, min_speed: Option, ) -> (axum::response::Response, StreamTail) { + let memory = request.memory.clone(); let readers = Arc::new(Mutex::new(Readers::default())); let readers_for_tail = Arc::clone(&readers); @@ -460,7 +464,7 @@ pub(crate) fn build_chat_pump( backoff = Some(wait); continue 'attempt; } - match deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client) { + match think_watch_common::resources::cpu(0, || deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client)).await { Ok(out) => { if !out.is_empty() { if last_frame.is_in(&out) { @@ -490,16 +494,27 @@ pub(crate) fn build_chat_pump( }; match item { Ok(chunk) => { + if !memory.lock().unwrap_or_else(|e| e.into_inner()).append(chunk.len()) { + metrics::counter!("gateway_memory_busy_total").increment(1); + let err = GatewayError::rate_limited("gateway_memory", 1); + let out = cut(&shaper, convert.as_mut(), client, &[], &err); + tell(&mut done_tx, StreamOutcome::UpstreamError { + error_type: err.error_tag().into(), + message: err.to_string(), status_code: err.status_code(), + }); + yield Ok(Bytes::from(out)); + return; + } if let Some(s) = speed.as_mut() { s.feed(&chunk, std::time::Instant::now()); } - let (sealed, said) = match readers.lock() { + let (sealed, said) = think_watch_common::resources::cpu(chunk.len(), || match readers.lock() { Ok(mut r) => { r.feed(&chunk); (r.sealed_refusal(), r.reported(false)) } Err(_) => (None, None), - }; + }).await; // The upstream cannot read the conversation's encrypted // history: the client is told why in plain words, in // place of the upstream's own error. @@ -538,7 +553,7 @@ pub(crate) fn build_chat_pump( hold.take().map(Opening::release).unwrap_or_default() } }; - match deliver(&chunk, &mut convert, &mut shaper, inspector.as_mut(), client) { + match think_watch_common::resources::cpu(chunk.len(), || deliver(&chunk, &mut convert, &mut shaper, inspector.as_mut(), client)).await { Ok(out) => { if !out.is_empty() { if last_frame.is_in(&out) { @@ -569,7 +584,7 @@ pub(crate) fn build_chat_pump( continue 'attempt; } // The opening it sent goes out before the error. - match deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client) { + match think_watch_common::resources::cpu(0, || deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client)).await { Ok(out) => { if !out.is_empty() { yield Ok(Bytes::from(out)); @@ -1513,6 +1528,7 @@ impl Surface for ChatCompletionSurface { }; if let Some(body) = body { let cached = crate::cache::Cached { + memory: Bytes::new(), body: body.clone(), prompt_tokens, completion_tokens, diff --git a/crates/gateway/src/proxy/body_capture.rs b/crates/gateway/src/proxy/body_capture.rs index c2889f8d..dc024ee0 100644 --- a/crates/gateway/src/proxy/body_capture.rs +++ b/crates/gateway/src/proxy/body_capture.rs @@ -198,7 +198,7 @@ async fn process_body( offloaded_flag: &mut bool, ) -> String { if redact_pii { - s = redaction.redact_blob(&s); + s = think_watch_common::resources::cpu(s.len(), || redaction.redact_blob(&s)).await; } if s.len() <= max_bytes { return s; diff --git a/crates/gateway/src/proxy/generate.rs b/crates/gateway/src/proxy/generate.rs index eb85bc3d..8321c059 100644 --- a/crates/gateway/src/proxy/generate.rs +++ b/crates/gateway/src/proxy/generate.rs @@ -32,6 +32,7 @@ //! are thinkwatch-core's, shared with the desktop gateway. use std::convert::Infallible; +use think_watch_common::resources; use crate::call_ctx::CallCtx; use crate::error::GatewayError; @@ -628,9 +629,12 @@ pub(crate) async fn send( ) -> Result<(reqwest::Response, Wire), GatewayError> { let upstream = &entry.upstream; let official = upstream.is_official(); - let first = outbound - .address(entry.protocol, model, official)? - .for_upstream(upstream, model); + let first = resources::cpu(usize::MAX, || { + outbound + .address(entry.protocol, model, official) + .map(|wire| wire.for_upstream(upstream, model)) + }) + .await?; // Only a request that can carry an effort asks the memory. let memory = if outbound.reasoning.any { effort::recall(redis, entry.provider_id, first.dialect, model).await @@ -724,9 +728,12 @@ pub(crate) async fn send( to = %protocol, "Upstream rejected the configured protocol — retrying with an alternate" ); - let wire = outbound - .address(*protocol, model, official)? - .for_upstream(upstream, model); + let wire = resources::cpu(usize::MAX, || { + outbound + .address(*protocol, model, official) + .map(|wire| wire.for_upstream(upstream, model)) + }) + .await?; let wire = if outbound.reasoning.any { let memory = effort::recall(redis, entry.provider_id, wire.dialect, model).await; wire.remembered(&memory) @@ -770,33 +777,46 @@ pub(crate) async fn read_whole( caller_model: &str, input_estimate: u64, ) -> Result { - let upstream = resp - .bytes() + let upstream = resources::read_response(resp, resources::policy().max_body_bytes) .await - .map_err(super::transport::transport_error)?; - - let mut sniffer = tw_dialect::usage::Sniffer::new(); - sniffer.feed(&upstream); - let (usage, usage_estimated) = - crate::usage_estimate::complete(sniffer.finish(), true, input_estimate, Some(&upstream)); - if usage_estimated { - metrics::counter!("gateway_usage_estimated_total").increment(1); - } + .map_err(|error| match error { + resources::ReadError::Busy => GatewayError::rate_limited("gateway_memory", 1), + resources::ReadError::Transport(error) => super::transport::transport_error(error), + resources::ReadError::TooLarge => { + GatewayError::ProviderInvalidResponse(error.to_string()) + } + })?; + let memory = upstream.clone(); + resources::cpu(upstream.len(), || { + let mut sniffer = tw_dialect::usage::Sniffer::new(); + sniffer.feed(&upstream); + let (usage, usage_estimated) = crate::usage_estimate::complete( + sniffer.finish(), + true, + input_estimate, + Some(&upstream), + ); + if usage_estimated { + metrics::counter!("gateway_usage_estimated_total").increment(1); + } - let body = match &wire.convert { - Some(session) => session.response(&upstream).ok_or_else(|| { - GatewayError::ProviderInvalidResponse( - "The upstream answered with something that is not JSON.".into(), - ) - })?, - None => upstream.to_vec(), - }; - Ok(Completed { - body: rewrite_model(&body, caller_model), - logged: None, - usage, - usage_estimated, + let body = match &wire.convert { + Some(session) => session.response(&upstream).ok_or_else(|| { + GatewayError::ProviderInvalidResponse( + "The upstream answered with something that is not JSON.".into(), + ) + })?, + None => upstream.to_vec(), + }; + Ok(Completed { + memory: Some(memory), + body: rewrite_model(&body, caller_model), + logged: None, + usage, + usage_estimated, + }) }) + .await } // ───────────────────────────────────────────── the pipeline @@ -860,11 +880,14 @@ async fn run( "(unknown)", request_started_at, ); - let raw: Value = serde_json::from_slice(&body).map_err(|_| { - early_ctx.emit(GatewayError::TransformError( - "The request body is not valid JSON.".into(), - )) - })?; + let original_body = body.clone(); + let raw: Value = resources::cpu(body.len(), || serde_json::from_slice(&body)) + .await + .map_err(|_| { + early_ctx.emit(GatewayError::TransformError( + "The request body is not valid JSON.".into(), + )) + })?; let (model, is_stream) = if surface.dialect == Dialect::Gemini { gemini_target(path).ok_or_else(|| { early_ctx.emit(GatewayError::TransformError(format!( @@ -927,7 +950,8 @@ async fn run( // the caller's 403, quoting their words (masked). The rules that // check answers run on the way back (the shaper for a stream, below // for a whole answer). - let screening = guards.content.screen(surface.dialect, &body); + let screening = + resources::cpu(body.len(), || guards.content.screen(surface.dialect, &body)).await; crate::content_filter::record(&state.audit, &caller, guards.content.mode, &screening); if let Some(hit) = screening.refusal() { return Err(ctx @@ -938,12 +962,14 @@ async fn run( // every hop, the audit row — the request is the changed one. let (body, raw) = match screening.body { Some(stripped) => { - let raw = serde_json::from_slice(&stripped).map_err(|_| { - ctx.emit(GatewayError::TransformError( + let raw = resources::cpu(stripped.len(), || serde_json::from_slice(&stripped)) + .await + .map_err(|_| { + ctx.emit(GatewayError::TransformError( "The request could not be read after the content filter changed text in it." .into(), )) - })?; + })?; (stripped, raw) } None => (body, raw), @@ -954,8 +980,10 @@ async fn run( // conversation kept on OpenAI's servers (`previous_response_id`), or // carries a compaction only OpenAI can read — can still be forwarded // to an upstream of its own format, so that is where it goes (step 8). - let decoded = - tw_dialect::convert::decode(surface.dialect, &raw, path, internal_query(surface.dialect)); + let decoded = resources::cpu(body.len(), || { + tw_dialect::convert::decode(surface.dialect, &raw, path, internal_query(surface.dialect)) + }) + .await; // 5. Outbound redaction: the text written in the request — located by // the structure of the caller's format, never its ids, types or @@ -968,19 +996,25 @@ async fn run( // The audit row keeps what the caller wrote (their at-rest // redaction setting applies on top). let shape = Shape::Generate(surface.dialect); - let (findings, ledger) = guards.redaction.look(&body, shape); + let (findings, ledger) = + resources::cpu(body.len(), || guards.redaction.look(&body, shape)).await; crate::redaction::record(&state.audit, &caller, &guards.redaction, &findings); - let request_for_audit = body.to_vec(); - let outbound_body = if ledger.is_empty() { - raw - } else { - let (replaced, _) = guards.redaction.replace(body.to_vec(), &ledger, shape); - serde_json::from_slice(&replaced).map_err(|_| { - ctx.emit(GatewayError::TransformError( - "The request could not be read after redaction.".into(), - )) - })? - }; + let request_for_audit = resources::retain(body.clone(), original_body); + let outbound_body = resources::cpu(body.len(), || { + if ledger.is_empty() { + Ok::<_, GatewayErrorResponse>(raw) + } else { + let (replaced, _) = guards.redaction.replace(body.to_vec(), &ledger, shape); + serde_json::from_slice(&replaced) + .map_err(|_| { + ctx.emit(GatewayError::TransformError( + "The request could not be read after redaction.".into(), + )) + }) + .map_err(Into::into) + } + }) + .await?; // 6. The model's settings. The output cap is applied to each hop as it // is addressed (`Outbound::cap_output`): whether to fill one in @@ -1005,22 +1039,27 @@ async fn run( // budgets like a real call would — the caller received them, and // otherwise repeating a cached prompt would get round every token // limit. - let cache_fingerprint = if surface.caches { - ResponseCache::fingerprint(&outbound_body).map(|mut fp| { - // The cap and the reasoning setting are applied after this, per - // hop: an answer made under one must not be served under - // another. - if let Some(cap) = max_output_tokens { - fp.extend_from_slice(format!("\nmax_output_tokens={cap}").as_bytes()); - } - if let Some(setting) = reasoning.setting { - fp.extend_from_slice(format!("\nreasoning_effort={}", setting.slug()).as_bytes()); - } - fp - }) - } else { - None - }; + let cache_fingerprint = resources::cpu(body.len(), || { + if surface.caches { + ResponseCache::fingerprint(&outbound_body).map(|mut fp| { + // The cap and the reasoning setting are applied after this, per + // hop: an answer made under one must not be served under + // another. + if let Some(cap) = max_output_tokens { + fp.extend_from_slice(format!("\nmax_output_tokens={cap}").as_bytes()); + } + if let Some(setting) = reasoning.setting { + fp.extend_from_slice( + format!("\nreasoning_effort={}", setting.slug()).as_bytes(), + ); + } + fp + }) + } else { + None + } + }) + .await; if let Some(fp) = &cache_fingerprint && let Some(cached) = state.cache.get(fp).await { @@ -1028,23 +1067,33 @@ async fn run( // With this caller's values in it: what they would run. A stored // answer passed the content rules and the inspection in force when // it was stored, not necessarily the ones in force now. - let restored = crate::redaction::restore_body(&ledger, &cached.body); - let restored = crate::content_filter::whole( - &guards.content, - &state.audit, - &caller, - "cache", - surface.dialect, - restored, - ); - if let Some(e) = crate::tool_inspection::check_whole( - &guards.tools, - &guards.redaction, - &state.audit, - &caller, - "cache", - &restored, - ) { + let restored = resources::cpu(cached.body.len(), || { + crate::redaction::restore_body(&ledger, &cached.body) + }) + .await; + let restored = resources::cpu(restored.len(), || { + crate::content_filter::whole( + &guards.content, + &state.audit, + &caller, + "cache", + surface.dialect, + restored, + ) + }) + .await; + if let Some(e) = resources::cpu(restored.len(), || { + crate::tool_inspection::check_whole( + &guards.tools, + &guards.redaction, + &state.audit, + &caller, + "cache", + &restored, + ) + }) + .await + { return Err(ctx.emit(e).into()); } // The tokens the stored answer records, weighted for this model, @@ -1110,12 +1159,13 @@ async fn run( // client gets the framing it asked for. let body = String::from_utf8_lossy(&restored).into_owned(); let events = async_stream::stream! { + let _memory = cached.memory; yield Ok::<_, Infallible>(axum::response::sse::Event::default().data(body)); yield Ok::<_, Infallible>(axum::response::sse::Event::default().data("[DONE]")); }; axum::response::sse::Sse::new(events).into_response() } else { - json_response(restored) + json_response(resources::retain(restored.into(), cached.memory)) }; response .headers_mut() @@ -1213,6 +1263,7 @@ async fn run( state: state.clone(), guards: guards.clone(), request: crate::lifecycle::ChatRequestSnapshot { + memory: resources::stream_budget(), identity: identity.clone(), trace_id: metadata.request_id.clone(), session_id: session_id.clone(), @@ -1361,28 +1412,39 @@ async fn run( // The answer as the caller will receive it: their values restored, // then its text through the content rules that check answers. - let restored = crate::redaction::restore_body(&answer_ledger, &completed.body); - let restored = crate::content_filter::whole( - &guards.content, - &state.audit, - &caller, - &entry.provider_name, - surface.dialect, - restored, - ); + let response_memory = completed.memory.clone(); + let restored = resources::cpu(completed.body.len(), || { + crate::redaction::restore_body(&answer_ledger, &completed.body) + }) + .await; + let restored = resources::cpu(restored.len(), || { + crate::content_filter::whole( + &guards.content, + &state.audit, + &caller, + &entry.provider_name, + surface.dialect, + restored, + ) + }) + .await; // Tool calls, on the whole answer before any of it has gone out — in // the form the caller would run them. A refusal is the gateway's // policy, not the upstream failing, so the route's health counts it as // a success; the answer is neither cached nor billed. - if let Some(e) = crate::tool_inspection::check_whole( - &guards.tools, - &guards.redaction, - &state.audit, - &caller, - &entry.provider_name, - &restored, - ) { + if let Some(e) = resources::cpu(restored.len(), || { + crate::tool_inspection::check_whole( + &guards.tools, + &guards.redaction, + &state.audit, + &caller, + &entry.provider_name, + &restored, + ) + }) + .await + { finalize_health(&state, &sel_record, true).await; return Err(ctx.emit(e).into()); } @@ -1408,7 +1470,10 @@ async fn run( "Audit log: request completed" ); - let mut response = json_response(restored); + let mut response = json_response(resources::retain( + restored.into(), + response_memory.unwrap_or_default(), + )); response .headers_mut() .insert("X-Cache", HeaderValue::from_static("MISS")); @@ -1475,13 +1540,13 @@ fn dialect_headers(headers: &HeaderMap) -> Vec<(String, String)> { .collect() } -fn json_response(body: Vec) -> axum::response::Response { +fn json_response(body: impl Into) -> axum::response::Response { ( [( header::CONTENT_TYPE, HeaderValue::from_static("application/json"), )], - body, + body.into(), ) .into_response() } diff --git a/crates/gateway/src/proxy/transport.rs b/crates/gateway/src/proxy/transport.rs index 84431279..8d301b4f 100644 --- a/crates/gateway/src/proxy/transport.rs +++ b/crates/gateway/src/proxy/transport.rs @@ -304,7 +304,10 @@ async fn check_status( } if !status.is_success() { let retry_after = retry_after(resp.headers()); - let body = resp.text().await.unwrap_or_default(); + let body = think_watch_common::resources::read_response(resp, 64 * 1024) + .await + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .unwrap_or_else(|error| error.to_string()); tracing::warn!(provider = label, status = %status, body = %body, "upstream returned non-2xx"); const CLIENT_MAX: usize = 512; let shown = if body.len() > CLIENT_MAX { diff --git a/crates/mcp-gateway/src/cache.rs b/crates/mcp-gateway/src/cache.rs index a34ad8aa..50bb818d 100644 --- a/crates/mcp-gateway/src/cache.rs +++ b/crates/mcp-gateway/src/cache.rs @@ -105,16 +105,26 @@ impl McpResponseCache { ) -> Option { let scope_label = scope_metric_label(caller); let key = Self::cache_key(server_id, caller, request); - let cached: Option = self.redis.get(&key).await.ok().flatten(); + let cached = think_watch_common::resources::cached_body(&self.redis, &key).await; let hit = cached.is_some(); - let parsed = cached.and_then(|json| { - serde_json::from_str::(&json) - .map_err(|e| { - tracing::warn!("Failed to deserialize cached MCP response: {e}"); - e - }) - .ok() - }); + let parsed = if let Some(json) = cached { + let memory = json.clone(); + think_watch_common::resources::cpu(json.len(), || { + serde_json::from_slice::(&json) + .map_err(|e| { + tracing::warn!("Failed to deserialize cached MCP response: {e}"); + e + }) + .ok() + .map(|mut response| { + response.memory = Some(memory); + response + }) + }) + .await + } else { + None + }; // Three counters per lookup, each labelled by scope so the // dashboard can break hit-rate down by global vs per-caller: // * mcp_cache_total — denominator diff --git a/crates/mcp-gateway/src/lifecycle/mod.rs b/crates/mcp-gateway/src/lifecycle/mod.rs index 095f343f..435f0e7e 100644 --- a/crates/mcp-gateway/src/lifecycle/mod.rs +++ b/crates/mcp-gateway/src/lifecycle/mod.rs @@ -150,6 +150,10 @@ impl Surface for McpSurface { // shared response-code classifier so a single user's bad // INVALID_PARAMS doesn't trip the breaker for everyone. match &invoked.view { + CapturedView::Streaming { + outcome: StreamOutcome::UpstreamError { error_type, .. }, + .. + } if error_type == "gateway_memory" => {} CapturedView::Streaming { outcome: StreamOutcome::UpstreamError { .. }, .. @@ -240,6 +244,8 @@ impl Surface for McpSurface { /// shaped); the rest are surface-specific identifiers the pipeline /// is generic over, so they cannot live on `Invoked` itself. pub struct McpPostInvokeDeps { + /// Retains stream capture admission until cache/audit hooks finish. + pub stream_memory: think_watch_common::resources::SharedStreamBudget, pub proxy: McpProxy, pub server_id: Uuid, pub server_name: String, @@ -278,6 +284,19 @@ fn response_for_hooks(invoked: &Invoked, deps: &McpPostInvokeDeps) - match &invoked.view { CapturedView::Buffered(r) => r.clone(), CapturedView::Streaming { outcome, captured } => { + if let StreamOutcome::UpstreamError { + error_type, + message, + .. + } = outcome + && error_type == "gateway_memory" + { + return err_response( + deps.original_request_id.clone(), + crate::proxy::RESOURCE_BUSY, + message.clone(), + ); + } pick_response_envelope(captured, deps.original_request_id.as_ref()).unwrap_or_else( || { let msg = match outcome { diff --git a/crates/mcp-gateway/src/pool.rs b/crates/mcp-gateway/src/pool.rs index 9c5b74e2..d95892e4 100644 --- a/crates/mcp-gateway/src/pool.rs +++ b/crates/mcp-gateway/src/pool.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::sync::Arc; use reqwest::Client; +use think_watch_common::resources; use tokio::sync::RwLock; use uuid::Uuid; @@ -55,6 +56,9 @@ pub enum PoolError { #[error("Failed to parse upstream JSON-RPC response: {0}")] ParseError(String), + + #[error(transparent)] + Buffering(#[from] resources::ReadError), } /// Manages a pool of `McpConnection`s keyed by server ID. @@ -219,7 +223,10 @@ impl ConnectionPool { let status = resp.status(); if !status.is_success() { - let body = resp.text().await.unwrap_or_default(); + let body = resources::read_response(resp, 64 * 1024) + .await + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .unwrap_or_else(|error| error.to_string()); return Err(PoolError::UpstreamError { status: status.as_u16(), body, @@ -237,28 +244,30 @@ impl ConnectionPool { .unwrap_or("") .to_lowercase(); - let (json_resp, stream_audit_body): (JsonRpcResponse, Option) = + let data = resources::read_response(resp, resources::policy().max_body_bytes).await?; + let (mut json_resp, stream_audit_body): (JsonRpcResponse, Option) = if content_type.contains("text/event-stream") { - let text = resp - .text() - .await - .map_err(|e| PoolError::ParseError(e.to_string()))?; + let text = String::from_utf8_lossy(&data); // Pick the response envelope by matching id against the // request id — without this, a tool that emits // `notifications/progress` before the final response // had its first progress notification returned AS the // response (notifications have id=None which deserializes // into JsonRpcResponse just fine, masking the real reply). - let (matched, full_audit) = parse_sse_json_rpc(&text, request.id.as_ref())?; + let (matched, full_audit) = resources::cpu(data.len(), || { + parse_sse_json_rpc(&text, request.id.as_ref()) + }) + .await?; (matched, Some(full_audit)) } else { - let envelope: JsonRpcResponse = resp - .json() - .await - .map_err(|e| PoolError::ParseError(e.to_string()))?; + let envelope: JsonRpcResponse = + resources::cpu(data.len(), || serde_json::from_slice(&data)) + .await + .map_err(|e| PoolError::ParseError(e.to_string()))?; (envelope, None) }; + json_resp.memory = Some(data); // Validate JSON-RPC version if json_resp.jsonrpc != "2.0" { return Err(PoolError::ParseError("Invalid JSON-RPC version".into())); @@ -309,7 +318,10 @@ impl ConnectionPool { // Same error-surfacing contract as send_request — buffer // the small error body so the caller can see the upstream // diagnostic instead of an opaque "non-2xx". - let body = resp.text().await.unwrap_or_default(); + let body = resources::read_response(resp, 64 * 1024) + .await + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .unwrap_or_else(|error| error.to_string()); return Err(PoolError::UpstreamError { status: status.as_u16(), body, @@ -477,6 +489,27 @@ fn parse_sse_json_rpc( Ok((json_resp, audit_body)) } +pub(crate) async fn read_json_response( + response: reqwest::Response, +) -> Result { + let bytes = resources::read_response(response, resources::policy().max_body_bytes).await?; + let mut result: JsonRpcResponse = + resources::cpu(bytes.len(), || serde_json::from_slice(&bytes)) + .await + .map_err(|e| PoolError::ParseError(e.to_string()))?; + result.memory = Some(bytes); + Ok(result) +} + +pub(crate) fn error_response(id: Option, error: &PoolError) -> JsonRpcResponse { + let code = if matches!(error, PoolError::Buffering(resources::ReadError::Busy)) { + crate::proxy::RESOURCE_BUSY + } else { + crate::proxy::INTERNAL_ERROR + }; + crate::proxy::err_response(id, code, format!("Upstream server error: {error}")) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/mcp-gateway/src/proxy.rs b/crates/mcp-gateway/src/proxy.rs index 762c358f..1210d6ad 100644 --- a/crates/mcp-gateway/src/proxy.rs +++ b/crates/mcp-gateway/src/proxy.rs @@ -23,7 +23,7 @@ mod streaming; use jsonrpc::ok_response; pub use jsonrpc::{ INTERNAL_ERROR, INVALID_PARAMS, INVALID_REQUEST, JsonRpcError, JsonRpcRequest, JsonRpcResponse, - METHOD_NOT_FOUND, NEEDS_USER_CREDENTIALS, err_response, + METHOD_NOT_FOUND, NEEDS_USER_CREDENTIALS, RESOURCE_BUSY, err_response, }; use streaming::build_replay_payload; pub(crate) use streaming::pick_response_envelope; @@ -794,6 +794,7 @@ impl McpProxy { ), }; return HandleOutcome::Buffered(JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id: request.id, result: None, @@ -824,6 +825,7 @@ impl McpProxy { CredentialOwner::AdminShared => "/admin/mcp/servers", }; return HandleOutcome::Buffered(JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id: request.id, result: None, @@ -882,6 +884,7 @@ impl McpProxy { // to `run_post_invoke` per branch. let logged_arguments = params.get("arguments").cloned(); let mut deps = crate::lifecycle::McpPostInvokeDeps { + stream_memory: think_watch_common::resources::stream_budget(), proxy: self.clone(), server_id, server_name: server_name.clone(), @@ -958,6 +961,8 @@ impl McpProxy { self.build_mcp_pump( resp, crate::proxy::streaming::PumpContext { + memory: deps.stream_memory.clone(), + request_id: request.id.clone(), identity: auth_identity, trace_id: auth_trace_id, started_at: auth_started_at, @@ -968,13 +973,9 @@ impl McpProxy { ) } else { // Upstream chose application/json — buffer it. - let response = match resp.json::().await { + let response = match crate::pool::read_json_response(resp).await { Ok(r) => r, - Err(e) => err_response( - request.id.clone(), - INTERNAL_ERROR, - format!("Upstream server error: parse failed: {e}"), - ), + Err(e) => crate::pool::error_response(request.id.clone(), &e), }; Invocation::Buffered(Invoked { identity: auth_identity, @@ -993,11 +994,7 @@ impl McpProxy { error = %e, "streaming upstream tools/call failed" ); - let response = err_response( - request.id.clone(), - INTERNAL_ERROR, - format!("Upstream server error: {e}"), - ); + let response = crate::pool::error_response(request.id.clone(), &e); Invocation::Buffered(Invoked { identity: auth_identity, trace_id: auth_trace_id, @@ -1049,11 +1046,7 @@ impl McpProxy { error = %e, "upstream tools/call failed" ); - let response = err_response( - request.id.clone(), - INTERNAL_ERROR, - format!("Upstream server error: {e}"), - ); + let response = crate::pool::error_response(request.id.clone(), &e); Invocation::Buffered(Invoked { identity: auth_identity, trace_id: auth_trace_id, @@ -1077,7 +1070,12 @@ impl McpProxy { let response = emitted .response .expect("Invocation::Buffered always yields Emitted.response"); - if ctx.wants_streaming { + if ctx.wants_streaming + && !response + .error + .as_ref() + .is_some_and(|e| e.code == RESOURCE_BUSY) + { // Client signalled SSE but the upstream replied // buffered. Wrap the response (and any captured // upstream timeline) as a sequence of SSE events. @@ -1120,6 +1118,13 @@ impl McpProxy { server_name: &str, response: &JsonRpcResponse, ) { + if response + .error + .as_ref() + .is_some_and(|error| error.code == RESOURCE_BUSY) + { + return; + } let is_server_failure = response .error .as_ref() diff --git a/crates/mcp-gateway/src/proxy/jsonrpc.rs b/crates/mcp-gateway/src/proxy/jsonrpc.rs index be117652..b9da25f0 100644 --- a/crates/mcp-gateway/src/proxy/jsonrpc.rs +++ b/crates/mcp-gateway/src/proxy/jsonrpc.rs @@ -15,6 +15,9 @@ pub struct JsonRpcRequest { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct JsonRpcResponse { + /// Internal ownership only; excluded from the JSON-RPC wire format. + #[serde(skip)] + pub memory: Option, pub jsonrpc: String, pub id: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -36,6 +39,7 @@ pub const INVALID_REQUEST: i32 = -32600; pub const METHOD_NOT_FOUND: i32 = -32601; pub const INVALID_PARAMS: i32 = -32602; pub const INTERNAL_ERROR: i32 = -32603; +pub const RESOURCE_BUSY: i32 = -32005; /// ThinkWatch-specific: the calling user (or the calling API key's /// chosen account label) has no upstream credential for the requested /// MCP server. The `data` field carries `server_id` and an optional @@ -48,6 +52,7 @@ pub(crate) fn ok_response( result: serde_json::Value, ) -> JsonRpcResponse { JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id, result: Some(result), @@ -61,6 +66,7 @@ pub fn err_response( message: impl Into, ) -> JsonRpcResponse { JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id, result: None, diff --git a/crates/mcp-gateway/src/proxy/streaming.rs b/crates/mcp-gateway/src/proxy/streaming.rs index d69834c9..2312e252 100644 --- a/crates/mcp-gateway/src/proxy/streaming.rs +++ b/crates/mcp-gateway/src/proxy/streaming.rs @@ -108,6 +108,8 @@ pub(crate) fn pick_response_envelope( /// /// [`run_post_invoke`]: think_watch_common::lifecycle::stages::run_post_invoke pub(crate) struct PumpContext { + pub memory: think_watch_common::resources::SharedStreamBudget, + pub request_id: Option, pub identity: crate::lifecycle::McpIdentity, pub trace_id: String, pub started_at: std::time::Instant, @@ -151,7 +153,13 @@ impl McpProxy { let bytes_source = upstream_resp .bytes_stream() .map(|r| r.map_err(|e| e.to_string())); - let body = build_passthrough_body(bytes_source, events_buf, done_tx); + let body = build_passthrough_body( + bytes_source, + events_buf, + done_tx, + ctx.memory.clone(), + ctx.request_id, + ); let payload = StreamingPayload { body, @@ -205,6 +213,8 @@ pub(crate) fn build_passthrough_body( source: S, events_buf: std::sync::Arc>>, done_tx: tokio::sync::oneshot::Sender, + memory: think_watch_common::resources::SharedStreamBudget, + request_id: Option, ) -> std::pin::Pin< Box< dyn futures::stream::Stream< @@ -232,6 +242,23 @@ where while let Some(chunk) = source.next().await { match chunk { Ok(bytes) => { + let admitted = memory.lock().unwrap_or_else(|e| e.into_inner()).append(bytes.len()); + if !admitted { + metrics::counter!("gateway_memory_busy_total").increment(1); + metrics::counter!("mcp_memory_busy_total").increment(1); + if let Some(tx) = done_tx.take() { + let _ = tx.send(StreamOutcome::UpstreamError { + error_type: "gateway_memory".into(), + message: "Gateway buffering capacity is occupied; retry shortly".into(), + status_code: 429, + }); + } + let error = crate::proxy::err_response(request_id.clone(), crate::proxy::RESOURCE_BUSY, + "Gateway buffering capacity is occupied; retry shortly"); + yield Ok::(Event::default().data(serde_json::to_string(&error).unwrap())); + text_buf.clear(); + break; + } // Reject the chunk BEFORE the allocation when its // post-append size would push us past the cap. // Tearing down the stream is preferable to letting @@ -264,7 +291,7 @@ where // Best-effort JSON parse so non-JSON // `data:` payloads (rare) still land in // the timeline verbatim. - let parsed = serde_json::from_str::(&payload) + let parsed = think_watch_common::resources::cpu(payload.len(), || serde_json::from_str::(&payload)).await .unwrap_or_else(|_| { serde_json::Value::String(payload.clone()) }); @@ -303,7 +330,7 @@ where && let Some(payload) = extract_sse_data_payload(trailing) && !payload.is_empty() { - let parsed = serde_json::from_str::(&payload) + let parsed = think_watch_common::resources::cpu(payload.len(), || serde_json::from_str::(&payload)).await .unwrap_or_else(|_| serde_json::Value::String(payload.clone())); if let Ok(mut g) = events_buf.lock() { g.push(parsed); @@ -352,11 +379,12 @@ pub(crate) fn build_replay_payload( { events.push(s); } - let stream = futures::stream::iter( - events - .into_iter() - .map(|s| Ok::<_, Infallible>(Event::default().data(s))), - ); + let stream = async_stream::stream! { + let _memory = response.memory; + for event in events { + yield Ok::<_, Infallible>(Event::default().data(event)); + } + }; StreamingPayload { body: Box::pin(stream), new_session_id: None, @@ -368,8 +396,24 @@ mod replay_tests { use super::*; use futures::StreamExt; + #[tokio::test] + async fn replay_retains_buffer_admission_until_the_slow_consumer_drops() { + let budget = think_watch_common::resources::ByteBudget::new(100); + let mut lease = budget.reserve(); + assert!(lease.grow(100)); + let mut response = make_response(1, serde_json::json!({"ok":true})); + response.memory = Some(lease.wrap(bytes::Bytes::from_static(b"upstream"))); + let mut payload = build_replay_payload(response, None); + assert_eq!(budget.used(), 100); + let _ = payload.body.next().await.unwrap().unwrap(); + assert_eq!(budget.used(), 100); + drop(payload); + assert_eq!(budget.used(), 0); + } + fn make_response(id: u64, result: serde_json::Value) -> JsonRpcResponse { JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id: Some(serde_json::json!(id)), result: Some(result), @@ -537,13 +581,54 @@ mod passthrough_tests { use futures::stream::{self, StreamExt}; use std::sync::{Arc, Mutex}; + #[tokio::test] + async fn small_events_share_admission_and_a_detached_tail_keeps_it_after_cancel() { + use think_watch_common::resources::{BODY_HEADROOM, ByteBudget, StreamBudget}; + let chunk = Bytes::from_static(b"data: {\"result\":1}\n\n"); + let budget = ByteBudget::new(chunk.len() * BODY_HEADROOM); + let memory = Arc::new(Mutex::new(StreamBudget::new(&budget, 1024))); + let tail_memory = memory.clone(); + let events = Arc::new(Mutex::new(Vec::new())); + let (tx, rx) = tokio::sync::oneshot::channel(); + let mut body = build_passthrough_body( + stream::iter([Ok(chunk.clone()), Ok(chunk.clone())]), + events.clone(), + tx, + memory, + Some(serde_json::json!(42)), + ); + let _ = body.next().await.unwrap().unwrap(); + let busy = body.next().await.unwrap().unwrap(); + let busy = format!("{busy:?}"); + assert!(busy.contains("-32005"), "{busy}"); + assert!( + busy.contains("42"), + "the busy response must identify the request: {busy}" + ); + assert!( + matches!(rx.await.unwrap(), StreamOutcome::UpstreamError { error_type, status_code: 429, .. } if error_type == "gateway_memory") + ); + assert_eq!(events.lock().unwrap().len(), 1); + drop(body); + assert_eq!(budget.used(), chunk.len() * BODY_HEADROOM); + drop(events); + drop(tail_memory); + assert_eq!(budget.used(), 0); + } + async fn run_body( chunks: Vec>, ) -> (usize, Vec, Option<&'static str>) { let events_buf: Arc>> = Arc::new(Mutex::new(Vec::new())); let (done_tx, done_rx) = tokio::sync::oneshot::channel::(); let source = stream::iter(chunks); - let body = build_passthrough_body(source, events_buf.clone(), done_tx); + let body = build_passthrough_body( + source, + events_buf.clone(), + done_tx, + think_watch_common::resources::stream_budget(), + None, + ); let mut count = 0usize; let mut body = body; while let Some(_ev) = body.next().await { @@ -629,7 +714,13 @@ mod passthrough_tests { "data: {\"jsonrpc\":\"2.0\",\"method\":\"x\"}\n\n", ))]) .chain(stream::pending()); - let body = build_passthrough_body(source, events_buf.clone(), done_tx); + let body = build_passthrough_body( + source, + events_buf.clone(), + done_tx, + think_watch_common::resources::stream_budget(), + None, + ); let mut body = body; let _ = body.next().await; drop(body); diff --git a/crates/mcp-gateway/src/transport/streamable_http.rs b/crates/mcp-gateway/src/transport/streamable_http.rs index a0f86554..b2455ac7 100644 --- a/crates/mcp-gateway/src/transport/streamable_http.rs +++ b/crates/mcp-gateway/src/transport/streamable_http.rs @@ -108,7 +108,7 @@ pub async fn handle_post( State(state): State>, headers: HeaderMap, axum::Extension(identity): axum::Extension, - Json(request): Json, + think_watch_common::resources::http::ScheduledJson(request): think_watch_common::resources::http::ScheduledJson, ) -> impl IntoResponse { // --- Session management ------------------------------------------------ let sessions = &state.proxy.sessions; @@ -162,6 +162,16 @@ pub async fn handle_post( match outcome { HandleOutcome::Buffered(response) => { + let memory = response.memory.clone().unwrap_or_default(); + if response + .error + .as_ref() + .is_some_and(|e| e.code == crate::proxy::RESOURCE_BUSY) + { + resp_headers.insert(axum::http::header::RETRY_AFTER, "1".parse().unwrap()); + return (StatusCode::TOO_MANY_REQUESTS, resp_headers, Json(response)) + .into_response(); + } // Buffered shape happens for initialize / tools/list, for // tools/call cache hits / early rejections, AND for // tools/call when the upstream replied with plain @@ -172,13 +182,25 @@ pub async fn handle_post( // plain JSON-RPC clients. if sse_capable { let payload = serde_json::to_string(&response).unwrap_or_default(); - let event = Ok::<_, std::convert::Infallible>(Event::default().data(payload)); - let stream = futures::stream::iter(std::iter::once(event)); + let stream = async_stream::stream! { + let _memory = memory; + yield Ok::<_, std::convert::Infallible>(Event::default().data(payload)); + }; let mut sse_resp = Sse::new(stream).into_response(); sse_resp.headers_mut().extend(resp_headers); sse_resp } else { - (StatusCode::OK, resp_headers, Json(response)).into_response() + let bytes = serde_json::to_vec(&response).unwrap_or_default(); + resp_headers.insert( + axum::http::header::CONTENT_TYPE, + "application/json".parse().unwrap(), + ); + ( + StatusCode::OK, + resp_headers, + think_watch_common::resources::retain(bytes.into(), memory), + ) + .into_response() } } HandleOutcome::Streaming(payload) => { diff --git a/crates/mcp-gateway/src/user_token.rs b/crates/mcp-gateway/src/user_token.rs index 24e1df1e..c92f008a 100644 --- a/crates/mcp-gateway/src/user_token.rs +++ b/crates/mcp-gateway/src/user_token.rs @@ -825,7 +825,10 @@ impl UserTokenResolver { }; let status = resp.status(); - let text = resp.text().await.unwrap_or_default(); + let bytes = think_watch_common::resources::read_response(resp, 64 * 1024) + .await + .map_err(|e| OAuthRefreshFailure::Transient(format!("token response read: {e}")))?; + let text = String::from_utf8_lossy(&bytes); if !status.is_success() { // Only treat actually-irrecoverable grant rejections as // Permanent — the original "any 4xx → Permanent" rule diff --git a/crates/server/Cargo.toml b/crates/server/Cargo.toml index c14e4140..49796768 100644 --- a/crates/server/Cargo.toml +++ b/crates/server/Cargo.toml @@ -11,6 +11,8 @@ name = "think-watch-server" path = "src/main.rs" [dependencies] +mimalloc = { version = "0.1.52", features = ["no_thp"] } +libc = "0.2" tw-breaker = { workspace = true } think-watch-common = { workspace = true } think-watch-auth = { workspace = true } diff --git a/crates/server/src/app.rs b/crates/server/src/app.rs index e4e551dc..91349b97 100644 --- a/crates/server/src/app.rs +++ b/crates/server/src/app.rs @@ -440,6 +440,9 @@ pub async fn create_gateway_app(_config: &AppConfig, state: AppState) -> anyhow: .merge(ai_routes) .merge(key_routes) .merge(mcp_routes) + .layer(axum::middleware::from_fn( + think_watch_common::resources::http::admit, + )) .layer(RequestBodyLimitLayer::new(10 * 1024 * 1024)) // 10MB for large prompts // When the timeout below starts counting: a whole answer's retries // stop in time for its error to be written (see `gateway_proxy::retry`). @@ -1186,6 +1189,7 @@ pub fn create_console_app(config: &AppConfig, state: AppState) -> anyhow::Result // upgrade path bypasses tower-http body transforms); permessage- // deflate negotiation belongs at the reverse-proxy layer. .layer(CompressionLayer::new()) + .layer(axum::middleware::from_fn(think_watch_common::resources::http::admit)) .layer(RequestBodyLimitLayer::new(1024 * 1024)) // 1MB for console API .layer(TimeoutLayer::with_status_code( axum::http::StatusCode::REQUEST_TIMEOUT, diff --git a/crates/server/src/handlers/mcp_oauth.rs b/crates/server/src/handlers/mcp_oauth.rs index 88311499..f9297494 100644 --- a/crates/server/src/handlers/mcp_oauth.rs +++ b/crates/server/src/handlers/mcp_oauth.rs @@ -1340,14 +1340,16 @@ pub(crate) async fn resolve_upstream_subject( return Some(s); } let endpoint = userinfo_endpoint?; - let body = match http + let (body, _memory): (serde_json::Value, _) = match http .get(endpoint) .bearer_auth(access_token) .header("Accept", "application/json") .send() .await { - Ok(r) if r.status().is_success() => r.json::().await.ok()?, + Ok(r) if r.status().is_success() => think_watch_common::resources::read_json(r, 64 * 1024) + .await + .ok()?, Ok(r) => { tracing::debug!( endpoint = %endpoint, diff --git a/crates/server/src/handlers/mcp_oauth/discovery.rs b/crates/server/src/handlers/mcp_oauth/discovery.rs index 997a80d5..629bb533 100644 --- a/crates/server/src/handlers/mcp_oauth/discovery.rs +++ b/crates/server/src/handlers/mcp_oauth/discovery.rs @@ -77,13 +77,14 @@ pub async fn oauth_discover( for url in &candidates { match http.get(url).send().await { Ok(resp) if resp.status().is_success() => { - let body: serde_json::Value = match resp.json().await { - Ok(v) => v, - Err(e) => { - last_err = Some(format!("{url}: {e}")); - continue; - } - }; + let (body, _memory): (serde_json::Value, _) = + match think_watch_common::resources::read_json(resp, 64 * 1024).await { + Ok(v) => v, + Err(e) => { + last_err = Some(format!("{url}: {e}")); + continue; + } + }; return Ok(Json(parse_oauth_metadata(&body))); } Ok(resp) => { @@ -455,13 +456,14 @@ async fn fetch_protected_resource( return None; } }; - let body: serde_json::Value = match resp.json().await { - Ok(v) => v, - Err(e) => { - diag.push(format!("{url}: parse: {e}")); - return None; - } - }; + let (body, _memory): (serde_json::Value, _) = + match think_watch_common::resources::read_json(resp, 64 * 1024).await { + Ok(v) => v, + Err(e) => { + diag.push(format!("{url}: parse: {e}")); + return None; + } + }; let issuer = body .get("authorization_servers") .and_then(|v| v.as_array()) @@ -501,8 +503,10 @@ async fn fetch_authz_server_metadata( } match http.get(url).send().await { Ok(resp) if resp.status().is_success() => { - match resp.json::().await { - Ok(body) => { + match think_watch_common::resources::read_json::(resp, 64 * 1024) + .await + { + Ok((body, _memory)) => { diag.push(format!("found authz-server metadata at {url}")); return parse_authz_server_metadata(&body); } @@ -655,7 +659,14 @@ async fn register_dynamic_client( } }; let status = resp.status(); - let text = resp.text().await.unwrap_or_default(); + let bytes = match think_watch_common::resources::read_response(resp, 64 * 1024).await { + Ok(bytes) => bytes, + Err(error) => { + diag.push(format!("dynamic-registration response read: {error}")); + return None; + } + }; + let text = String::from_utf8_lossy(&bytes); if !status.is_success() { diag.push(format!( "dynamic-registration HTTP {status}: {}", diff --git a/crates/server/src/handlers/mcp_shared.rs b/crates/server/src/handlers/mcp_shared.rs index e8875ee0..6e044d21 100644 --- a/crates/server/src/handlers/mcp_shared.rs +++ b/crates/server/src/handlers/mcp_shared.rs @@ -111,35 +111,33 @@ pub async fn probe_mcp_endpoint( .unwrap_or("") .to_lowercase(); - let json: serde_json::Value = if content_type.contains("text/event-stream") { - match resp.text().await { - Ok(text) => crate::mcp_runtime::parse_sse_json(&text).unwrap_or_else(|e| { - tracing::warn!( - endpoint = %endpoint_url, - error = %e, - "MCP probe: failed to parse SSE JSON response" - ); - serde_json::Value::Null - }), - Err(e) => { - tracing::warn!( - endpoint = %endpoint_url, - error = %e, - "MCP probe: failed to read SSE response body" - ); - serde_json::Value::Null - } + let memory = match think_watch_common::resources::read_response( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await + { + Ok(bytes) => bytes, + Err(error) => { + return McpProbeOutcome { + success: false, + requires_auth: false, + message: format!("MCP response could not be read: {error}"), + latency_ms, + tools: vec![], + }; } - } else { - resp.json().await.unwrap_or_else(|e| { - tracing::warn!( - endpoint = %endpoint_url, - error = %e, - "MCP probe: failed to parse JSON response" - ); + }; + let json: serde_json::Value = think_watch_common::resources::cpu(memory.len(), || { + if content_type.contains("text/event-stream") { + crate::mcp_runtime::parse_sse_json(&String::from_utf8_lossy(&memory)) + } else { + serde_json::from_slice(&memory).map_err(Into::into) + }.unwrap_or_else(|e| { + tracing::warn!(endpoint = %endpoint_url, error = %e, "MCP probe: failed to parse response"); serde_json::Value::Null }) - }; + }).await; let result_field = json.get("result"); let tools: Vec = result_field diff --git a/crates/server/src/handlers/mcp_store.rs b/crates/server/src/handlers/mcp_store.rs index b57efbeb..cb79b150 100644 --- a/crates/server/src/handlers/mcp_store.rs +++ b/crates/server/src/handlers/mcp_store.rs @@ -274,10 +274,12 @@ pub async fn sync_registry( ))); } - let registry: RegistryResponse = resp - .json() - .await - .map_err(|e| AppError::BadRequest(format!("Invalid registry JSON: {e}")))?; + let (registry, _memory): (RegistryResponse, _) = think_watch_common::resources::read_json( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await + .map_err(|e| AppError::BadRequest(format!("Invalid registry JSON: {e}")))?; // Wrap the whole sync (upserts + post-loop cleanup DELETE) in a // single transaction. Previously each `INSERT ... ON CONFLICT diff --git a/crates/server/src/handlers/providers.rs b/crates/server/src/handlers/providers.rs index 72ad74e3..38b4dd19 100644 --- a/crates/server/src/handlers/providers.rs +++ b/crates/server/src/handlers/providers.rs @@ -806,7 +806,19 @@ async fn list_models(client: &reqwest::Client, url: &str, headers: &[(String, St }; let status = resp.status(); - let body: serde_json::Value = resp.json().await.unwrap_or(serde_json::Value::Null); + let memory = match think_watch_common::resources::read_response( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await + { + Ok(value) => value, + Err(error) => return Listing::Failed(error.to_string()), + }; + let body: serde_json::Value = think_watch_common::resources::cpu(memory.len(), || { + serde_json::from_slice(&memory).unwrap_or(serde_json::Value::Null) + }) + .await; if !status.is_success() { let message = body .get("error") diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 864b2658..2d66e50c 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -5,8 +5,40 @@ use think_watch_common::db; use think_watch_server::{app, handlers, init, tracing_ch}; use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; -#[tokio::main] -async fn main() -> anyhow::Result<()> { +#[global_allocator] +static ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; + +fn main() -> anyhow::Result<()> { + #[cfg(target_os = "linux")] + { + // A process-local policy, inherited by subsequently created threads. + // Sparse small-object heaps should not pin entire 2 MiB huge pages. + // SAFETY: PR_SET_THP_DISABLE only consumes integer arguments. + unsafe { + libc::prctl( + libc::PR_SET_THP_DISABLE, + 1 as libc::c_ulong, + 0 as libc::c_ulong, + 0 as libc::c_ulong, + 0 as libc::c_ulong, + ) + }; + } + let policy = think_watch_common::resources::policy(); + let workers = match std::env::var("TOKIO_WORKER_THREADS") { + Ok(value) => value.parse::()?.get(), + Err(std::env::VarError::NotPresent) => policy.io_workers, + Err(error) => return Err(error.into()), + }; + tokio::runtime::Builder::new_multi_thread() + .worker_threads(workers) + .thread_keep_alive(think_watch_common::resources::IDLE_TIMEOUT) + .enable_all() + .build()? + .block_on(serve()) +} + +async fn serve() -> anyhow::Result<()> { // Container healthcheck. The runtime image is distroless/static — // no shell, no curl — so the only executable available to Docker's // healthcheck is this binary. `think-watch-server healthcheck` @@ -29,6 +61,8 @@ async fn main() -> anyhow::Result<()> { .with(ch_layer) .init(); + tracing::info!(policy = ?think_watch_common::resources::policy(), allocator = "mimalloc", "Automatic resource admission enabled"); + let config = AppConfig::from_env()?; if let Err(e) = config.validate() { tracing::error!("Configuration validation failed: {e}"); diff --git a/crates/server/src/mcp_runtime.rs b/crates/server/src/mcp_runtime.rs index 1274c69e..a8084073 100644 --- a/crates/server/src/mcp_runtime.rs +++ b/crates/server/src/mcp_runtime.rs @@ -441,12 +441,19 @@ async fn try_discover_and_persist_tools_with_auth( .unwrap_or("") .to_lowercase(); - let json: serde_json::Value = if content_type.contains("text/event-stream") { - let text = resp.text().await?; - parse_sse_json(&text)? - } else { - resp.json().await? - }; + let memory = think_watch_common::resources::read_response( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await?; + let json: serde_json::Value = think_watch_common::resources::cpu(memory.len(), || { + if content_type.contains("text/event-stream") { + parse_sse_json(&String::from_utf8_lossy(&memory)) + } else { + serde_json::from_slice(&memory).map_err(Into::into) + } + }) + .await?; // Strict: require a `result` field so malformed responses are // caught instead of silently treated as success. let result = json @@ -569,12 +576,19 @@ pub async fn discover_user_tools( .and_then(|v| v.to_str().ok()) .unwrap_or("") .to_lowercase(); - let json: serde_json::Value = if content_type.contains("text/event-stream") { - let text = resp.text().await?; - parse_sse_json(&text)? - } else { - resp.json().await? - }; + let memory = think_watch_common::resources::read_response( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await?; + let json: serde_json::Value = think_watch_common::resources::cpu(memory.len(), || { + if content_type.contains("text/event-stream") { + parse_sse_json(&String::from_utf8_lossy(&memory)) + } else { + serde_json::from_slice(&memory).map_err(Into::into) + } + }) + .await?; let result = json .get("result") .ok_or_else(|| anyhow::anyhow!("tools/list response missing `result` field"))? diff --git a/crates/server/src/protocol_probe.rs b/crates/server/src/protocol_probe.rs index 1d42d3ea..c02d5675 100644 --- a/crates/server/src/protocol_probe.rs +++ b/crates/server/src/protocol_probe.rs @@ -247,7 +247,11 @@ async fn probe_one(materials: &ProviderMaterials, upstream_model: &str) -> Probe match attempt { Ok(Ok(resp)) => { // Drain it so the connection goes back to the pool. - let _ = resp.bytes().await; + let _ = tokio::time::timeout( + PROBE_TIMEOUT, + think_watch_common::resources::read_response(resp, 64 * 1024), + ) + .await; return Probe::Verdict(Verdict::Ok(candidate)); } Ok(Err(GatewayError::UpstreamAuthError { status, message })) => { diff --git a/crates/test-support/tests/resources.rs b/crates/test-support/tests/resources.rs new file mode 100644 index 00000000..9fe69211 --- /dev/null +++ b/crates/test-support/tests/resources.rs @@ -0,0 +1,73 @@ +//! Real Redis decoding and HTTP admission/recovery through the assembled app. + +use fred::prelude::*; +use think_watch_common::resources; +use think_watch_test_support::prelude::*; + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn redis_body_reads_preserve_bytes_and_reservations_until_the_last_clone() { + let app = TestApp::spawn().await; + let key = unique_name("resource-body"); + let data = b"\0\xff\"cached JSON\"\xe4\xb8\xad"; + let _: () = app + .state + .redis + .set(&key, data.as_slice(), None, None, false) + .await + .unwrap(); + let budget = resources::body_budget(); + let baseline = budget.used(); + let body = resources::cached_body(&app.state.redis, &key) + .await + .unwrap(); + assert_eq!(&body[..], data); + assert_eq!( + budget.used(), + baseline + data.len() * resources::BODY_HEADROOM + ); + let copy = body.clone(); + drop(body); + assert!(budget.used() > baseline); + drop(copy); + assert_eq!(budget.used(), baseline); + + let mut occupied = budget.reserve(); + assert!(occupied.grow(resources::policy().body_bytes - baseline)); + assert!( + resources::cached_body(&app.state.redis, &key) + .await + .is_none() + ); + drop(occupied); + assert!( + resources::cached_body(&app.state.redis, &key) + .await + .is_some() + ); + let _: () = app.state.redis.del(&key).await.unwrap(); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn http_body_pressure_is_retryable_and_releases_without_affecting_health() { + let app = TestApp::spawn().await; + let budget = resources::body_budget(); + let mut occupied = budget.reserve(); + assert!(occupied.grow(resources::policy().body_bytes - budget.used())); + let gateway = app.gateway_client(); + let response = gateway + .post("/v1/chat/completions", json!({"model":"test"})) + .await + .unwrap(); + response.assert_status(429); + assert_eq!(response.headers["retry-after"], "1"); + drop(occupied); + let response = gateway + .post("/v1/chat/completions", json!({"model":"test"})) + .await + .unwrap(); + response.assert_status(401); + let health = gateway.get("/health/live").await.unwrap(); + health.assert_ok(); +} diff --git a/docs/operations/resources.md b/docs/operations/resources.md new file mode 100644 index 00000000..57a3e49d --- /dev/null +++ b/docs/operations/resources.md @@ -0,0 +1,66 @@ +# Automatic process resources + +The server sizes its execution and buffering policy at startup. No configuration +change is required. CPU availability includes the process affinity and container +CPU quota. Memory uses the smaller of physical RAM and the cgroup v1/v2 memory +limit, including ancestor limits. + +## Execution + +The I/O runtime keeps at most eight workers. Large or concurrent JSON decoding, +protocol conversion, redaction and content inspection use elastic blocking +execution. CPU admission is shared with password hashing: at most one job per +available CPU, further limited to one job per 128 MiB of memory allowance. +An 80 CPU server with sufficient memory can execute 80 jobs concurrently. +Extra execution threads retire after ten idle seconds. + +Argon2 also has its own memory admission: one concurrent hash per 32 MiB within +one eighth of the memory allowance, with a minimum of one. Cancelling an HTTP +request does not release a hash's admission until that hash actually finishes. +The existing `TOKIO_WORKER_THREADS` override remains supported. + +The server uses mimalloc with transparent huge pages disabled. On Linux it also +applies a process-local THP policy before creating runtime threads. This needs +no host sysctl or environment changes. + +## Buffered bodies and streams + +Requests, upstream AI/MCP responses, Redis cache reads and retained streaming +timelines share one byte budget: one quarter of the memory allowance, between +8 MiB and 8 GiB. Reservations allow 32 times the wire size for decoded JSON, +conversion, inspection and capture copies. A single response or stream timeline +is limited to the smaller of 64 MiB and one thirty-second of that budget. +The existing gateway 10 MiB and console 1 MiB request limits still apply. +Error bodies and credential endpoint responses have a separate per-read 64 KiB +limit and use the same shared budget. MCP tool discovery, provider and registry +catalogs, OAuth discovery, audit forwarder errors and S3 reads also use bounded +reads, so administrative/background requests cannot bypass admission on ingress. + +Admission is incremental, including chunked responses. Readers reject growth +before appending rather than waiting while holding partial bodies. Reservations +follow shared buffers through transformations, slow clients, cancellation and +detached stream audit work. Cache reads reserve before fetching and use bounded +Redis reads even if a key changes between the size check and fetch. + +Memory pressure produces HTTP 429 with `Retry-After: 1` before streaming starts. +Already-started streams terminate with an error event (MCP code `-32005`). +Local pressure does not count as an upstream failure for circuit breakers. +Oversized requests return 413; oversized buffered upstream responses fail rather +than allocating without a bound. Cache reads under pressure become misses. + +## Audit backlog + +Queued audit entries and retained ClickHouse retry batches share a separate byte +budget: one sixteenth of the memory allowance, between 2 MiB and 128 MiB. The +estimate includes owned strings, JSON allocations and room for temporary flush +copies. The existing 100,000 entry queue limit is a secondary bound. +When full, the logger drops new entries and emits a throttled error. + +Monitor `audit_log_buffered_bytes`, `audit_log_memory_dropped_total`, +`audit_log_dropped_total` and `gateway_memory_busy_total` alongside the existing +queue-depth metrics. A drop requires operational attention, especially when +ClickHouse is unavailable for an extended period. + +These budgets estimate application allocations; they are not a hard RSS cap. +Connection state, database pools, configuration, stacks and allocator overhead +also consume memory. Deployment memory and CPU limits remain authoritative.