I am implementing a websocket chat application where I want to gracefully shutdown all clients when the server stop because of the ctrl+c signal.
I am listening to incoming events using mio poll and tokens. Any new socket connection is registered with mio poll and any event received on the socket is successfully captured on the polling.
My initial idea was to use tokio::select with listen events and shutdown as 2 branches. But, I guess this design needs some modification to enable graceful shutdown.
// get_events function
fn get_events(poll, conn) -> Option<mio::Event>{
let shutdown = tokio::signal::ctrl_c();
let res = tokio::select!{
res = get_poll_events(poll) => { // returns events asynchronously, its an async method
// ... handler when events are received
res
},
_ = shutdown => {
// ... handler when ctrl+c signal is received
// send close connection message to TcpStream
return None;
}
};
Some(res)
}
// main function
let poll = Poll::new();
let shared_poll = Arc::new(Mutex::new(poll));
let conn: Arc<Mutex<HashMap<Token, WebSocketClient>>> = Arc::new(Mutex::new(HashMap::new()));
let (shutdown_notifier, shutdown_receiver) = mpsc::channel(10);
loop {
let res = WebSocketServer::get_events(shared_poll.clone(), conn.clone()).await;
if let None = res {
drop(shutdown_notifier); // drop the original mpsc::Sender, cloned in all clients.
let _ = shutdown_receiver.recv().await; // mpsc::Receiver only returns error when all clients are dropped, thus dropping the senders inside them.
break; // stop the program
}
// use event to register a new client with mpsc::sender
// if readable; accept incoming message
// broadcast message to other subscribers
// ... other tasks
...
}
On Ctrl+c signal received the shutdown handler gets activated. Post that, close connection message is sent on the TcpStream to all active clients. At this point, everything is as expected. All the clients receive close messages. But the real issue starts here, since the shutdown handler executes, the get_poll_events branch of tokio::select gets dropped. Meaning, no more polling for the incoming events on the registered sources (TcpStreams).
Ideally, the server should only close when it has received back the close message response from all the clients. The clients communicate back on the TcpStream, but since there is no active mechanism to listen to these events, I am unable to capture them and hence not able to drop the clients. However, instead of sending close message to clients if I dropped all the clients manually, then there was no need to listen to close message response and things worked albeit an incorrect implementation.
Tokio::select isn't an ideal choice as far as I can guess, but I am unable to come up with a solution on how to implement this case, where the server is still active to listen to all the close message responses from clients. It can then close when all the clients have gone out of scope in the received close messages.
What would be a way to achieve this functionality? TIA.
