Add log-init-function and websocket retry

emgauwa-controller/src/main.rs

@@ -1,19 +1,20 @@
-use std::str;
+use std::time::Duration;
 
 use crate::relay_loop::run_relay_loop;
 use crate::settings::Settings;
-use emgauwa_lib::db;
+use emgauwa_lib::constants::WEBSOCKET_RETRY_TIMEOUT;
 use emgauwa_lib::db::{DbController, DbRelay};
 use emgauwa_lib::handlers::v1::ws::controllers::ControllerWsAction;
 use emgauwa_lib::models::{Controller, FromDbModel};
 use emgauwa_lib::types::ControllerUid;
-use futures::channel::mpsc;
-use futures::{future, pin_mut, SinkExt, StreamExt};
+use emgauwa_lib::{db, utils};
+use futures::{SinkExt, StreamExt};
 use sqlx::pool::PoolConnection;
 use sqlx::Sqlite;
-use tokio::io::AsyncReadExt;
+use tokio::time;
 use tokio_tungstenite::tungstenite::Error;
 use tokio_tungstenite::{connect_async, tungstenite::protocol::Message};
+use utils::init_logging;
 
 mod driver;
 mod relay_loop;
@@ -51,6 +52,7 @@ async fn create_this_relay(
 #[tokio::main]
 async fn main() {
 	let settings = settings::init();
+	init_logging(&settings.logging.level);
 
 	let pool = db::init(&settings.database).await;
 
@@ -87,46 +89,47 @@ async fn main() {
 		settings.core.host, settings.core.port
 	);
 
-	let (stdin_tx, stdin_rx) = mpsc::unbounded();
-	tokio::spawn(read_stdin(stdin_tx));
-
-	let (ws_stream, _) = connect_async(url).await.expect("Failed to connect");
-
-	let (mut write, read) = ws_stream.split();
-
-	let ws_action = ControllerWsAction::Register(this);
-	println!("Sending action: {:?}", ws_action);
-	let ws_action_json = serde_json::to_string(&ws_action).unwrap();
-	println!("Sending json: {}", ws_action_json);
-	write.send(Message::text(ws_action_json)).await.unwrap();
-	let ws_to_stdout = read.for_each(handle_message);
-	let stdin_to_ws = stdin_rx.map(Ok).forward(write);
-
 	tokio::spawn(run_relay_loop(settings));
 
-	pin_mut!(stdin_to_ws, ws_to_stdout);
-	future::select(stdin_to_ws, ws_to_stdout).await;
-}
-
-// Our helper method which will read data from stdin and send it along the
-// sender provided.
-async fn read_stdin(tx: mpsc::UnboundedSender<Message>) {
-	let mut stdin = tokio::io::stdin();
 	loop {
-		let mut buf = vec![0; 1024];
-		let n = match stdin.read(&mut buf).await {
-			Err(_) | Ok(0) => break,
-			Ok(n) => n,
-		};
-		buf.truncate(n);
-		tx.unbounded_send(Message::text(str::from_utf8(&buf).unwrap()))
-			.unwrap();
+		time::sleep(WEBSOCKET_RETRY_TIMEOUT).await;
+
+		let connect_result = connect_async(&url).await;
+		if let Err(err) = connect_result {
+			log::warn!(
+				"Failed to connect to websocket: {}. Retrying in {} seconds...",
+				err,
+				WEBSOCKET_RETRY_TIMEOUT.as_secs()
+			);
+			continue;
+		}
+		let (ws_stream, _) = connect_result.unwrap();
+
+		let (mut write, read) = ws_stream.split();
+
+		let ws_action = ControllerWsAction::Register(this.clone());
+
+		let ws_action_json = serde_json::to_string(&ws_action).unwrap();
+		write.send(Message::text(ws_action_json)).await.unwrap();
+
+		let read_handler = read.for_each(handle_message);
+
+		read_handler.await;
+
+		log::warn!(
+			"Lost connection to websocket. Retrying in {} seconds...",
+			WEBSOCKET_RETRY_TIMEOUT.as_secs()
+		);
 	}
 }
 
 pub async fn handle_message(message_result: Result<Message, Error>) {
 	match message_result {
-		Ok(message) => println!("{}", message.into_text().unwrap()),
-		Err(err) => println!("Error: {}", err),
+		Ok(message) => {
+			if let Message::Text(msg_text) = message {
+				log::debug!("{}", msg_text)
+			}
+		}
+		Err(err) => log::debug!("Error: {}", err),
 	}
 }