astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-10-31 21:02:01 +03:00
Code Issues Packages Projects Releases Wiki Activity

Process-wide concurrency + Cancellation ownership fixes

Browse Source
This commit is contained in:
Alexey
2026-09-23 22:30:51 +03:00
parent baa9bfbb01
commit f1107c21d9
47 changed files with 2177 additions and 762 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/proxy/middle_relay/session.rs
+21 -9
View File
@@ -2,11 +2,15 @@ use super::*;
// Bounded C2ME sender and downstream writer tasks.
mod tasks;
// Child-task ownership aborts relay tasks when the parent future is cancelled.
mod children;
// Conntrack close classification.
mod close_reason;
use children::RelayChildTasks;
use close_reason::classify_conntrack_close_reason;
use tasks::{run_c2me_sender, run_me_writer};
struct RelayConnLease {
connection: Option<ConnLease>,
conn_id: u64,
@@ -185,7 +189,7 @@ where
let c2me_byte_semaphore = Arc::new(Semaphore::new(c2me_byte_budget));
let (c2me_tx, c2me_rx) = mpsc::channel::<C2MeCommand>(c2me_channel_capacity);
let me_pool_c2me = me_pool.clone();
let mut c2me_sender = tokio::spawn(run_c2me_sender(
let c2me_sender = AbortOnDropHandle::new(tokio::spawn(run_c2me_sender(
c2me_rx,
me_pool_c2me,
conn_id,
@@ -193,7 +197,7 @@ where
peer,
translated_local_addr,
effective_tag_array,
));
)));
let (stop_tx, stop_rx) = oneshot::channel::<()>();
let flow_cancel = CancellationToken::new();
@@ -208,7 +212,7 @@ where
let last_downstream_activity_ms_clone = last_downstream_activity_ms.clone();
let bytes_me2c_clone = bytes_me2c.clone();
let d2c_flush_policy = MeD2cFlushPolicy::from_config(&config);
let mut me_writer = tokio::spawn(run_me_writer(
let me_writer = AbortOnDropHandle::new(tokio::spawn(run_me_writer(
crypto_writer,
me_rx_task,
stats_clone,
@@ -226,7 +230,13 @@ where
session_started_at,
conn_id,
stop_rx,
));
)));
let mut child_tasks = RelayChildTasks {
c2me_sender,
me_writer,
flow_cancel: flow_cancel.clone(),
stop_tx: Some(stop_tx),
};
let mut main_result: Result<()> = Ok(());
let mut client_closed = false;
@@ -454,28 +464,30 @@ where
}
drop(c2me_tx);
let c2me_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut c2me_sender).await {
let c2me_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut child_tasks.c2me_sender).await {
Ok(joined) => {
joined.unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME sender join error: {e}"))))
}
Err(_) => {
stats.increment_me_child_join_timeout_total();
stats.increment_me_child_abort_total();
c2me_sender.abort();
child_tasks.c2me_sender.abort();
Err(ProxyError::Proxy("ME sender join timeout".into()))
}
};
flow_cancel.cancel();
let _ = stop_tx.send(());
let mut writer_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut me_writer).await {
if let Some(stop_tx) = child_tasks.stop_tx.take() {
let _ = stop_tx.send(());
}
let mut writer_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut child_tasks.me_writer).await {
Ok(joined) => {
joined.unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME writer join error: {e}"))))
}
Err(_) => {
stats.increment_me_child_join_timeout_total();
stats.increment_me_child_abort_total();
me_writer.abort();
child_tasks.me_writer.abort();
Err(ProxyError::Proxy("ME writer join timeout".into()))
}
};