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:
Riz Ashraf committed 2026-09-20 23:08:40 +01:00
1 parent acf5ea8ebe
commit b9e44e1fa3
1 file changed
+37 -14
+37 -14
View File
@@ -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)]