fix: abort all socket tasks when any task exits
This commit is contained in:
parent
2924759b01
commit
93c4093620
1 changed files with 22 additions and 4 deletions
|
@ -44,7 +44,7 @@ pub async fn handle_socket(socket: WebSocket, state: State, channel_id: usize) {
|
|||
let (sender, receiver) = mpsc::channel::<String>();
|
||||
|
||||
// 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() {
|
||||
if message == "Close" {
|
||||
let _ = sink.close().await;
|
||||
|
@ -60,7 +60,7 @@ pub async fn handle_socket(socket: WebSocket, state: State, channel_id: usize) {
|
|||
|
||||
// ping
|
||||
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));
|
||||
|
||||
loop {
|
||||
|
@ -189,10 +189,28 @@ pub async fn handle_socket(socket: WebSocket, state: State, channel_id: usize) {
|
|||
// ...
|
||||
let close_sender = sender.clone();
|
||||
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) => {
|
||||
let _ = close_sender.send("Close".to_string());
|
||||
send_task.abort()
|
||||
send_task.abort();
|
||||
forward_task.abort();
|
||||
heartbeat_task.abort();
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue