wacker: gracefully shutdown the server

wacker-cli/src/main.rs

@@ -9,9 +9,7 @@ use wacker_api::config::SOCK_PATH;
 
 #[derive(Parser)]
 #[command(name = "wacker")]
-#[command(author = "ia")]
-#[command(version = "0.1.1")]
-#[command(about = "wacker client", long_about = None)]
+#[command(author, version, about, long_about = None)]
 struct Wacker {
     #[clap(subcommand)]
     subcommand: Subcommand,
@@ -35,7 +33,7 @@ enum Subcommand {
 
 impl Wacker {
     /// Executes the command.
-    pub async fn execute(self) -> Result<()> {
+    async fn execute(self) -> Result<()> {
         let home_dir = dirs::home_dir().expect("Can't get home dir");
         let path = home_dir.join(SOCK_PATH);
 
@@ -59,5 +57,5 @@ impl Wacker {
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    return Wacker::parse().execute().await;
+    Wacker::parse().execute().await
 }

wacker/Cargo.toml

@@ -19,15 +19,16 @@ path = "src/main.rs"
 [dependencies]
 wacker-api = { path = "../wacker-api", version = "0.1.1" }
 anyhow = "1.0.75"
-tokio = { version = "1.34.0", features = ["rt", "rt-multi-thread", "macros"] }
+tokio = { version = "1.34.0", features = ["rt", "rt-multi-thread", "macros", "signal"] }
 tokio-stream = { version = "0.1.14", features = ["net"] }
-wasi-common = "15.0.0"
+wasi-common = "15.0.1"
 wasmtime = { version = "15.0.1", features = ["cranelift"] }
-wasmtime-wasi = { version = "15.0.0", features = ["tokio"] }
+wasmtime-wasi = { version = "15.0.1", features = ["tokio"] }
 tonic = "0.10.2"
 dirs = "5.0.1"
 log = "0.4.20"
 env_logger = "0.10.1"
 cap-std = "2.0.0"
 rand = "0.8.5"
 sled = "0.34.7"
+clap = { version = "4.4.11", features = ["derive"] }

wacker/src/main.rs

@@ -3,42 +3,65 @@ mod run;
 
 use crate::module::Service;
 use anyhow::Result;
+use clap::Parser;
 use log::info;
 use std::fs;
 use tokio::net::UnixListener;
+use tokio::signal;
 use tokio_stream::wrappers::UnixListenerStream;
 use tonic::transport::Server;
-use wacker_api::{config::SOCK_PATH, modules_server::ModulesServer};
+use wacker_api::{
+    config::{DB_PATH, SOCK_PATH},
+    modules_server::ModulesServer,
+};
 
-#[tokio::main]
-async fn main() -> Result<()> {
-    let home_dir = dirs::home_dir().expect("Can't get home dir");
-    let binding = home_dir.join(SOCK_PATH);
-    let path = binding.as_path();
-    let parent_path = path.parent().unwrap();
+#[derive(Parser)]
+#[command(name = "wackerd")]
+#[command(author, version, about, long_about = None)]
+struct WackerDaemon {}
 
-    if !parent_path.exists() {
-        fs::create_dir_all(parent_path)?;
-    }
-    if path.exists() {
-        fs::remove_file(path).expect("Failed to remove existing socket file");
-    }
+impl WackerDaemon {
+    async fn execute(self) -> Result<()> {
+        let home_dir = dirs::home_dir().expect("Can't get home dir");
+        let binding = home_dir.join(SOCK_PATH);
+        let path = binding.as_path();
+        let parent_path = path.parent().unwrap();
+
+        if !parent_path.exists() {
+            fs::create_dir_all(parent_path)?;
+        }
+        if path.exists() {
+            fs::remove_file(path)?;
+        }
 
-    let uds = UnixListener::bind(path)?;
-    let uds_stream = UnixListenerStream::new(uds);
+        let uds = UnixListener::bind(path)?;
+        let uds_stream = UnixListenerStream::new(uds);
 
-    let inner = Service::new(home_dir).await?;
+        let db = sled::open(home_dir.join(DB_PATH))?;
+        let inner = Service::new(home_dir, db.clone()).await?;
 
-    let env = env_logger::Env::default()
-        .filter_or("LOG_LEVEL", "info")
-        .write_style_or("LOG_STYLE", "never");
-    env_logger::init_from_env(env);
+        let env = env_logger::Env::default()
+            .filter_or("LOG_LEVEL", "info")
+            .write_style_or("LOG_STYLE", "never");
+        env_logger::init_from_env(env);
 
-    info!("server listening on {:?}", path);
-    Server::builder()
-        .add_service(ModulesServer::new(inner))
-        .serve_with_incoming(uds_stream)
-        .await?;
+        info!("server listening on {:?}", path);
+        Server::builder()
+            .add_service(ModulesServer::new(inner))
+            .serve_with_incoming_shutdown(uds_stream, async {
+                signal::ctrl_c().await.expect("failed to listen for event");
+                println!();
+                info!("Shutting down the server");
+                fs::remove_file(path).expect("failed to remove existing socket file");
+                db.flush_async().await.expect("failed to flush the db");
+            })
+            .await?;
 
-    Ok(())
+        Ok(())
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    WackerDaemon::parse().execute().await
 }

wacker/src/module.rs

@@ -13,7 +13,7 @@ use tokio::{
     task,
 };
 use tonic::{Request, Response, Status};
-use wacker_api::config::{DB_PATH, LOGS_DIR};
+use wacker_api::config::LOGS_DIR;
 
 pub struct Service {
     db: Db,
@@ -31,14 +31,13 @@ struct InnerModule {
 }
 
 impl Service {
-    pub async fn new(home_dir: PathBuf) -> Result<Self, Error> {
+    pub async fn new(home_dir: PathBuf, db: Db) -> Result<Self, Error> {
         if let Err(e) = create_dir(home_dir.join(LOGS_DIR)) {
             if e.kind() != ErrorKind::AlreadyExists {
                 bail!("create logs dir failed: {}", e);
             }
         }
 
-        let db = sled::open(home_dir.join(DB_PATH))?;
         // Create an environment shared by all wasm execution. This contains
         // the `Engine` we are executing.
         let env = Environment::new()?;