opentally-server/src/web/websocket.rs

87 lines
2.3 KiB
Rust

use crate::TallyState;
use actix::prelude::*;
use futures::sync::mpsc::Sender;
use std::collections::HashMap;
#[derive(Message)]
pub struct Message(pub String);
#[derive(Message)]
pub struct Connect {
pub addr: Recipient<Message>,
}
#[derive(Message)]
pub struct Disconnect {
pub addr: Recipient<Message>,
}
#[derive(Message)]
pub struct NewState(Result<TallyState, ()>);
// Basically just convenience so you don't have to construct a NewState with an Ok all the time
impl From<TallyState> for NewState {
fn from(state: TallyState) -> Self {
NewState(Ok(state))
}
}
pub struct WSServer {
state_tx_tx: std::sync::mpsc::Sender<Sender<NewState>>,
clients: HashMap<Recipient<Message>, ()>,
}
// This lets the WSServer act as a futures channel reciever, with incoming NewStates being handled like actix messages
impl StreamHandler<NewState, ()> for WSServer {
fn handle(&mut self, item: NewState, _: &mut Context<WSServer>) {
if let Ok(state) = item.0 {
self.broadcast(serde_json::to_string(&state).unwrap())
}
}
}
impl Actor for WSServer {
type Context = Context<Self>;
fn started(&mut self, ctx: &mut Self::Context) {
log::info!("Websocket started");
let (tx, rx) = futures::sync::mpsc::channel::<NewState>(10);
Self::add_stream(rx, ctx);
self.state_tx_tx.send(tx).unwrap();
}
}
impl WSServer {
pub fn new(state_tx_tx: std::sync::mpsc::Sender<Sender<NewState>>) -> WSServer {
WSServer {
state_tx_tx,
clients: HashMap::new(),
}
}
fn broadcast(&mut self, msg: String) {
for c in self.clients.iter_mut() {
c.0.do_send(Message(msg.to_owned()))
.unwrap_or_else(|e| log::error!("Error sending message to client: {}", e));
}
}
}
impl Handler<Connect> for WSServer {
type Result = ();
fn handle(&mut self, msg: Connect, _: &mut Context<Self>) -> Self::Result {
log::info!("Client connected!");
self.clients.insert(msg.addr, ());
}
}
impl Handler<Disconnect> for WSServer {
type Result = ();
fn handle(&mut self, msg: Disconnect, _: &mut Context<Self>) -> Self::Result {
log::info!("Disconnected");
self.clients.remove(&msg.addr);
}
}