1 files changed, 114 insertions, 40 deletions

diff --git a/src/main.rs b/src/main.rs
index f755db0..51b1e38 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,22 +2,25 @@ use std::fs::File;
 use std::io::{Read, Write};
 use std::path::PathBuf;
 use std::thread;
-use std::time::Duration;
 
-use log::{debug, error, info};
+use slog::{debug, error, info, o, Drain, Logger};
 use structopt::clap::AppSettings;
 use structopt::StructOpt;
+#[cfg(unix)]
+use tokio::signal::unix::*;
 use tsclientlib::Identity;
 
 mod audio_player;
 mod bot;
 mod command;
+mod log_bridge;
 mod playlist;
 mod teamspeak;
 mod web_server;
 mod youtube_dl;
 
-use bot::{MasterArgs, MasterBot, MusicBot, MusicBotArgs};
+use bot::{MasterArgs, MasterBot, MusicBot, MusicBotArgs, Quit};
+use log_bridge::LogBridge;
 
 #[derive(StructOpt, Debug)]
 #[structopt(global_settings = &[AppSettings::ColoredHelp])]
@@ -51,6 +54,10 @@ pub struct Args {
         help = "The channel the master bot should connect to"
     )]
     master_channel: Option<String>,
+    // 0. Print nothing
+    // 1. Print command string
+    // 2. Print packets
+    // 3. Print udp packets
     #[structopt(
         short = "v",
         long = "verbose",
@@ -58,25 +65,44 @@ pub struct Args {
         parse(from_occurrences)
     )]
     verbose: u8,
