feat(cache): implement flow node caching (#4808)

* feat(cache): implement flow node caching * feat(cache): implement script caching * feat(cache): improve cache
2026-09-07 08:02:40 +00:00 · 2024-11-29 09:35:00 +01:00
parent f7908682d3
commit 3fbb2bfc8a
12 changed files with 471 additions and 88 deletions
@@ -0,0 +1,34 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT lock AS \"lock: String\", code AS \"code: String\", flow::text AS \"flow: Box<str>\" FROM flow_node WHERE id = $1 LIMIT 1",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "lock: String",
+        "type_info": "Text"
+      },
+      {
+        "ordinal": 1,
+        "name": "code: String",
+        "type_info": "Text"
+      },
+      {
+        "ordinal": 2,
+        "name": "flow: Box<str>",
+        "type_info": "Text"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Int8"
+      ]
+    },
+    "nullable": [
+      true,
+      true,
+      null
+    ]
+  },
+  "hash": "c57ed2d91de46d7de88e20b94b7afbafb622528864f2b23c8b7278bd506d967f"
+}
@@ -0,0 +1,72 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "SELECT lock AS \"lock: String\", content AS \"code!: String\",\n                        language AS \"language: Option<ScriptLang>\", envs AS \"envs: Vec<String>\", codebase AS \"codebase: String\" FROM script WHERE hash = $1 AND workspace_id = $2 LIMIT 1",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "lock: String",
+        "type_info": "Text"
+      },
+      {
+        "ordinal": 1,
+        "name": "code!: String",
+        "type_info": "Text"
+      },
+      {
+        "ordinal": 2,
+        "name": "language: Option<ScriptLang>",
+        "type_info": {
+          "Custom": {
+            "name": "script_lang",
+            "kind": {
+              "Enum": [
+                "python3",
+                "deno",
+                "go",
+                "bash",
+                "postgresql",
+                "nativets",
+                "bun",
+                "mysql",
+                "bigquery",
+                "snowflake",
+                "graphql",
+                "powershell",
+                "mssql",
+                "php",
+                "bunnative",
+                "rust",
+                "ansible"
+              ]
+            }
+          }
+        }
+      },
+      {
+        "ordinal": 3,
+        "name": "envs: Vec<String>",
+        "type_info": "VarcharArray"
+      },
+      {
+        "ordinal": 4,
+        "name": "codebase: String",
+        "type_info": "Varchar"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Int8",
+        "Text"
+      ]
+    },
+    "nullable": [
+      true,
+      false,
+      false,
+      true,
+      true
+    ]
+  },
+  "hash": "df52a71d59eb84a2b08133d25f0a8bba7f2b56625fdb7a7c0e10b51377eeb1d4"
+}
@@ -10625,6 +10625,7 @@ dependencies = [
 "mail-send",
 "object_store",
 "prometheus",
+ "quick_cache",
 "rand 0.8.5",
 "regex",
 "reqwest 0.12.9",
@@ -61,6 +61,7 @@ crc.workspace = true
 windmill-macros.workspace = true
 semver.workspace = true
 croner = "2.0.6"
+quick_cache.workspace = true

 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 tikv-jemalloc-ctl = { optional = true, workspace = true }
@@ -0,0 +1,327 @@
+use crate::error;
+
+use std::path::{Path, PathBuf};
+
+use quick_cache::sync::Cache;
+use sqlx::PgExecutor;
+
+/// Cache directory for windmill server/worker(s).
+pub const CACHE_DIR: &str = "/tmp/windmill/cache/";
+
+pub mod flow {
+    use super::*;
+    use crate::flows::{FlowNodeId, FlowValue};
+
+    /// Cache directory for windmill server/worker(s) flow nodes.
+    pub const CACHE_DIR: &str = const_format::concatcp!(super::CACHE_DIR, "flow");
+
+    lazy_static::lazy_static! {
+        /// Flow node cache.
+        /// FIXME: This should be a static but [`Cache`] does not have a const constructor.
+        /// FIXME: Use `Arc<Val>` for cheap cloning.
+        static ref CACHE: Cache<FlowNodeId, Val> = Cache::new(1000);
+    }
+
+    /// Flow node cache value.
+    #[derive(Debug, Clone, Default)]
+    struct Val {
+        lock: Option<String>,
+        code: Option<String>,
+        flow: Option<FlowValue>,
+    }
+
+    /// Fetch the flow node script referenced by `node` from the cache.
+    /// If not present, import from the file-system cache or fetch it from the database and write
+    /// it to the file system and cache.
+    /// This should be preferred over fetching the database directly.
+    pub async fn fetch_script(e: impl PgExecutor<'_>, node: FlowNodeId)
+        -> error::Result<(Option<String>, String)>
+    {
+        fetch(e, node).await.and_then(|Val { lock, code, .. }| Ok((lock, code.ok_or_else(|| {
+            error::Error::InternalErr(format!("Flow node ({:x}) isn't a script node.", node.0))
+        })?)))
+    }
+
+    /// Fetch the flow node flow value referenced by `node` from the cache.
+    /// If not present, import from the file-system cache or fetch it from the database and write
+    /// it to the file system and cache.
+    /// This should be preferred over fetching the database directly.
+    pub async fn fetch_flow(e: impl PgExecutor<'_>, node: FlowNodeId)
+        -> error::Result<FlowValue>
+    {
+        fetch(e, node).await.and_then(|Val { flow, .. }| flow.ok_or_else(|| {
+            error::Error::InternalErr(format!("Flow node ({:x}) isn't a flow value node.", node.0))
+        }))
+    }
+
+    /// Fetch the flow node referenced by `node` from the cache.
+    /// If not present, import from the file-system cache or fetch it from the database and write
+    /// it to the file system and cache.
+    /// This should be preferred over fetching the database directly.
+    async fn fetch(e: impl PgExecutor<'_>, node: FlowNodeId) -> error::Result<Val> {
+        // If not present, `get_or_insert_async` will lock the key until the future completes,
+        // so only one thread will be able to fetch the data from the database and write it to
+        // the file system and cache, hence no race on the file system.
+        CACHE.get_or_insert_async(
+            &node,
+            fs::import_or_insert_with(CACHE_DIR, node.0 as u64, async {
+                sqlx::query!(
+                    "SELECT \
+                        lock AS \"lock: String\", \
+                        code AS \"code: String\", \
+                        flow::text AS \"flow: Box<str>\" \
+                    FROM flow_node WHERE id = $1 LIMIT 1",
+                    node.0,
+                )
+                .fetch_one(e)
+                .await
+                .map_err(Into::into)
+                .and_then(|r| Ok(Val {
+                    lock: r.lock.and_then(|x| if x.is_empty() { None } else { Some(x) }),
+                    code: r.code,
+                    flow: match r.flow {
+                        None => None,
+                        Some(flow) => serde_json::from_str(&flow).map_err(|err| {
+                            error::Error::InternalErr(format!("Unable to parse flow value: {err:?}"))
+                        })?,
+                    }
+                }))
+            })
+        ).await
+    }
+
+    // ----------------------------------------------------------------------------------------------
+    // impl `fs::Bundle` for `Val`.
+
+    #[derive(Copy, Clone)]
+    enum Item {
+        Lock,
+        Code,
+        Flow,
+    }
+
+    impl fs::Item for Item {
+        fn path(&self, root: &Path) -> PathBuf {
+            match self {
+                Self::Lock => root.join("lock.txt"),
+                Self::Code => root.join("code.txt"),
+                Self::Flow => root.join("flow.json"),
+            }
+        }
+    }
+
+    impl fs::Bundle for Val {
+        type Item = Item;
+
+        fn items() -> &'static [Self::Item] {
+            &[Item::Lock, Item::Code, Item::Flow]
+        }
+
+        fn import(&mut self, item: Self::Item, data: Vec<u8>) -> error::Result<()> {
+            match item {
+                Item::Lock => self.lock = Some(String::from_utf8(data)?),
+                Item::Code => self.code = Some(String::from_utf8(data)?),
+                Item::Flow => self.flow = Some(serde_json::from_slice(&data)?),
+            }
+            Ok(())
+        }
+
+        fn export(&self, item: Self::Item) -> error::Result<Option<Vec<u8>>> {
+            match item {
+                Item::Lock => Ok(self.lock.as_ref().map(|s| s.as_bytes().to_vec())),
+                Item::Code => Ok(self.code.as_ref().map(|s| s.as_bytes().to_vec())),
+                Item::Flow => Ok(self.flow.as_ref().map(|f| serde_json::to_vec(f)).transpose()?),
+            }
+        }
+    }
+}
+
+pub mod script {
+    use super::*;
+    use crate::scripts::{ScriptHash, ScriptLang};
+
+    /// Cache directory for windmill server/worker(s) scripts.
+    pub const CACHE_DIR: &str = const_format::concatcp!(super::CACHE_DIR, "script");
+
+    lazy_static::lazy_static! {
+        /// Scripts cache.
+        /// FIXME: This should be a static but [`Cache`] does not have a const constructor.
+        /// FIXME: Use `Arc<Val>` for cheap cloning.
+        static ref CACHE: Cache<ScriptHash, Val> = Cache::new(1000);
+    }
+
+    /// Script cache value.
+    #[derive(Debug, Clone, Default)]
+    pub struct Val {
+        pub lock: Option<String>,
+        pub code: String,
+        pub language: Option<ScriptLang>,
+        pub envs: Option<Vec<String>>,
+        pub codebase: Option<String>,
+    }
+
+    /// Fetch the script referenced by `hash` from the cache.
+    /// If not present, import from the file-system cache or fetch it from the database and write
+    /// it to the file system and cache.
+    /// This should be preferred over fetching the database directly.
+    pub async fn fetch(e: impl PgExecutor<'_>, hash: ScriptHash, workspace_id: &str)
+        -> error::Result<Val>
+    {
+        // If not present, `get_or_insert_async` will lock the key until the future completes,
+        // so only one thread will be able to fetch the data from the database and write it to
+        // the file system and cache, hence no race on the file system.
+        CACHE.get_or_insert_async(
+            &hash,
+            fs::import_or_insert_with(CACHE_DIR, hash.0 as u64, async {
+                sqlx::query!(
+                    "SELECT \
+                        lock AS \"lock: String\", \
+                        content AS \"code!: String\",
+                        language AS \"language: Option<ScriptLang>\", \
+                        envs AS \"envs: Vec<String>\", \
+                        codebase AS \"codebase: String\" \
+                    FROM script WHERE hash = $1 AND workspace_id = $2 LIMIT 1",
+                    hash.0,
+                    workspace_id,
+                )
+                .fetch_one(e)
+                .await
+                .map_err(Into::into)
+                .map(|r| Val {
+                    lock: r.lock.and_then(|x| if x.is_empty() { None } else { Some(x) }),
+                    code: r.code,
+                    language: r.language,
+                    envs: r.envs,
+                    codebase: r.codebase,
+                })
+            })
+        )
+        .await
+    }
+
+    // ----------------------------------------------------------------------------------------------
+    // impl `fs::Bundle` for `Val`.
+
+    #[derive(Copy, Clone)]
+    pub enum Item {
+        Lock,
+        Code,
+        Info,
+    }
+
+    impl fs::Item for Item {
+        fn path(&self, root: &Path) -> PathBuf {
+            match self {
+                Item::Lock => root.join("lock.txt"),
+                Item::Code => root.join("code.txt"),
+                Item::Info => root.join("info.json"),
+            }
+        }
+    }
+
+    impl fs::Bundle for Val {
+        type Item = Item;
+
+        fn items() -> &'static [Self::Item] {
+            &[Item::Lock, Item::Code, Item::Info]
+        }
+
+        fn import(&mut self, item: Self::Item, data: Vec<u8>) -> error::Result<()> {
+            match item {
+                Item::Lock => self.lock = Some(String::from_utf8(data)?),
+                Item::Code => self.code = String::from_utf8(data)?,
+                Item::Info => (self.language, self.envs, self.codebase) = serde_json::from_slice(&data)?,
+            }
+            Ok(())
+        }
+
+        fn export(&self, item: Self::Item) -> error::Result<Option<Vec<u8>>> {
+            match item {
+                Item::Lock => Ok(self.lock.as_ref().map(|s| s.as_bytes().to_vec())),
+                Item::Code => Ok(Some(self.code.as_bytes().to_vec())),
+                Item::Info => Ok(Some(serde_json::to_vec(&(&self.language, &self.envs, &self.codebase))?)),
+            }
+        }
+    }
+}
+
+mod fs {
+    use super::*;
+
+    use std::future::Future;
+
+    use std::fs::{self, OpenOptions};
+    use std::io::{Read, Write};
+
+    /// A bundle of items that can be imported/exported from/into the file-system.
+    pub trait Bundle: Default {
+        /// Item type of the bundle.
+        type Item: Item;
+        /// Returns a slice of all items than **can** exists within the bundle.
+        fn items() -> &'static [Self::Item];
+        /// Import the given `data` into the `item`.
+        fn import(&mut self, item: Self::Item, data: Vec<u8>) -> error::Result<()>;
+        /// Export the `item` into a `Vec<u8>`.
+        fn export(&self, item: Self::Item) -> error::Result<Option<Vec<u8>>>;
+    }
+
+    /// An item that can be imported/exported from/into the file-system.
+    pub trait Item: Copy + 'static {
+        /// Returns the path of the item within the given `root` path.
+        fn path(&self, root: &Path) -> PathBuf;
+    }
+
+    /// Import or insert a bundle within the given combination of `{root}/{key}/`.
+    pub async fn import_or_insert_with<T, F>(root: &str, key: u64, f: F)
+        -> error::Result<T>
+    where
+        T: Bundle,
+        F: Future<Output = error::Result<T>>,
+    {
+        // Generate the file path from `root` path and `key`.
+        let path = Path::new(root).join(format!("{:016x}", key));
+        // Retrieve the data from the cache directory or the database.
+        if fs::metadata(&path).is_ok() {
+            // Cache path exists, read its contents.
+            let import = || -> error::Result<T> {
+                let mut data = T::default();
+                for item in T::items() {
+                    let mut buf = vec![];
+                    let Ok(mut file) = OpenOptions::new().read(true).open(item.path(&path))
+                    else { continue };
+                    file.read_to_end(&mut buf)?;
+                    data.import(*item, buf)?;
+                }
+                tracing::debug!("Imported from file-system: {:?}", path);
+                Ok(data)
+            };
+            match import() {
+                Ok(data) => return Ok(data),
+                Err(err) => tracing::warn!(
+                    "Failed to import from file-system, fetch source..: {path:?}: {err:?}"
+                )
+            }
+        }
+        // Cache path doesn't exist or import failed, generate the content.
+        let data = f.await?;
+        let export = |data: &T| -> error::Result<()> {
+            fs::create_dir_all(&path)?;
+            // Write the generated data to the file.
+            for item in T::items() {
+                let Some(buf) = data.export(*item)?
+                else { continue };
+                let mut file = OpenOptions::new().write(true).create(true).open(item.path(&path))?;
+                file.write_all(&buf)?;
+            }
+            tracing::debug!("Exported to file-system: {:?}", path);
+            Ok(())
+        };
+        // Try to export data to the file-system.
+        // If failed, remove the directory but still return the data.
+        if let Err(err) = export(&data) {
+            tracing::warn!("Failed to export to file-system: {path:?}: {err:?}");
+            let _ = fs::remove_dir_all(&path);
+        }
+        Ok(data)
+    }
+}
@@ -62,6 +62,10 @@ pub enum Error {
    AiError(String),
    #[error("{0}")]
    AlreadyCompleted(String),
+    #[error("{0}")]
+    Utf8(#[from] std::string::FromUtf8Error),
+    #[error("Encoding/decoding error: {0}")]
+    SerdeJson(#[from] serde_json::Error),
 }

 impl Error {
@@ -18,6 +18,7 @@ use sqlx::types::Json;
 use sqlx::types::JsonRawValue;

 use crate::{
+    cache,
    error::Error,
    more_serde::{default_empty_string, default_id, default_null, default_true, is_default},
    scripts::{Schema, ScriptHash, ScriptLang},
@@ -402,7 +403,7 @@ pub enum InputTransform {
 }

 /// Id in the `flow_node` table.
-#[derive(Serialize, Deserialize, Debug, Copy, Clone, Hash)]
+#[derive(Serialize, Deserialize, Debug, Copy, Clone, Hash, Eq, PartialEq)]
 #[serde(transparent)]
 pub struct FlowNodeId(pub i64);

@@ -759,11 +760,7 @@ pub async fn resolve_module(
            let (lock, content) = if !with_code {
                (Some("...".to_string()), "...".to_string())
            } else {
-                sqlx::query!("SELECT lock, code AS \"code!: String\" FROM flow_node WHERE id = $1", id.0)
-                    .fetch_one(e)
-                    .await
-                    .map_err(Error::SqlErr)
-                    .map(|record| (record.lock, record.code))?
+                cache::flow::fetch_script(e, id).await?
            };
            val = RawScript {
                input_transforms, content, lock, path: None, tag, language, custom_concurrency_key,
@@ -799,31 +796,12 @@ pub async fn resolve_modules(
 ) -> Result<(), Error> {
    // Replace the `modules_node` with the actual modules.
    if let Some(id) = modules_node {
-        *modules = load_flow_modules(e, id).await?;
+        *modules = cache::flow::fetch_flow(e, id)
+            .await
+            .map(|flow| flow.modules)?;
    }
    for module in modules.iter_mut() {
        Box::pin(resolve_module(e, workspace_id, &mut module.value, with_code)).await?;
    }
    Ok(())
 }
-
-pub async fn load_flow_modules(
-    e: &sqlx::PgPool,
-    id: FlowNodeId,
-) -> Result<Vec<FlowModule>, Error> {
-    #[derive(Deserialize)]
-    struct FlowModulesOnly { modules: Vec<FlowModule> }
-
-    sqlx::query_scalar!(
-        "SELECT flow AS \"flow!: Json<Box<JsonRawValue>>\" FROM flow_node WHERE id = $1 LIMIT 1",
-        id.0
-    )
-    .fetch_one(e)
-    .await
-    .map_err(Error::SqlErr)
-    .and_then(|value| {
-        serde_json::from_str::<FlowModulesOnly>(value.get())
-            .map(|x| x.modules)
-            .map_err(|err| Error::InternalErr(format!("Failed to parse flow node value: {}", err)))
-    })
-}
@@ -20,6 +20,7 @@ pub mod apps;
 pub mod auth;
 #[cfg(feature = "benchmark")]
 pub mod bench;
+pub mod cache;
 pub mod db;
 pub mod ee;
 pub mod email_ee;
@@ -71,7 +71,7 @@ impl ScriptLang {
    }
 }

-#[derive(PartialEq, Debug, Hash, Clone, Copy, sqlx::Type)]
+#[derive(Eq, PartialEq, Debug, Hash, Clone, Copy, sqlx::Type)]
 #[sqlx(transparent)]
 pub struct ScriptHash(pub i64);

@@ -39,6 +39,7 @@ use windmill_audit::audit_ee::{audit_log, AuditAuthor};
 use windmill_audit::ActionKind;

 use windmill_common::{
+    cache,
    auth::{fetch_authed_from_permissioned_as, permissioned_as_to_username},
    db::{Authed, UserDB},
    error::{self, to_anyhow, Error},
@@ -3209,16 +3210,7 @@ pub async fn push<'c, 'd>(
            None,
        ),
        JobPayload::FlowNode { id, path } => {
-            let flow_value = sqlx::query_scalar!(
-                "SELECT flow as \"flow!: sqlx::types::Json<Box<RawValue>>\" FROM flow_node WHERE id = $1 LIMIT 1",
-                id.0
-            ).fetch_one(_db)
-            .await?;
-            let value = serde_json::from_str::<FlowValue>(flow_value.get()).map_err(|err| {
-                Error::InternalErr(format!(
-                    "could not convert json to flow for node={}: {err:?}", id.0
-                ))
-            })?;
+            let value = cache::flow::fetch_flow(_db, id).await?;
            let status = Some(FlowStatus::new(&value));
            (
                Some(id.0),
@@ -16,6 +16,7 @@ use windmill_common::error::Error;
 use windmill_common::flows::FlowValue;
 use windmill_common::worker::WORKER_CONFIG;
 use windmill_common::{
+    cache,
    error,
    flows::{FlowModule, FlowModuleValue},
    jobs::QueuedJob,
@@ -393,18 +394,14 @@ async fn spawn_dedicated_workers_for_flow(
                    }
                }
                FlowModuleValue::FlowScript { id, language, .. } => {
-                    let spawn = sqlx::query!(
-                        "SELECT lock, code AS \"code!: String\" FROM flow_node WHERE id = $1 LIMIT 1",
-                        id.0
-                    )
-                    .fetch_one(db)
-                    .await
-                    .map(|record| SpawnWorker::RawScript {
-                        path: "".to_string(),
-                        content: record.code,
-                        lock: record.lock,
-                        lang: language.clone(),
-                    });
+                    let spawn = cache::flow::fetch_script(db, *id)
+                        .await
+                        .map(|(lock, content)| SpawnWorker::RawScript {
+                            path: "".to_string(),
+                            content,
+                            lock,
+                            lang: language.clone(),
+                        });
                    match spawn {
                        Ok(spawn) => {
                            if let Some(dedi_w) = spawn_dedicated_worker(
@@ -46,7 +46,9 @@ use std::{
 use uuid::Uuid;

 use windmill_common::{
+    cache,
    error::{self, to_anyhow, Error},
+    flows::FlowNodeId,
    get_latest_deployed_hash_for_path,
    jobs::{JobKind, QueuedJob},
    scripts::{get_full_hub_script_by_path, ScriptHash, ScriptLang, PREVIEW_IS_CODEBASE_HASH},
@@ -2218,38 +2220,16 @@ pub async fn get_script_content_by_hash(
    w_id: &str,
    db: &DB,
 ) -> error::Result<ContentReqLangEnvs> {
-    let r = sqlx::query_as::<
-        _,
-        (
-            String,
-            Option<String>,
-            Option<ScriptLang>,
-            Option<Vec<String>>,
-            Option<bool>
-        ),
-    >(
-        "SELECT content, lock, language, envs, codebase LIKE '%.tar' as codebase FROM script WHERE hash = $1 AND workspace_id = $2",
-    )
-    .bind(script_hash.0)
-    .bind(w_id)
-    .fetch_optional(db)
-    .await?
-    .ok_or_else(|| Error::InternalErr(format!("expected content and lock")))?;
+    let script = cache::script::fetch(db, *script_hash, w_id).await?;
    Ok(ContentReqLangEnvs {
-        content: r.0,
-        lockfile: r.1,
-        language: r.2,
-        envs: r.3,
-        codebase: if r.4.is_some() {
-            let b = r.4.unwrap();
-            let sh = script_hash.to_string();
-            if b {
-                Some(format!("{sh}.tar"))
-            } else {
-                Some(sh)
-            }
-        } else {
-            None
+        content: script.code,
+        lockfile: script.lock,
+        language: script.language,
+        envs: script.envs,
+        codebase: match script.codebase {
+            None => None,
+            Some(x) if x.ends_with(".tar") => Some(format!("{}.tar", script_hash)),
+            Some(_) => Some(script_hash.to_string()),
        },
    })
 }
@@ -2305,13 +2285,9 @@ async fn handle_code_execution_job(
            .await?
        }
        JobKind::FlowScript => {
-            let (lockfile, content) = sqlx::query!(
-                "SELECT lock, code AS \"code!: String\" FROM flow_node WHERE id = $1 LIMIT 1",
+            let (lockfile, content) = cache::flow::fetch_script(db, FlowNodeId(
                job.script_hash.unwrap_or(ScriptHash(0)).0
-            )
-            .fetch_one(db)
-            .await
-            .map(|record| (record.lock, record.code))?;
+            )).await?;
            ContentReqLangEnvs {
                content,
                lockfile,