From 501c254c4710fe6e9322135e5f54f39bbd8a744c Mon Sep 17 00:00:00 2001 From: Taoda Date: Sun, 17 May 2026 17:31:22 +0800 Subject: [PATCH 1/2] feat(server): graceful shutdown --- Cargo.toml | 1 + .../chat-server-graceful-shutdown/Cargo.toml | 12 ++ .../chat-server-graceful-shutdown/src/main.rs | 183 ++++++++++++++++++ src/server.rs | 29 ++- 4 files changed, 221 insertions(+), 4 deletions(-) create mode 100644 examples/chat-server-graceful-shutdown/Cargo.toml create mode 100644 examples/chat-server-graceful-shutdown/src/main.rs diff --git a/Cargo.toml b/Cargo.toml index 20320db..b42e260 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -82,6 +82,7 @@ members = [ "examples/chat-client-wasm", "examples/chat-server", "examples/chat-server-axum", + "examples/chat-server-graceful-shutdown", "examples/echo-server", "examples/echo-server-native-tls", "examples/simple-client", diff --git a/examples/chat-server-graceful-shutdown/Cargo.toml b/examples/chat-server-graceful-shutdown/Cargo.toml new file mode 100644 index 0000000..b7a253f --- /dev/null +++ b/examples/chat-server-graceful-shutdown/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "ezsockets-chat-server-graceful-shutdown" +version = "0.1.0" +edition = "2021" + +[dependencies] +async-trait = "0.1.52" +axum = "0.8.1" +ezsockets = { path = "../../", features = ["axum"] } +tokio = { version = "1.17.0", features = ["full"] } +tracing = "0.1.32" +tracing-subscriber = "0.3.9" diff --git a/examples/chat-server-graceful-shutdown/src/main.rs b/examples/chat-server-graceful-shutdown/src/main.rs new file mode 100644 index 0000000..5924f89 --- /dev/null +++ b/examples/chat-server-graceful-shutdown/src/main.rs @@ -0,0 +1,183 @@ +use async_trait::async_trait; +use axum::extract::Extension; +use axum::extract::Query; +use axum::response::IntoResponse; +use axum::routing::get; +use axum::Router; +use ezsockets::axum::Upgrade; +use ezsockets::axum_tungstenite::rejection::WebSocketUpgradeRejection; +use ezsockets::CloseFrame; +use ezsockets::Error; +use ezsockets::Server; +use std::collections::HashMap; +use std::net::SocketAddr; +use tokio::net::TcpListener; +use tokio::io::AsyncBufReadExt; +use tokio::io::BufReader; + +type SessionID = u16; +type Session = ezsockets::Session; + +#[derive(Debug)] +enum ChatMessage { + Send { from: SessionID, text: String }, +} + +struct ChatServer { + sessions: HashMap, + handle: Server, +} + +#[async_trait] +impl ezsockets::ServerExt for ChatServer { + type Session = ChatSession; + type Call = ChatMessage; + + async fn on_connect( + &mut self, + socket: ezsockets::Socket, + _request: ezsockets::Request, + _address: SocketAddr, + ) -> Result> { + let id = (0..).find(|i| !self.sessions.contains_key(i)).unwrap_or(0); + let session = Session::create( + |_| ChatSession { + id, + server: self.handle.clone(), + }, + id, + socket, + ); + self.sessions.insert(id, session.clone()); + Ok(session) + } + + async fn on_disconnect( + &mut self, + id: ::ID, + _reason: Result, Error>, + ) -> Result<(), Error> { + assert!(self.sessions.remove(&id).is_some()); + Ok(()) + } + + async fn on_call(&mut self, call: Self::Call) -> Result<(), Error> { + match call { + ChatMessage::Send { text, from } => { + let sessions = self.sessions.iter().filter(|(id, _)| from != **id); + let text = format!("from {from}: {text}"); + for (id, handle) in sessions { + tracing::info!("sending {text} to {id}"); + handle.text(text.clone()).unwrap(); + } + } + }; + Ok(()) + } +} + +struct ChatSession { + id: SessionID, + server: Server, +} + +#[async_trait] +impl ezsockets::SessionExt for ChatSession { + type ID = SessionID; + type Call = (); + + fn id(&self) -> &Self::ID { + &self.id + } + async fn on_text(&mut self, text: ezsockets::Utf8Bytes) -> Result<(), Error> { + tracing::info!("received: {text}"); + self.server + .call(ChatMessage::Send { + from: self.id, + text: text.to_string(), + }) + .unwrap(); + Ok(()) + } + + async fn on_binary(&mut self, _bytes: ezsockets::Bytes) -> Result<(), Error> { + unimplemented!() + } + + async fn on_call(&mut self, call: Self::Call) -> Result<(), Error> { + let () = call; + Ok(()) + } +} + +#[tokio::main] +async fn main() { + tracing_subscriber::fmt::init(); + let (server, joiner) = Server::create(|handle| ChatServer { + sessions: HashMap::new(), + handle, + }); + + let app = Router::new() + .route("/websocket", get(websocket_handler)) + .layer(Extension(server.clone())); + + let address = SocketAddr::from(([127, 0, 0, 1], 8080)); + + tokio::spawn(async move { + tracing::info!("listening on {}", address); + let listener = TcpListener::bind(address).await.unwrap(); + axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .await + .unwrap(); + }); + + let stdin = tokio::io::stdin(); + let stdin = BufReader::new(stdin); + let mut lines = stdin.lines(); + loop { + tokio::select! { + _ = tokio::signal::ctrl_c() => { + tracing::info!("ctrl-c is received. \ + An additional ENTER is required to quit the process."); + // The reason is explained in + // in [here](https://docs.rs/tokio/latest/tokio/io/fn.stdin.html). + server.graceful_shutdown().await.unwrap(); + break; + } + line = lines.next_line() => { + let line = line.unwrap().unwrap(); + server + .call(ChatMessage::Send { + text: line, + from: SessionID::MAX, // reserve some ID for the server + }) + .unwrap(); + } + } + } + + joiner.await.unwrap(); +} + +async fn websocket_handler( + Extension(server): Extension>, + Query(query): Query>, + ezsocket: Result, +) -> impl IntoResponse { + let ezsocket = ezsocket.expect("Upgrade should not fail"); + tracing::info!("received upgrade request for {:?}", ezsocket.address()); + let kick_me = query.get("kick_me"); + let kick_me = kick_me.map(|s| s.as_str()); + if matches!(kick_me, Some("Yes")) { + return ( + axum::http::StatusCode::BAD_REQUEST, + "we won't accept you because of kick_me query parameter", + ) + .into_response(); + } + ezsocket.on_upgrade(server) +} diff --git a/src/server.rs b/src/server.rs index 6556f14..eaa25d3 100644 --- a/src/server.rs +++ b/src/server.rs @@ -115,6 +115,7 @@ struct ServerActor { connection_receiver: mpsc::UnboundedReceiver, disconnection_receiver: mpsc::UnboundedReceiver>, server_call_receiver: mpsc::UnboundedReceiver, + shutdown_receiver: mpsc::Receiver<()>, server: Server, extension: E, } @@ -127,8 +128,12 @@ where async fn run(mut self) { tracing::info!("starting websocket server"); loop { - if let Err(err) = async { + let one_event = async { tokio::select! { + _ = self.shutdown_receiver.recv() => { + tracing::info!("The websocket server is going to shutdown."); + return Ok(None); + } Some(NewConnection{socket, address, request}) = self.connection_receiver.recv() => { let socket_sink = socket.sink.clone(); match self.extension.on_connect(socket, request, address).await { @@ -168,10 +173,17 @@ where } else => return Err("server actor branches broken".into()), } - Ok::<_, Error>(()) + Ok::<_, Error>(Some(())) } - .await { - tracing::warn!("error when processing: {err:?}"); + .await; + match one_event { + Err(err) => { + tracing::warn!("error when processing: {err:?}"); + } + Ok(None) => { + break; + } + _ => {} } } } @@ -214,6 +226,7 @@ pub struct Server { connection_sender: mpsc::UnboundedSender, disconnection_sender: mpsc::UnboundedSender>, server_call_sender: mpsc::UnboundedSender, + shutdown_sender: mpsc::Sender<()>, } impl From> for mpsc::UnboundedSender { @@ -227,16 +240,19 @@ impl Server { let (connection_sender, connection_receiver) = mpsc::unbounded_channel(); let (disconnection_sender, disconnection_receiver) = mpsc::unbounded_channel(); let (server_call_sender, server_call_receiver) = mpsc::unbounded_channel(); + let (shutdown_sender, shutdown_receiver) = mpsc::channel(1); let handle = Self { connection_sender, server_call_sender, disconnection_sender, + shutdown_sender, }; let extension = create(handle.clone()); let actor = ServerActor { connection_receiver, disconnection_receiver, server_call_receiver, + shutdown_receiver, extension, server: handle.clone(), }; @@ -295,6 +311,10 @@ impl Server { }; Some(result) } + + pub async fn graceful_shutdown(&self) -> Result<(), mpsc::error::SendError<()>> { + self.shutdown_sender.send(()).await + } } impl std::clone::Clone for Server { @@ -303,6 +323,7 @@ impl std::clone::Clone for Server { connection_sender: self.connection_sender.clone(), disconnection_sender: self.disconnection_sender.clone(), server_call_sender: self.server_call_sender.clone(), + shutdown_sender: self.shutdown_sender.clone(), } } } From 3f371a48c55c9572b4ec8c04e794670c33d52289 Mon Sep 17 00:00:00 2001 From: Taoda Date: Mon, 18 May 2026 14:37:10 +0800 Subject: [PATCH 2/2] style: run `cargo fmt` --- examples/chat-server-graceful-shutdown/src/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/chat-server-graceful-shutdown/src/main.rs b/examples/chat-server-graceful-shutdown/src/main.rs index 5924f89..4d06e17 100644 --- a/examples/chat-server-graceful-shutdown/src/main.rs +++ b/examples/chat-server-graceful-shutdown/src/main.rs @@ -11,9 +11,9 @@ use ezsockets::Error; use ezsockets::Server; use std::collections::HashMap; use std::net::SocketAddr; -use tokio::net::TcpListener; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; +use tokio::net::TcpListener; type SessionID = u16; type Session = ezsockets::Session;