mirror of
https://github.com/actix/examples
synced 2025-06-27 09:29:02 +02:00
reduce boilerplate on response channels
This commit is contained in:
@ -10,11 +10,44 @@ use std::{
|
||||
};
|
||||
|
||||
use rand::{thread_rng, Rng as _};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
use crate::{Command, ConnId, Msg, RoomId};
|
||||
use crate::{ConnId, Msg, RoomId};
|
||||
|
||||
/// A command received by the [`ChatServer`].
|
||||
#[derive(Debug)]
|
||||
enum Command {
|
||||
Connect {
|
||||
conn_tx: mpsc::UnboundedSender<Msg>,
|
||||
res_tx: oneshot::Sender<ConnId>,
|
||||
},
|
||||
|
||||
Disconnect {
|
||||
conn: ConnId,
|
||||
},
|
||||
|
||||
List {
|
||||
res_tx: oneshot::Sender<Vec<RoomId>>,
|
||||
},
|
||||
|
||||
Join {
|
||||
conn: ConnId,
|
||||
room: RoomId,
|
||||
res_tx: oneshot::Sender<()>,
|
||||
},
|
||||
|
||||
Message {
|
||||
msg: Msg,
|
||||
conn: ConnId,
|
||||
res_tx: oneshot::Sender<()>,
|
||||
},
|
||||
}
|
||||
|
||||
/// A multi-room chat server.
|
||||
///
|
||||
/// Contains the logic of how connections chat with each other plus room management.
|
||||
///
|
||||
/// Call and spawn [`run`](Self::run) to start processing commands.
|
||||
#[derive(Debug)]
|
||||
pub struct ChatServer {
|
||||
/// Map of connection IDs to their message receivers.
|
||||
@ -27,41 +60,39 @@ pub struct ChatServer {
|
||||
visitor_count: Arc<AtomicUsize>,
|
||||
|
||||
/// Command receiver.
|
||||
rx: mpsc::UnboundedReceiver<Command>,
|
||||
cmd_rx: mpsc::UnboundedReceiver<Command>,
|
||||
}
|
||||
|
||||
impl ChatServer {
|
||||
pub fn new() -> (Self, mpsc::UnboundedSender<Command>) {
|
||||
pub fn new() -> (Self, ChatServerHandle) {
|
||||
// create empty server
|
||||
let mut rooms = HashMap::with_capacity(4);
|
||||
|
||||
// create default room
|
||||
rooms.insert("main".to_owned(), HashSet::new());
|
||||
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
|
||||
|
||||
(
|
||||
Self {
|
||||
sessions: HashMap::new(),
|
||||
rooms,
|
||||
visitor_count: Arc::new(AtomicUsize::new(0)),
|
||||
rx,
|
||||
cmd_rx,
|
||||
},
|
||||
tx,
|
||||
ChatServerHandle { cmd_tx },
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatServer {
|
||||
/// Send message to all users in the room.
|
||||
/// Send message to users in a room.
|
||||
///
|
||||
/// `skip_id` is used to prevent messages send by a connection also being received by it.
|
||||
async fn send_message(&self, room: &str, msg: impl Into<String>, skip_id: ConnId) {
|
||||
/// `skip` is used to prevent messages triggered by a connection also being received by it.
|
||||
async fn send_system_message(&self, room: &str, skip: ConnId, msg: impl Into<String>) {
|
||||
if let Some(sessions) = self.rooms.get(room) {
|
||||
let msg = msg.into();
|
||||
|
||||
for conn_id in sessions {
|
||||
if *conn_id != skip_id {
|
||||
if *conn_id != skip {
|
||||
if let Some(tx) = self.sessions.get(conn_id) {
|
||||
// errors if client disconnected abruptly and hasn't been timed-out yet
|
||||
let _ = tx.send(msg.clone());
|
||||
@ -71,14 +102,26 @@ impl ChatServer {
|
||||
}
|
||||
}
|
||||
|
||||
/// Handler for Connect message.
|
||||
/// Send message to all other users in current room.
|
||||
///
|
||||
/// Register new session and assign unique id to this session
|
||||
/// `conn` is used to find current room and prevent messages sent by a connection also being
|
||||
/// received by it.
|
||||
async fn send_message(&self, conn: ConnId, msg: impl Into<String>) {
|
||||
if let Some(room) = self
|
||||
.rooms
|
||||
.iter()
|
||||
.find_map(|(room, participants)| participants.contains(&conn).then_some(room))
|
||||
{
|
||||
self.send_system_message(&room, conn, msg).await;
|
||||
};
|
||||
}
|
||||
|
||||
/// Register new session and assign unique ID to this session
|
||||
async fn connect(&mut self, tx: mpsc::UnboundedSender<Msg>) -> ConnId {
|
||||
log::info!("Someone joined");
|
||||
|
||||
// notify all users in same room
|
||||
self.send_message("main", "Someone joined", 0).await;
|
||||
self.send_system_message("main", 0, "Someone joined").await;
|
||||
|
||||
// register session with random connection ID
|
||||
let id = thread_rng().gen::<usize>();
|
||||
@ -91,14 +134,14 @@ impl ChatServer {
|
||||
.insert(id);
|
||||
|
||||
let count = self.visitor_count.fetch_add(1, Ordering::SeqCst);
|
||||
self.send_message("main", format!("Total visitors {count}"), 0)
|
||||
self.send_system_message("main", 0, format!("Total visitors {count}"))
|
||||
.await;
|
||||
|
||||
// send id back
|
||||
id
|
||||
}
|
||||
|
||||
/// Handler for Disconnect message.
|
||||
/// Unregister connection from room map and broadcast disconnection message.
|
||||
async fn disconnect(&mut self, conn_id: ConnId) {
|
||||
println!("Someone disconnected");
|
||||
|
||||
@ -116,19 +159,14 @@ impl ChatServer {
|
||||
|
||||
// send message to other users
|
||||
for room in rooms {
|
||||
self.send_message(&room, "Someone disconnected", 0).await;
|
||||
self.send_system_message(&room, 0, "Someone disconnected")
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Handler for `ListRooms` message.
|
||||
/// Returns list of created room names.
|
||||
fn list_rooms(&mut self) -> Vec<String> {
|
||||
let mut rooms = Vec::new();
|
||||
|
||||
for key in self.rooms.keys() {
|
||||
rooms.push(key.to_owned())
|
||||
}
|
||||
|
||||
rooms
|
||||
self.rooms.keys().cloned().collect()
|
||||
}
|
||||
|
||||
/// Join room, send disconnect message to old room send join message to new room.
|
||||
@ -143,7 +181,8 @@ impl ChatServer {
|
||||
}
|
||||
// send message to other users
|
||||
for room in rooms {
|
||||
self.send_message(&room, "Someone disconnected", 0).await;
|
||||
self.send_system_message(&room, 0, "Someone disconnected")
|
||||
.await;
|
||||
}
|
||||
|
||||
self.rooms
|
||||
@ -151,12 +190,13 @@ impl ChatServer {
|
||||
.or_insert_with(HashSet::new)
|
||||
.insert(conn_id);
|
||||
|
||||
self.send_message(&room, "Someone connected", conn_id).await;
|
||||
self.send_system_message(&room, conn_id, "Someone connected")
|
||||
.await;
|
||||
}
|
||||
|
||||
pub async fn run(mut self) -> io::Result<()> {
|
||||
loop {
|
||||
let cmd = match self.rx.recv().await {
|
||||
let cmd = match self.cmd_rx.recv().await {
|
||||
Some(cmd) => cmd,
|
||||
None => break,
|
||||
};
|
||||
@ -180,13 +220,8 @@ impl ChatServer {
|
||||
let _ = res_tx.send(());
|
||||
}
|
||||
|
||||
Command::Message {
|
||||
room,
|
||||
msg,
|
||||
skip,
|
||||
res_tx,
|
||||
} => {
|
||||
self.send_message(&room, msg, skip).await;
|
||||
Command::Message { conn, msg, res_tx } => {
|
||||
self.send_message(conn, msg).await;
|
||||
let _ = res_tx.send(());
|
||||
}
|
||||
}
|
||||
@ -195,3 +230,77 @@ impl ChatServer {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle and command sender for chat server.
|
||||
///
|
||||
/// Reduces boilerplate of setting up response channels in WebSocket handlers.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChatServerHandle {
|
||||
cmd_tx: mpsc::UnboundedSender<Command>,
|
||||
}
|
||||
|
||||
impl ChatServerHandle {
|
||||
/// Register client message sender and obtain connection ID.
|
||||
pub async fn connect(&self, conn_tx: mpsc::UnboundedSender<String>) -> ConnId {
|
||||
let (res_tx, res_rx) = oneshot::channel();
|
||||
|
||||
// unwrap: chat server should not have been dropped
|
||||
self.cmd_tx
|
||||
.send(Command::Connect { conn_tx, res_tx })
|
||||
.unwrap();
|
||||
|
||||
// unwrap: chat server does not drop out response channel
|
||||
res_rx.await.unwrap()
|
||||
}
|
||||
|
||||
/// List all created rooms.
|
||||
pub async fn list_rooms(&self) -> Vec<String> {
|
||||
let (res_tx, res_rx) = oneshot::channel();
|
||||
|
||||
// unwrap: chat server should not have been dropped
|
||||
self.cmd_tx.send(Command::List { res_tx }).unwrap();
|
||||
|
||||
// unwrap: chat server does not drop our response channel
|
||||
res_rx.await.unwrap()
|
||||
}
|
||||
|
||||
/// Join `room`, creating it if it does not exist.
|
||||
pub async fn join_room(&self, conn: ConnId, room: impl Into<String>) {
|
||||
let (res_tx, res_rx) = oneshot::channel();
|
||||
|
||||
// unwrap: chat server should not have been dropped
|
||||
self.cmd_tx
|
||||
.send(Command::Join {
|
||||
conn,
|
||||
room: room.into(),
|
||||
res_tx,
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// unwrap: chat server does not drop our response channel
|
||||
res_rx.await.unwrap();
|
||||
}
|
||||
|
||||
/// Broadcast message to current room.
|
||||
pub async fn send_message(&self, conn: ConnId, msg: impl Into<String>) {
|
||||
let (res_tx, res_rx) = oneshot::channel();
|
||||
|
||||
// unwrap: chat server should not have been dropped
|
||||
self.cmd_tx
|
||||
.send(Command::Message {
|
||||
msg: msg.into(),
|
||||
conn,
|
||||
res_tx,
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
// unwrap: chat server does not drop our response channel
|
||||
res_rx.await.unwrap();
|
||||
}
|
||||
|
||||
/// Unregister message sender and broadcast disconnection message to current room.
|
||||
pub fn disconnect(&self, conn: ConnId) {
|
||||
// unwrap: chat server should not have been dropped
|
||||
self.cmd_tx.send(Command::Disconnect { conn }).unwrap();
|
||||
}
|
||||
}
|
||||
|
Reference in New Issue
Block a user