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>();
|
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();
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue