Implement graceful shutdown support

Cargo.toml

@@ -57,3 +57,4 @@ tracing = "0.1.44"
 tracing-subscriber = { version = "0.3.22", features = ["env-filter"] }
 uuid = { version = "1.19.0", features = ["v4"] }
 walkdir = "2.5.0"
+tokio-util = { version = "0.7", features = ["rt"] }

java_wrapper/wrapper/src/main/java/dev/quozul/BungeeCordPlugin.java

@@ -8,6 +8,8 @@
 
 public class BungeeCordPlugin extends Plugin {
 
+    private PicoLimboRunner worker;
+
     @Override
     public void onEnable() {
         Path dataDirectory = getDataFolder().toPath();
@@ -21,8 +23,14 @@ public void onEnable() {
         }
 
         Path configurationFile = dataDirectory.resolve("server.toml");
-        PicoLimboRunner worker = new PicoLimboRunner(configurationFile);
+        this.worker = new PicoLimboRunner(configurationFile);
 
         getProxy().getScheduler().runAsync(this, worker);
     }
+
+    @Override
+    public void onDisable() {
+        if (worker != null)
+            worker.stop();
+    }
 }

java_wrapper/wrapper/src/main/java/dev/quozul/PicoLimboRunner.java

@@ -5,17 +5,35 @@
 public class PicoLimboRunner implements Runnable {
 
     private final Path configurationPath;
+    private volatile Standalone.RustLib lib;
+    private volatile boolean running = false;
 
     public PicoLimboRunner(Path configurationPath) {
         this.configurationPath = configurationPath;
     }
 
     @Override
     public void run() {
-        String[] args = {"--config", configurationPath.toString()};
-        Standalone.main(args);
+        try {
+            lib = Standalone.loadLib();
+            running = true;
+
+            String[] args = {
+                    "pico_limbo_java_wrapper",
+                    "--config",
+                    configurationPath.toString()
+            };
+
+            lib.start_app(args.length, args);
+        } catch (Exception e) {
+            e.printStackTrace();
+        } finally {
+            running = false;
+        }
     }
 
     public void stop() {
+        if (lib != null && running)
+            lib.stop_app();
     }
 }

java_wrapper/wrapper/src/main/java/dev/quozul/Standalone.java

