Implement service startup/shutdown timeout

src/config.rs

@@ -4,7 +4,6 @@ use std::{
     fmt::{Display, Formatter},
     fs, io,
     path::PathBuf,
-    time::Duration,
 };
 use thiserror::Error;
 
@@ -38,23 +37,16 @@ fn discord_token_default() -> String {
     String::from("Please provide a token")
 }
 
-fn discord_timeout_default() -> Duration {
-    Duration::from_secs(10)
-}
-
 #[derive(Debug, PartialEq, PartialOrd, Serialize, Deserialize, Clone)]
 pub struct Config {
     #[serde(rename = "discordToken", default = "discord_token_default")]
     pub discord_token: String,
-    #[serde(rename = "discordTimeout", default = "discord_timeout_default")]
-    pub discord_timeout: Duration,
 }
 
 impl Default for Config {
     fn default() -> Self {
         Config {
             discord_token: discord_token_default(),
-            discord_timeout: discord_timeout_default(),
         }
     }
 }

src/main.rs

@@ -54,8 +54,7 @@ fn initialize_services(config: &Config) -> Vec<Arc<RwLock<dyn Service>>> {
     //TODO: Add services
     //...
 
-    let discord_service =
-        DiscordService::new(config.discord_token.as_str(), config.discord_timeout);
+    let discord_service = DiscordService::new(config.discord_token.as_str());
 
     vec![Arc::new(RwLock::new(discord_service))]
 }

src/service.rs

@@ -10,8 +10,9 @@ use std::{
     mem,
     pin::Pin,
     sync::Arc,
+    time::Duration,
 };
-use tokio::sync::RwLock;
+use tokio::{sync::RwLock, time::timeout};
 
 use crate::setlock::SetLock;
 
@@ -290,16 +291,35 @@ impl ServiceManager {
         drop(status);
 
         let service_manager = Arc::clone(self.arc.read().await.unwrap());
-        match service.start(service_manager).await {
-            Ok(()) => {
-                info!("Started service: {}", service.info().name);
-                service.info().set_status(Status::Started).await;
-            }
+
+        let start = service.start(service_manager);
+
+        let duration = Duration::from_secs(10); //TODO: Add to config instead of hardcoding
+        let timeout_result = timeout(duration, start);
+
+        match timeout_result.await {
+            Ok(start_result) => match start_result {
+                Ok(()) => {
+                    info!("Started service: {}", service.info().name);
+                    service.info().set_status(Status::Started).await;
+                }
+                Err(error) => {
+                    error!("Failed to start service {}: {}", service.info().name, error);
+                    service
+                        .info()
+                        .set_status(Status::FailedToStart(error))
+                        .await;
+                }
+            },
             Err(error) => {
-                error!("Failed to start service {}: {}", service.info().name, error);
+                error!(
+                    "Failed to start service {}: Timeout of {} seconds reached.",
+                    service.info().name,
+                    duration.as_secs()
+                );
                 service
                     .info()
-                    .set_status(Status::FailedToStart(error))
+                    .set_status(Status::FailedToStart(Box::new(error)))
                     .await;
             }
         }
@@ -329,14 +349,32 @@ impl ServiceManager {
         *status = Status::Stopping;
         drop(status);
 
-        match service.stop().await {
-            Ok(()) => {
-                info!("Stopped service: {}", service.info().name);
-                service.info().set_status(Status::Stopped).await;
-            }
+        let stop = service.stop();
+
+        let duration = Duration::from_secs(10); //TODO: Add to config instead of hardcoding
+        let timeout_result = timeout(duration, stop);
+
+        match timeout_result.await {
+            Ok(stop_result) => match stop_result {
+                Ok(()) => {
+                    info!("Stopped service: {}", service.info().name);
+                    service.info().set_status(Status::Stopped).await;
+                }
+                Err(error) => {
+                    error!("Failed to stop service {}: {}", service.info().name, error);
+                    service.info().set_status(Status::FailedToStop(error)).await;
+                }
+            },
             Err(error) => {
-                error!("Failed to stop service {}: {}", service.info().name, error);
-                service.info().set_status(Status::FailedToStop(error)).await;
+                error!(
+                    "Failed to stop service {}: Timeout of {} seconds reached.",
+                    service.info().name,
+                    duration.as_secs()
+                );
+                service
+                    .info()
+                    .set_status(Status::FailedToStop(Box::new(error)))
+                    .await;
             }
         }
     }

src/service/discord.rs

@@ -14,16 +14,15 @@ use serenity::{
 };
 use std::{sync::Arc, time::Duration};
 use tokio::{
-    spawn,
+    select, spawn,
     sync::{Mutex, Notify, RwLock},
     task::JoinHandle,
-    time::{sleep, timeout},
+    time::sleep,
 };
 
 pub struct DiscordService {
     info: ServiceInfo,
     discord_token: String,
-    connection_timeout: Duration,
     pub ready: Arc<RwLock<SetLock<Ready>>>,
     client_handle: Option<JoinHandle<Result<(), Error>>>,
     pub cache: SetLock<Arc<Cache>>,
@@ -35,11 +34,10 @@ pub struct DiscordService {
 }
 
 impl DiscordService {
-    pub fn new(discord_token: &str, connection_timeout: Duration) -> Self {
+    pub fn new(discord_token: &str) -> Self {
         Self {
             info: ServiceInfo::new("lum_builtin_discord", "Discord", Priority::Essential),
             discord_token: discord_token.to_string(),
-            connection_timeout,
             ready: Arc::new(RwLock::new(SetLock::new())),
             client_handle: None,
             cache: SetLock::new(),
@@ -101,29 +99,16 @@ impl Service for DiscordService {
             info!("Connecting to Discord");
             let client_handle = spawn(async move { client.start().await });
 
-            // This prevents waiting for the timeout if the client fails immediately
-            // TODO: Optimize this, as it will currently add 1000mqs to the startup time
-            sleep(Duration::from_secs(1)).await;
+            select! {
+                _ = client_ready_notify.notified() => {},
+                _ = sleep(Duration::from_secs(2)) => {},
+            }
+
             if client_handle.is_finished() {
                 client_handle.await??;
-                return Err("Discord client stopped unexpectedly and with no error".into());
+                return Err("Discord client stopped unexpectedly".into());
             }
 
-            if timeout(self.connection_timeout, client_ready_notify.notified())
-                .await
-                .is_err()
-            {
-                client_handle.abort();
-                let result = convert_thread_result(client_handle).await;
-                result?;
-
-                return Err(format!(
-                    "Discord client failed to connect within {} seconds",
-                    self.connection_timeout.as_secs()
-                )
-                .into());
-            };
-
             self.client_handle = Some(client_handle);
             Ok(())
         })