feat(server): graceful shutdown

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",

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"

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::io::AsyncBufReadExt;
+use tokio::io::BufReader;
+use tokio::net::TcpListener;
+
+type SessionID = u16;
+type Session = ezsockets::Session<SessionID, ()>;
+
+#[derive(Debug)]
+enum ChatMessage {
+    Send { from: SessionID, text: String },
+}
+
+struct ChatServer {
+    sessions: HashMap<SessionID, Session>,
+    handle: Server<Self>,
+}
+
+#[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<Session, Option<CloseFrame>> {
+        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: <Self::Session as ezsockets::SessionExt>::ID,
+        _reason: Result<Option<CloseFrame>, 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<ChatServer>,
+}
+
+#[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::<SocketAddr>(),
+        )
+        .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<Server<ChatServer>>,
+    Query(query): Query<HashMap<String, String>>,
+    ezsocket: Result<Upgrade, WebSocketUpgradeRejection>,
+) -> 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)
+}

src/server.rs

@@ -115,6 +115,7 @@ struct ServerActor<E: ServerExt> {
     connection_receiver: mpsc::UnboundedReceiver<NewConnection>,
     disconnection_receiver: mpsc::UnboundedReceiver<Disconnected<E>>,
     server_call_receiver: mpsc::UnboundedReceiver<E::Call>,
+    shutdown_receiver: mpsc::Receiver<()>,
     server: Server<E>,
     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<E: ServerExt> {
     connection_sender: mpsc::UnboundedSender<NewConnection>,
     disconnection_sender: mpsc::UnboundedSender<Disconnected<E>>,
     server_call_sender: mpsc::UnboundedSender<E::Call>,
+    shutdown_sender: mpsc::Sender<()>,
 }
 
 impl<E: ServerExt> From<Server<E>> for mpsc::UnboundedSender<E::Call> {
@@ -227,16 +240,19 @@ impl<E: ServerExt + 'static> Server<E> {
         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<E: ServerExt> Server<E> {
         };
         Some(result)
     }
+
+    pub async fn graceful_shutdown(&self) -> Result<(), mpsc::error::SendError<()>> {
+        self.shutdown_sender.send(()).await
+    }
 }
 
 impl<E: ServerExt> std::clone::Clone for Server<E> {
@@ -303,6 +323,7 @@ impl<E: ServerExt> std::clone::Clone for Server<E> {
             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(),
         }
     }
 }