-    // 0. Print nothing
-    // 1. Print command string
-    // 2. Print packets
-    // 3. Print udp packets
 }
 
 #[tokio::main]
 async fn main() {
-    if let Err(e) = run().await {
-        println!("Error: {}", e);
+    let root_logger = {
+        let config = log4rs::load_config_file("log4rs.yml", Default::default()).unwrap();
+        let drain = LogBridge(log4rs::Logger::new(config)).fuse();
+        // slog_async adds a channel because log4rs if not unwind safe
+        let drain = slog_async::Async::new(drain).build().fuse();
+
+        Logger::root(drain, o!())
+    };
+
+    let scope_guard = slog_scope::set_global_logger(root_logger.clone());
+    // On SIGTERM the logger resets for some reason which makes the bot panic
+    // if it tries to log anything
+    scope_guard.cancel_reset();
+
+    slog_stdlog::init().unwrap();
+
+    if let Err(e) = run(root_logger.clone()).await {
+        error!(root_logger, "{}", e);
     }
 }
 
-async fn run() -> Result<(), Box<dyn std::error::Error>> {
-    log4rs::init_file("log4rs.yml", Default::default()).unwrap();
-
+async fn run(root_logger: Logger) -> Result<(), Box<dyn std::error::Error>> {
     // Parse command line options
     let args = Args::from_args();
 
+    // Set up signal handlers
+    let ctrl_c = tokio::task::spawn(tokio::signal::ctrl_c());
+    #[cfg(unix)]
+    let (sighup, sigterm, sigquit) = (
+        tokio::task::spawn(hangup()),
+        tokio::task::spawn(terminate()),
+        tokio::task::spawn(quit()),
+    );
+
     let mut file = File::open(&args.config_path)?;
     let mut toml = String::new();
     file.read_to_string(&mut toml)?;
@@ -107,14 +133,14 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
 
     if let Some(level) = args.wanted_level {
         if let Some(id) = &mut config.id {
-            info!("Upgrading master identity");
+            info!(root_logger, "Upgrading master identity");
             id.upgrade_level(level).expect("can upgrade level");
         }
 
         if let Some(ids) = &mut config.ids {
             let len = ids.len();
             for (i, id) in ids.iter_mut().enumerate() {
-                info!("Upgrading bot identity {}/{}", i + 1, len);
+                info!(root_logger, "Upgrading bot identity"; "current" => i + 1, "amount" => len);
                 id.upgrade_level(level).expect("can upgrade level");
             }
         }
@@ -127,53 +153,101 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
     }
 
     if config.id.is_none() || config.ids.is_none() {
-        error!("Failed to find required identites, try running with `-g`");
+        error!(
+            root_logger,
+            "Failed to find required identites, try running with `-g`"
+        );
         return Ok(());
     }
 
+    let local = args.local;
     let bot_args = config.merge(args);
 
-    info!("Starting PokeBot!");
-    debug!("Received CLI arguments: {:?}", std::env::args());
+    info!(root_logger, "Starting PokeBot!");
+    debug!(root_logger, "Received CLI arguments"; "args" => ?std::env::args());
 
-    if bot_args.local {
+    if local {
         let name = bot_args.names[0].clone();
-        let id = bot_args.ids.expect("identies should exists")[0].clone();
-
-        let disconnect_cb = Box::new(move |_, _, _| {});
+        let identity = bot_args.ids.expect("identies should exists")[0].clone();
 
         let bot_args = MusicBotArgs {
             name,
-            name_index: 0,
-            id_index: 0,
+            master: None,
             local: true,
             address: bot_args.address.clone(),
-            id,
+            identity,
             channel: String::from("local"),
             verbose: bot_args.verbose,
-            disconnect_cb,
+            logger: root_logger,
         };
-        MusicBot::new(bot_args).await.1.await;
+        MusicBot::spawn(bot_args).await;
+
+        ctrl_c.await??;
     } else {
         let domain = bot_args.domain.clone();
         let bind_address = bot_args.bind_address.clone();
-        let (bot, fut) = MasterBot::new(bot_args).await;
-
-        thread::spawn(|| {
-            let web_args = web_server::WebServerArgs {
-                domain,
-                bind_address,
-                bot,
-            };
-            if let Err(e) = web_server::start(web_args) {
-                error!("Error in web server: {}", e);
+        let bot_name = bot_args.master_name.clone();
+        let bot_logger = root_logger.new(o!("master" => bot_name.clone()));
+        let bot = MasterBot::spawn(bot_args, bot_logger).await;
+
+        let web_args = web_server::WebServerArgs {
+            domain,
+            bind_address,
+            bot: bot.downgrade(),
+        };
+        spawn_web_server(web_args, root_logger.new(o!("webserver" => bot_name)));
+
+        #[cfg(unix)]
+        tokio::select! {
+            res = ctrl_c => {
+                res??;
+                info!(root_logger, "Received signal, shutting down"; "signal" => "SIGINT");
             }
-        });
+            _ = sigterm => {
+                info!(root_logger, "Received signal, shutting down"; "signal" => "SIGTERM");
+            }
+            _ = sighup => {
+                info!(root_logger, "Received signal, shutting down"; "signal" => "SIGHUP");
+            }
+            _ = sigquit => {
+                info!(root_logger, "Received signal, shutting down"; "signal" => "SIGQUIT");
+            }
+        };
+
+        #[cfg(windows)]
+        ctrl_c.await??;
 
-        fut.await;
-        // Keep tokio running while the bot disconnects
-        tokio::time::delay_for(Duration::from_secs(1)).await;
+        bot.send(Quit(String::from("Stopping")))
+            .await
+            .unwrap()
+            .unwrap();
     }
 
     Ok(())
 }
+
+pub fn spawn_web_server(args: web_server::WebServerArgs, logger: Logger) {
+    thread::spawn(move || {
+        if let Err(e) = web_server::start(args, logger.clone()) {
+            error!(logger, "Error in web server"; "error" => %e);
+        }
+    });
+}
+
+#[cfg(unix)]
+pub async fn terminate() -> std::io::Result<()> {
+    signal(SignalKind::terminate())?.recv().await;
+    Ok(())
+}
+
+#[cfg(unix)]
+pub async fn hangup() -> std::io::Result<()> {
+    signal(SignalKind::hangup())?.recv().await;
+    Ok(())
+}
+
+#[cfg(unix)]
+pub async fn quit() -> std::io::Result<()> {
+    signal(SignalKind::quit())?.recv().await;
+    Ok(())
+}