From 63b2f8f7739870d3f0d5a71ffe86f989c536fe17 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Tue, 15 Sep 2026 16:53:23 +0200 Subject: [PATCH] fix: key a cached workflow-as-code task on its name and arguments --- backend/tests/bun_jobs.rs | 69 ++++++++++++--------- backend/tests/fixtures/wac_flow_script.sql | 20 ++++++ backend/windmill-common/src/wac.rs | 33 ++++++++++ backend/windmill-worker/src/bun_executor.rs | 2 + backend/windmill-worker/src/common.rs | 53 +++++++++------- 5 files changed, 125 insertions(+), 52 deletions(-) diff --git a/backend/tests/bun_jobs.rs b/backend/tests/bun_jobs.rs index 63d5341ba3..67391a6134 100644 --- a/backend/tests/bun_jobs.rs +++ b/backend/tests/bun_jobs.rs @@ -1257,8 +1257,10 @@ async fn test_bun_wac_task_dispatch_from_flow_script(db: Pool) -> anyh } /// `task(fn, { cache_ttl })` on an inline task of a deployed flow's step: the child runs -/// the parent's code with the parent's arguments, so its result-cache key carries its -/// step key, or the parent and every sibling would read its result back as their own. +/// the parent's code with the parent's arguments, so its result is keyed on the task +/// and the arguments it was called with, or the parent and every sibling would read +/// its result back as their own, and a task fed by an earlier step would be served a +/// result computed for another input. #[sqlx::test(fixtures("base", "wac_flow_script"))] async fn test_bun_wac_inline_task_cache_is_per_task(db: Pool) -> anyhow::Result<()> { use windmill_common::flows::FlowNodeId; @@ -1266,40 +1268,49 @@ async fn test_bun_wac_inline_task_cache_is_per_task(db: Pool) -> anyho let server = ApiServer::start(db.clone()).await?; let port = server.addr.port(); - let payload = || JobPayload::FlowScript { - id: FlowNodeId(3000000000000012), - path: "f/system/wac_flow_script/a".to_string(), - language: ScriptLang::Bun, - cache_ttl: None, - cache_ignore_s3_path: None, - dedicated_worker: None, - concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default(), - }; - - let mut children_from_cache = Vec::new(); - for _ in 0..2 { - let job = RunJob::from(payload()) - .arg("n", serde_json::json!(5)) - .run_until_complete(&db, false, port) - .await; - assert_eq!( - job.json_result().unwrap(), - serde_json::json!({"doubled": 10, "tripled": 15}) - ); + // Runs the step of node `node` with n=5; returns its result and how many of its + // task children were served from the result cache. + async fn run( + db: &Pool, + port: u16, + node: i64, + ) -> anyhow::Result<(serde_json::Value, i64)> { + let job = RunJob::from(JobPayload::FlowScript { + id: FlowNodeId(node), + path: "f/system/wac_flow_script/a".to_string(), + language: ScriptLang::Bun, + cache_ttl: None, + cache_ignore_s3_path: None, + dedicated_worker: None, + concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default( + ), + }) + .arg("n", serde_json::json!(5)) + .run_until_complete(db, false, port) + .await; let from_cache: i64 = sqlx::query_scalar( "SELECT count(*) FROM job_logs l JOIN v2_job j ON j.id = l.job_id \ WHERE j.parent_job = $1 AND l.logs LIKE '%found in cache%'", ) .bind(job.id) - .fetch_one(&db) + .fetch_one(db) .await?; - children_from_cache.push(from_cache); + Ok((job.json_result().unwrap(), from_cache)) + } + + // Two cached tasks on the same input: the second run serves each from its own entry. + for expected_from_cache in [0, 2] { + let (result, from_cache) = run(&db, port, 3000000000000012).await?; + assert_eq!(result, serde_json::json!({"doubled": 10, "tripled": 15})); + assert_eq!(from_cache, expected_from_cache); + } + + // A cached task called with a fresh value each run is never served a stale result. + for _ in 0..2 { + let (result, from_cache) = run(&db, port, 3000000000000013).await?; + assert_eq!(result, serde_json::json!({"fresh": true})); + assert_eq!(from_cache, 0); } - assert_eq!( - children_from_cache, - vec![0, 2], - "the second run serves each task from its own cache entry" - ); Ok(()) } diff --git a/backend/tests/fixtures/wac_flow_script.sql b/backend/tests/fixtures/wac_flow_script.sql index 1b780c6cf4..da50a6d35d 100644 --- a/backend/tests/fixtures/wac_flow_script.sql +++ b/backend/tests/fixtures/wac_flow_script.sql @@ -51,3 +51,23 @@ export const main = workflow(async (n: number) => { return { doubled: d, tripled: t }; });' ); + +-- The same flow's step with a cached task fed by a value that differs on every run. +INSERT INTO public.flow_node(id, workspace_id, path, hash_v2, lock, code) VALUES ( +3000000000000013, +'test-workspace', +'f/system/wac_flow_script', +'0000000000000000000000000000000000000000000000000000000000000013', +NULL, +E'import { workflow, task, step } from "windmill-client"; + +const double = task(async (n: number) => { + return n * 2; +}, { cache_ttl: 60 }); + +export const main = workflow(async (n: number) => { + const r = await step("pick", async () => Math.random()); + const d = await double(r); + return { fresh: d === r * 2 }; +});' +); diff --git a/backend/windmill-common/src/wac.rs b/backend/windmill-common/src/wac.rs index 05c343b427..fe9b6c154f 100644 --- a/backend/windmill-common/src/wac.rs +++ b/backend/windmill-common/src/wac.rs @@ -40,6 +40,14 @@ pub struct WacCheckpoint { #[serde(skip_serializing_if = "Option::is_none")] #[serde(default)] pub _executing_key: Option, + /// With `_executing_key`: the task the child runs and `task_args_hash` of the + /// arguments it was called with, the identity its cached result is keyed on. + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(default)] + pub _executing_task: Option, + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(default)] + pub _executing_args_hash: Option, /// `resume_job.id` values already consumed by earlier approval steps (the /// row primary key, not the distinct integer `resume_id` column). Rows are /// never deleted, so a workflow with several sequential wait_for_approval() @@ -72,9 +80,34 @@ pub fn approval_resume_id(step_key: &str) -> u32 { u32::from_be_bytes([digest[0], digest[1], digest[2], digest[3]]) } +/// Hash of a task call's arguments, with the keys in a fixed order so the order a +/// caller passes them in cannot split its cached results. +pub fn task_args_hash(args: &serde_json::Map) -> String { + use sha2::{Digest, Sha256}; + let ordered: std::collections::BTreeMap<&String, &Value> = args.iter().collect(); + let json = serde_json::to_string(&ordered).unwrap_or_default(); + format!("{:x}", Sha256::digest(json.as_bytes())) +} + #[cfg(test)] mod tests { use super::approval_resume_id; + use super::task_args_hash; + + #[test] + fn task_args_hash_ignores_key_order_and_sees_values() { + let parse = |s: &str| { + serde_json::from_str::>(s).unwrap() + }; + assert_eq!( + task_args_hash(&parse(r#"{"a":1,"b":2}"#)), + task_args_hash(&parse(r#"{"b":2,"a":1}"#)) + ); + assert_ne!( + task_args_hash(&parse(r#"{"a":1,"b":2}"#)), + task_args_hash(&parse(r#"{"a":1,"b":3}"#)) + ); + } /// Golden values: worker and API must agree on this mapping, and they can run /// different builds during a rolling deploy. Changing it strands every resume diff --git a/backend/windmill-worker/src/bun_executor.rs b/backend/windmill-worker/src/bun_executor.rs index 117fbda086..d4aa7b4fb8 100644 --- a/backend/windmill-worker/src/bun_executor.rs +++ b/backend/windmill-worker/src/bun_executor.rs @@ -3237,6 +3237,8 @@ pub async fn handle_wac_v2_output( let child_checkpoint_json = serde_json::json!({ "completed_steps": &checkpoint.completed_steps, "_executing_key": &step.key, + "_executing_task": &step.name, + "_executing_args_hash": windmill_common::wac::task_args_hash(&step.args), }); sqlx::query( "INSERT INTO v2_job_status (id, workflow_as_code_status) diff --git a/backend/windmill-worker/src/common.rs b/backend/windmill-worker/src/common.rs index 9c012fd117..748cafe776 100644 --- a/backend/windmill-worker/src/common.rs +++ b/backend/windmill-worker/src/common.rs @@ -1574,43 +1574,50 @@ pub async fn cached_result_path( _ => {} } } - // A workflow-as-code task child runs its parent's code with the parent's - // arguments; the step it executes is what tells its result from the parent's - // and from its siblings'. - if let Some(step_key) = wac_executing_key(db, job).await? { - hasher.update(b"wac_step:"); - hasher.update(step_key.as_bytes()); + // A workflow-as-code task child carries its parent's arguments, which say + // nothing about the task's own inputs; its result is keyed on the task it + // runs and the arguments that task was called with. + if let Some((task, args_hash)) = wac_task_identity(db, job).await? { + hasher.update(b"wac_task:"); + hasher.update(task.as_bytes()); + hasher.update(b":"); + hasher.update(args_hash.as_bytes()); + } else { + hash_args( + db, + client, + &job.workspace_id, + &job.args, + &mut hasher, + &job.id, + job.cache_ignore_s3_path.unwrap_or(false), + ) + .await; } - hash_args( - db, - client, - &job.workspace_id, - &job.args, - &mut hasher, - &job.id, - job.cache_ignore_s3_path.unwrap_or(false), - ) - .await; Ok(format!("g/results/{:064x}", hasher.finalize())) } -/// The checkpoint step key a workflow-as-code parent seeded for this child at push -/// time; `None` for any job that is not such a child. -async fn wac_executing_key( +/// The task name and arguments hash a workflow-as-code parent seeded in this +/// child's checkpoint at push time; `None` for any job that is not such a child. +async fn wac_task_identity( db: &DB, job: &MiniPulledJob, -) -> windmill_common::error::Result> { +) -> windmill_common::error::Result> { if job.parent_job.is_none() || job.flow_step_id.is_some() { return Ok(None); } - let key: Option> = sqlx::query_scalar( - "SELECT workflow_as_code_status->'_checkpoint'->>'_executing_key' \ + let identity: Option<(Option, Option)> = sqlx::query_as( + "SELECT workflow_as_code_status->'_checkpoint'->>'_executing_task', \ + workflow_as_code_status->'_checkpoint'->>'_executing_args_hash' \ FROM v2_job_status WHERE id = $1", ) .bind(job.id) .fetch_optional(db) .await?; - Ok(key.flatten()) + Ok(match identity { + Some((Some(task), Some(args_hash))) => Some((task, args_hash)), + _ => None, + }) } #[cfg(feature = "parquet")]