fix: abort all socket tasks when any task exits

This commit is contained in:
trisua 2025-04-28 00:57:21 -04:00
parent 2924759b01
commit 93c4093620

View file

@ -44,7 +44,7 @@ pub async fn handle_socket(socket: WebSocket, state: State, channel_id: usize) {
let (sender, receiver) = mpsc::channel::<String>(); let (sender, receiver) = mpsc::channel::<String>();
// forward messages from mpsc to the sink // forward messages from mpsc to the sink
tokio::spawn(async move { let mut forward_task = tokio::spawn(async move {
while let Ok(message) = receiver.recv() { while let Ok(message) = receiver.recv() {
if message == "Close" { if message == "Close" {
let _ = sink.close().await; let _ = sink.close().await;
@ -60,7 +60,7 @@ pub async fn handle_socket(socket: WebSocket, state: State, channel_id: usize) {
// ping // ping
let ping_sender = sender.clone(); let ping_sender = sender.clone();
tokio::spawn(async move { let mut heartbeat_task = tokio::spawn(async move {
let mut heartbeat = tokio::time::interval(Duration::from_secs(30)); let mut heartbeat = tokio::time::interval(Duration::from_secs(30));
loop { loop {
@ -189,10 +189,28 @@ pub async fn handle_socket(socket: WebSocket, state: State, channel_id: usize) {
// ... // ...
let close_sender = sender.clone(); let close_sender = sender.clone();
tokio::select! { tokio::select! {
_ = (&mut send_task) => recv_task.abort(), _ = (&mut heartbeat_task) => {
let _ = close_sender.send("Close".to_string());
forward_task.abort();
recv_task.abort();
send_task.abort();
}
_ = (&mut forward_task) => {
send_task.abort();
recv_task.abort();
heartbeat_task.abort();
}
_ = (&mut send_task) => {
let _ = close_sender.send("Close".to_string());
forward_task.abort();
recv_task.abort();
heartbeat_task.abort();
},
_ = (&mut recv_task) => { _ = (&mut recv_task) => {
let _ = close_sender.send("Close".to_string()); let _ = close_sender.send("Close".to_string());
send_task.abort() send_task.abort();
forward_task.abort();
heartbeat_task.abort();
}, },
}; };
} }