Refactor websocket
This commit is contained in:
parent
ebac452a86
commit
f3d367e479
GPG key ID:
13AD60C237A28DFE
22 changed files with 924 additions and 773 deletions
333
src/server/server.rs
Normal file
333
src/server/server.rs
Normal file
|
|
@ -0,0 +1,333 @@
|
|||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use futures::executor::block_on;
|
||||
use rand::random;
|
||||
use sqlx::{Pool, Sqlite};
|
||||
use tokio::signal;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::mpsc::UnboundedSender;
|
||||
use emgauwa_common::db::{DbController, DbJunctionRelaySchedule, DbRelay, DbSchedule};
|
||||
use emgauwa_common::errors::{DatabaseError, EmgauwaError};
|
||||
use emgauwa_common::models::{convert_db_list, Controller, FromDbModel, Relay};
|
||||
use emgauwa_common::types::{ControllersWsAction, EmgauwaUid, RelayStates};
|
||||
use emgauwa_common::utils;
|
||||
use crate::server::{WsConnId, WsControllerAction, WsRelayAction, WsServerAction};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct WsServer {
|
||||
/// Map of connection IDs to their message receivers.
|
||||
controller_sessions: HashMap<WsConnId, (mpsc::UnboundedSender<WsControllerAction>, Option<EmgauwaUid>)>,
|
||||
controllers: HashMap<EmgauwaUid, (Controller, WsConnId)>,
|
||||
|
||||
relay_sessions: HashMap<WsConnId, mpsc::UnboundedSender<WsRelayAction>>,
|
||||
|
||||
/// Command receiver.
|
||||
cmd_rx: mpsc::UnboundedReceiver<WsServerAction>,
|
||||
|
||||
pub pool: Pool<Sqlite>,
|
||||
}
|
||||
|
||||
impl WsServer {
|
||||
pub fn new(pool: Pool<Sqlite>) -> (Self, mpsc::UnboundedSender<WsServerAction>) {
|
||||
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
|
||||
|
||||
let new_self = Self {
|
||||
controller_sessions: HashMap::new(),
|
||||
controllers: HashMap::new(),
|
||||
relay_sessions: HashMap::new(),
|
||||
cmd_rx,
|
||||
pool,
|
||||
};
|
||||
(new_self, cmd_tx)
|
||||
}
|
||||
|
||||
/// Register new session and assign unique ID to this session
|
||||
async fn connect_controller(&mut self, tx: mpsc::UnboundedSender<WsControllerAction>) -> WsConnId {
|
||||
log::info!("Controller connected");
|
||||
|
||||
// register a session with random connection ID
|
||||
let id = random::<WsConnId>();
|
||||
self.controller_sessions.insert(id, (tx, None));
|
||||
|
||||
// send id back
|
||||
id
|
||||
}
|
||||
|
||||
async fn connect_relay(&mut self, tx: mpsc::UnboundedSender<WsRelayAction>) -> Result<WsConnId, EmgauwaError> {
|
||||
log::debug!("Relay client connected");
|
||||
|
||||
let relays_json = serde_json::to_string(&self.get_relays()?)?;
|
||||
tx.send(WsRelayAction::Forward(relays_json.clone()))?;
|
||||
|
||||
// register a session with random connection ID
|
||||
let id = random::<WsConnId>();
|
||||
self.relay_sessions.insert(id, tx);
|
||||
|
||||
// send id back
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
async fn disconnect_controller(&mut self, sender: UnboundedSender<WsControllerAction>, emgauwa_uid: EmgauwaUid, update_db: bool) -> Result<(), EmgauwaError> {
|
||||
if let Some((controller, _)) = self.controllers.remove(&emgauwa_uid) {
|
||||
if update_db {
|
||||
let mut pool_conn = block_on(self.pool.acquire())?;
|
||||
log::debug!("Marking controller {} as inactive", controller.c.uid);
|
||||
if let Err(err) = block_on(controller.c.update_active(&mut pool_conn, false)) {
|
||||
log::error!(
|
||||
"Failed to mark controller {} as inactive: {:?}",
|
||||
controller.c.uid,
|
||||
err
|
||||
);
|
||||
return Err(EmgauwaError::from(err));
|
||||
}
|
||||
}
|
||||
log::debug!("Telling websocket server to disconnect a controller");
|
||||
sender.send(WsControllerAction::Disconnect)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Unregister the connection from the controller map and broadcast a disconnection message.
|
||||
async fn disconnect_controller_conn(&mut self, conn_id: WsConnId, update_db: bool) -> Result<(), EmgauwaError> {
|
||||
log::debug!("Trying to find controller connection for disconnect: {}", conn_id);
|
||||
if let Some((sender, Some(emgauwa_uid))) = self.controller_sessions.remove(&conn_id) {
|
||||
self.disconnect_controller(sender, emgauwa_uid, update_db).await?;
|
||||
}
|
||||
self.notify_relay_clients()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn disconnect_relay_conn(&mut self, conn_id: WsConnId) -> Result<(), EmgauwaError> {
|
||||
if let Some(relay_tx) = self.relay_sessions.remove(&conn_id) {
|
||||
relay_tx.send(WsRelayAction::Disconnect)?;
|
||||
}
|
||||
self.notify_relay_clients()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_relays(&mut self) -> Result<Vec<Relay>, EmgauwaError> {
|
||||
let mut pool_conn = block_on(self.pool.acquire())?;
|
||||
let db_controllers = block_on(DbController::get_all(&mut pool_conn))?;
|
||||
let mut controllers: Vec<Controller> = convert_db_list(&mut pool_conn, db_controllers)?;
|
||||
|
||||
controllers.iter_mut().for_each(|c| {
|
||||
if let Some((cc, _)) = self.controllers.get(&c.c.uid) {
|
||||
c.apply_relay_states(&cc.get_relay_states())
|
||||
}
|
||||
});
|
||||
|
||||
let mut relays: Vec<Relay> = Vec::new();
|
||||
controllers.iter().for_each(|c| {
|
||||
relays.extend(c.relays.clone());
|
||||
});
|
||||
|
||||
Ok(relays)
|
||||
}
|
||||
|
||||
fn notify_relay_clients(&mut self) -> Result<(), EmgauwaError> {
|
||||
let relays_json = serde_json::to_string(&self.get_relays()?)?;
|
||||
|
||||
self.relay_sessions.retain(|_, tx| {
|
||||
tx.send(WsRelayAction::Forward(relays_json.clone())).is_ok()
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// You may want to log errors rather than return them in this function.
|
||||
pub async fn run(self) -> io::Result<()> {
|
||||
tokio::select! {
|
||||
output = self.real_run() => output,
|
||||
_ = signal::ctrl_c() => Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn real_run(mut self) -> io::Result<()> {
|
||||
while let Some(cmd) = self.cmd_rx.recv().await {
|
||||
match cmd {
|
||||
WsServerAction::ConnectController(conn_tx, res_tx) => {
|
||||
let conn_id = self.connect_controller(conn_tx).await;
|
||||
let _ = res_tx.send(conn_id);
|
||||
}
|
||||
|
||||
WsServerAction::DisconnectControllerConn(conn, update_db) => {
|
||||
if let Err(err) = self.disconnect_controller_conn(conn, update_db).await {
|
||||
log::error!("Error disconnecting controller: {:?}", err);
|
||||
}
|
||||
}
|
||||
WsServerAction::DisconnectController(emgauwa_uid, update_db) => {
|
||||
if let Some((_, conn_id)) = self.controllers.get(&emgauwa_uid) {
|
||||
if let Err(err ) = self.disconnect_controller_conn(*conn_id, update_db).await {
|
||||
log::error!("Error disconnecting controller: {:?}", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
WsServerAction::ControllerMessage(conn, msg) => {
|
||||
match serde_json::from_str(&msg) {
|
||||
Ok(action) => {
|
||||
self.handle_action(conn, action);
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!("Error deserializing action: {:?}", e);
|
||||
self.text_to_controller_conn(
|
||||
conn,
|
||||
serde_json::to_string(&EmgauwaError::Serialization(e))
|
||||
.unwrap_or(String::from("Error in deserializing action")),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
WsServerAction::ForwardToController(uid, action) => {
|
||||
log::debug!("Forwarding action: {:?}", action);
|
||||
if let Some((_, address)) = self.controllers.get(&uid) {
|
||||
self.text_to_controller_conn(*address, serde_json::to_string(&action)?);
|
||||
}
|
||||
}
|
||||
WsServerAction::GetRelays(tx) => {
|
||||
tx.send(self.get_relays()).map_err(|_| io::Error::other("Error sending relays"))?;
|
||||
}
|
||||
WsServerAction::ConnectRelay(conn_tx, res_tx) => {
|
||||
if let Ok(conn_id) = self.connect_relay(conn_tx).await {
|
||||
if let Err(e) = res_tx.send(conn_id) {
|
||||
log::error!("Error sending relay conn id: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
WsServerAction::DisconnectRelayConn(conn) => {
|
||||
if let Err(err) = self.disconnect_relay_conn(conn).await {
|
||||
log::error!("Error disconnecting relay conn: {:?}", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn handle_action(
|
||||
&mut self,
|
||||
conn: WsConnId,
|
||||
action: ControllersWsAction,
|
||||
) {
|
||||
let action_res = match action {
|
||||
ControllersWsAction::Register(controller) => self.handle_register(conn, controller),
|
||||
ControllersWsAction::RelayStates((controller_uid, relay_states)) => {
|
||||
self.handle_relay_states(controller_uid, relay_states)
|
||||
}
|
||||
_ => Ok(()),
|
||||
};
|
||||
if let Err(e) = action_res {
|
||||
log::error!("Error handling action: {:?}", e);
|
||||
self.text_to_controller_conn(
|
||||
conn,
|
||||
serde_json::to_string(&e).unwrap_or(format!("Error in handling action: {:?}", e)),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn text_to_controller_conn(&self, target: WsConnId, msg: String) {
|
||||
if let Some((tx, _)) = self.controller_sessions.get(&target) {
|
||||
log::debug!("Sending message to controller conn {}", target);
|
||||
// errors if a client has disconnected abruptly and hasn't been timed-out yet
|
||||
let _ = tx.send(WsControllerAction::Forward(msg.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
pub fn handle_register(
|
||||
&mut self,
|
||||
conn: WsConnId,
|
||||
controller: Controller,
|
||||
) -> Result<(), EmgauwaError> {
|
||||
log::info!(
|
||||
"Registering controller: {} ({})",
|
||||
controller.c.name,
|
||||
controller.c.uid
|
||||
);
|
||||
let mut pool_conn = block_on(self.pool.acquire())?;
|
||||
|
||||
let c = &controller.c;
|
||||
let controller_db = block_on(DbController::get_by_uid_or_create(
|
||||
&mut pool_conn,
|
||||
&c.uid,
|
||||
&c.name,
|
||||
c.relay_count,
|
||||
))?;
|
||||
block_on(controller_db.update_active(&mut pool_conn, true))?;
|
||||
// update only the relay count
|
||||
block_on(controller_db.update(&mut pool_conn, &controller_db.name, c.relay_count))?;
|
||||
|
||||
for relay in &controller.relays {
|
||||
log::debug!(
|
||||
"Registering relay: {} ({})",
|
||||
relay.r.name,
|
||||
match relay.is_on {
|
||||
Some(true) => "+",
|
||||
Some(false) => "-",
|
||||
None => "?",
|
||||
}
|
||||
);
|
||||
let (new_relay, created) = block_on(DbRelay::get_by_controller_and_num_or_create(
|
||||
&mut pool_conn,
|
||||
&controller_db,
|
||||
relay.r.number,
|
||||
&relay.r.name,
|
||||
))?;
|
||||
if created {
|
||||
let mut relay_schedules = Vec::new();
|
||||
for schedule in &relay.schedules {
|
||||
let (new_schedule, _) = block_on(DbSchedule::get_by_uid_or_create(
|
||||
&mut pool_conn,
|
||||
schedule.uid.clone(),
|
||||
&schedule.name,
|
||||
&schedule.periods,
|
||||
))?;
|
||||
relay_schedules.push(new_schedule);
|
||||
}
|
||||
|
||||
block_on(DbJunctionRelaySchedule::set_schedules(
|
||||
&mut pool_conn,
|
||||
&new_relay,
|
||||
relay_schedules.iter().collect(),
|
||||
))?;
|
||||
}
|
||||
}
|
||||
|
||||
let relay_states = controller.get_relay_states();
|
||||
|
||||
let controller_uid = &controller.c.uid;
|
||||
let controller_db = block_on(DbController::get_by_uid(&mut pool_conn, controller_uid))?
|
||||
.ok_or(DatabaseError::InsertGetError)?;
|
||||
let mut controller = Controller::from_db_model(&mut pool_conn, controller_db)?;
|
||||
controller.apply_relay_states(&relay_states);
|
||||
|
||||
self.controllers.insert(controller_uid.clone(), (controller, conn));
|
||||
|
||||
// Update uid in controller session
|
||||
if let Some((tx, _)) = self.controller_sessions.remove(&conn) {
|
||||
log::debug!("Updating controller uid for conn {}", conn);
|
||||
self.controller_sessions.insert(conn, (tx, Some(controller_uid.clone())));
|
||||
}
|
||||
|
||||
log::debug!("Done registering controller");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn handle_relay_states(
|
||||
&mut self,
|
||||
controller_uid: EmgauwaUid,
|
||||
relay_states: RelayStates,
|
||||
) -> Result<(), EmgauwaError> {
|
||||
log::debug!(
|
||||
"Received relay states: {} for {}",
|
||||
utils::printable_relay_states(&relay_states),
|
||||
controller_uid
|
||||
);
|
||||
|
||||
if let Some((controller, _)) = self.controllers.get_mut(&controller_uid) {
|
||||
controller.apply_relay_states(&relay_states);
|
||||
}
|
||||
self.notify_relay_clients()
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue