mirror of
https://github.com/telemt/telemt.git
synced 2026-09-24 03:25:58 +03:00
Split oversized runtime modules + Async tests hardened
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
@@ -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();
|
||||
|
||||
@@ -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
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}));
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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(
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user