fix(cache): reuse decrypted db across wal-only updates (#58)

src/daemon/cache.rs

@@ -30,6 +30,7 @@ struct CacheEntry {
 pub struct DbCache {
     db_dir: PathBuf,
     cache_dir: PathBuf,
+    mtime_file: PathBuf,
     all_keys: HashMap<String, String>, // rel_key -> enc_key(hex)
     inner: Arc<Mutex<HashMap<String, CacheEntry>>>,
 }
@@ -39,15 +40,24 @@ impl DbCache {
         db_dir: PathBuf,
         all_keys: HashMap<String, String>,
     ) -> Result<Self> {
-        let cache_dir = config::cache_dir();
+        Self::with_dirs(db_dir, config::cache_dir(), config::mtime_file(), all_keys).await
+    }
+
+    /// 注入 `cache_dir` / `mtime_file`（测试用 + 生产 `new()` 复用）
+    pub(crate) async fn with_dirs(
+        db_dir: PathBuf,
+        cache_dir: PathBuf,
+        mtime_file: PathBuf,
+        all_keys: HashMap<String, String>,
+    ) -> Result<Self> {
         tokio::fs::create_dir_all(&cache_dir).await?;
 
-        let inner: HashMap<String, CacheEntry> = HashMap::new();
         let cache = DbCache {
             db_dir,
             cache_dir,
+            mtime_file,
             all_keys,
-            inner: Arc::new(Mutex::new(inner)),
+            inner: Arc::new(Mutex::new(HashMap::new())),
         };
 
         cache.load_persistent().await;
@@ -67,7 +77,7 @@ impl DbCache {
 
     /// 从持久化文件加载 mtime 记录，复用未过期的解密文件
     async fn load_persistent(&self) {
-        let mtime_file = config::mtime_file();
+        let mtime_file = &self.mtime_file;
         let content = match tokio::fs::read_to_string(&mtime_file).await {
             Ok(c) => c,
             Err(_) => return,
@@ -88,12 +98,17 @@ impl DbCache {
             let wal_path = wal_path_for(&db_path);
 
             let db_mt = mtime_nanos(&db_path);
-            let wal_mt = if wal_path.exists() { mtime_nanos(&wal_path) } else { 0 };
+            let _wal_mt = if wal_path.exists() { mtime_nanos(&wal_path) } else { 0 };
 
-            if db_mt == entry.db_mt && wal_mt == entry.wal_mt {
+            // 只要主 .db 没变，就把 cached 产物载回来。
+            // 如果 WAL mtime 变了，后续 `get()` 会自动走 Path 2：在已有 cached DB 上增量 apply_wal，
+            // 而不是 daemon 重启后第一条请求又退回全量解密。
+            if db_mt == entry.db_mt {
                 inner.insert(rel_key.clone(), CacheEntry {
                     db_mtime: db_mt,
-                    wal_mtime: wal_mt,
+                    // 保留"cached 产物构建时看到的 wal_mtime"，让 `get()` 去比较当前 WAL
+                    // 是否发生了变化，从而决定 exact-hit 还是 WAL 增量。
+                    wal_mtime: entry.wal_mt,
                     decrypted_path: dec_path,
                 });
                 reused += 1;
@@ -106,7 +121,7 @@ impl DbCache {
 
     /// 持久化 mtime 记录
     async fn save_persistent(&self) {
-        let mtime_file = config::mtime_file();
+        let mtime_file = &self.mtime_file;
         let inner = self.inner.lock().await;
         let data: HashMap<String, MtimeEntry> = inner.iter().map(|(k, v)| {
             (k.clone(), MtimeEntry {
@@ -124,7 +139,14 @@ impl DbCache {
 
     /// 获取解密后的数据库路径
     ///
-    /// 如果 mtime 未变，直接返回缓存路径；否则重新解密
+    /// 三种命中路径：
+    /// 1. 主 `.db` 和 WAL mtime 都未变 → 直接返回缓存路径
+    /// 2. 主 `.db` 未变、WAL mtime 变了 → 在已有 cached 产物上**增量** `apply_wal`
+    ///    （apply_wal 是幂等的：旧帧 redo 同样的 page 写入，新帧追加生效；不重新 full_decrypt）
+    /// 3. 主 `.db` mtime 变了 → 重新 `full_decrypt` + `apply_wal`
+    ///
+    /// WeChat 在写消息时只 append WAL（除非触发 checkpoint），因此 path 2 是常态；
+    /// 这条路径把"每次请求都全量解密 ~1.8GB DB（~120s）"压到"只解 WAL 帧（典型 < 10s）"。
     pub async fn get(&self, rel_key: &str) -> Result<Option<PathBuf>> {
         let enc_key_hex = match self.all_keys.get(rel_key) {
             Some(k) => k.clone(),
@@ -140,28 +162,53 @@ impl DbCache {
         }
 
         let wal_path = wal_path_for(&db_path);
-
         let db_mt = mtime_nanos(&db_path);
         let wal_mt = if wal_path.exists() { mtime_nanos(&wal_path) } else { 0 };
 
-        // 检查缓存
+        let cached = {
-        {
             let inner = self.inner.lock().await;
-            if let Some(entry) = inner.get(rel_key) {
+            inner.get(rel_key).cloned()
-                if entry.db_mtime == db_mt
+        };
-                    && entry.wal_mtime == wal_mt
-                    && entry.decrypted_path.exists()
-                {
-                    return Ok(Some(entry.decrypted_path.clone()));
-                }
-            }
-        }
 
-        // 需要重新解密
-        let out_path = self.cache_file_path(rel_key);
         let enc_key_bytes = hex_to_32bytes(&enc_key_hex)
             .with_context(|| format!("密钥格式错误: {}", rel_key))?;
 
+        // Path 1 / Path 2：主 .db mtime 未变且 cached 产物仍在
+        if let Some(entry) = cached.as_ref() {
+            if entry.db_mtime == db_mt && entry.decrypted_path.exists() {
+                if entry.wal_mtime == wal_mt {
+                    return Ok(Some(entry.decrypted_path.clone()));
+                }
+
+                // Path 2: WAL-only 变化 → 在 cached 产物上重新 apply_wal
+                // 不存在的 WAL 也要更新 wal_mtime=0（虽然 SQLite 不会自发"主库不变 + WAL 清空"）
+                let out_path = entry.decrypted_path.clone();
+                let t0 = std::time::Instant::now();
+                if wal_path.exists() {
+                    let out_path2 = out_path.clone();
+                    let wal_path2 = wal_path.clone();
+                    let key_copy = enc_key_bytes;
+                    tokio::task::spawn_blocking(move || {
+                        wal::apply_wal(&wal_path2, &out_path2, &key_copy)
+                    }).await??;
+                }
+                eprintln!("[cache] WAL 增量 {} ({}ms)", rel_key, t0.elapsed().as_millis());
+
+                {
+                    let mut inner = self.inner.lock().await;
+                    inner.insert(rel_key.to_string(), CacheEntry {
+                        db_mtime: db_mt,
+                        wal_mtime: wal_mt,
+                        decrypted_path: out_path.clone(),
+                    });
+                }
+                self.save_persistent().await;
+                return Ok(Some(out_path));
+            }
+        }
+
+        // Path 3: 主 .db 变了 / 缓存 miss → 全量解密
+        let out_path = self.cache_file_path(rel_key);
         let t0 = std::time::Instant::now();
         let db_path2 = db_path.clone();
         let out_path2 = out_path.clone();
@@ -170,7 +217,6 @@ impl DbCache {
             crypto::full_decrypt(&db_path2, &out_path2, &key_copy)
         }).await??;
 
-        // 应用 WAL
         if wal_path.exists() {
             let out_path3 = out_path.clone();
             let wal_path3 = wal_path.clone();
@@ -180,10 +226,8 @@ impl DbCache {
             }).await??;
         }
 
-        let elapsed_ms = t0.elapsed().as_millis();
+        eprintln!("[cache] 全量解密 {} ({}ms)", rel_key, t0.elapsed().as_millis());
-        eprintln!("[cache] 解密 {} ({}ms)", rel_key, elapsed_ms);
 
-        // 更新内存缓存
         {
             let mut inner = self.inner.lock().await;
             inner.insert(rel_key.to_string(), CacheEntry {
@@ -223,3 +267,231 @@ fn hex_to_32bytes(s: &str) -> Result<[u8; 32]> {
     }
     Ok(out)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// 64 字符 hex（不需要是真 SQLCipher key — 仅用来证明"是否触发了 full_decrypt"）
+    const FAKE_KEY_HEX: &str =
+        "0000000000000000000000000000000000000000000000000000000000000000";
+
+    /// 路径区分约定：
+    /// - 完全 hit / WAL 增量 → `decrypted_path` **内容不变**
+    /// - 全量解密 → `crypto::full_decrypt` 把 cached file **重写为 PAGE_SZ 倍数**
+    ///   （fake key 解出 4096 字节垃圾，但仍写入 — 不验证内容合法性）
+    /// 因此用 cached file 的"size 是否被改"来判断走了哪条路径。
+    const ORIGINAL_CACHED_BYTES: &[u8] = b"original cached contents";
+
+    fn unique_tmpdir(tag: &str) -> PathBuf {
+        let pid = std::process::id();
+        let nanos = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        let p = std::env::temp_dir().join(format!("wx-cli-cache-test-{}-{}-{}", tag, pid, nanos));
+        std::fs::create_dir_all(&p).unwrap();
+        p
+    }
+
+    /// 准备一份 "DbCache 已经 reuse 了 cached 解密产物" 的初始状态。
+    /// 返回 (cache, db_path, decrypted_path, mtime_file, rel_key)。
+    async fn setup_seeded_cache(tag: &str) -> (DbCache, PathBuf, PathBuf, PathBuf, String) {
+        let root = unique_tmpdir(tag);
+        let db_dir = root.join("db_storage");
+        let cache_dir = root.join("cache");
+        std::fs::create_dir_all(&db_dir).unwrap();
+        std::fs::create_dir_all(&cache_dir).unwrap();
+
+        let rel_key = "message_0.db".to_string();
+        let db_path = db_dir.join(&rel_key);
+        std::fs::write(&db_path, b"fake encrypted db").unwrap();
+
+        let cached_hash = format!("{:x}", md5::compute(rel_key.as_bytes()));
+        let decrypted_path = cache_dir.join(format!("{}.db", cached_hash));
+        std::fs::write(&decrypted_path, ORIGINAL_CACHED_BYTES).unwrap();
+
+        let db_mt = mtime_nanos(&db_path);
+        let mtime_file = cache_dir.join("_mtimes.json");
+        let payload = serde_json::to_string(&serde_json::json!({
+            &rel_key: {
+                "db_mt": db_mt,
+                "wal_mt": 0u64,
+                "path": decrypted_path.display().to_string(),
+            }
+        }))
+        .unwrap();
+        std::fs::write(&mtime_file, payload).unwrap();
+
+        let mut all_keys = HashMap::new();
+        all_keys.insert(rel_key.clone(), FAKE_KEY_HEX.to_string());
+        let cache = DbCache::with_dirs(db_dir, cache_dir, mtime_file.clone(), all_keys)
+            .await
+            .unwrap();
+
+        (cache, db_path, decrypted_path, mtime_file, rel_key)
+    }
+
+    #[tokio::test]
+    async fn exact_mtime_hit_skips_decrypt() {
+        let (cache, _db_path, decrypted_path, _mtime_file, rel_key) =
+            setup_seeded_cache("exact").await;
+
+        let p = cache.get(&rel_key).await.unwrap().expect("cache should hit");
+        assert_eq!(p, decrypted_path);
+
+        // 完全 hit → cached file 内容不应被改
+        let body = std::fs::read(&decrypted_path).unwrap();
+        assert_eq!(body, ORIGINAL_CACHED_BYTES);
+    }
+
+    #[tokio::test]
+    async fn wal_only_change_uses_incremental_path() {
+        // 自己构造（不走 setup_seeded_cache）以便初始 mtime.json 同时写 db_mt 和 wal_mt
+        let root = unique_tmpdir("walonly");
+        let db_dir = root.join("db_storage");
+        let cache_dir = root.join("cache");
+        std::fs::create_dir_all(&db_dir).unwrap();
+        std::fs::create_dir_all(&cache_dir).unwrap();
+
+        let rel_key = "message_0.db".to_string();
+        let db_path = db_dir.join(&rel_key);
+        std::fs::write(&db_path, b"fake encrypted db").unwrap();
+
+        let wal_path = wal_path_for(&db_path);
+        std::fs::write(&wal_path, [0u8; 31]).unwrap(); // ≤ WAL_HDR_SZ=32 → apply_wal noop
+
+        let cached_hash = format!("{:x}", md5::compute(rel_key.as_bytes()));
+        let decrypted_path = cache_dir.join(format!("{}.db", cached_hash));
+        std::fs::write(&decrypted_path, ORIGINAL_CACHED_BYTES).unwrap();
+
+        let db_mt = mtime_nanos(&db_path);
+        let wal_mt0 = mtime_nanos(&wal_path);
+        let mtime_file = cache_dir.join("_mtimes.json");
+        let payload = serde_json::to_string(&serde_json::json!({
+            &rel_key: {
+                "db_mt": db_mt,
+                "wal_mt": wal_mt0,
+                "path": decrypted_path.display().to_string(),
+            }
+        }))
+        .unwrap();
+        std::fs::write(&mtime_file, payload).unwrap();
+
+        let mut all_keys = HashMap::new();
+        all_keys.insert(rel_key.clone(), FAKE_KEY_HEX.to_string());
+        let cache = DbCache::with_dirs(db_dir, cache_dir, mtime_file, all_keys)
+            .await
+            .unwrap();
+
+        // 第一次：完全 hit
+        let p1 = cache.get(&rel_key).await.unwrap().expect("first get hits");
+        assert_eq!(p1, decrypted_path);
+        assert_eq!(std::fs::read(&decrypted_path).unwrap(), ORIGINAL_CACHED_BYTES);
+
+        // bump WAL mtime（重写仍 31 bytes，apply_wal 仍 noop）
+        std::thread::sleep(std::time::Duration::from_millis(20));
+        std::fs::write(&wal_path, [0xffu8; 31]).unwrap();
+        let wal_mt1 = mtime_nanos(&wal_path);
+        assert_ne!(wal_mt0, wal_mt1, "rewriting WAL should bump mtime");
+
+        // 第二次：WAL 增量路径
+        // 如果错误地走 full_decrypt → cached file 大小会被重写为 ≥ PAGE_SZ
+        let p2 = cache
+            .get(&rel_key)
+            .await
+            .unwrap()
+            .expect("WAL-incremental path should produce path");
+        assert_eq!(p2, decrypted_path);
+
+        let body = std::fs::read(&decrypted_path).unwrap();
+        assert_eq!(
+            body, ORIGINAL_CACHED_BYTES,
+            "WAL-incremental should NOT rewrite cached file"
+        );
+    }
+
+    #[tokio::test]
+    async fn db_mtime_change_triggers_full_decrypt() {
+        let (cache, db_path, decrypted_path, _mtime_file, rel_key) =
+            setup_seeded_cache("dbchange").await;
+
+        // bump 主 .db 的 mtime（重写一份不同 bytes）
+        std::thread::sleep(std::time::Duration::from_millis(20));
+        std::fs::write(&db_path, b"different fake encrypted bytes").unwrap();
+        assert_ne!(
+            mtime_nanos(&db_path),
+            cache.inner.lock().await.get(&rel_key).unwrap().db_mtime,
+            "rewriting db file should bump mtime"
+        );
+
+        // 走 full_decrypt 路径 → fake key 不会让 full_decrypt 失败（它不验证内容），
+        // 但会把 cached file 重写为 PAGE_SZ 倍数。原始内容是 24 bytes，重写后应该 ≥ 4096 bytes。
+        let p = cache
+            .get(&rel_key)
+            .await
+            .unwrap()
+            .expect("cache should produce path");
+        assert_eq!(p, decrypted_path);
+
+        let new_size = std::fs::metadata(&decrypted_path).unwrap().len() as usize;
+        assert!(
+            new_size >= crate::crypto::PAGE_SZ,
+            "expected full_decrypt to rewrite cached file to PAGE_SZ multiple, got size={}",
+            new_size,
+        );
+    }
+
+    #[tokio::test]
+    async fn restart_with_wal_change_still_reuses_cached_db_then_applies_wal() {
+        let root = unique_tmpdir("restart-wal");
+        let db_dir = root.join("db_storage");
+        let cache_dir = root.join("cache");
+        std::fs::create_dir_all(&db_dir).unwrap();
+        std::fs::create_dir_all(&cache_dir).unwrap();
+
+        let rel_key = "message_0.db".to_string();
+        let db_path = db_dir.join(&rel_key);
+        std::fs::write(&db_path, b"fake encrypted db").unwrap();
+
+        let wal_path = wal_path_for(&db_path);
+        std::fs::write(&wal_path, [0u8; 31]).unwrap(); // WAL 增量仍是 noop
+
+        let cached_hash = format!("{:x}", md5::compute(rel_key.as_bytes()));
+        let decrypted_path = cache_dir.join(format!("{}.db", cached_hash));
+        std::fs::write(&decrypted_path, ORIGINAL_CACHED_BYTES).unwrap();
+
+        let db_mt = mtime_nanos(&db_path);
+        let wal_mt0 = mtime_nanos(&wal_path);
+        let mtime_file = cache_dir.join("_mtimes.json");
+        let payload = serde_json::to_string(&serde_json::json!({
+            &rel_key: {
+                "db_mt": db_mt,
+                "wal_mt": wal_mt0,
+                "path": decrypted_path.display().to_string(),
+            }
+        }))
+        .unwrap();
+        std::fs::write(&mtime_file, payload).unwrap();
+
+        // 模拟 daemon 重启前又有新消息写入 WAL
+        std::thread::sleep(std::time::Duration::from_millis(20));
+        std::fs::write(&wal_path, [0xffu8; 31]).unwrap();
+        let wal_mt1 = mtime_nanos(&wal_path);
+        assert_ne!(wal_mt0, wal_mt1);
+
+        let mut all_keys = HashMap::new();
+        all_keys.insert(rel_key.clone(), FAKE_KEY_HEX.to_string());
+        let cache = DbCache::with_dirs(db_dir, cache_dir, mtime_file, all_keys)
+            .await
+            .unwrap();
+
+        let p = cache.get(&rel_key).await.unwrap().expect("cache should reuse persisted DB");
+        assert_eq!(p, decrypted_path);
+        let body = std::fs::read(&decrypted_path).unwrap();
+        assert_eq!(
+            body, ORIGINAL_CACHED_BYTES,
+            "restart + WAL-only change should still reuse cached DB and avoid full_decrypt"
+        );
+    }
+}