fix(server): Fix memory leak in handle_socket on abrupt disconnects
- Used a SessionCleanup Drop guard (RAII pattern) to guarantee that the websocket session is evicted from the global clients map and background read/write tasks are forcibly aborted even when Axum terminates the handler mid-execution due to a TCP connection loss.
This commit is contained in:
1 parent
acf5ea8ebe
commit
b9e44e1fa3
1 file changed
+37
-14
+37
-14
@@ -557,22 +557,45 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
|
||||
);
|
||||
});
|
||||
|
||||
tokio::select! {
|
||||
_ = (&mut send_task) => {
|
||||
tracing::info!("Websocket send task finished for session {}", session_id);
|
||||
recv_task.abort();
|
||||
},
|
||||
_ = (&mut recv_task) => {
|
||||
tracing::info!("Websocket recv task finished for session {}", session_id);
|
||||
send_task.abort();
|
||||
},
|
||||
struct SessionCleanup {
|
||||
session_id: String,
|
||||
state: Arc<AppState>,
|
||||
send_task: Option<tokio::task::JoinHandle<()>>,
|
||||
recv_task: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl Drop for SessionCleanup {
|
||||
fn drop(&mut self) {
|
||||
self.state.clients.write().unwrap().remove(&self.session_id);
|
||||
if let Some(task) = self.send_task.take() {
|
||||
task.abort();
|
||||
}
|
||||
if let Some(task) = self.recv_task.take() {
|
||||
task.abort();
|
||||
}
|
||||
tracing::info!(
|
||||
"Websocket session {} closed and cleaned up",
|
||||
self.session_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let mut cleanup = SessionCleanup {
|
||||
session_id: session_id.clone(),
|
||||
state: Arc::clone(&state),
|
||||
send_task: Some(send_task),
|
||||
recv_task: Some(recv_task),
|
||||
};
|
||||
|
||||
state.clients.write().unwrap().remove(&session_id);
|
||||
tracing::info!(
|
||||
"Websocket session {} closed and removed from state",
|
||||
session_id
|
||||
);
|
||||
tokio::select! {
|
||||
_ = cleanup.send_task.as_mut().unwrap() => {
|
||||
tracing::info!("Websocket send task finished for session {}", session_id);
|
||||
},
|
||||
_ = cleanup.recv_task.as_mut().unwrap() => {
|
||||
tracing::info!("Websocket recv task finished for session {}", session_id);
|
||||
},
|
||||
};
|
||||
// Drop guard automatically handles removal and aborts the other task.
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, serde::Serialize, Debug)]
|
||||
|
||||
Reference in new issue
Block a user