world: put cache option

2 files changed, 44 insertions, 13 deletions

diff --git a/shared/src/store.rs b/shared/src/store.rs
index 548a402..ddcee2c 100644
--- a/shared/src/store.rs
+++ b/shared/src/store.rs
@@ -37,13 +37,15 @@ pub enum ResourceStore {
 }
 impl ResourceStore {
     pub fn new_env() -> Result<Self> {
-        if var("WEARECHAT_RES_CACHE_USE_REDB").is_ok() {
-            ResourceStore::new_redb(
+        if var("WEARECHAT_RES_CACHE_USE_MEMORY").is_ok() {
+            Ok(Self::new_memory())
+        } else if var("WEARECHAT_RES_CACHE_USE_REDB").is_ok() {
+            Self::new_redb(
                 &xdg::BaseDirectories::with_prefix("wearechat")?
                     .place_cache_file("resources.db")?,
             )
         } else {
-            ResourceStore::new_filesystem(
+            Self::new_filesystem(
                 &xdg::BaseDirectories::with_prefix("wearechat")?
                     .create_cache_directory("resources")?,
             )
@@ -112,19 +114,38 @@ impl ResourceStore {
             }
             ResourceStore::Filesystem(root) => {
                 let path = fs_cache_path(&root, key);
-                let path_temp = path.with_extension("part");
-                File::create(&path_temp)?.write_all(value)?;
-                rename(path_temp, path)?;
+                if !path.exists() {
+                    let path_temp = path.with_extension("part");
+                    File::create(&path_temp)?.write_all(value)?;
+                    rename(path_temp, path)?;
+                }
             }
         }
         Ok(key)
     }
-    pub fn iter(&self, mut cb: impl FnMut(Resource, &[u8])) -> Result<()> {
+    pub fn iter(&self, mut cb: impl FnMut(Resource, usize)) -> Result<()> {
         match self {
             ResourceStore::Redb(_database) => todo!(),
-            ResourceStore::Filesystem(_root) => todo!(),
+            ResourceStore::Filesystem(_root) => {
+                for e in _root.read_dir()? {
+                    let e = e?;
+                    let k = e.file_name();
+                    let k = k.to_string_lossy();
+                    let m = e.metadata()?.len();
+                    let mut r = [0; 32];
+                    for i in 0..32 {
+                        r[i] = u8::from_str_radix(&k[i * 2..(i + 1) * 2], 16)?
+                    }
+                    cb(Resource(r, PhantomData), m as usize)
+                }
+                Ok(())
+            }
             ResourceStore::Memory(mutex) => {
-                mutex.lock().unwrap().iter().for_each(|(k, v)| cb(*k, v));
+                mutex
+                    .lock()
+                    .unwrap()
+                    .iter()
+                    .for_each(|(k, v)| cb(*k, v.len()));
                 Ok(())
             }
         }
diff --git a/world/src/main.rs b/world/src/main.rs
index b5ce1a4..d46f4d1 100644
--- a/world/src/main.rs
+++ b/world/src/main.rs
@@ -81,6 +81,9 @@ pub struct Args {
     dry_run: bool,
 
     #[arg(long)]
+    put_cache: bool,
+
+    #[arg(long)]
     debug_light: bool,
 }
 
@@ -88,7 +91,11 @@ fn main() -> Result<()> {
     env_logger::init_from_env("LOG");
     let args = Args::parse();
 
-    let store = ResourceStore::new_memory();
+    let store = if args.put_cache {
+        ResourceStore::new_env()?
+    } else {
+        ResourceStore::new_memory()
+    };
 
     let mut prefabs = Vec::new();
     let texture_cache = Arc::new(Mutex::new(HashMap::new()));
@@ -249,7 +256,7 @@ fn main() -> Result<()> {
     }
 
     let mut size = 0;
-    store.iter(|_k, d| size += d.len()).unwrap();
+    store.iter(|_k, len| size += len).unwrap();
     info!(
         "prefab has network size of {}",
         humansize::format_size(size, BINARY)
@@ -293,9 +300,12 @@ fn main() -> Result<()> {
         });
     }
 
+    if args.put_cache {
+        return Ok(());
+    }
     if args.push {
-        store.iter(|k, v| {
-            Packet::RespondResource(k, Data(v.to_vec()))
+        store.iter(|k, _| {
+            Packet::RespondResource(k, Data(store.get_raw(k).unwrap().unwrap()))
                 .write(&mut sock)
                 .unwrap();
         })?;