feat: persist cache using SQLite & support passing additional args

Cargo.toml

@@ -48,6 +48,11 @@ serde_yml = "0.0.12"
 twox-hash = "2.1.1"
 vite_workspace = { path = "crates/vite_workspace" }
 wax = "0.6.0"
+bstr = "1.12.0"
+clap = "4.5.40"
+edit = "0.1.5"
+rusqlite = "0.36.0"
+shell-escape = "0.1.5"
 
 [profile.dev]
 # Disabling debug info speeds up local and CI builds,

crates/vite_task/Cargo.toml

@@ -14,15 +14,20 @@ workspace = true
 [dependencies]
 anyhow = { workspace = true }
 bincode = { workspace = true, features = ["derive"] }
+bstr = { workspace = true }
+clap = { workspace = true, features = ["derive"] }
 compact_str = { workspace = true, features = ["serde"] }
 dashmap = { workspace = true }
 diff-struct = { workspace = true }
+edit = { workspace = true }
 itertools = { workspace = true }
 petgraph = { workspace = true }
 rayon = { workspace = true }
 relative-path = { workspace = true }
+rusqlite = { workspace = true, features = ["bundled"] }
 serde = { workspace = true, features = ["derive", "rc"] }
 serde_json = { workspace = true }
+shell-escape = { workspace = true }
 twox-hash = { workspace = true }
 vite_workspace = { workspace = true }
 wax = { workspace = true }

crates/vite_task/src/cache.rs

@@ -1,20 +1,18 @@
-use crate::collections::HashMap;
-use std::fs::{self, File};
-use std::io::{BufReader, BufWriter};
-use std::path::{Path, PathBuf};
-use std::sync::Arc;
+use std::path::Path;
+use std::sync::{Arc, Mutex};
 
 // use bincode::config::{Configuration, standard};
-use bincode::{Decode, Encode};
-use serde::{Deserialize, Serialize};
+use bincode::{Decode, Encode, decode_from_slice, encode_to_vec};
+use rusqlite::{Connection, OptionalExtension as _};
+use serde::Serialize;
 
 use crate::config::ResolvedTask;
 use crate::execute::{ExecutedTask, StdOutput};
 use crate::fingerprint::{FingerprintMismatch, TaskFingerprint};
 use crate::fs::FileSystem;
 use crate::str::Str;
 
