diff --git a/src/cli.rs b/src/cli.rs index cbeb33f..55f9347 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -113,6 +113,10 @@ pub fn start_cli(server: Arc, plugin_loader: PluginLoader) { "ping" => { LOGGER.info("pong"); } + "shutdown" => { + server.shutdown(); + break; + } _ => LOGGER.error(format!("Command not found: {}", args[0])), } } diff --git a/src/plugin/loader.rs b/src/plugin/loader.rs index 26d1f4c..92c6b15 100644 --- a/src/plugin/loader.rs +++ b/src/plugin/loader.rs @@ -33,11 +33,13 @@ impl PluginLoader { let plugin_json: PluginJson = serde_json::from_str(&json_string).unwrap(); LOGGER.info(format!("Loading {}", plugin_json.id)); - Command::new(plugin_json.file) - .args(plugin_json.args) - .current_dir(path) - .spawn() - .unwrap(); + let child = Arc::new(Mutex::new( + Command::new(plugin_json.file) + .args(plugin_json.args) + .current_dir(path) + .spawn() + .unwrap(), + )); while self .plugin_clients @@ -59,6 +61,7 @@ impl PluginLoader { stream: a.try_clone().unwrap(), reader: BufReader::new(a), id: plugin_json.id, + child, } } @@ -88,6 +91,7 @@ impl Clone for Plugin { stream: self.stream.try_clone().unwrap(), reader: BufReader::new(self.stream.try_clone().unwrap()), id: self.id.clone(), + child: self.child.clone(), } } } diff --git a/src/plugin/mod.rs b/src/plugin/mod.rs index 399e57c..7d5ca0e 100644 --- a/src/plugin/mod.rs +++ b/src/plugin/mod.rs @@ -4,7 +4,8 @@ pub mod types; use std::{ io::{BufRead, BufReader, Write}, net::TcpStream, - sync::Arc, + process::Child, + sync::{Arc, Mutex}, }; use crate::{ @@ -17,6 +18,7 @@ pub struct Plugin { stream: TcpStream, reader: BufReader, id: String, + child: Arc>, } impl Plugin { @@ -60,4 +62,10 @@ impl Plugin { } } } + + pub fn stop(&mut self) { + self.send(&LoaderMessage::Shutdown).unwrap(); + self.stream.shutdown(std::net::Shutdown::Both).unwrap(); + self.child.lock().unwrap().kill().unwrap(); + } } diff --git a/src/plugin/types.rs b/src/plugin/types.rs index 169ce4b..7c7d076 100644 --- a/src/plugin/types.rs +++ b/src/plugin/types.rs @@ -31,6 +31,8 @@ pub enum LoaderMessage { user_id: String, msg: Message, }, + + Shutdown, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src/server.rs b/src/server.rs index 3f285b4..705b3d6 100644 --- a/src/server.rs +++ b/src/server.rs @@ -3,7 +3,10 @@ use std::{ fs::{self}, net::{TcpListener, TcpStream}, path::{Path, PathBuf}, - sync::{Arc, Mutex}, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, }; use crate::{ @@ -22,13 +25,13 @@ pub struct ServerConfig { pub channels: Vec, } -#[allow(dead_code)] pub struct Server { pub root: PathBuf, pub config: ServerConfig, pub clients: Mutex>, pub plugins: Mutex>, pub db: utils::database::Database, + pub shutting_down: AtomicBool, } impl Default for ServerConfig { @@ -63,6 +66,7 @@ impl Server { config, clients: Mutex::new(HashSet::new()), plugins: Mutex::new(Vec::new()), + shutting_down: AtomicBool::new(false), }) } @@ -182,7 +186,7 @@ impl Server { fn handle_client(self: &Arc, client: &Client) -> crate::Result<()> { // The main req/res loop - loop { + while !self.shutting_down.load(Ordering::SeqCst) { let req = client.read()?; if let Some(r) = &req { self.send_plugin_message(&LoaderMessage::Request { @@ -192,6 +196,7 @@ impl Server { self.wrap_err(&client, self.call_request(r, &client))?; } } + Ok(()) } /// When there is a error it removes the client @@ -217,4 +222,29 @@ impl Server { } Ok(()) } + + pub fn shutdown(self: &Arc) { + Self::LOGGER.info("Server shutting down..."); + + // Signal shutdown + self.shutting_down.store(true, Ordering::SeqCst); + + // Disconnect clients + let clients = self.clients.lock().unwrap(); + for client in clients.iter() { + let _ = client.send(types::message::ServerMessage::Shutdown { + message: format!("Server shutting down... we'll be back shortly"), + }); + let _ = client.close(); + } + + // Stop plugins + for plugin in self.plugins.lock().unwrap().iter_mut() { + let _ = plugin.stop(); + } + + Self::LOGGER.info("Shutdown complete"); + Self::LOGGER.info("Exiting process.."); + std::process::exit(0); + } } diff --git a/src/types.rs b/src/types.rs index a28b3bd..644cee4 100644 --- a/src/types.rs +++ b/src/types.rs @@ -115,6 +115,10 @@ pub mod message { user_id: Author, channel_id: String, }, + + Shutdown { + message: String, + }, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src/utils/client.rs b/src/utils/client.rs index d7b77aa..fa8812d 100644 --- a/src/utils/client.rs +++ b/src/utils/client.rs @@ -346,6 +346,11 @@ impl Client { pub fn set_uuid(&mut self, uuid: &str) { self.1 = Some(uuid.to_string()) } + + pub fn close(&self) -> crate::Result<()> { + self.0.shutdown(std::net::Shutdown::Both)?; + Ok(()) + } } impl Clone for Client {