@@ -13,31 +13,35 @@ public class Standalone {
     public interface RustLib extends Library {
         // void start_app(int argc, char** argv);
         byte start_app(int argc, String[] argv);
+
+        void stop_app();
     }
 
-    public static void main(String[] args) {
-        try {
-            String libName = BuildConstants.LIB_NAME;
-            String resourcePath = getResourcePath(libName);
-
-            String extension = resourcePath.substring(resourcePath.lastIndexOf('.'));
-            File tempLib = File.createTempFile(libName, extension);
-            tempLib.deleteOnExit();
-
-            try (InputStream in = Standalone.class.getResourceAsStream(resourcePath)) {
-                if (in == null) {
-                    throw new RuntimeException("Library file not found in JAR: " + resourcePath);
-                }
-                Files.copy(in, tempLib.toPath(), StandardCopyOption.REPLACE_EXISTING);
+    public static RustLib loadLib() throws Exception {
+        String libName = BuildConstants.LIB_NAME;
+        String resourcePath = getResourcePath(libName);
+
+        String extension = resourcePath.substring(resourcePath.lastIndexOf('.'));
+        File tempLib = File.createTempFile(libName, extension);
+        tempLib.deleteOnExit();
+
+        try (InputStream in = Standalone.class.getResourceAsStream(resourcePath)) {
+            if (in == null) {
+                throw new RuntimeException("Library file not found in JAR: " + resourcePath);
             }
+            Files.copy(in, tempLib.toPath(), StandardCopyOption.REPLACE_EXISTING);
+        }
 
-            RustLib lib = Native.load(tempLib.getAbsolutePath(), RustLib.class);
+        return Native.load(tempLib.getAbsolutePath(), RustLib.class);
+    }
 
+    public static void main(String[] args) {
+        try {
             String[] effectiveArgs = new String[args.length + 1];
             effectiveArgs[0] = "pico_limbo_java_wrapper";
             System.arraycopy(args, 0, effectiveArgs, 1, args.length);
 
-            byte exitCode = lib.start_app(effectiveArgs.length, effectiveArgs);
+            byte exitCode = loadLib().start_app(effectiveArgs.length, effectiveArgs);
             System.exit(exitCode);
 
         } catch (Exception e) {

pico_limbo/Cargo.toml

@@ -41,3 +41,4 @@ toml = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 uuid = { workspace = true }
+tokio-util = { workspace = true }

pico_limbo/src/lib.rs

@@ -10,6 +10,10 @@ use crate::cli::Cli;
 use clap::Parser;
 use std::ffi::{CStr, c_char, c_int};
 use std::slice;
+use std::sync::OnceLock;
+use tokio_util::sync::CancellationToken;
+
+static CANCEL_TOKEN: OnceLock<CancellationToken> = OnceLock::new();
 
 /// Some docs
 ///
@@ -42,6 +46,8 @@ pub unsafe extern "C" fn start_app(argc: c_int, argv: *const *const c_char) {
 
     match Cli::try_parse_from(&rust_args) {
         Ok(cli) => {
+            let token = CANCEL_TOKEN.get_or_init(CancellationToken::new).clone();
+
             let rt = tokio::runtime::Builder::new_current_thread()
                 .enable_all()
                 .build()
@@ -50,10 +56,23 @@ pub unsafe extern "C" fn start_app(argc: c_int, argv: *const *const c_char) {
             let _ = rt.block_on(server::start_server::start_server(
                 cli.config_path,
                 cli.verbose,
+                token,
             ));
         }
         Err(e) => {
             e.print().expect("Failed to print error");
         }
     }
 }
+
+/// Some docs
+///
+/// # Safety
+///
+/// Pretty safe actually
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn stop_app() {
+    if let Some(token) = CANCEL_TOKEN.get() {
+        token.cancel();
+    }
+}

pico_limbo/src/main.rs

@@ -10,9 +10,10 @@ mod server_state;
 use crate::cli::Cli;
 use clap::Parser;
 use std::process::ExitCode;
+use tokio_util::sync::CancellationToken;
 
 #[tokio::main]
 async fn main() -> ExitCode {
     let cli = Cli::parse();
-    server::start_server::start_server(cli.config_path, cli.verbose).await
+    server::start_server::start_server(cli.config_path, cli.verbose, CancellationToken::new()).await
 }

pico_limbo/src/server/start_server.rs

@@ -8,12 +8,17 @@ use crate::server::network::Server;
 use crate::server_state::{ServerState, ServerStateBuilderError};
 use std::path::PathBuf;
 use std::process::ExitCode;
+use tokio_util::sync::CancellationToken;
 use tracing::{Level, debug, error};
 use tracing_subscriber::EnvFilter;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
 
-pub async fn start_server(config_path: PathBuf, logging_level: u8) -> ExitCode {
+pub async fn start_server(
+    config_path: PathBuf,
+    logging_level: u8,
+    token: CancellationToken,
+) -> ExitCode {
     enable_logging(logging_level);
     let Some(cfg) = load_configuration(&config_path) else {
         return ExitCode::FAILURE;
@@ -23,7 +28,12 @@ pub async fn start_server(config_path: PathBuf, logging_level: u8) -> ExitCode {
 
     match build_state(cfg) {
         Ok(server_state) => {
-            Server::new(&bind, server_state).run().await;
+            tokio::select! {
+                () = Server::new(&bind, server_state).run() => {}
+                () = token.cancelled() => {
+                    tracing::info!("Shutdown signal received, stopping server...");
+                }
+            }
             ExitCode::SUCCESS
         }
         Err(err) => {