-#[derive(Debug, Encode, Decode, Serialize, Deserialize)]
+#[derive(Debug, Encode, Decode, Serialize)]
 pub struct CachedTask {
     pub fingerprint: TaskFingerprint,
     pub std_outputs: Arc<[StdOutput]>,
@@ -33,11 +31,16 @@ impl CachedTask {
 }
 
 pub struct TaskCache {
-    cached_tasks_by_name: HashMap<Str, CachedTask>,
-    path: PathBuf,
+    conn: Mutex<Connection>,
 }
 
-// const BINCODE_CONFIG: Configuration = standard();
+#[derive(Debug, Hash, Encode, Decode, Serialize)]
+pub struct TaskCacheKey {
+    pub task_name: Str,
+    pub args: Arc<[Str]>,
+}
+
+const BINCODE_CONFIG: bincode::config::Configuration = bincode::config::standard();
 
 #[derive(Debug)]
 pub enum CacheMiss {
@@ -48,55 +51,95 @@ pub enum CacheMiss {
 impl TaskCache {
     pub fn load_from_file(path: impl AsRef<Path>) -> anyhow::Result<Self> {
         let path = path.as_ref();
-        let cached_tasks_by_name: HashMap<Str, CachedTask> = match File::open(path) {
-            Ok(file) => {
-                let reader = BufReader::new(file);
-                // Using json for easy debugging
-                // Will switch to bincode for better performance
-                serde_json::from_reader(reader)?
-                // bincode::decode_from_std_read(&mut reader, BINCODE_CONFIG)?
-            }
-            Err(err) => {
-                if err.kind() == std::io::ErrorKind::NotFound {
-                    HashMap::new()
-                } else {
-                    return Err(err.into());
+        let conn = Connection::open(path)?;
+        conn.execute_batch("PRAGMA journal_mode=WAL")?;
+        conn.execute("BEGIN EXCLUSIVE", ())?;
+        loop {
+            let user_version: u32 = conn.query_one("PRAGMA user_version", (), |row| row.get(0))?;
+            match user_version {
+                0 => {
+                    // fresh new db
+                    conn.execute("CREATE TABLE tasks (key BLOB PRIMARY KEY, value BLOB);", ())?;
+                    conn.execute("PRAGMA user_version = 1", ())?;
                 }
+                // Migration done here
+                1 => break,
+                2.. => anyhow::bail!("Unrecognized cache db version: {user_version}"),
             }
-        };
-        Ok(Self { cached_tasks_by_name, path: path.to_path_buf() })
-    }
-    pub fn save(&self) -> anyhow::Result<()> {
-        let path = self.path.as_path();
-        if let Some(parent) = path.parent() {
-            fs::create_dir_all(parent)?;
         }
-        let file = File::create(path)?;
-        let mut writer = BufWriter::new(file);
-        serde_json::to_writer_pretty(&mut writer, &self.cached_tasks_by_name)?;
-        // bincode::encode_into_std_write(&self.cached_tasks_by_name, &mut writer, BINCODE_CONFIG)?;
-        writer.into_inner()?.sync_data()?;
+        Ok(Self { conn: Mutex::new(conn) })
+    }
+    pub fn save(self) -> anyhow::Result<()> {
+        let conn = self.conn.into_inner().unwrap();
+        conn.execute("COMMIT", ())?;
         Ok(())
     }
 
-    pub fn update(&mut self, task_name: Str, cached_task: CachedTask) -> anyhow::Result<()> {
-        self.cached_tasks_by_name.insert(task_name, cached_task);
+    pub fn update(
+        &mut self,
+        task_name: Str,
+        args: Arc<[Str]>,
+        cached_task: CachedTask,
+    ) -> anyhow::Result<()> {
+        let key = TaskCacheKey { task_name, args };
+        let conn = self.conn.lock().unwrap();
+        let key_blob = encode_to_vec(&key, BINCODE_CONFIG)?;
+        let value_blob = encode_to_vec(&cached_task, BINCODE_CONFIG)?;
+        let mut update_stmt = conn.prepare_cached(
+            "INSERT INTO tasks (key, value) VALUES (?1, ?2) ON CONFLICT(key) DO UPDATE SET value=?2"
+        )?;
+        update_stmt.execute([key_blob, value_blob])?;
+        Ok(())
+    }
+
+    pub fn get_cache(
+        &self,
+        task_name: Str,
+        args: Arc<[Str]>,
+    ) -> anyhow::Result<Option<CachedTask>> {
+        let conn = self.conn.lock().unwrap();
+        let mut select_stmt = conn.prepare_cached("SELECT value FROM tasks WHERE key=?")?;
+        let key_blob = encode_to_vec(&TaskCacheKey { task_name, args }, BINCODE_CONFIG)?;
+        let Some(value_blob) =
+            select_stmt.query_row::<Vec<u8>, _, _>([key_blob], |row| row.get(0)).optional()?
+        else {
+            return Ok(None);
+        };
+        let (cached_task, _) = decode_from_slice::<CachedTask, _>(&value_blob, BINCODE_CONFIG)?;
+        Ok(Some(cached_task))
+    }
+
+    pub fn list_cache(
+        &self,
+        mut f: impl FnMut(TaskCacheKey, CachedTask) -> anyhow::Result<()>,
+    ) -> anyhow::Result<()> {
+        let conn = self.conn.lock().unwrap();
+        let mut select_stmt = conn.prepare_cached("SELECT key, value FROM tasks")?;
+        let cache_list = select_stmt.query_and_then((), |row| {
+            let key_blob: Vec<u8> = row.get(0)?;
+            let value_blob: Vec<u8> = row.get(1)?;
+            let (key, _) = decode_from_slice::<TaskCacheKey, _>(&key_blob, BINCODE_CONFIG)?;
+            let (cached_task, _) = decode_from_slice::<CachedTask, _>(&value_blob, BINCODE_CONFIG)?;
+            anyhow::Ok((key, cached_task))
+        })?;
+        for cache in cache_list {
+            let (key, cached_task) = cache?;
+            f(key, cached_task)?;
+        }
         Ok(())
     }
 
     /// Tries to get the task cache if the fingerprint matches, otherwise returns why the cache misses
-    pub fn try_hit<'me>(
-        &'me self,
+    pub fn try_hit(
+        &self,
         task: &ResolvedTask,
         fs: &impl FileSystem,
         base_dir: &Path,
-    ) -> anyhow::Result<Result<&'me CachedTask, CacheMiss>> {
-        let Some(cached_task) = self.cached_tasks_by_name.get(&task.name) else {
+    ) -> anyhow::Result<Result<CachedTask, CacheMiss>> {
+        let Some(cached_task) = self.get_cache(task.name.clone(), task.args.clone())? else {
             return Ok(Err(CacheMiss::NotFound));
         };
-        if let Some(fingerprint_mismatch) =
-            cached_task.fingerprint.validate(&task.config, fs, base_dir)?
-        {
+        if let Some(fingerprint_mismatch) = cached_task.fingerprint.validate(task, fs, base_dir)? {
             return Ok(Err(CacheMiss::FingerprintMismatch(fingerprint_mismatch)));
         }
         Ok(Ok(cached_task))

crates/vite_task/src/config.rs

@@ -1,7 +1,9 @@
 use std::{
+    collections::BTreeSet,
+    ffi::OsStr,
     fs::File,
     io::BufReader,
-    iter::once,
+    iter::{self},
     path::{Path, PathBuf},
     sync::Arc,
 };
@@ -16,9 +18,10 @@ use crate::{
 use anyhow::Context;
 
 use bincode::{Decode, Encode};
-use compact_str::CompactString;
 use diff::Diff;
+use itertools::Itertools;
 use petgraph::{graph::NodeIndex, stable_graph::StableDiGraph};
+use relative_path::RelativePath;
 use serde::{Deserialize, Serialize};
 use vite_workspace::PackageInfo;
 
@@ -66,9 +69,52 @@ pub struct Workspace {
 #[derive(Debug)]
 pub struct ResolvedTask {
     pub name: Str,
+    pub args: Arc<[Str]>,
+    pub resolved_config: ResolvedTaskConfig,
+    pub resolved_command: ResolvedTaskCommand,
+}
+
+#[derive(Encode, Decode, Debug, Serialize, PartialEq, Eq, Diff)]
+#[diff(attr(#[derive(Debug)]))]
+pub struct ResolvedTaskConfig {
     pub config_dir: Str,
     pub config: TaskConfig,
-    pub envs: TaskEnvs,
+}
+
+impl ResolvedTaskConfig {
+    fn resolve_command(&self, task_args: &[Str]) -> anyhow::Result<ResolvedTaskCommand> {
+        let cwd = RelativePath::new(&self.config_dir).join(self.config.cwd.as_str());
+        let command_line = iter::once(self.config.command.clone())
+            .chain(
+                task_args
+                    .iter()
+                    .map(|arg| shell_escape::escape(arg.as_str().into()).as_ref().into()),
+            )
+            .join(" ");
+        let task_envs = TaskEnvs::resolve(&self.config)?;
+        Ok(ResolvedTaskCommand {
+            fingerprint: CommandFingerprint {
+                cwd: cwd.as_str().into(),
+                command_line: command_line.as_str().into(),
+                envs_without_pass_through: task_envs.envs_without_pass_through,
+            },
+            all_envs: task_envs.all_envs,
+        })
+    }
+}
+
+#[derive(Debug)]
+pub struct ResolvedTaskCommand {
+    pub fingerprint: CommandFingerprint,
+    pub all_envs: HashMap<Str, Arc<OsStr>>,
+}
+
+#[derive(Encode, Decode, Debug, Serialize, PartialEq, Eq, Diff)]
+#[diff(attr(#[derive(Debug)]))]
+pub struct CommandFingerprint {
+    pub cwd: Str,
+    pub command_line: Str,
+    pub envs_without_pass_through: HashMap<Str, Str>,
 }
 
 impl Workspace {
@@ -91,30 +137,33 @@ impl Workspace {
                     Err(err) => {
                         if err.kind() == std::io::ErrorKind::NotFound {
                             continue;
-                        } else {
-                            return Err(err.into());
                         }
+                        return Err(err.into());
                     }
                 }))?;
             vite_task_jsons.push((vite_task_json, pkg));
         }
 
-        let cache_path = dir.join("node_modules/.vite/task-cache.json");
+        let cache_path = dir.join("node_modules/.vite/task-cache.db");
         let task_cache = TaskCache::load_from_file(&cache_path)?;
 
         Ok(Self { vite_task_jsons, dir, task_cache, fs: CachedFileSystem::default() })
     }
+    pub const fn cache(&self) -> &TaskCache {
+        &self.task_cache
+    }
 
     pub fn unload(self) -> anyhow::Result<()> {
         self.task_cache.save()?;
         Ok(())
     }
 
-    pub fn to_task_graph(
+    pub fn resolve_tasks(
         &self,
-        mut task_names: Vec<Str>,
+        task_names: &[Str],
+        task_args: Arc<[Str]>,
     ) -> anyhow::Result<StableDiGraph<ResolvedTask, ()>> {
-        let mut tasks_by_full_name: HashMap<Str, (TaskConfigWithDeps, PackageInfo)> =
+        let mut task_configs_by_full_name: HashMap<Str, (TaskConfigWithDeps, PackageInfo)> =
             HashMap::new();
         for (task_json, package_info) in &self.vite_task_jsons {
             for (task_name, task_config_json) in &task_json.tasks {
@@ -123,38 +172,46 @@ impl Workspace {
                 } else {
                     format!("{}#{}", &package_info.name, task_name).as_str().into()
                 };
-                if tasks_by_full_name
+                if task_configs_by_full_name
                     .insert(full_name.clone(), (task_config_json.clone(), package_info.clone()))
                     .is_some()
                 {
                     anyhow::bail!("Duplicated task name '{}'", &full_name)
                 }
             }
         }
-        // let mut vite_task_json = self.vite_task_json.clone();
-        // let capacity = vite_task_json.tasks.len();
+
+        let mut task_names: BTreeSet<Str> = task_names.iter().cloned().collect();
+
         let mut task_graph = StableDiGraph::<ResolvedTask, ()>::new();
         let mut ids_by_task_name = HashMap::<Str, NodeIndex>::new();
         let mut edges = Vec::<(Str, Str)>::new();
 
-        while let Some(task_name) = task_names.pop() {
-            let (task_config, package_info) = tasks_by_full_name
+        while let Some(task_name) = task_names.pop_first() {
+            let (task_config_with_deps, package_info) = task_configs_by_full_name
                 .remove(&task_name)
                 .with_context(|| format!("Task '{}' not found", &task_name))?;
 
+            let resolved_config = ResolvedTaskConfig {
+                config_dir: package_info.path.as_str().into(),
+                config: task_config_with_deps.config,
+            };
+
+            let resolved_command = resolved_config.resolve_command(&task_args)?;
+
             let id = task_graph.add_node(ResolvedTask {
                 name: task_name.clone(),
-                config_dir: package_info.path.as_str().into(),
-                envs: TaskEnvs::resolve(&task_config.config)?,
-                config: task_config.config,
+                args: task_args.clone(),
+                resolved_command,
+                resolved_config,
             });
             if ids_by_task_name.insert(task_name.clone(), id).is_some() {
                 anyhow::bail!("Duplicated task name '{}'", &task_name)
             }
 
-            for dep in task_config.depends_on {
+            for dep in task_config_with_deps.depends_on {
                 edges.push((dep.clone(), task_name.clone()));
-                task_names.push(dep);
+                task_names.insert(dep);
             }
         }