2 files changed, 50 insertions, 46 deletions

diff --git a/src/main.rs b/src/main.rs
index 2559a2a..81b5bec 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,9 @@ use std::fs::File;
 use std::io::{Read, Write};
 use std::path::PathBuf;
 use std::thread;
+use std::time::{Duration, Instant};
 
+use futures::compat::Future01CompatExt;
 use futures::future::{FutureExt, TryFutureExt};
 use log::{debug, error, info};
 use structopt::clap::AppSettings;
@@ -97,49 +99,56 @@ fn run() -> Result<(), Box<dyn std::error::Error>> {
     info!("Starting PokeBot!");
     debug!("Received CLI arguments: {:?}", std::env::args());
 
-    tokio::run(
-        async {
-            if bot_args.local {
-                let name = bot_args.names[0].clone();
-                let id = bot_args.ids[0].clone();
-
-                let disconnect_cb = Box::new(move |_, _, _| {});
-
-                let bot_args = MusicBotArgs {
-                    name: name,
-                    name_index: 0,
-                    id_index: 0,
-                    local: true,
-                    address: bot_args.address.clone(),
-                    id,
-                    channel: String::from("local"),
-                    verbose: bot_args.verbose,
-                    disconnect_cb,
-                };
-                MusicBot::new(bot_args).await.1.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,
+    tokio::runtime::Runtime::new()?
+        .block_on(
+            async {
+                if bot_args.local {
+                    let name = bot_args.names[0].clone();
+                    let id = bot_args.ids[0].clone();
+
+                    let disconnect_cb = Box::new(move |_, _, _| {});
+
+                    let bot_args = MusicBotArgs {
+                        name: name,
+                        name_index: 0,
+                        id_index: 0,
+                        local: true,
+                        address: bot_args.address.clone(),
+                        id,
+                        channel: String::from("local"),
+                        verbose: bot_args.verbose,
+                        disconnect_cb,
                     };
-                    if let Err(e) = web_server::start(web_args) {
-                        error!("Error in web server: {}", e);
-                    }
-                });
-
-                fut.await;
+                    MusicBot::new(bot_args).await.1.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);
+                        }
+                    });
+
+                    fut.await;
+                    // Keep tokio running while the bot disconnects
+                    tokio::timer::Delay::new(Instant::now() + Duration::from_secs(1))
+                        .compat()
+                        .await
+                        .expect("Failed to wait for delay");
+                }
             }
-        }
-        .unit_error()
-        .boxed()
-        .compat(),
-    );
+            .unit_error()
+            .boxed()
+            .compat(),
+        )
+        .expect("Runtime exited on an error");
 
     Ok(())
 }
diff --git a/src/teamspeak/mod.rs b/src/teamspeak/mod.rs
index 7551e77..4d31f2e 100644
--- a/src/teamspeak/mod.rs
+++ b/src/teamspeak/mod.rs
@@ -1,5 +1,4 @@
 use std::sync::{Arc, RwLock};
-use std::time::{Duration, Instant};
 
 use futures::compat::Future01CompatExt;
 use futures01::{future::Future, sink::Sink};
@@ -224,9 +223,5 @@ impl TeamSpeakConnection {
                 .disconnect(opt)
                 .map_err(|e| error!("Failed to send message: {}", e)),
         );
-        // Might or might not be required to keep tokio running while the bot disconnects
-        tokio::spawn(
-            tokio::timer::Delay::new(Instant::now() + Duration::from_secs(1)).map_err(|_| ()),
-        );
     }
 }