use actix::Addr; use emgauwa_common::constants::WEBSOCKET_RETRY_TIMEOUT; use emgauwa_common::db::{DbController, DbJunctionRelaySchedule, DbRelay, DbSchedule}; use emgauwa_common::errors::{DatabaseError, EmgauwaError}; use emgauwa_common::models::{Controller, Relay}; use emgauwa_common::types::{ControllerWsAction, ScheduleUid}; use futures::{future, pin_mut, SinkExt, StreamExt}; use sqlx::pool::PoolConnection; use sqlx::{Pool, Sqlite}; use tokio::time; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::{connect_async, tungstenite}; use crate::app_state::AppState; use crate::utils; use crate::utils::{app_state_get_relay_notifier, app_state_get_this}; pub async fn run_ws_loop(pool: Pool<Sqlite>, app_state: Addr<AppState>, url: String) { log::debug!("Spawned ws loop"); loop { let run_result = run_websocket(pool.clone(), &app_state, &url).await; if let Err(err) = run_result { log::error!("Error running websocket: {}", err); } log::info!( "Retrying to connect in {} seconds...", WEBSOCKET_RETRY_TIMEOUT.as_secs() ); time::sleep(WEBSOCKET_RETRY_TIMEOUT).await; } } async fn run_websocket( pool: Pool<Sqlite>, app_state: &Addr<AppState>, url: &str, ) -> Result<(), EmgauwaError> { log::debug!("Trying to connect to {}", url); match connect_async(url).await { Ok(connection) => { log::info!("Websocket connected"); let (ws_stream, _) = connection; let (mut write, read) = ws_stream.split(); let ws_action = ControllerWsAction::Register(app_state_get_this(app_state).await?); let ws_action_json = serde_json::to_string(&ws_action)?; if let Err(err) = write.send(Message::text(ws_action_json)).await { log::error!("Failed to register at websocket: {}", err); return Ok(()); } let (app_state_tx, app_state_rx) = futures_channel::mpsc::unbounded::<Message>(); tokio::spawn(read_app_state(app_state.clone(), app_state_tx)); let app_state_to_ws = app_state_rx.map(Ok).forward(write); let read_handler = read.for_each(|msg| handle_message(pool.clone(), app_state, msg)); pin_mut!(app_state_to_ws, read_handler); future::select(app_state_to_ws, read_handler).await; log::warn!("Lost connection to websocket"); } Err(err) => { log::warn!("Failed to connect to websocket: {}", err,); } } Ok(()) } async fn read_app_state( app_state: Addr<AppState>, tx: futures_channel::mpsc::UnboundedSender<Message>, ) -> Result<(), EmgauwaError> { let notifier = &*app_state_get_relay_notifier(&app_state).await?; loop { notifier.notified().await; log::debug!("Relay change detected"); let this = app_state_get_this(&app_state).await?; let relay_states = this.get_relay_states(); let ws_action = ControllerWsAction::RelayStates((this.c.uid, relay_states)); let ws_action_json = serde_json::to_string(&ws_action)?; tx.unbounded_send(Message::text(ws_action_json)) .map_err(|_| { EmgauwaError::Other(String::from( "Failed to forward message from app state to websocket", )) })?; } } async fn handle_message( pool: Pool<Sqlite>, app_state: &Addr<AppState>, message_result: Result<Message, tungstenite::Error>, ) { let msg = match message_result { Ok(msg) => msg, Err(err) => { log::error!("Error reading message: {}", err); return; } }; if let Message::Text(text) = msg { match serde_json::from_str(&text) { Ok(action) => { log::debug!("Received action: {:?}", action); let mut pool_conn = match pool.acquire().await { Ok(conn) => conn, Err(err) => { log::error!("Failed to acquire database connection: {:?}", err); return; } }; let action_res = handle_action(&mut pool_conn, app_state, action).await; if let Err(e) = action_res { log::error!("Error handling action: {:?}", e); } } Err(e) => { log::error!("Error deserializing action: {:?}", e); } } } } pub async fn handle_action( conn: &mut PoolConnection<Sqlite>, app_state: &Addr<AppState>, action: ControllerWsAction, ) -> Result<(), EmgauwaError> { let this = app_state_get_this(app_state).await?; match action { ControllerWsAction::Controller(controller) => { handle_controller(conn, &this, controller).await? } ControllerWsAction::Relays(relays) => handle_relays(conn, &this, relays).await?, ControllerWsAction::Schedules(schedules) => handle_schedules(conn, schedules).await?, ControllerWsAction::RelayPulse((relay_num, duration)) => { handle_relay_pulse(app_state, relay_num, duration).await? } _ => return Ok(()), }; utils::app_state_reload(app_state).await } async fn handle_controller( conn: &mut PoolConnection<Sqlite>, this: &Controller, controller: Controller, ) -> Result<(), EmgauwaError> { if controller.c.uid != this.c.uid { return Err(EmgauwaError::Other(String::from( "Controller UID mismatch during update", ))); } DbController::get_by_uid(conn, &controller.c.uid) .await? .ok_or(DatabaseError::NotFound)? .update(conn, controller.c.name.as_str(), this.c.relay_count) .await?; Ok(()) } async fn handle_schedules( conn: &mut PoolConnection<Sqlite>, schedules: Vec<DbSchedule>, ) -> Result<(), EmgauwaError> { let mut handled_uids = vec![ // on and off schedules are always present and should not be updated ScheduleUid::On, ScheduleUid::Off, ]; for schedule in schedules { if handled_uids.contains(&schedule.uid) { continue; } handled_uids.push(schedule.uid.clone()); log::debug!("Handling schedule: {:?}", schedule); let schedule_db = DbSchedule::get_by_uid(conn, &schedule.uid).await?; if let Some(schedule_db) = schedule_db { schedule_db .update(conn, schedule.name.as_str(), &schedule.periods) .await?; } else { DbSchedule::create( conn, schedule.uid.clone(), schedule.name.as_str(), &schedule.periods, ) .await?; } } Ok(()) } async fn handle_relays( conn: &mut PoolConnection<Sqlite>, this: &Controller, relays: Vec<Relay>, ) -> Result<(), EmgauwaError> { for relay in relays { if relay.controller.uid != this.c.uid { return Err(EmgauwaError::Other(String::from( "Controller UID mismatch during relay update", ))); } let db_relay = DbRelay::get_by_controller_and_num(conn, &this.c, relay.r.number) .await? .ok_or(DatabaseError::NotFound)?; db_relay.update(conn, relay.r.name.as_str()).await?; handle_schedules(conn, relay.schedules.clone()).await?; let mut schedules = Vec::new(); // We need to get the schedules from the database to have the right IDs for schedule in relay.schedules { schedules.push( DbSchedule::get_by_uid(conn, &schedule.uid) .await? .ok_or(DatabaseError::NotFound)?, ); } DbJunctionRelaySchedule::set_schedules(conn, &db_relay, schedules.iter().collect()).await?; } Ok(()) } async fn handle_relay_pulse( app_state: &Addr<AppState>, relay_num: i64, duration: Option<u32>, ) -> Result<(), EmgauwaError> { utils::app_state_relay_pulse(app_state, relay_num, duration).await }