1
0
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:
Rob Ede
2022-07-12 02:33:54 +01:00
parent 25368e6b65
commit 6c39b4b138
4 changed files with 176 additions and 143 deletions

View File

@ -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();
}
}