mod handlers; use actix::Addr; use emgauwa_common::constants::WEBSOCKET_RETRY_TIMEOUT; use emgauwa_common::errors::EmgauwaError; use emgauwa_common::types::ControllerWsAction; use futures::{future, pin_mut, SinkExt, StreamExt}; use sqlx::{Pool, Sqlite}; use tokio::time; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; use crate::app_state::AppState; use crate::utils::{app_state_get_relay_notifier, app_state_get_this}; use crate::ws::handlers::handle_message; 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", )) })?; } }