@@ -449,9 +449,107 @@ Here's some important information to get you started:
449449This project involves a lot of concepts that are probably new to you, and you
450450probably will get stuck somewhere along the way. That's okay! Try to get as far
451451as you can, and consult resources when you get stuck. Once you've given it a
452- good try,
453- [ click here to open the completed example code] ( ../example/src/main.rs ) .
452+ good try, you can compare your code to the example below.
454453
455454Good luck! If you get stuck along the way, look online for other resources or
456455examples. The axum repository has some great examples showcasing different
457456features, not just WebSockets.
457+
458+ ## Example Code
459+
460+ Here's an example implementation of this project. You can use it as a reference
461+ for your own code, but please give the project a try before consulting the
462+ example! The best way to learn is by spending some time struggling with the
463+ problem.
464+
465+ ``` rust
466+ use axum :: extract :: {
467+ ws :: {Message , WebSocket },
468+ State , WebSocketUpgrade ,
469+ };
470+ use axum :: response :: IntoResponse ;
471+ use axum :: routing :: any;
472+ use axum :: Router ;
473+ use std :: sync :: {Arc , Mutex };
474+ use tokio :: net :: TcpListener ;
475+ use tokio :: sync :: broadcast;
476+ use tower_http :: services :: ServeDir ;
477+
478+ #[derive(Clone )]
479+ struct ServerState {
480+ /// The list of every message that this server has received.
481+ message_history : Arc <Mutex <Vec <String >>>,
482+ /// A channel to share messages between all connected clients.
483+ message_channel : broadcast :: Sender <String >,
484+ }
485+
486+ #[tokio:: main]
487+ async fn main () {
488+ let message_history = Arc :: new (Mutex :: new (Vec :: new ()));
489+ let (tx , _rx ) = broadcast :: channel (32 );
490+
491+ let state = ServerState {
492+ // This is syntax sugar for message_history: message_history.
493+ message_history ,
494+ message_channel : tx ,
495+ };
496+
497+ let app = Router :: new ()
498+ . route (" /socket" , any (ws_handler ))
499+ . fallback_service (ServeDir :: new (" public" ))
500+ . with_state (state );
501+
502+ let listener = TcpListener :: bind (" localhost:3000" ). await . unwrap ();
503+ axum :: serve (listener , app ). await . unwrap ();
504+ }
505+
506+ async fn ws_handler (ws : WebSocketUpgrade , State (state ): State <ServerState >) -> impl IntoResponse {
507+ ws . on_upgrade (move | socket | handle_socket (socket , state ))
508+ }
509+
510+ async fn handle_socket (mut socket : WebSocket , state : ServerState ) {
511+ // Send over every message to the client.
512+ let messages = state . message_history. lock (). unwrap (). clone ();
513+
514+ for msg in messages {
515+ if let Err (err ) = socket . send (Message :: text (msg )). await {
516+ eprintln! (" Error while sending initial data: {err:?}" );
517+ return ;
518+ }
519+ }
520+
521+ let mut recv = state . message_channel. subscribe ();
522+
523+ // Receive loop.
524+ loop {
525+ tokio :: select! {
526+ val = recv . recv () => {
527+ let Ok (val ) = val else {
528+ break ;
529+ };
530+
531+ if let Err (err ) = socket . send (Message :: text (val )). await {
532+ eprintln! (" Error while sending message to socket: {err:?}" );
533+ return ;
534+ }
535+ }
536+ val = socket . recv () => {
537+ let Some (Ok (val )) = val else {
538+ break ;
539+ };
540+
541+ // Only handle text messages.
542+ if let Message :: Text (text ) = val {
543+ // Send message to all connected clients.
544+ if state . message_channel. send (text . to_string ()). is_err () {
545+ break ;
546+ }
547+
548+ // Save message to history.
549+ state . message_history. lock (). unwrap (). push (text . to_string ());
550+ }
551+ }
552+ }
553+ }
554+ }
555+ ```
0 commit comments