mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-13 06:30:20 +08:00
fix: resolve concurrency hardening lint failures
Use typed connection admission errors, group HTTP limits, and make test lock lifetimes explicit. Handle fixture reads and remove unnecessary cloning and manual divisibility checks.
This commit is contained in:
@@ -1242,7 +1242,7 @@ mod tests {
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut request = [0_u8; 4096];
|
||||
socket.read(&mut request).await.unwrap();
|
||||
assert!(socket.read(&mut request).await.unwrap() > 0);
|
||||
let response = if content_type == "application/json" {
|
||||
format!("HTTP/1.1 200 OK\r\ncontent-type: {content_type}\r\ncontent-length: 1024\r\n\r\n{first_chunk}")
|
||||
} else {
|
||||
|
||||
@@ -1879,39 +1879,38 @@ fn gateway_listeners(
|
||||
Ok(listeners)
|
||||
}
|
||||
|
||||
async fn serve_gateway_router(
|
||||
listeners: Vec<tokio::net::TcpListener>,
|
||||
router: axum::Router,
|
||||
connection_budget: Arc<HttpConnectionBudget>,
|
||||
#[derive(Clone, Copy)]
|
||||
struct GatewayHttpLimits {
|
||||
http2_max_concurrent_streams: u32,
|
||||
http_header_read_timeout_ms: u64,
|
||||
http_header_max_bytes: usize,
|
||||
http_max_headers: usize,
|
||||
}
|
||||
|
||||
async fn serve_gateway_router(
|
||||
listeners: Vec<tokio::net::TcpListener>,
|
||||
router: axum::Router,
|
||||
connection_budget: Arc<HttpConnectionBudget>,
|
||||
limits: GatewayHttpLimits,
|
||||
shutdown: CancellationToken,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let http2_max_concurrent_streams =
|
||||
gateway_http2_max_concurrent_streams(http2_max_concurrent_streams);
|
||||
let http_header_read_timeout_ms =
|
||||
gateway_http_header_read_timeout_ms(http_header_read_timeout_ms);
|
||||
let http_header_max_bytes = gateway_http_header_max_bytes(http_header_max_bytes);
|
||||
let http_max_headers = gateway_http_max_headers(http_max_headers);
|
||||
let limits = GatewayHttpLimits {
|
||||
http2_max_concurrent_streams: gateway_http2_max_concurrent_streams(
|
||||
limits.http2_max_concurrent_streams,
|
||||
),
|
||||
http_header_read_timeout_ms: gateway_http_header_read_timeout_ms(
|
||||
limits.http_header_read_timeout_ms,
|
||||
),
|
||||
http_header_max_bytes: gateway_http_header_max_bytes(limits.http_header_max_bytes),
|
||||
http_max_headers: gateway_http_max_headers(limits.http_max_headers),
|
||||
};
|
||||
let mut servers = tokio::task::JoinSet::new();
|
||||
for listener in listeners {
|
||||
let router = router.clone();
|
||||
let connection_budget = Arc::clone(&connection_budget);
|
||||
let shutdown = shutdown.clone();
|
||||
servers.spawn(async move {
|
||||
serve_gateway_listener(
|
||||
listener,
|
||||
router,
|
||||
connection_budget,
|
||||
http2_max_concurrent_streams,
|
||||
http_header_read_timeout_ms,
|
||||
http_header_max_bytes,
|
||||
http_max_headers,
|
||||
shutdown,
|
||||
)
|
||||
.await
|
||||
serve_gateway_listener(listener, router, connection_budget, limits, shutdown).await
|
||||
});
|
||||
}
|
||||
let mut failure = None;
|
||||
@@ -1941,12 +1940,15 @@ async fn serve_gateway_listener(
|
||||
listener: tokio::net::TcpListener,
|
||||
router: axum::Router,
|
||||
connection_budget: Arc<HttpConnectionBudget>,
|
||||
http2_max_concurrent_streams: u32,
|
||||
http_header_read_timeout_ms: u64,
|
||||
http_header_max_bytes: usize,
|
||||
http_max_headers: usize,
|
||||
limits: GatewayHttpLimits,
|
||||
shutdown: CancellationToken,
|
||||
) -> Result<(), std::io::Error> {
|
||||
let GatewayHttpLimits {
|
||||
http2_max_concurrent_streams,
|
||||
http_header_read_timeout_ms,
|
||||
http_header_max_bytes,
|
||||
http_max_headers,
|
||||
} = limits;
|
||||
let mut make_service = router.into_make_service_with_connect_info::<std::net::SocketAddr>();
|
||||
let mut connections = tokio::task::JoinSet::new();
|
||||
loop {
|
||||
@@ -2594,10 +2596,12 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
listeners,
|
||||
router,
|
||||
Arc::clone(&http_connection_budget),
|
||||
args.http2_max_concurrent_streams,
|
||||
args.http_header_read_timeout_ms,
|
||||
args.http_header_max_bytes,
|
||||
args.http_max_headers,
|
||||
GatewayHttpLimits {
|
||||
http2_max_concurrent_streams: args.http2_max_concurrent_streams,
|
||||
http_header_read_timeout_ms: args.http_header_read_timeout_ms,
|
||||
http_header_max_bytes: args.http_header_max_bytes,
|
||||
http_max_headers: args.http_max_headers,
|
||||
},
|
||||
shutdown.clone(),
|
||||
);
|
||||
tokio::pin!(server);
|
||||
|
||||
@@ -7,7 +7,7 @@ use tokio::net::{TcpListener, TcpStream};
|
||||
use tokio::sync::Notify;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::super::{serve_gateway_router, HttpConnectionBudget};
|
||||
use super::super::{serve_gateway_router, GatewayHttpLimits, HttpConnectionBudget};
|
||||
|
||||
async fn start(
|
||||
router: Router,
|
||||
@@ -28,10 +28,12 @@ async fn start(
|
||||
vec![listener],
|
||||
router,
|
||||
shared,
|
||||
16,
|
||||
10_000,
|
||||
32_768,
|
||||
100,
|
||||
GatewayHttpLimits {
|
||||
http2_max_concurrent_streams: 16,
|
||||
http_header_read_timeout_ms: 10_000,
|
||||
http_header_max_bytes: 32_768,
|
||||
http_max_headers: 100,
|
||||
},
|
||||
stop,
|
||||
)
|
||||
.await
|
||||
|
||||
@@ -9,7 +9,7 @@ use std::time::Duration;
|
||||
|
||||
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError};
|
||||
use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned};
|
||||
|
||||
const MAX_HTTP_CONNECTIONS: usize = 65_536;
|
||||
@@ -71,10 +71,12 @@ impl HttpConnectionBudget {
|
||||
|
||||
/// 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);
|
||||
})?;
|
||||
pub fn try_admit<T>(self: &Arc<Self>, io: T) -> Result<AdmittedConnection<T>, TryAcquireError> {
|
||||
let permit = Arc::clone(&self.permits)
|
||||
.try_acquire_owned()
|
||||
.inspect_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 {
|
||||
|
||||
@@ -173,7 +173,10 @@ async fn assert_transfer_source_unchanged(
|
||||
source: &str,
|
||||
entry: &RuntimeQueueEntry,
|
||||
) {
|
||||
assert_eq!(transfer_entries(admin, source).await, [entry.clone()]);
|
||||
assert_eq!(
|
||||
transfer_entries(admin, source).await.as_slice(),
|
||||
std::slice::from_ref(&entry)
|
||||
);
|
||||
let pending = transfer_pending(admin, source).await;
|
||||
assert_eq!(pending.len(), 1);
|
||||
assert_eq!(pending[0].0, entry.id);
|
||||
|
||||
@@ -100,7 +100,7 @@ async fn copy_and_commit(
|
||||
if plan.first() != Some(&2) {
|
||||
return Ok(Some((plan, false)));
|
||||
}
|
||||
if plan.len() < 4 || (plan.len() - 1) % 3 != 0 {
|
||||
if plan.len() < 4 || !(plan.len() - 1).is_multiple_of(3) {
|
||||
return Err(DataLayerError::UnexpectedValue(
|
||||
"invalid usage window copy plan".to_string(),
|
||||
));
|
||||
|
||||
@@ -8681,12 +8681,13 @@ mod tests {
|
||||
);
|
||||
}
|
||||
assert_eq!(store.enrichment_calls.load(Ordering::Acquire), 2);
|
||||
let records = store.records.lock().expect("records lock");
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].total_cost_usd, Some(0.456));
|
||||
assert_eq!(records[0].actual_total_cost_usd, Some(0.123));
|
||||
assert_eq!(records[0].total_tokens, Some(12));
|
||||
drop(records);
|
||||
{
|
||||
let records = store.records.lock().expect("records lock");
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].total_cost_usd, Some(0.456));
|
||||
assert_eq!(records[0].actual_total_cost_usd, Some(0.123));
|
||||
assert_eq!(records[0].total_tokens, Some(12));
|
||||
}
|
||||
{
|
||||
let coalescer = &runtime.lifecycle_coalescer;
|
||||
let entries = coalescer.shards[coalescer.shard_index(request_id)]
|
||||
|
||||
@@ -1670,20 +1670,22 @@ mod tests {
|
||||
.await
|
||||
.expect("pending entry should become reclaimable");
|
||||
|
||||
let records = store.records.lock().expect("records lock");
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].total_cost_usd, Some(0.456));
|
||||
assert_eq!(records[0].actual_total_cost_usd, Some(0.123));
|
||||
assert_eq!(records[0].total_tokens, Some(10));
|
||||
drop(records);
|
||||
let reconciliations = store.reconciliations.lock().expect("reconciliations lock");
|
||||
assert_eq!(reconciliations.len(), 1);
|
||||
assert_eq!(reconciliations[0].actual_cost_units, 12_300_000);
|
||||
assert_eq!(
|
||||
reconciliations[0].reservation_token,
|
||||
"pricing-retry-reservation"
|
||||
);
|
||||
drop(reconciliations);
|
||||
{
|
||||
let records = store.records.lock().expect("records lock");
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].total_cost_usd, Some(0.456));
|
||||
assert_eq!(records[0].actual_total_cost_usd, Some(0.123));
|
||||
assert_eq!(records[0].total_tokens, Some(10));
|
||||
}
|
||||
{
|
||||
let reconciliations = store.reconciliations.lock().expect("reconciliations lock");
|
||||
assert_eq!(reconciliations.len(), 1);
|
||||
assert_eq!(reconciliations[0].actual_cost_units, 12_300_000);
|
||||
assert_eq!(
|
||||
reconciliations[0].reservation_token,
|
||||
"pricing-retry-reservation"
|
||||
);
|
||||
}
|
||||
assert_eq!(store.settlements.lock().expect("settlements lock").len(), 1);
|
||||
assert_eq!(
|
||||
store.enrich_calls.lock().expect("enrich calls lock").len(),
|
||||
@@ -2400,11 +2402,12 @@ mod tests {
|
||||
.await
|
||||
.expect("worker should stop")
|
||||
.expect("worker task");
|
||||
let records = store.records.lock().expect("records lock");
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].response_body, event.data.response_body);
|
||||
assert_eq!(records[0].total_tokens, Some(10));
|
||||
drop(records);
|
||||
{
|
||||
let records = store.records.lock().expect("records lock");
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].response_body, event.data.response_body);
|
||||
assert_eq!(records[0].total_tokens, Some(10));
|
||||
}
|
||||
let snapshot = budget.snapshot();
|
||||
assert_eq!(snapshot.reserved_bytes, 0);
|
||||
assert_eq!(snapshot.oversized_entries_total, 1);
|
||||
|
||||
Reference in New Issue
Block a user