Split oversized runtime modules + Async tests hardened

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-08-30 09:39:25 +03:00
parent 281f63f940
commit 66f2b8889f
154 changed files with 25561 additions and 24915 deletions
+8 -8
View File
@@ -129,12 +129,10 @@ impl ProcessControlPlane {
if self.inner.shutdown_completed.load(Ordering::Acquire) {
return true;
}
let registrations_stopped = tokio::time::timeout_at(
deadline,
self.inner.admission.wait_for_registrations(),
)
.await
.is_ok();
let registrations_stopped =
tokio::time::timeout_at(deadline, self.inner.admission.wait_for_registrations())
.await
.is_ok();
let tasks_stopped = tokio::time::timeout_at(deadline, self.inner.tasks.wait())
.await
.is_ok();
@@ -188,7 +186,8 @@ mod tests {
let first = tokio::spawn(async move { first_scope.shutdown(Duration::from_secs(1)).await });
tokio::task::yield_now().await;
let second_scope = scope.clone();
let second = tokio::spawn(async move { second_scope.shutdown(Duration::from_secs(1)).await });
let second =
tokio::spawn(async move { second_scope.shutdown(Duration::from_secs(1)).await });
tokio::task::yield_now().await;
assert!(!first.is_finished());
@@ -204,7 +203,8 @@ mod tests {
let scope = ProcessControlPlane::new();
let registration = scope.inner.admission.try_register().unwrap();
let first_scope = scope.clone();
let first = tokio::spawn(async move { first_scope.shutdown(Duration::from_secs(30)).await });
let first =
tokio::spawn(async move { first_scope.shutdown(Duration::from_secs(30)).await });
tokio::task::yield_now().await;
first.abort();
+26 -30
View File
@@ -213,11 +213,9 @@ async fn run_accept_loop(
continue;
}
if web_runtime.is_shutdown() {
web_runtime
.telemetry()
.record_rejection(
crate::web::telemetry::WebRejectionReason::RuntimeClosed,
);
web_runtime.telemetry().record_rejection(
crate::web::telemetry::WebRejectionReason::RuntimeClosed,
);
drop(stream);
continue;
}
@@ -233,9 +231,8 @@ async fn run_accept_loop(
Err(HttpConnectionAdmissionError::AtCapacity) => {
let config = web_runtime.active_generation().config();
let action = config.web.http_connection_capacity_action;
let phase_timeout = Duration::from_millis(
config.web.timeouts.http_overload_timeout_ms,
);
let phase_timeout =
Duration::from_millis(config.web.timeouts.http_overload_timeout_ms);
drop(config);
if action == crate::config::WebHttpConnectionCapacityAction::Drop {
web_runtime.telemetry().record_rejection(
@@ -247,30 +244,29 @@ async fn run_accept_loop(
drop(stream);
continue;
}
let overload_permit =
match web_runtime.try_http_overload_connection() {
Ok(permit) => permit,
Err(HttpConnectionAdmissionError::Closed) => {
web_runtime.telemetry().record_rejection(
crate::web::telemetry::WebRejectionReason::RuntimeClosed,
);
web_runtime.telemetry().record_overload(
WebHttpConnectionOverloadOutcome::ShutdownDrop,
);
drop(stream);
continue;
}
Err(HttpConnectionAdmissionError::AtCapacity) => {
web_runtime.telemetry().record_rejection(
let overload_permit = match web_runtime.try_http_overload_connection() {
Ok(permit) => permit,
Err(HttpConnectionAdmissionError::Closed) => {
web_runtime.telemetry().record_rejection(
crate::web::telemetry::WebRejectionReason::RuntimeClosed,
);
web_runtime.telemetry().record_overload(
WebHttpConnectionOverloadOutcome::ShutdownDrop,
);
drop(stream);
continue;
}
Err(HttpConnectionAdmissionError::AtCapacity) => {
web_runtime.telemetry().record_rejection(
crate::web::telemetry::WebRejectionReason::HttpConnectionCapacity,
);
web_runtime.telemetry().record_overload(
WebHttpConnectionOverloadOutcome::OverflowCapacityDrop,
);
drop(stream);
continue;
}
};
web_runtime.telemetry().record_overload(
WebHttpConnectionOverloadOutcome::OverflowCapacityDrop,
);
drop(stream);
continue;
}
};
connections.spawn(web_overload::serve(
stream,
peer_addr,
+6 -115
View File
@@ -22,57 +22,11 @@ use crate::transport::middle_proxy::MePool;
use super::generation::RuntimeTaskScope;
use super::helpers::load_startup_proxy_config_snapshot;
async fn supervise_me_task<F, Fut>(task_name: &'static str, mut task: F)
where
F: FnMut() -> Fut,
Fut: Future<Output = ()> + Send + 'static,
{
loop {
let result = AbortOnDropHandle::new(tokio::spawn(task())).await;
match result {
Ok(()) => warn!(
task = task_name,
"Middle-End supervisor task exited unexpectedly, restarting"
),
Err(error) => {
error!(task = task_name, error = %error, "Middle-End supervisor task panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
}
fn spawn_me_supervisors(
task_scope: RuntimeTaskScope,
pool: Arc<MePool>,
rng: Arc<SecureRandom>,
min_connections: usize,
) {
let health_pool = pool.clone();
let health_rng = rng;
task_scope.spawn(supervise_me_task("health_monitor", move || {
let pool = health_pool.clone();
let rng = health_rng.clone();
async move {
crate::transport::middle_proxy::me_health_monitor(pool, rng, min_connections).await;
}
}));
let drain_pool = pool.clone();
task_scope.spawn(supervise_me_task("drain_timeout_enforcer", move || {
let pool = drain_pool.clone();
async move {
crate::transport::middle_proxy::me_drain_timeout_enforcer(pool).await;
}
}));
task_scope.spawn(supervise_me_task("zombie_writer_watchdog", move || {
let pool = pool.clone();
async move {
crate::transport::middle_proxy::me_zombie_writer_watchdog(pool).await;
}
}));
}
// Restarting supervisors for long-lived ME maintenance tasks.
mod supervisor;
use supervisor::spawn_me_supervisors;
#[cfg(test)]
use supervisor::supervise_me_task;
pub(crate) async fn initialize_me_pool(
use_middle_proxy: bool,
@@ -587,67 +541,4 @@ pub(crate) async fn initialize_me_pool(
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::Notify;
struct DropSignal(Arc<Notify>);
impl Drop for DropSignal {
fn drop(&mut self) {
self.0.notify_one();
}
}
#[tokio::test]
async fn scoped_supervisor_aborts_its_current_child() {
let scope = RuntimeTaskScope::new();
let dropped = Arc::new(Notify::new());
let dropped_for_task = dropped.clone();
scope.spawn(supervise_me_task("test", move || {
let dropped = dropped_for_task.clone();
async move {
let _signal = DropSignal(dropped);
std::future::pending::<()>().await;
}
}));
tokio::task::yield_now().await;
scope.stop().await;
tokio::time::timeout(Duration::from_secs(1), dropped.notified())
.await
.unwrap();
}
#[tokio::test]
async fn supervisor_restarts_exited_child_and_stops_with_runtime_scope() {
let scope = RuntimeTaskScope::new();
let starts = Arc::new(AtomicUsize::new(0));
let restarted = Arc::new(Notify::new());
let starts_task = starts.clone();
let restarted_task = restarted.clone();
scope.spawn(supervise_me_task("restart_test", move || {
let starts = starts_task.clone();
let restarted = restarted_task.clone();
async move {
if starts.fetch_add(1, Ordering::AcqRel) + 1 >= 3 {
restarted.notify_one();
}
}
}));
tokio::time::timeout(Duration::from_secs(1), restarted.notified())
.await
.unwrap();
scope.stop().await;
let stopped_at = starts.load(Ordering::Acquire);
for _ in 0..100 {
tokio::task::yield_now().await;
}
assert!(stopped_at >= 3);
assert_eq!(starts.load(Ordering::Acquire), stopped_at);
}
}
mod tests;
+53
View File
@@ -0,0 +1,53 @@
use super::*;
pub(super) async fn supervise_me_task<F, Fut>(task_name: &'static str, mut task: F)
where
F: FnMut() -> Fut,
Fut: Future<Output = ()> + Send + 'static,
{
loop {
let result = AbortOnDropHandle::new(tokio::spawn(task())).await;
match result {
Ok(()) => warn!(
task = task_name,
"Middle-End supervisor task exited unexpectedly, restarting"
),
Err(error) => {
error!(task = task_name, error = %error, "Middle-End supervisor task panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
}
pub(super) fn spawn_me_supervisors(
task_scope: RuntimeTaskScope,
pool: Arc<MePool>,
rng: Arc<SecureRandom>,
min_connections: usize,
) {
let health_pool = pool.clone();
let health_rng = rng;
task_scope.spawn(supervise_me_task("health_monitor", move || {
let pool = health_pool.clone();
let rng = health_rng.clone();
async move {
crate::transport::middle_proxy::me_health_monitor(pool, rng, min_connections).await;
}
}));
let drain_pool = pool.clone();
task_scope.spawn(supervise_me_task("drain_timeout_enforcer", move || {
let pool = drain_pool.clone();
async move {
crate::transport::middle_proxy::me_drain_timeout_enforcer(pool).await;
}
}));
task_scope.spawn(supervise_me_task("zombie_writer_watchdog", move || {
let pool = pool.clone();
async move {
crate::transport::middle_proxy::me_zombie_writer_watchdog(pool).await;
}
}));
}
+62
View File
@@ -0,0 +1,62 @@
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::Notify;
struct DropSignal(Arc<Notify>);
impl Drop for DropSignal {
fn drop(&mut self) {
self.0.notify_one();
}
}
#[tokio::test]
async fn scoped_supervisor_aborts_its_current_child() {
let scope = RuntimeTaskScope::new();
let dropped = Arc::new(Notify::new());
let dropped_for_task = dropped.clone();
scope.spawn(supervise_me_task("test", move || {
let dropped = dropped_for_task.clone();
async move {
let _signal = DropSignal(dropped);
std::future::pending::<()>().await;
}
}));
tokio::task::yield_now().await;
scope.stop().await;
tokio::time::timeout(Duration::from_secs(1), dropped.notified())
.await
.unwrap();
}
#[tokio::test]
async fn supervisor_restarts_exited_child_and_stops_with_runtime_scope() {
let scope = RuntimeTaskScope::new();
let starts = Arc::new(AtomicUsize::new(0));
let restarted = Arc::new(Notify::new());
let starts_task = starts.clone();
let restarted_task = restarted.clone();
scope.spawn(supervise_me_task("restart_test", move || {
let starts = starts_task.clone();
let restarted = restarted_task.clone();
async move {
if starts.fetch_add(1, Ordering::AcqRel) + 1 >= 3 {
restarted.notify_one();
}
}
}));
tokio::time::timeout(Duration::from_secs(1), restarted.notified())
.await
.unwrap();
scope.stop().await;
let stopped_at = starts.load(Ordering::Acquire);
for _ in 0..100 {
tokio::task::yield_now().await;
}
assert!(stopped_at >= 3);
assert_eq!(starts.load(Ordering::Acquire), stopped_at);
}
+3 -10
View File
@@ -52,16 +52,9 @@ pub(super) async fn run_telemt_core(
let runtime_task_scope = generation::RuntimeTaskScope::new();
stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
let quota_state_path = config.general.quota_state_path.clone();
let quota_state = crate::quota_state::QuotaStateOwner::new(
quota_state_path,
quota_store.clone(),
);
let configured_quota_users = config
.access
.users
.keys()
.cloned()
.collect::<BTreeSet<_>>();
let quota_state =
crate::quota_state::QuotaStateOwner::new(quota_state_path, quota_store.clone());
let configured_quota_users = config.access.users.keys().cloned().collect::<BTreeSet<_>>();
quota_state.load(&configured_quota_users).await;
let upstream_manager = Arc::new(
+2 -5
View File
@@ -23,9 +23,9 @@ use super::control_plane::ProcessControlPlane;
use super::generation::RuntimeGeneration;
use super::helpers::{format_uptime, unit_label};
use super::reload_supervisor::ReloadSupervisorHandle;
use crate::quota_state::QuotaStateOwner;
use crate::stats::Stats;
use crate::synlimit_control;
use crate::quota_state::QuotaStateOwner;
/// Signal that triggered shutdown.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -130,10 +130,7 @@ async fn perform_shutdown(
warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
}
if !process_control_plane
.shutdown(Duration::from_secs(5))
.await
{
if !process_control_plane.shutdown(Duration::from_secs(5)).await {
warn!("Process control-plane task shutdown deadline expired");
}