From 17bcc048005dd1f1b8fb223b4160863660e24e60 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 3 Sep 2026 10:05:35 -0700 Subject: [PATCH] perf(rust): bound provider connection pools and adopt mimalloc in the bridge The four shared reqwest clients (core messages/chat_completions/ audio_transcription, ai-gateway) now set pool_max_idle_per_host(64) and tcp_keepalive(60s): - reqwest defaults pool_max_idle_per_host to usize::MAX, so a burst of N concurrent calls to one provider host leaves N idle TLS connections in the pool after completion, each held until its idle timeout; under sustained bursts this grows file descriptors without bound. 64 keeps a warm set large enough for handshake-free reuse. - tcp_keepalive lets silent half-open peers (NAT/LB reaping, vanished upstreams) be detected during long streaming responses instead of blocking until the 600s full-request timeout. The python-bridge cdylib now uses mimalloc as its Rust global allocator (the pydantic-core/orjson pattern). It only governs allocations made by Rust code in the extension - CPython's pymalloc heap is untouched - but Rust-side request/response buffering in a long-lived interpreter process gets mimalloc's sharded free lists and eager page purging instead of glibc malloc, which favors RSS growth under churn. No behavior change beyond the socket/allocator characteristics; timeouts and request paths are untouched. --- litellm-rust/Cargo.lock | 19 +++++++++++++++++++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/ai-gateway/src/client.rs | 4 ++++ .../core/src/audio_transcription/client.rs | 6 +++++- .../core/src/chat_completions/client.rs | 7 ++++++- litellm-rust/crates/core/src/constants.rs | 12 ++++++++++++ .../crates/core/src/messages/client.rs | 7 ++++++- litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 3 +++ 9 files changed, 57 insertions(+), 3 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b3dac5ca935..90536e66753 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1404,6 +1404,15 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libmimalloc-sys" +version = "0.1.49" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" +dependencies = [ + "cc", +] + [[package]] name = "litellm-ai-gateway" version = "0.1.0" @@ -1454,6 +1463,7 @@ dependencies = [ "litellm-ai-gateway", "litellm-core", "litellm-python-interop", + "mimalloc", "pyo3", "pyo3-async-runtimes", "serde", @@ -1505,6 +1515,15 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[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" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a13dd4c04b0..9a7fbe72c35 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -35,6 +35,7 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" +mimalloc = "0.1" [profile.release] opt-level = 3 diff --git a/litellm-rust/crates/ai-gateway/src/client.rs b/litellm-rust/crates/ai-gateway/src/client.rs index ff2606f0229..16a22733fac 100644 --- a/litellm-rust/crates/ai-gateway/src/client.rs +++ b/litellm-rust/crates/ai-gateway/src/client.rs @@ -2,12 +2,16 @@ use std::sync::OnceLock; use std::time::Duration; const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600; +const HTTP_CLIENT_POOL_MAX_IDLE_PER_HOST: usize = 64; +const HTTP_CLIENT_TCP_KEEPALIVE_SECS: u64 = 60; pub(crate) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { reqwest::Client::builder() .timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS)) + .pool_max_idle_per_host(HTTP_CLIENT_POOL_MAX_IDLE_PER_HOST) + .tcp_keepalive(Some(Duration::from_secs(HTTP_CLIENT_TCP_KEEPALIVE_SECS))) .build() .expect("failed to build reqwest client") }) diff --git a/litellm-rust/crates/core/src/audio_transcription/client.rs b/litellm-rust/crates/core/src/audio_transcription/client.rs index 0e612628dc6..86edf504923 100644 --- a/litellm-rust/crates/core/src/audio_transcription/client.rs +++ b/litellm-rust/crates/core/src/audio_transcription/client.rs @@ -1,13 +1,17 @@ use std::sync::OnceLock; use std::time::Duration; -use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS; +use crate::constants::{ + AUDIO_TRANSCRIPTION_TIMEOUT_SECS, PROVIDER_POOL_MAX_IDLE_PER_HOST, PROVIDER_TCP_KEEPALIVE_SECS, +}; pub(super) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { reqwest::Client::builder() .timeout(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)) + .pool_max_idle_per_host(PROVIDER_POOL_MAX_IDLE_PER_HOST) + .tcp_keepalive(Some(Duration::from_secs(PROVIDER_TCP_KEEPALIVE_SECS))) .build() .unwrap_or_else(|_| reqwest::Client::new()) }) diff --git a/litellm-rust/crates/core/src/chat_completions/client.rs b/litellm-rust/crates/core/src/chat_completions/client.rs index f2ef73ed030..636bf31a048 100644 --- a/litellm-rust/crates/core/src/chat_completions/client.rs +++ b/litellm-rust/crates/core/src/chat_completions/client.rs @@ -1,7 +1,10 @@ use std::sync::OnceLock; use std::time::Duration; -use crate::constants::{CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS}; +use crate::constants::{ + CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS, + PROVIDER_POOL_MAX_IDLE_PER_HOST, PROVIDER_TCP_KEEPALIVE_SECS, +}; pub(super) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); @@ -9,6 +12,8 @@ pub(super) fn http_client() -> &'static reqwest::Client { reqwest::Client::builder() .timeout(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)) .connect_timeout(Duration::from_secs(CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS)) + .pool_max_idle_per_host(PROVIDER_POOL_MAX_IDLE_PER_HOST) + .tcp_keepalive(Some(Duration::from_secs(PROVIDER_TCP_KEEPALIVE_SECS))) .build() .unwrap_or_else(|_| reqwest::Client::new()) }) diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index a73961060eb..f24bf9db2a9 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -32,6 +32,18 @@ pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10; pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600; +/// Idle connections retained per host in the shared provider clients. +/// reqwest defaults to `usize::MAX`, so a burst of concurrent calls to one +/// provider host leaves every connection idle in the pool until its timeout; +/// this bounds file descriptors while keeping enough warm connections for +/// TLS-handshake-free reuse. +pub(crate) const PROVIDER_POOL_MAX_IDLE_PER_HOST: usize = 64; + +/// TCP keep-alive probe interval for provider sockets. Detects half-open +/// peers (NAT/LB reaping, vanished upstreams) during long-lived streaming +/// responses that otherwise sit silent for the whole request timeout. +pub(crate) const PROVIDER_TCP_KEEPALIVE_SECS: u64 = 60; + /// `object` field every non-streaming chat completion response carries. pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion"; diff --git a/litellm-rust/crates/core/src/messages/client.rs b/litellm-rust/crates/core/src/messages/client.rs index 6281270b964..6d377f476d0 100644 --- a/litellm-rust/crates/core/src/messages/client.rs +++ b/litellm-rust/crates/core/src/messages/client.rs @@ -1,7 +1,10 @@ use std::sync::OnceLock; use std::time::Duration; -use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS}; +use crate::constants::{ + MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS, PROVIDER_POOL_MAX_IDLE_PER_HOST, + PROVIDER_TCP_KEEPALIVE_SECS, +}; pub(super) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); @@ -9,6 +12,8 @@ pub(super) fn http_client() -> &'static reqwest::Client { reqwest::Client::builder() .timeout(Duration::from_secs(MESSAGES_TIMEOUT_SECS)) .connect_timeout(Duration::from_secs(MESSAGES_CONNECT_TIMEOUT_SECS)) + .pool_max_idle_per_host(PROVIDER_POOL_MAX_IDLE_PER_HOST) + .tcp_keepalive(Some(Duration::from_secs(PROVIDER_TCP_KEEPALIVE_SECS))) .build() .unwrap_or_else(|_| reqwest::Client::new()) }) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 637e5580170..c4be72ab8f5 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -17,6 +17,7 @@ panic-test = [] [dependencies] futures-util.workspace = true +mimalloc.workspace = true tracing.workspace = true tracing-subscriber.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 5f36a22370a..7e160a4c7bf 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -11,6 +11,9 @@ use pyo3::prelude::*; use pyo3::types::PyAny; use serde_json::Value; +#[global_allocator] +static ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; + use crate::errors::core_error_to_pyerr; use crate::marshal::{marshal_headers, optional_timeout};