update lru_cache

kfatyuip · Aug 1, 2024 · e2bba15 · e2bba15
1 parent 84098e3
commit e2bba15
Show file tree

Hide file tree

Showing 4 changed files with 119 additions and 58 deletions.
diff --git a/README.md b/README.md
@@ -21,6 +21,9 @@ server:
   info: "Powered by Rust"
   root: .
   error_page: 404.html # optional
+  cache: # optional
+    index_capacity: 16
+    file_capacity: 32
 
 allowlist: # optional
   - 127.0.0.1

diff --git a/src/config.rs b/src/config.rs
@@ -15,6 +15,28 @@ pub struct Config {
     pub logging: Option<LoggingConfig>,
 }
 
+impl Default for Config {
+    fn default() -> Self {
+        Config {
+            bind: BindConfig {
+                addr: "0.0.0.0".to_owned(),
+                listen: 8080,
+            },
+            server: ServerConfig {
+                info: "Powered by Rust".to_owned(),
+                root: current_dir().unwrap_or(".".into()),
+                error_page: Some("404.html".to_owned().into()),
+                cache: None,
+            },
+            allowlist: None,
+            blocklist: None,
+            rate_limit: None,
+            locations: None,
+            logging: None,
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize, Clone)]
 pub struct BindConfig {
     pub addr: String,
@@ -26,14 +48,30 @@ pub struct ServerConfig {
     pub info: String,
     pub root: PathBuf,
     pub error_page: Option<PathBuf>,
+    pub cache: Option<CacheConfig>,
+}
+
+#[derive(Serialize, Deserialize, Clone)]
+pub struct CacheConfig {
+    pub index_capacity: Option<usize>,
+    pub file_capacity: Option<usize>,
+}
+
+impl Default for CacheConfig {
+    fn default() -> Self {
+        CacheConfig {
+            index_capacity: Some(16),
+            file_capacity: Some(32),
+        }
+    }
 }
 
 #[derive(Serialize, Deserialize, Clone)]
 pub struct RateLimitConfig {
     pub max_requests: usize,
 }
 
-#[derive(Serialize, Deserialize, Clone)]
+#[derive(Serialize, Deserialize, Clone, Default)]
 pub struct LocationConfig {
     pub auto_index: Option<bool>,
     pub index: Option<PathBuf>,
@@ -45,27 +83,6 @@ pub struct LoggingConfig {
     pub error_log: Option<String>,
 }
 
-impl Default for Config {
-    fn default() -> Self {
-        Config {
-            bind: BindConfig {
-                addr: "0.0.0.0".to_owned(),
-                listen: 8080,
-            },
-            server: ServerConfig {
-                info: "Powered by Rust".to_owned(),
-                root: current_dir().unwrap_or(".".into()),
-                error_page: Some("404.html".to_owned().into()),
-            },
-            allowlist: None,
-            blocklist: None,
-            rate_limit: None,
-            locations: None,
-            logging: None,
-        }
-    }
-}
-
 #[derive(Parser)]
 #[command(version, about, long_about = None)]
 pub struct Args {

diff --git a/src/main.rs b/src/main.rs
@@ -1,5 +1,5 @@
 use zest::{
-    config::{Config, ARGS, CONFIG, CONFIG_PATH},
+    config::{ARGS, CONFIG, CONFIG_PATH},
     route::{location_index, mime_match, root_relative, status_page},
 };
 
@@ -10,7 +10,14 @@ use std::{
 };
 
 #[cfg(feature = "log")]
-use log::logger;
+use {
+    log::logger,
+    log4rs::{
+        append::{console::ConsoleAppender, file::FileAppender},
+        config::{Appender, Logger, Root},
+        encode::pattern::PatternEncoder,
+    },
+};
 
 #[cfg(target_os = "android")]
 use std::os::android::fs::MetadataExt;
@@ -40,8 +47,34 @@ const LOG_FORMAT: &str = "[{d(%Y-%m-%dT%H:%M:%SZ)} {h({l})}  zest] {m}\n";
 
 #[cfg(feature = "lru_cache")]
 lazy_static! {
-    static ref CACHE: Mutex<LruCache<String, String>> = {
-        let cache = LruCache::new(NonZeroUsize::new(64).unwrap());
+    static ref INDEX_CACHE: Mutex<LruCache<String, String>> = {
+        let cache = LruCache::new(
+            NonZeroUsize::new(
+                CONFIG
+                    .server
+                    .cache
+                    .clone()
+                    .unwrap_or_default()
+                    .index_capacity
+                    .unwrap_or_default(),
+            )
+            .unwrap(),
+        );
+        Mutex::new(cache)
+    };
+    static ref FILE_CACHE: Mutex<LruCache<String, Vec<u8>>> = {
+        let cache = LruCache::new(
+            NonZeroUsize::new(
+                CONFIG
+                    .server
+                    .cache
+                    .clone()
+                    .unwrap_or_default()
+                    .file_capacity
+                    .unwrap_or_default(),
+            )
+            .unwrap(),
+        );
         Mutex::new(cache)
     };
 }
@@ -152,15 +185,15 @@ where
             let mut html: String = String::new();
             #[cfg(feature = "lru_cache")]
             {
-                let mut cache = CACHE.lock().await;
+                let mut cache = INDEX_CACHE.lock().await;
                 if let Some(ctx) = cache.get(&location) {
                     html.clone_from(ctx);
                 } else if let Ok(index) = location_index(path, &location).await {
                     cache
-                        .push(location.clone(), index)
+                        .push(location.clone().into(), index)
                         .to_owned()
                         .unwrap_or_default();
-                    html.clone_from(cache.get(&location).unwrap());
+                    html.clone_from(cache.get(&location).unwrap().into());
                 } else {
                     response.status_code = 301;
                 }
@@ -181,6 +214,22 @@ where
                 Ok(f) => {
                     let mut file = f;
                     mime_type = mime_match(path.to_str().unwrap());
+
+                    #[cfg(feature = "lru_cache")]
+                    {
+                        let mut cache = FILE_CACHE.lock().await;
+                        if let Some(content) = cache.get(&location) {
+                            buffer = content.to_vec();
+                        } else {
+                            file.read_to_end(&mut buffer).await?;
+                            cache
+                                .push(location.clone(), buffer.clone())
+                                .to_owned()
+                                .unwrap_or_default();
+                        }
+                    }
+
+                    #[cfg(not(feature = "lru_cache"))]
                     file.read_to_end(&mut buffer).await?;
 
                     response.send_header(
@@ -215,15 +264,9 @@ where
 }
 
 #[inline]
-fn init_logger(config: Config) {
-    use log4rs::{
-        append::{console::ConsoleAppender, file::FileAppender},
-        config::{Appender, Logger, Root},
-        encode::pattern::PatternEncoder,
-        Config,
-    };
-
-    let mut builder = Config::builder();
+#[cfg(feature = "log")]
+fn init_logger(config: zest::config::Config) {
+    let mut builder = log4rs::Config::builder();
 
     let stdout = ConsoleAppender::builder()
         .encoder(Box::new(PatternEncoder::new(LOG_FORMAT)))

diff --git a/src/route.rs b/src/route.rs
@@ -17,28 +17,26 @@ pub fn root_relative(p: &str) -> &str {
 pub async fn location_index(path: PathBuf, location: &str) -> Result<String> {
     let config = CONFIG.deref();
 
-    if let Some(locations) = &config.locations {
-        for (s, v) in locations {
-            if root_relative(s) == location.trim_end_matches('/') {
-                match from_value::<LocationConfig>(v.clone()) {
-                    Ok(_location) => {
-                        if let Some(index) = _location.index {
-                            return fs::read_to_string(root_relative(
-                                PathBuf::from(location)
-                                    .join(index.clone())
-                                    .as_path()
-                                    .to_str()
-                                    .unwrap(),
-                            ))
-                            .await;
-                        } else if _location.auto_index.is_none() || !_location.auto_index.unwrap() {
-                            return Err(ErrorKind::Unsupported.into());
-                        }
-                    }
-                    _ => {
-                        continue;
+    for (s, v) in &config.locations.clone().unwrap_or_default() {
+        if root_relative(s) == location.trim_end_matches('/') {
+            match from_value::<LocationConfig>(v.clone()) {
+                Ok(_location) => {
+                    if let Some(index) = _location.index {
+                        return fs::read_to_string(root_relative(
+                            PathBuf::from(location)
+                                .join(index.clone())
+                                .as_path()
+                                .to_str()
+                                .unwrap(),
+                        ))
+                        .await;
+                    } else if _location.auto_index.is_none() || !_location.auto_index.unwrap() {
+                        return Err(ErrorKind::Unsupported.into());
                     }
                 }
+                _ => {
+                    continue;
+                }
             }
         }
     }