add pid lock

kfatyuip · Aug 22, 2024 · 49626f1 · 49626f1
1 parent 4830aea
commit 49626f1
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 6 deletions.
diff --git a/src/init.rs b/src/init.rs
@@ -4,9 +4,15 @@ use async_rwlock::RwLock;
 use lazy_static::lazy_static;
 use log4rs::Handle;
 use lru::LruCache;
+use signal_hook::consts::SIGINT;
 use signal_hook::{consts::SIGHUP, iterator::Signals};
-use std::num::NonZero;
-use std::{env::set_current_dir, io, num::NonZeroUsize, sync::Arc, thread};
+use std::{
+    fs,
+    num::NonZero,
+    path::PathBuf,
+    process,
+    {env::set_current_dir, io, num::NonZeroUsize, sync::Arc, thread},
+};
 
 #[cfg(feature = "log")]
 use {
@@ -20,6 +26,7 @@ use {
 };
 
 lazy_static! {
+    pub static ref PID_FILE: Mutex<Option<PathBuf>> = Mutex::new(None);
     pub static ref T: Arc<RwLock<Option<i32>>> = Arc::new(RwLock::new(None));
     pub static ref LOGGER_HANDLE: Mutex<Option<Handle>> = Mutex::new(None);
 }
@@ -146,7 +153,7 @@ where
 }
 
 pub async fn init_signal() -> io::Result<()> {
-    let mut signals = Signals::new([SIGHUP])?;
+    let mut signals = Signals::new([SIGHUP, SIGINT])?;
 
     tokio::spawn(async move {
         for sig in signals.forever() {
@@ -182,6 +189,9 @@ pub async fn init_signal() -> io::Result<()> {
                 let mut t = T.write().await;
                 *t = None;
                 drop(t);
+            } else if sig == SIGINT {
+                fs::remove_file(PID_FILE.try_lock().unwrap().clone().unwrap().as_path()).unwrap();
+                process::exit(0);
             }
         }
     });

diff --git a/src/server.rs b/src/server.rs
@@ -1,14 +1,20 @@
 use crate::{
     config::{Config, ARGS, CONFIG, CONFIG_PATH, DEFAULT_CONFIG, DEFAULT_INTERVAL},
-    init::{init_cache, init_signal, DATE_FORMAT, FILE_CACHE, INDEX_CACHE, T},
+    init::{init_cache, init_signal, DATE_FORMAT, FILE_CACHE, INDEX_CACHE, PID_FILE, T},
     route::{location_index, mime_match, root_relative, status_page},
 };
 
 use anyhow::{Context, Result};
 use chrono::{DateTime, Utc};
 use mime::Mime;
 use std::{
-    collections::HashMap, env::set_current_dir, error::Error, ops::Deref, path::Path, sync::Arc,
+    collections::HashMap,
+    env::{self, set_current_dir},
+    error::Error,
+    fs,
+    ops::Deref,
+    path::Path,
+    sync::Arc,
 };
 
 #[cfg(feature = "log")]
@@ -313,12 +319,30 @@ pub async fn zest_main() -> Result<(), Box<dyn Error>> {
 
     set_current_dir(config.clone().server.root)?;
 
+    let runtime_dir = env::temp_dir();
+    let zest_pid = runtime_dir.join("zest.pid");
+    fs::create_dir_all(zest_pid.clone()).with_context(|| {
+        format!(
+            "failed to create dir {}",
+            zest_pid.as_path().to_str().unwrap()
+        )
+    })?;
+    let pid_file = zest_pid.clone().join(std::process::id().to_string());
+    *PID_FILE.try_lock().unwrap() = Some(pid_file.clone());
+
+    File::create(pid_file.clone()).await.with_context(|| {
+        format!(
+            "failed to create file {}",
+            pid_file.as_path().to_str().unwrap()
+        )
+    })?;
+
     #[cfg(feature = "log")]
     init_logger(&config.clone())
         .await
         .context("failed to init logger")?;
 
-    init_signal().await.context("failed to init signal")?;
+    init_signal().await.context("failed to init signal hook")?;
 
     #[cfg(feature = "lru_cache")]
     init_cache().await.context("failed to init lru cache")?;