mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 16:37:46 +08:00
fix: harden concurrency limits and high-RPM runtime paths
Bound request, stream, queue, and shutdown resource lifetimes. Reduce scheduler and Redis hot-path work and isolate database maintenance. Include regression coverage, load probes, and concurrency audit results.
This commit is contained in:
elky
@@ -10,13 +10,18 @@ description = "HTTP frontdoor middleware and request lifecycle primitives for Ae
|
||||
aether-ai-formats.workspace = true
|
||||
axum.workspace = true
|
||||
bytes.workspace = true
|
||||
futures-util.workspace = true
|
||||
http.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio = { workspace = true, features = ["io-util"] }
|
||||
tokio-util = { workspace = true, features = ["rt"] }
|
||||
tracing.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
futures-util.workspace = true
|
||||
http-body-util = "0.1"
|
||||
hyper = { version = "1", features = ["client", "server", "http1", "http2"] }
|
||||
hyper-util = { version = "0.1", features = ["tokio"] }
|
||||
serde_json.workspace = true
|
||||
tokio = { workspace = true, features = ["test-util"] }
|
||||
tower = { version = "0.5", features = ["util"] }
|
||||
tracing-subscriber.workspace = true
|
||||
|
||||
@@ -1,14 +1,13 @@
|
||||
//! Bounded request-body buffering for frontdoor adapters.
|
||||
//!
|
||||
//! The policy reserves weighted memory before reading a body and holds the
|
||||
//! reservation through the caller's normalization callback. This keeps body
|
||||
//! buffering independent from gateway business routing while preventing a
|
||||
//! burst of compressed requests from bypassing the memory budget.
|
||||
//! The policy grows weighted reservations as bytes arrive and holds them through
|
||||
//! normalization. Growth never waits while retaining a partial buffer, so
|
||||
//! concurrent uploads cannot deadlock while competing for the remaining budget.
|
||||
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::body::Body;
|
||||
use bytes::Bytes;
|
||||
use futures_util::StreamExt;
|
||||
use http::{header, HeaderMap, StatusCode};
|
||||
use std::error::Error as StdError;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
@@ -127,7 +126,12 @@ impl BodyBufferPolicy {
|
||||
}
|
||||
|
||||
pub fn reservation_bytes(&self, headers: &HeaderMap) -> usize {
|
||||
reservation_bytes(headers, self.max_bytes, self.budget_bytes)
|
||||
reservation_bytes(
|
||||
headers,
|
||||
self.max_bytes,
|
||||
self.budget_bytes,
|
||||
self.permit_bytes,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn reservation_permits(&self, reservation_bytes: usize) -> u32 {
|
||||
@@ -168,67 +172,134 @@ impl BodyBufferPolicy {
|
||||
};
|
||||
|
||||
Ok(BodyBufferReservation {
|
||||
permit,
|
||||
memory: BodyBufferBudget {
|
||||
permit,
|
||||
budget: Arc::clone(&self.budget),
|
||||
budget_bytes: self.budget_bytes,
|
||||
permit_bytes: self.permit_bytes,
|
||||
requested_bytes,
|
||||
},
|
||||
max_bytes: effective_max_bytes,
|
||||
read_timeout: self.read_timeout,
|
||||
requested_bytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct BodyBufferReservation {
|
||||
permit: OwnedSemaphorePermit,
|
||||
memory: BodyBufferBudget,
|
||||
max_bytes: u64,
|
||||
read_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct BodyBufferBudget {
|
||||
permit: OwnedSemaphorePermit,
|
||||
budget: Arc<Semaphore>,
|
||||
budget_bytes: usize,
|
||||
permit_bytes: usize,
|
||||
requested_bytes: usize,
|
||||
}
|
||||
|
||||
impl BodyBufferBudget {
|
||||
/// Reserve a new high-water mark before retaining or decoding more bytes.
|
||||
/// Never queue for growth while another partial request may hold the rest.
|
||||
pub fn try_reserve_bytes(&mut self, requested_bytes: usize) -> Result<(), BodyBufferError> {
|
||||
if requested_bytes > self.budget_bytes {
|
||||
return Err(self.overloaded(requested_bytes));
|
||||
}
|
||||
let permits = reservation_permits(requested_bytes, self.permit_bytes) as usize;
|
||||
let additional = permits.saturating_sub(self.permit.num_permits());
|
||||
if additional > 0 {
|
||||
let permit = Arc::clone(&self.budget)
|
||||
.try_acquire_many_owned(additional as u32)
|
||||
.map_err(|_| self.overloaded(requested_bytes))?;
|
||||
self.permit.merge(permit);
|
||||
}
|
||||
self.requested_bytes = self.requested_bytes.max(requested_bytes);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn overloaded(&self, requested_bytes: usize) -> BodyBufferError {
|
||||
BodyBufferError::Overloaded {
|
||||
requested_bytes,
|
||||
budget_bytes: self.budget_bytes,
|
||||
timeout_ms: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BodyBufferReservation {
|
||||
pub fn requested_bytes(&self) -> usize {
|
||||
self.requested_bytes
|
||||
self.memory.requested_bytes
|
||||
}
|
||||
|
||||
pub async fn collect(self, body: Body) -> Result<BufferedBody, BodyBufferError> {
|
||||
let Self {
|
||||
permit,
|
||||
mut memory,
|
||||
max_bytes,
|
||||
read_timeout,
|
||||
requested_bytes,
|
||||
} = self;
|
||||
let started_at = Instant::now();
|
||||
let body_limit = usize::try_from(max_bytes).unwrap_or(usize::MAX);
|
||||
let collected = match read_timeout {
|
||||
Some(read_timeout) => {
|
||||
match tokio::time::timeout(read_timeout, to_bytes(body, body_limit)).await {
|
||||
Ok(result) => result,
|
||||
Err(_) => {
|
||||
return Err(BodyBufferError::Timeout {
|
||||
timeout_ms: duration_millis(read_timeout),
|
||||
});
|
||||
let collect = async {
|
||||
let mut stream = body.into_data_stream();
|
||||
let mut bytes = Vec::new();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(|error| BodyBufferError::ReadFailed {
|
||||
message: error.to_string(),
|
||||
})?;
|
||||
let length =
|
||||
bytes
|
||||
.len()
|
||||
.checked_add(chunk.len())
|
||||
.ok_or(BodyBufferError::TooLarge {
|
||||
limit_bytes: max_bytes,
|
||||
})?;
|
||||
if length as u64 > max_bytes {
|
||||
return Err(BodyBufferError::TooLarge {
|
||||
limit_bytes: max_bytes,
|
||||
});
|
||||
}
|
||||
if length > bytes.capacity() {
|
||||
let mut capacity = if length <= DEFAULT_BODY_BUFFER_PERMIT_BYTES {
|
||||
length
|
||||
} else {
|
||||
bytes
|
||||
.capacity()
|
||||
.saturating_mul(2)
|
||||
.max(length)
|
||||
.min(usize::try_from(max_bytes).unwrap_or(usize::MAX))
|
||||
};
|
||||
if memory.try_reserve_bytes(capacity).is_err() {
|
||||
memory.try_reserve_bytes(length)?;
|
||||
capacity = length;
|
||||
}
|
||||
// Account for geometric growth, falling back to the bytes needed under load.
|
||||
bytes
|
||||
.try_reserve_exact(capacity - bytes.len())
|
||||
.map_err(|error| BodyBufferError::ReadFailed {
|
||||
message: error.to_string(),
|
||||
})?;
|
||||
if bytes.capacity() > capacity {
|
||||
memory.try_reserve_bytes(bytes.capacity())?;
|
||||
}
|
||||
}
|
||||
bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
None => to_bytes(body, body_limit).await,
|
||||
};
|
||||
let bytes = match collected {
|
||||
Ok(bytes) => bytes,
|
||||
Err(error) if collection_exceeded_limit(&error) => {
|
||||
return Err(BodyBufferError::TooLarge {
|
||||
limit_bytes: max_bytes,
|
||||
});
|
||||
}
|
||||
Err(error) => {
|
||||
return Err(BodyBufferError::ReadFailed {
|
||||
message: error.to_string(),
|
||||
});
|
||||
}
|
||||
Ok(Bytes::from(bytes))
|
||||
};
|
||||
let bytes = match read_timeout {
|
||||
Some(read_timeout) => tokio::time::timeout(read_timeout, collect)
|
||||
.await
|
||||
.map_err(|_| BodyBufferError::Timeout {
|
||||
timeout_ms: duration_millis(read_timeout),
|
||||
}),
|
||||
None => Ok(collect.await),
|
||||
}??;
|
||||
|
||||
Ok(BufferedBody {
|
||||
bytes,
|
||||
permit: Some(permit),
|
||||
requested_bytes,
|
||||
memory,
|
||||
elapsed: started_at.elapsed(),
|
||||
})
|
||||
}
|
||||
@@ -237,8 +308,7 @@ impl BodyBufferReservation {
|
||||
#[derive(Debug)]
|
||||
pub struct BufferedBody {
|
||||
bytes: Bytes,
|
||||
permit: Option<OwnedSemaphorePermit>,
|
||||
requested_bytes: usize,
|
||||
memory: BodyBufferBudget,
|
||||
elapsed: Duration,
|
||||
}
|
||||
|
||||
@@ -248,19 +318,28 @@ impl BufferedBody {
|
||||
}
|
||||
|
||||
pub fn requested_bytes(&self) -> usize {
|
||||
self.requested_bytes
|
||||
self.memory.requested_bytes
|
||||
}
|
||||
|
||||
pub fn elapsed(&self) -> Duration {
|
||||
self.elapsed
|
||||
}
|
||||
|
||||
/// Apply normalization while retaining the memory permit until the
|
||||
/// callback completes.
|
||||
/// Retain the permit for a callback that does not grow the buffered payload.
|
||||
pub fn try_map<T, E>(self, map: impl FnOnce(Bytes) -> Result<T, E>) -> Result<T, E> {
|
||||
let Self { bytes, permit, .. } = self;
|
||||
let result = map(bytes);
|
||||
drop(permit);
|
||||
self.try_map_with_budget(|bytes, _| map(bytes))
|
||||
}
|
||||
|
||||
/// The callback must account for decoded buffers before allocating them.
|
||||
pub fn try_map_with_budget<T, E>(
|
||||
self,
|
||||
map: impl FnOnce(Bytes, &mut BodyBufferBudget) -> Result<T, E>,
|
||||
) -> Result<T, E> {
|
||||
let Self {
|
||||
bytes, mut memory, ..
|
||||
} = self;
|
||||
let result = map(bytes, &mut memory);
|
||||
drop(memory);
|
||||
result
|
||||
}
|
||||
}
|
||||
@@ -389,23 +468,15 @@ fn invalid_body_headers(message: &str) -> BodyBufferError {
|
||||
}
|
||||
}
|
||||
|
||||
fn reservation_bytes(headers: &HeaderMap, max_bytes: u64, budget_bytes: usize) -> usize {
|
||||
fn reservation_bytes(
|
||||
headers: &HeaderMap,
|
||||
max_bytes: u64,
|
||||
budget_bytes: usize,
|
||||
permit_bytes: usize,
|
||||
) -> usize {
|
||||
let reservation_ceiling = usize::try_from(max_bytes)
|
||||
.unwrap_or(usize::MAX)
|
||||
.min(budget_bytes);
|
||||
let encoded = headers
|
||||
.get_all(header::CONTENT_ENCODING)
|
||||
.iter()
|
||||
.any(|value| {
|
||||
value.to_str().map_or(true, |value| {
|
||||
value.split(',').map(str::trim).any(|encoding| {
|
||||
!encoding.is_empty() && !encoding.eq_ignore_ascii_case("identity")
|
||||
})
|
||||
})
|
||||
});
|
||||
if encoded {
|
||||
return reservation_ceiling;
|
||||
}
|
||||
declared_content_length(headers)
|
||||
.ok()
|
||||
.flatten()
|
||||
@@ -414,7 +485,7 @@ fn reservation_bytes(headers: &HeaderMap, max_bytes: u64, budget_bytes: usize) -
|
||||
.unwrap_or(usize::MAX)
|
||||
.min(reservation_ceiling)
|
||||
})
|
||||
.unwrap_or(reservation_ceiling)
|
||||
.unwrap_or_else(|| permit_bytes.min(reservation_ceiling))
|
||||
}
|
||||
|
||||
fn reservation_permits(reservation_bytes: usize, permit_bytes: usize) -> u32 {
|
||||
@@ -429,17 +500,6 @@ fn duration_millis(duration: Duration) -> u64 {
|
||||
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
|
||||
}
|
||||
|
||||
fn collection_exceeded_limit(error: &(dyn StdError + 'static)) -> bool {
|
||||
let mut current = Some(error);
|
||||
while let Some(error) = current {
|
||||
if error.to_string().contains("length limit exceeded") {
|
||||
return true;
|
||||
}
|
||||
current = error.source();
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{BodyBufferError, BodyBufferPolicy, DEFAULT_BODY_BUFFER_PERMIT_BYTES};
|
||||
@@ -579,9 +639,9 @@ mod tests {
|
||||
let reservation = policy
|
||||
.reserve(&headers)
|
||||
.await
|
||||
.expect("encoded unlimited body should reserve the available budget");
|
||||
assert_eq!(reservation.requested_bytes(), 4);
|
||||
assert_eq!(budget.available_permits(), 0);
|
||||
.expect("encoded unlimited body should reserve its initial chunk");
|
||||
assert_eq!(reservation.requested_bytes(), 1);
|
||||
assert_eq!(budget.available_permits(), 3);
|
||||
|
||||
let error = reservation
|
||||
.collect(Body::from(Bytes::from_static(b"01234")))
|
||||
@@ -707,4 +767,172 @@ mod tests {
|
||||
.expect_err("exhausted budget should fail closed");
|
||||
assert!(matches!(error, BodyBufferError::Overloaded { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn small_compressed_and_unknown_length_requests_share_the_budget() {
|
||||
let budget_bytes = 256 * 1024 * 1024;
|
||||
let budget = Arc::new(Semaphore::new(
|
||||
budget_bytes / DEFAULT_BODY_BUFFER_PERMIT_BYTES,
|
||||
));
|
||||
let policy = policy(
|
||||
budget_bytes as u64,
|
||||
Duration::from_secs(1),
|
||||
Arc::clone(&budget),
|
||||
);
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(header::CONTENT_ENCODING, HeaderValue::from_static("gzip"));
|
||||
headers.insert(header::CONTENT_LENGTH, HeaderValue::from_static("1024"));
|
||||
let compressed = policy.reserve(&headers).await.unwrap();
|
||||
let unknown = policy.reserve(&HeaderMap::new()).await.unwrap();
|
||||
assert_eq!(compressed.requested_bytes(), 1024);
|
||||
assert_eq!(unknown.requested_bytes(), DEFAULT_BODY_BUFFER_PERMIT_BYTES);
|
||||
assert_eq!(budget.available_permits(), 4094);
|
||||
drop((compressed, unknown));
|
||||
assert_eq!(budget.available_permits(), 4096);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_length_body_grows_its_reservation_and_holds_it_until_normalized() {
|
||||
let budget = Arc::new(Semaphore::new(8));
|
||||
let policy = BodyBufferPolicy::with_permit_bytes(
|
||||
8,
|
||||
Duration::from_secs(1),
|
||||
Duration::from_secs(1),
|
||||
8,
|
||||
1,
|
||||
Arc::clone(&budget),
|
||||
);
|
||||
let reservation = policy.reserve(&HeaderMap::new()).await.unwrap();
|
||||
assert_eq!(budget.available_permits(), 7);
|
||||
let body = Body::from_stream(stream::iter([
|
||||
Ok::<_, std::io::Error>(Bytes::from_static(b"ab")),
|
||||
Ok(Bytes::from_static(b"cd")),
|
||||
]));
|
||||
let buffered = reservation.collect(body).await.unwrap();
|
||||
assert_eq!(buffered.bytes().as_ref(), b"abcd");
|
||||
assert_eq!(budget.available_permits(), 4);
|
||||
buffered
|
||||
.try_map_with_budget(|_, memory| {
|
||||
memory.try_reserve_bytes(7)?;
|
||||
assert_eq!(budget.available_permits(), 1);
|
||||
Ok::<_, BodyBufferError>(())
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(budget.available_permits(), 8);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn partial_upload_growth_rejects_without_waiting_and_releases_its_budget() {
|
||||
let budget = Arc::new(Semaphore::new(2));
|
||||
let policy = BodyBufferPolicy::with_permit_bytes(
|
||||
2,
|
||||
Duration::from_secs(60),
|
||||
Duration::from_secs(60),
|
||||
2,
|
||||
1,
|
||||
Arc::clone(&budget),
|
||||
);
|
||||
let first = policy.reserve(&HeaderMap::new()).await.unwrap();
|
||||
let second = policy.reserve(&HeaderMap::new()).await.unwrap();
|
||||
let error =
|
||||
tokio::time::timeout(Duration::from_millis(100), first.collect(Body::from("ab")))
|
||||
.await
|
||||
.expect("growth must not wait while holding a partial reservation")
|
||||
.unwrap_err();
|
||||
assert_eq!(
|
||||
error,
|
||||
BodyBufferError::Overloaded {
|
||||
requested_bytes: 2,
|
||||
budget_bytes: 2,
|
||||
timeout_ms: 0,
|
||||
}
|
||||
);
|
||||
assert_eq!(budget.available_permits(), 1);
|
||||
let buffered = second.collect(Body::from("ab")).await.unwrap();
|
||||
assert_eq!(budget.available_permits(), 0);
|
||||
drop(buffered);
|
||||
assert_eq!(budget.available_permits(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn normalization_budget_failure_releases_all_upload_permits() {
|
||||
let budget = Arc::new(Semaphore::new(4));
|
||||
let policy = BodyBufferPolicy::with_permit_bytes(
|
||||
4,
|
||||
Duration::from_secs(1),
|
||||
Duration::from_secs(1),
|
||||
4,
|
||||
1,
|
||||
Arc::clone(&budget),
|
||||
);
|
||||
let buffered = policy
|
||||
.reserve(&HeaderMap::new())
|
||||
.await
|
||||
.unwrap()
|
||||
.collect(Body::from("ab"))
|
||||
.await
|
||||
.unwrap();
|
||||
let result = buffered.try_map_with_budget(|_, memory| memory.try_reserve_bytes(5));
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(BodyBufferError::Overloaded {
|
||||
requested_bytes: 5,
|
||||
..
|
||||
})
|
||||
));
|
||||
assert_eq!(budget.available_permits(), 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upload_growth_uses_available_budget_without_requiring_geometric_headroom() {
|
||||
let budget = Arc::new(Semaphore::new(128));
|
||||
let held = Arc::clone(&budget).acquire_many_owned(32).await.unwrap();
|
||||
let policy = BodyBufferPolicy::with_permit_bytes(
|
||||
128 * 1024,
|
||||
Duration::from_secs(1),
|
||||
Duration::from_secs(1),
|
||||
128 * 1024,
|
||||
1024,
|
||||
Arc::clone(&budget),
|
||||
);
|
||||
let body = Body::from_stream(stream::iter([
|
||||
Ok::<_, std::io::Error>(Bytes::from(vec![b'a'; 70_000])),
|
||||
Ok(Bytes::from(vec![b'b'; 20_000])),
|
||||
]));
|
||||
let buffered = policy
|
||||
.reserve(&HeaderMap::new())
|
||||
.await
|
||||
.unwrap()
|
||||
.collect(body)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(buffered.bytes().len(), 90_000);
|
||||
assert_eq!(buffered.requested_bytes(), 90_000);
|
||||
drop((buffered, held));
|
||||
assert_eq!(budget.available_permits(), 128);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancellation_releases_partial_upload_budget() {
|
||||
let budget = Arc::new(Semaphore::new(4));
|
||||
let policy = BodyBufferPolicy::with_permit_bytes(
|
||||
4,
|
||||
Duration::from_secs(60),
|
||||
Duration::from_secs(1),
|
||||
4,
|
||||
1,
|
||||
Arc::clone(&budget),
|
||||
);
|
||||
let reservation = policy.reserve(&HeaderMap::new()).await.unwrap();
|
||||
let body = Body::from_stream(
|
||||
stream::once(async { Ok::<_, std::io::Error>(Bytes::from_static(b"ab")) })
|
||||
.chain(stream::pending()),
|
||||
);
|
||||
assert!(
|
||||
tokio::time::timeout(Duration::from_millis(10), reservation.collect(body))
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(budget.available_permits(), 4);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,229 @@
|
||||
use std::future::Future;
|
||||
use std::io;
|
||||
use std::net::SocketAddr;
|
||||
use std::pin::Pin;
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned};
|
||||
|
||||
const MAX_HTTP_CONNECTIONS: usize = 65_536;
|
||||
const FD_RESERVE: usize = 256;
|
||||
|
||||
/// Selects a per-server incoming TCP limit. The FD allowance leaves room for
|
||||
/// upstream sockets and process infrastructure; it is not a complete FD budget.
|
||||
pub fn http_connection_limit(
|
||||
configured: Option<usize>,
|
||||
request_limit: usize,
|
||||
websocket_limit: usize,
|
||||
fd_soft_limit: Option<usize>,
|
||||
) -> usize {
|
||||
let configured = configured
|
||||
.filter(|limit| *limit > 0)
|
||||
.unwrap_or_else(|| request_limit.saturating_add(websocket_limit))
|
||||
.clamp(1, MAX_HTTP_CONNECTIONS);
|
||||
let fd_allowance = fd_soft_limit
|
||||
.map(|limit| (limit.saturating_sub(FD_RESERVE) / 2).max(1))
|
||||
.unwrap_or(MAX_HTTP_CONNECTIONS);
|
||||
configured.min(fd_allowance)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct HttpConnectionBudgetSnapshot {
|
||||
pub limit: usize,
|
||||
pub in_flight: usize,
|
||||
pub high_watermark: usize,
|
||||
pub rejected_total: u64,
|
||||
pub accept_errors_total: u64,
|
||||
}
|
||||
|
||||
/// Share one budget across listeners. Admission belongs to the underlying IO,
|
||||
/// so HTTP/1 upgrades keep their permit and HTTP/2 streams share one permit.
|
||||
#[derive(Debug)]
|
||||
pub struct HttpConnectionBudget {
|
||||
limit: usize,
|
||||
permits: Arc<Semaphore>,
|
||||
in_flight: AtomicUsize,
|
||||
high_watermark: AtomicUsize,
|
||||
rejected_total: AtomicU64,
|
||||
accept_errors_total: AtomicU64,
|
||||
shutdown: CancellationToken,
|
||||
}
|
||||
|
||||
impl HttpConnectionBudget {
|
||||
pub fn new(limit: usize) -> Self {
|
||||
let limit = limit.clamp(1, MAX_HTTP_CONNECTIONS.min(Semaphore::MAX_PERMITS));
|
||||
Self {
|
||||
limit,
|
||||
permits: Arc::new(Semaphore::new(limit)),
|
||||
in_flight: AtomicUsize::new(0),
|
||||
high_watermark: AtomicUsize::new(0),
|
||||
rejected_total: AtomicU64::new(0),
|
||||
accept_errors_total: AtomicU64::new(0),
|
||||
shutdown: CancellationToken::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Admit after accepting. Waiting for a permit before accept can let idle
|
||||
/// reuseport listeners monopolize permits needed by a busy listener.
|
||||
pub fn try_admit<T>(self: &Arc<Self>, io: T) -> Result<AdmittedConnection<T>, ()> {
|
||||
let permit = Arc::clone(&self.permits).try_acquire_owned().map_err(|_| {
|
||||
self.rejected_total.fetch_add(1, Ordering::Relaxed);
|
||||
})?;
|
||||
let in_flight = self.in_flight.fetch_add(1, Ordering::Relaxed) + 1;
|
||||
self.high_watermark.fetch_max(in_flight, Ordering::Relaxed);
|
||||
Ok(AdmittedConnection {
|
||||
io,
|
||||
read_shutdown: Box::pin(self.shutdown.clone().cancelled_owned()),
|
||||
write_shutdown: Box::pin(self.shutdown.clone().cancelled_owned()),
|
||||
_permit: ConnectionPermit {
|
||||
budget: Arc::clone(self),
|
||||
_permit: permit,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> HttpConnectionBudgetSnapshot {
|
||||
HttpConnectionBudgetSnapshot {
|
||||
limit: self.limit,
|
||||
in_flight: self.in_flight.load(Ordering::Relaxed),
|
||||
high_watermark: self.high_watermark.load(Ordering::Relaxed),
|
||||
rejected_total: self.rejected_total.load(Ordering::Relaxed),
|
||||
accept_errors_total: self.accept_errors_total.load(Ordering::Relaxed),
|
||||
}
|
||||
}
|
||||
|
||||
/// End the drain deadline for all sockets, including upgraded connections.
|
||||
pub fn force_close(&self) {
|
||||
self.permits.close();
|
||||
self.shutdown.cancel();
|
||||
}
|
||||
|
||||
pub async fn wait_for_forced_close(&self) {
|
||||
self.shutdown.cancelled().await;
|
||||
}
|
||||
|
||||
pub async fn accept(&self, listener: &TcpListener) -> (TcpStream, SocketAddr) {
|
||||
self.accept_with(|| listener.accept()).await
|
||||
}
|
||||
|
||||
async fn accept_with<T, F, A>(&self, mut accept: A) -> T
|
||||
where
|
||||
A: FnMut() -> F,
|
||||
F: Future<Output = io::Result<T>>,
|
||||
{
|
||||
loop {
|
||||
match accept().await {
|
||||
Ok(connection) => return connection,
|
||||
Err(error) => {
|
||||
self.accept_errors_total.fetch_add(1, Ordering::Relaxed);
|
||||
// Match Axum's listener behavior: failed peers can be retried
|
||||
// immediately; resource failures such as EMFILE need backoff.
|
||||
if matches!(
|
||||
error.kind(),
|
||||
io::ErrorKind::ConnectionRefused
|
||||
| io::ErrorKind::ConnectionAborted
|
||||
| io::ErrorKind::ConnectionReset
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
tracing::error!(
|
||||
event_name = "http_connection_accept_failed",
|
||||
error = %error,
|
||||
"HTTP listener accept failed; retrying after one second"
|
||||
);
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ConnectionPermit {
|
||||
budget: Arc<HttpConnectionBudget>,
|
||||
_permit: OwnedSemaphorePermit,
|
||||
}
|
||||
|
||||
impl Drop for ConnectionPermit {
|
||||
fn drop(&mut self) {
|
||||
self.budget.in_flight.fetch_sub(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AdmittedConnection<T> {
|
||||
// Close the socket before returning its permit, including upgrade teardown.
|
||||
io: T,
|
||||
read_shutdown: Pin<Box<WaitForCancellationFutureOwned>>,
|
||||
write_shutdown: Pin<Box<WaitForCancellationFutureOwned>>,
|
||||
_permit: ConnectionPermit,
|
||||
}
|
||||
|
||||
impl<T: AsyncRead + Unpin> AsyncRead for AdmittedConnection<T> {
|
||||
fn poll_read(
|
||||
mut self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &mut ReadBuf<'_>,
|
||||
) -> Poll<io::Result<()>> {
|
||||
if self.read_shutdown.as_mut().poll(cx).is_ready() {
|
||||
return Poll::Ready(Err(shutdown_error()));
|
||||
}
|
||||
Pin::new(&mut self.io).poll_read(cx, buf)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AsyncWrite + Unpin> AsyncWrite for AdmittedConnection<T> {
|
||||
fn poll_write(
|
||||
mut self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
if self.write_shutdown.as_mut().poll(cx).is_ready() {
|
||||
return Poll::Ready(Err(shutdown_error()));
|
||||
}
|
||||
Pin::new(&mut self.io).poll_write(cx, buf)
|
||||
}
|
||||
|
||||
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
if self.write_shutdown.as_mut().poll(cx).is_ready() {
|
||||
return Poll::Ready(Err(shutdown_error()));
|
||||
}
|
||||
Pin::new(&mut self.io).poll_flush(cx)
|
||||
}
|
||||
|
||||
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
Pin::new(&mut self.io).poll_shutdown(cx)
|
||||
}
|
||||
|
||||
fn is_write_vectored(&self) -> bool {
|
||||
self.io.is_write_vectored()
|
||||
}
|
||||
|
||||
fn poll_write_vectored(
|
||||
mut self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
bufs: &[io::IoSlice<'_>],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
if self.write_shutdown.as_mut().poll(cx).is_ready() {
|
||||
return Poll::Ready(Err(shutdown_error()));
|
||||
}
|
||||
Pin::new(&mut self.io).poll_write_vectored(cx, bufs)
|
||||
}
|
||||
}
|
||||
|
||||
fn shutdown_error() -> io::Error {
|
||||
io::Error::new(
|
||||
io::ErrorKind::ConnectionAborted,
|
||||
"gateway shutdown deadline reached",
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "connection_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,411 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::convert::Infallible;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
|
||||
use bytes::Bytes;
|
||||
use http::{Request, Response, StatusCode};
|
||||
use http_body_util::{BodyExt, Empty};
|
||||
use hyper::body::Incoming;
|
||||
use hyper::service::service_fn;
|
||||
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
|
||||
use super::*;
|
||||
|
||||
async fn within<T>(future: impl Future<Output = T>) -> T {
|
||||
tokio::time::timeout(Duration::from_secs(3), future)
|
||||
.await
|
||||
.expect("connection test exceeded its deadline")
|
||||
}
|
||||
|
||||
async fn tcp_pair() -> (TcpStream, TcpStream) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let client = TcpStream::connect(listener.local_addr().unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
let (server, _) = listener.accept().await.unwrap();
|
||||
(client, server)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forced_shutdown_wakes_independent_read_and_write_waiters() {
|
||||
let budget = Arc::new(HttpConnectionBudget::new(2));
|
||||
let (io, _peer) = tokio::io::duplex(1);
|
||||
let io = budget.try_admit(io).unwrap();
|
||||
let (mut reader, mut writer) = tokio::io::split(io);
|
||||
writer.write_all(b"a").await.unwrap();
|
||||
let reading = tokio::spawn(async move { reader.read_u8().await });
|
||||
let writing = tokio::spawn(async move { writer.write_all(b"b").await });
|
||||
tokio::task::yield_now().await;
|
||||
budget.force_close();
|
||||
assert_eq!(
|
||||
within(reading).await.unwrap().unwrap_err().kind(),
|
||||
io::ErrorKind::ConnectionAborted
|
||||
);
|
||||
assert_eq!(
|
||||
within(writing).await.unwrap().unwrap_err().kind(),
|
||||
io::ErrorKind::ConnectionAborted
|
||||
);
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
assert!(budget.try_admit(tokio::io::empty()).is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forced_shutdown_before_first_poll_closes_read_write_and_flush() {
|
||||
let budget = Arc::new(HttpConnectionBudget::new(1));
|
||||
let (io, _peer) = tokio::io::duplex(16);
|
||||
let mut io = budget.try_admit(io).unwrap();
|
||||
budget.force_close();
|
||||
assert_eq!(
|
||||
io.read_u8().await.unwrap_err().kind(),
|
||||
io::ErrorKind::ConnectionAborted
|
||||
);
|
||||
assert_eq!(
|
||||
io.write_all(b"a").await.unwrap_err().kind(),
|
||||
io::ErrorKind::ConnectionAborted
|
||||
);
|
||||
assert_eq!(
|
||||
io.flush().await.unwrap_err().kind(),
|
||||
io::ErrorKind::ConnectionAborted
|
||||
);
|
||||
let buffers = [io::IoSlice::new(b"a")];
|
||||
assert_eq!(
|
||||
io.write_vectored(&buffers).await.unwrap_err().kind(),
|
||||
io::ErrorKind::ConnectionAborted
|
||||
);
|
||||
io.shutdown().await.unwrap();
|
||||
drop(io);
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_connection_limits_apply_to_auto_explicit_zero_and_fd_bounds() {
|
||||
for (configured, requests, websockets, fd_limit, expected) in [
|
||||
(None, 0, 0, None, 1),
|
||||
(None, 3, 5, None, 8),
|
||||
(Some(0), 3, 5, None, 8),
|
||||
(Some(9), 3, 5, None, 9),
|
||||
(Some(usize::MAX), 0, 0, None, MAX_HTTP_CONNECTIONS),
|
||||
(None, usize::MAX, usize::MAX, None, MAX_HTTP_CONNECTIONS),
|
||||
(None, 512, 512, Some(1_024), 384),
|
||||
(Some(500), 3, 5, Some(1_024), 384),
|
||||
(Some(32), 3, 5, Some(1_024), 32),
|
||||
(Some(500), 3, 5, Some(260), 2),
|
||||
(Some(500), 3, 5, Some(258), 1),
|
||||
(Some(500), 3, 5, Some(0), 1),
|
||||
] {
|
||||
assert_eq!(
|
||||
http_connection_limit(configured, requests, websockets, fd_limit),
|
||||
expected,
|
||||
);
|
||||
}
|
||||
assert_eq!(HttpConnectionBudget::new(0).snapshot().limit, 1);
|
||||
assert_eq!(
|
||||
HttpConnectionBudget::new(usize::MAX).snapshot().limit,
|
||||
MAX_HTTP_CONNECTIONS,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_connection_budget_drops_io_before_returning_capacity() {
|
||||
struct DropProbe {
|
||||
budget: Arc<HttpConnectionBudget>,
|
||||
observed: Arc<AtomicBool>,
|
||||
}
|
||||
impl Drop for DropProbe {
|
||||
fn drop(&mut self) {
|
||||
self.observed
|
||||
.store(self.budget.snapshot().in_flight == 1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
let budget = Arc::new(HttpConnectionBudget::new(1));
|
||||
let admitted_dropped = Arc::new(AtomicBool::new(false));
|
||||
let rejected_dropped = Arc::new(AtomicBool::new(false));
|
||||
let admitted = budget
|
||||
.try_admit(DropProbe {
|
||||
budget: Arc::clone(&budget),
|
||||
observed: Arc::clone(&admitted_dropped),
|
||||
})
|
||||
.unwrap();
|
||||
assert!(budget
|
||||
.try_admit(DropProbe {
|
||||
budget: Arc::clone(&budget),
|
||||
observed: Arc::clone(&rejected_dropped),
|
||||
})
|
||||
.is_err());
|
||||
assert!(rejected_dropped.load(Ordering::Relaxed));
|
||||
assert_eq!(budget.snapshot().in_flight, 1);
|
||||
drop(admitted);
|
||||
assert!(admitted_dropped.load(Ordering::Relaxed));
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
assert_eq!(budget.snapshot().high_watermark, 1);
|
||||
assert_eq!(budget.snapshot().rejected_total, 1);
|
||||
drop(budget.try_admit(()).unwrap());
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_connection_budget_forwards_read_write_vectored_and_half_close() {
|
||||
let budget = Arc::new(HttpConnectionBudget::new(1));
|
||||
let (mut peer, io) = tokio::io::duplex(32);
|
||||
let vectored = io.is_write_vectored();
|
||||
let mut admitted = budget.try_admit(io).unwrap();
|
||||
assert_eq!(admitted.is_write_vectored(), vectored);
|
||||
let written = admitted
|
||||
.write_vectored(&[io::IoSlice::new(b"ab"), io::IoSlice::new(b"cd")])
|
||||
.await
|
||||
.unwrap();
|
||||
assert!((1..=4).contains(&written));
|
||||
admitted.write_all(&b"abcd"[written..]).await.unwrap();
|
||||
admitted.flush().await.unwrap();
|
||||
let mut message = [0; 4];
|
||||
peer.read_exact(&mut message).await.unwrap();
|
||||
assert_eq!(&message, b"abcd");
|
||||
|
||||
admitted.shutdown().await.unwrap();
|
||||
assert_eq!(peer.read(&mut [0; 1]).await.unwrap(), 0);
|
||||
assert_eq!(budget.snapshot().in_flight, 1);
|
||||
peer.write_all(b"reply").await.unwrap();
|
||||
let mut reply = [0; 5];
|
||||
admitted.read_exact(&mut reply).await.unwrap();
|
||||
assert_eq!(&reply, b"reply");
|
||||
drop(admitted);
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_connection_budget_two_tcp_listeners_share_capacity_and_recover() {
|
||||
let first_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let second_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let budget = Arc::new(HttpConnectionBudget::new(1));
|
||||
let first_client = TcpStream::connect(first_listener.local_addr().unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
let (first_io, _) = within(budget.accept(&first_listener)).await;
|
||||
let first = budget.try_admit(first_io).unwrap();
|
||||
|
||||
let mut rejected_client = TcpStream::connect(second_listener.local_addr().unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
let (second_io, _) = within(budget.accept(&second_listener)).await;
|
||||
assert!(Arc::clone(&budget).try_admit(second_io).is_err());
|
||||
let rejected_read = within(rejected_client.read(&mut [0; 1])).await;
|
||||
assert!(
|
||||
matches!(rejected_read, Ok(0))
|
||||
|| matches!(rejected_read, Err(ref error) if matches!(
|
||||
error.kind(), io::ErrorKind::ConnectionReset | io::ErrorKind::ConnectionAborted
|
||||
))
|
||||
);
|
||||
assert_eq!(budget.snapshot().in_flight, 1);
|
||||
assert_eq!(budget.snapshot().rejected_total, 1);
|
||||
|
||||
drop((first, first_client, rejected_client));
|
||||
let replacement_client = TcpStream::connect(second_listener.local_addr().unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
let (replacement_io, _) = within(budget.accept(&second_listener)).await;
|
||||
let replacement = budget.try_admit(replacement_io).unwrap();
|
||||
assert_eq!(budget.snapshot().in_flight, 1);
|
||||
assert_eq!(budget.snapshot().high_watermark, 1);
|
||||
drop((replacement, replacement_client));
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_connection_budget_task_cancelled_before_first_poll_returns_permit() {
|
||||
let (client, server) = tcp_pair().await;
|
||||
let budget = Arc::new(HttpConnectionBudget::new(1));
|
||||
let admitted = budget.try_admit(server).unwrap();
|
||||
let polled = Arc::new(AtomicBool::new(false));
|
||||
let polled_by_task = Arc::clone(&polled);
|
||||
let task = tokio::spawn(async move {
|
||||
let _io = admitted;
|
||||
polled_by_task.store(true, Ordering::Relaxed);
|
||||
std::future::pending::<()>().await;
|
||||
});
|
||||
task.abort();
|
||||
assert!(within(task).await.unwrap_err().is_cancelled());
|
||||
assert!(!polled.load(Ordering::Relaxed));
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
drop(client);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_connection_budget_header_timeout_and_parse_failure_return_permit() {
|
||||
for malformed in [false, true] {
|
||||
let (mut client, server_io) = tcp_pair().await;
|
||||
let budget = Arc::new(HttpConnectionBudget::new(1));
|
||||
let admitted = budget.try_admit(server_io).unwrap();
|
||||
let requests = Arc::new(AtomicUsize::new(0));
|
||||
let requests_seen = Arc::clone(&requests);
|
||||
let server = tokio::spawn(async move {
|
||||
let service = service_fn(move |_: Request<Incoming>| {
|
||||
requests_seen.fetch_add(1, Ordering::Relaxed);
|
||||
async { Ok::<_, Infallible>(Response::new(Empty::<Bytes>::new())) }
|
||||
});
|
||||
hyper::server::conn::http1::Builder::new()
|
||||
.timer(TokioTimer::new())
|
||||
.header_read_timeout(Duration::from_millis(20))
|
||||
.serve_connection(TokioIo::new(admitted), service)
|
||||
.await
|
||||
});
|
||||
if malformed {
|
||||
client.write_all(b"invalid request\r\n\r\n").await.unwrap();
|
||||
}
|
||||
let _connection_result = within(server).await.unwrap();
|
||||
assert_eq!(requests.load(Ordering::Relaxed), 0);
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
drop(client);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_connection_budget_h1_upgrade_keeps_permit_after_connection_future_finishes() {
|
||||
let (mut client, server_io) = tcp_pair().await;
|
||||
let budget = Arc::new(HttpConnectionBudget::new(1));
|
||||
let admitted = budget.try_admit(server_io).unwrap();
|
||||
let (upgrade_finished_tx, upgrade_finished_rx) = tokio::sync::oneshot::channel();
|
||||
let upgrade_finished_tx = Arc::new(std::sync::Mutex::new(Some(upgrade_finished_tx)));
|
||||
let server = tokio::spawn(async move {
|
||||
let service = service_fn(move |mut request: Request<Incoming>| {
|
||||
let on_upgrade = hyper::upgrade::on(&mut request);
|
||||
let finished = upgrade_finished_tx.lock().unwrap().take().unwrap();
|
||||
tokio::spawn(async move {
|
||||
let mut upgraded = TokioIo::new(on_upgrade.await.unwrap());
|
||||
let mut message = [0; 4];
|
||||
upgraded.read_exact(&mut message).await.unwrap();
|
||||
upgraded.write_all(&message).await.unwrap();
|
||||
assert_eq!(upgraded.read(&mut [0; 1]).await.unwrap(), 0);
|
||||
drop(upgraded);
|
||||
let _ = finished.send(());
|
||||
});
|
||||
async {
|
||||
Ok::<_, Infallible>(
|
||||
Response::builder()
|
||||
.status(StatusCode::SWITCHING_PROTOCOLS)
|
||||
.header("connection", "upgrade")
|
||||
.header("upgrade", "echo")
|
||||
.body(Empty::<Bytes>::new())
|
||||
.unwrap(),
|
||||
)
|
||||
}
|
||||
});
|
||||
hyper::server::conn::http1::Builder::new()
|
||||
.serve_connection(TokioIo::new(admitted), service)
|
||||
.with_upgrades()
|
||||
.await
|
||||
});
|
||||
client
|
||||
.write_all(
|
||||
b"GET / HTTP/1.1\r\nHost: localhost\r\nConnection: upgrade\r\nUpgrade: echo\r\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let response = within(async {
|
||||
let mut headers = Vec::new();
|
||||
while !headers.ends_with(b"\r\n\r\n") {
|
||||
assert!(headers.len() < 4096);
|
||||
headers.push(client.read_u8().await.unwrap());
|
||||
}
|
||||
headers
|
||||
})
|
||||
.await;
|
||||
assert!(response.starts_with(b"HTTP/1.1 101"));
|
||||
within(server).await.unwrap().unwrap();
|
||||
assert_eq!(budget.snapshot().in_flight, 1);
|
||||
assert!(budget.try_admit(()).is_err());
|
||||
client.write_all(b"ping").await.unwrap();
|
||||
let mut echoed = [0; 4];
|
||||
within(client.read_exact(&mut echoed)).await.unwrap();
|
||||
assert_eq!(&echoed, b"ping");
|
||||
drop(client);
|
||||
within(upgrade_finished_rx).await.unwrap();
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_connection_budget_h2_parallel_streams_share_one_socket_permit() {
|
||||
let (client_io, server_io) = tcp_pair().await;
|
||||
let budget = Arc::new(HttpConnectionBudget::new(1));
|
||||
let admitted = budget.try_admit(server_io).unwrap();
|
||||
let requests = Arc::new(AtomicUsize::new(0));
|
||||
let concurrent = Arc::new(tokio::sync::Barrier::new(3));
|
||||
let server_requests = Arc::clone(&requests);
|
||||
let server_concurrent = Arc::clone(&concurrent);
|
||||
let server = tokio::spawn(async move {
|
||||
let service = service_fn(move |_: Request<Incoming>| {
|
||||
server_requests.fetch_add(1, Ordering::Relaxed);
|
||||
let concurrent = Arc::clone(&server_concurrent);
|
||||
async move {
|
||||
concurrent.wait().await;
|
||||
Ok::<_, Infallible>(Response::new(Empty::<Bytes>::new()))
|
||||
}
|
||||
});
|
||||
hyper::server::conn::http2::Builder::new(TokioExecutor::new())
|
||||
.max_concurrent_streams(2)
|
||||
.serve_connection(TokioIo::new(admitted), service)
|
||||
.await
|
||||
});
|
||||
let (sender, connection) = within(
|
||||
hyper::client::conn::http2::Builder::new(TokioExecutor::new())
|
||||
.handshake::<_, Empty<Bytes>>(TokioIo::new(client_io)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let client_driver = tokio::spawn(connection);
|
||||
let mut requests_in_flight = tokio::task::JoinSet::new();
|
||||
for path in ["first", "second"] {
|
||||
let mut sender = sender.clone();
|
||||
requests_in_flight.spawn(async move {
|
||||
let response = sender
|
||||
.send_request(
|
||||
Request::builder()
|
||||
.uri(format!("http://localhost/{path}"))
|
||||
.body(Empty::<Bytes>::new())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
response.into_body().collect().await.unwrap();
|
||||
});
|
||||
}
|
||||
within(concurrent.wait()).await;
|
||||
assert_eq!(requests.load(Ordering::Relaxed), 2);
|
||||
assert_eq!(budget.snapshot().in_flight, 1);
|
||||
assert_eq!(budget.snapshot().high_watermark, 1);
|
||||
within(async {
|
||||
while let Some(result) = requests_in_flight.join_next().await {
|
||||
result.unwrap();
|
||||
}
|
||||
})
|
||||
.await;
|
||||
assert_eq!(budget.snapshot().in_flight, 1);
|
||||
drop(sender);
|
||||
client_driver.abort();
|
||||
let _ = within(client_driver).await;
|
||||
let _ = within(server).await.unwrap();
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn http_connection_accept_retries_peer_errors_and_backs_off_resource_errors() {
|
||||
let budget = HttpConnectionBudget::new(1);
|
||||
let mut attempts = VecDeque::from([
|
||||
Err(io::Error::from(io::ErrorKind::ConnectionAborted)),
|
||||
Err(io::Error::from(io::ErrorKind::ConnectionReset)),
|
||||
Err(io::Error::other("injected file descriptor exhaustion")),
|
||||
Err(io::Error::other("injected temporary accept failure")),
|
||||
Ok(42),
|
||||
]);
|
||||
let started = tokio::time::Instant::now();
|
||||
let accepted = budget
|
||||
.accept_with(|| std::future::ready(attempts.pop_front().unwrap()))
|
||||
.await;
|
||||
assert_eq!(accepted, 42);
|
||||
assert!(attempts.is_empty());
|
||||
assert_eq!(started.elapsed(), Duration::from_secs(2));
|
||||
assert_eq!(budget.snapshot().accept_errors_total, 4);
|
||||
assert_eq!(budget.snapshot().in_flight, 0);
|
||||
assert_eq!(budget.snapshot().rejected_total, 0);
|
||||
}
|
||||
@@ -1,11 +1,15 @@
|
||||
pub mod body;
|
||||
mod connection;
|
||||
pub mod middleware;
|
||||
mod request_id;
|
||||
|
||||
pub use body::{
|
||||
BodyBufferError, BodyBufferPolicy, BodyBufferReservation, BufferedBody,
|
||||
BodyBufferBudget, BodyBufferError, BodyBufferPolicy, BodyBufferReservation, BufferedBody,
|
||||
DEFAULT_BODY_BUFFER_PERMIT_BYTES,
|
||||
};
|
||||
pub use connection::{
|
||||
http_connection_limit, AdmittedConnection, HttpConnectionBudget, HttpConnectionBudgetSnapshot,
|
||||
};
|
||||
|
||||
pub use middleware::access_log::{
|
||||
access_log_middleware, sanitize_access_log_path, should_downgrade_access_log,
|
||||
|
||||
Reference in New Issue
Block a user