kalugin66/windmill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: add inline-persist fast path for WAC v2 step() (#8807)

Browse Source
This commit is contained in:
Ruben Fiszel
2026-04-13 12:49:53 -04:00
committed by GitHub
parent 3f5841f84d
commit b3ef4bc26c
15 changed files with 762 additions and 385 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
backend/.sqlx/query-15862647eb83c6a3116aa4a1e59480f7fec559c88bc384c5475dba5f03526337.json generated Normal file
View File

@@ -0,0 +1,22 @@
{
"db_name": "PostgreSQL",
"query": "SELECT workflow_as_code_status as \"v: serde_json::Value\"\n FROM v2_job_completed WHERE id = $1",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "v: serde_json::Value",
"type_info": "Jsonb"
}
],
"parameters": {
"Left": [
"Uuid"
]
},
"nullable": [
true
]
},
"hash": "15862647eb83c6a3116aa4a1e59480f7fec559c88bc384c5475dba5f03526337"
}

15
backend/.sqlx/query-50b25537dcb799cc233dbb06c76798a860ce977954a856335bb62ae00f615659.json generated
View File

@@ -1,15 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO folder SELECT name, $1, display_name, owners, extra_perms, summary, edited_at, created_by FROM folder WHERE workspace_id = $2",
"describe": {
"columns": [],
"parameters": {
"Left": [
"Varchar",
"Text"
]
},
"nullable": []
},
"hash": "50b25537dcb799cc233dbb06c76798a860ce977954a856335bb62ae00f615659"
}

2
backend/.sqlx/query-5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55.json generated
View File

@@ -15,7 +15,7 @@
]
},
"nullable": [
null
true
]
},
"hash": "5a219a2532517869578c4504ff3153c43903f929ae5d62fbba12610f89c36d55"

20
backend/.sqlx/query-5f3f1f1ca72b0392f227d22adef9ef795049e60a823e8ce754a3f9d91ef1b6bb.json generated
View File

@@ -1,20 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO folder (workspace_id, name, display_name, owners, extra_perms, summary, created_by, edited_at) VALUES ($1, $2, $3, $4, $5, $6, $7, now())",
"describe": {
"columns": [],
"parameters": {
"Left": [
"Varchar",
"Varchar",
"Varchar",
"VarcharArray",
"Jsonb",
"Text",
"Varchar"
]
},
"nullable": []
},
"hash": "5f3f1f1ca72b0392f227d22adef9ef795049e60a823e8ce754a3f9d91ef1b6bb"
}

66
backend/.sqlx/query-8830fc3736ec6dbbaee20f0ebeba5a87c2c94227798f4545fe0931af2818a6fa.json generated
View File

@@ -1,66 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "SELECT workspace_id, name, display_name, owners, extra_perms, summary, created_by, edited_at FROM folder WHERE workspace_id = $1 ORDER BY name asc LIMIT $2 OFFSET $3",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "workspace_id",
"type_info": "Varchar"
},
{
"ordinal": 1,
"name": "name",
"type_info": "Varchar"
},
{
"ordinal": 2,
"name": "display_name",
"type_info": "Varchar"
},
{
"ordinal": 3,
"name": "owners",
"type_info": "VarcharArray"
},
{
"ordinal": 4,
"name": "extra_perms",
"type_info": "Jsonb"
},
{
"ordinal": 5,
"name": "summary",
"type_info": "Text"
},
{
"ordinal": 6,
"name": "created_by",
"type_info": "Varchar"
},
{
"ordinal": 7,
"name": "edited_at",
"type_info": "Timestamptz"
}
],
"parameters": {
"Left": [
"Text",
"Int8",
"Int8"
]
},
"nullable": [
false,
false,
false,
false,
false,
true,
true,
true
]
},
"hash": "8830fc3736ec6dbbaee20f0ebeba5a87c2c94227798f4545fe0931af2818a6fa"
}

65
backend/.sqlx/query-cb20f04352364f112ec564617722354577dc5f77169fd4573317454e47361d65.json generated
View File

@@ -1,65 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "SELECT workspace_id, name, display_name, owners, extra_perms, summary, created_by, edited_at FROM folder WHERE name = $1 AND workspace_id = $2",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "workspace_id",
"type_info": "Varchar"
},
{
"ordinal": 1,
"name": "name",
"type_info": "Varchar"
},
{
"ordinal": 2,
"name": "display_name",
"type_info": "Varchar"
},
{
"ordinal": 3,
"name": "owners",
"type_info": "VarcharArray"
},
{
"ordinal": 4,
"name": "extra_perms",
"type_info": "Jsonb"
},
{
"ordinal": 5,
"name": "summary",
"type_info": "Text"
},
{
"ordinal": 6,
"name": "created_by",
"type_info": "Varchar"
},
{
"ordinal": 7,
"name": "edited_at",
"type_info": "Timestamptz"
}
],
"parameters": {
"Left": [
"Text",
"Text"
]
},
"nullable": [
false,
false,
false,
false,
false,
true,
true,
true
]
},
"hash": "cb20f04352364f112ec564617722354577dc5f77169fd4573317454e47361d65"
}

15
backend/.sqlx/query-dbd66bee283a7b4f892673f968f92adce1be8b897288b655505729c97c6e324c.json generated
View File

@@ -1,15 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO folder (workspace_id, name, display_name, owners, extra_perms, summary, edited_at, created_by)\n SELECT $2, name, display_name, owners, extra_perms, summary, edited_at, created_by\n FROM folder\n WHERE workspace_id = $1",
"describe": {
"columns": [],
"parameters": {
"Left": [
"Text",
"Varchar"
]
},
"nullable": []
},
"hash": "dbd66bee283a7b4f892673f968f92adce1be8b897288b655505729c97c6e324c"
}

150
backend/tests/python_jobs.rs
View File

@@ -1014,3 +1014,153 @@ async def main(item: str, qty: int, email: str):
.await;
Ok(())
}
/// End-to-end comparison between the legacy `step()` suspend-and-replay path
/// and the new SDK inline-persist fast path, toggled per-job via the
/// `WM_WAC_INLINE_FAST_PATH` env var which the Python script sets on its own
/// `os.environ` at the top so parallel tests can't race on a global env var.
///
/// Runs the same 5-step WAC v2 Python workflow twice, asserts both modes
/// produce the same final result and the same `completed_steps` map, and
/// prints a wall-clock benchmark line so CI and manual runs can track the
/// speedup. The fast path is expected to be faster because it avoids N-1
/// subprocess spawns and N-1 queue round-trips for a workflow with N
/// `step()` calls, but the exact ratio depends on the CI runner so we don't
/// assert a hard threshold — behavioral equivalence is the important check.
#[cfg(feature = "python")]
#[sqlx::test(fixtures("base"))]
async fn test_python_wac_v2_step_inline_fast_path(db: Pool<Postgres>) -> anyhow::Result<()> {
initialize_tracing().await;
let server = ApiServer::start(db.clone()).await?;
let port = server.addr.port();
fn workflow_content(fast_path_enabled: bool) -> String {
let flag = if fast_path_enabled { "1" } else { "0" };
// NB: the `os.environ[...] = ...` line must execute BEFORE the first
// `step()` call inside the workflow subprocess — setting it at module
// scope (before `from wmill import ...` finishes importing is still
// fine because the SDK reads it lazily at call time, not at import).
format!(
r#"import os
os.environ["WM_WAC_INLINE_FAST_PATH"] = "{flag}"
from wmill import workflow, step
@workflow
async def main(n: int):
a = await step("a", lambda: n)
b = await step("b", lambda: a + 1)
c = await step("c", lambda: b + 1)
d = await step("d", lambda: c + 1)
e = await step("e", lambda: d + 1)
return {{"a": a, "b": b, "c": c, "d": d, "e": e}}
"#
)
}
let db_ref = &db;
async fn run_once(
db: &Pool<Postgres>,
port: u16,
content: String,
) -> (serde_json::Value, serde_json::Value, std::time::Duration) {
let mut job_id_out: Option<sqlx::types::Uuid> = None;
let mut result_out: Option<serde_json::Value> = None;
let t0 = std::time::Instant::now();
in_test_worker(
db,
async {
let job = Box::pin(
RunJob::from(JobPayload::Code(RawCode {
language: ScriptLang::Python3,
content,
..RawCode::default()
}))
.arg("n", json!(1))
.run_until_complete(db, false, port),
)
.await;
result_out = Some(job.json_result().unwrap_or_else(|| {
panic!("job {} returned no result — raw job = {:?}", job.id, job)
}));
job_id_out = Some(job.id);
},
port,
)
.await;
let elapsed = t0.elapsed();
let job_id = job_id_out.expect("job id");
// Fetch the full workflow_as_code_status for diagnostics, then extract
// _checkpoint.completed_steps. A None here means the step() path never
// wrote the checkpoint, which is exactly the signal we want to surface
// clearly (instead of panicking with an opaque Option::unwrap() error).
let full_status: Option<serde_json::Value> = sqlx::query_scalar!(
r#"SELECT workflow_as_code_status as "v: serde_json::Value"
FROM v2_job_completed WHERE id = $1"#,
job_id
)
.fetch_one(db)
.await
.expect("v2_job_completed row fetch");
let ckpt = full_status
.as_ref()
.and_then(|s| s.get("_checkpoint"))
.and_then(|c| c.get("completed_steps"))
.cloned()
.unwrap_or_else(|| {
panic!(
"job {job_id} completed_steps missing — full workflow_as_code_status = {:?}, job_result = {:?}",
full_status, result_out
)
});
(result_out.unwrap(), ckpt, elapsed)
}
// --- Legacy path: worker-side suspend & replay ---
let (legacy_result, legacy_ckpt, legacy_elapsed) =
run_once(db_ref, port, workflow_content(false)).await;
// --- Fast path: SDK persists the delta via the new API endpoint ---
let (fast_result, fast_ckpt, fast_elapsed) =
run_once(db_ref, port, workflow_content(true)).await;
// Behavioral equivalence: same final result and same completed_steps.
assert_eq!(
legacy_result, fast_result,
"legacy and fast path produced different workflow results"
);
assert_eq!(
legacy_result,
json!({"a": 1, "b": 2, "c": 3, "d": 4, "e": 5}),
"unexpected workflow result"
);
assert_eq!(
legacy_ckpt, fast_ckpt,
"legacy and fast path stored different completed_steps in the checkpoint"
);
// Benchmark output. We log and print but do NOT assert a hard threshold:
// CI runners have variable noise floors and a "fast < legacy" check would
// flake. Behavioral equivalence above is the important invariant.
let legacy_ms = legacy_elapsed.as_millis();
let fast_ms = fast_elapsed.as_millis();
let speedup = if fast_ms > 0 {
legacy_ms as f64 / fast_ms as f64
} else {
f64::INFINITY
};
tracing::info!(
"WAC v2 step() benchmark: legacy={}ms fast={}ms speedup={:.2}x",
legacy_ms,
fast_ms,
speedup
);
println!(
"WAC v2 step() benchmark: legacy={}ms fast={}ms speedup={:.2}x",
legacy_ms, fast_ms, speedup
);
Ok(())
}

87
backend/windmill-api/src/jobs.rs
View File

@@ -63,7 +63,6 @@ use windmill_worker::get_worker_internal_server_inline_utils;
use windmill_common::variables::get_workspace_key;
#[cfg(feature = "run_inline")]
use crate::db::OptJobAuthed;
use crate::triggers::trigger_helpers::{FlowId, ScriptId};
use crate::{
@@ -161,6 +160,10 @@ pub fn workspaced_service() -> Router {
.layer(cors.clone())
.layer(ce_headers.clone()),
)
.route(
"/wac/inline_checkpoint/{job_id}",
post(wac_inline_checkpoint).layer(cors.clone()),
)
.route(
"/restart/f/{job_id}",
post(restart_flow).head(|| async { "" }).layer(cors.clone()),
@@ -4445,6 +4448,88 @@ pub async fn run_workflow_as_code(
Ok((StatusCode::CREATED, uuid.to_string()))
}
#[derive(Deserialize)]
pub struct WacInlineCheckpointPayload {
pub key: String,
pub result: serde_json::Value,
#[serde(default)]
pub started_at: Option<String>,
#[serde(default)]
pub duration_ms: Option<u64>,
}
/// Fast-path endpoint called by the WAC v2 SDKs to persist a single `step()`
/// checkpoint delta without unwinding the parent workflow subprocess.
///
/// Mirrors the worker-side `WacOutput::InlineCheckpoint` arm in
/// `bun_executor::handle_wac_v2_output` exactly — same `completed_steps`
/// entry, same `_step/<key>` timeline entry, same source-hash validation —
/// but does **not** touch `v2_job_queue`, because the parent subprocess is
/// still live and about to return the next chunk of script output.
///
/// Auth: requires the job's ephemeral token (the one the worker sets into
/// `WM_TOKEN` when spawning the subprocess), not just any workspace-scoped
/// `ApiAuthed`. Because WAC v2 replays steps from `completed_steps`, a forged
/// entry directly changes the workflow's observed return values — this is
/// execution state, not user-facing metadata. `OptJobAuthed.job_id` is set
/// only when the caller presents a JWT whose `job_id` claim matches the URL
/// path, so rejecting mismatches closes the workspace-wide privilege gap.
///
/// Any error here causes the SDK to fall back to raising `_StepSuspend`,
/// which then goes through the untouched worker-side path. Old SDKs that
/// never call this endpoint continue to work unchanged.
pub async fn wac_inline_checkpoint(
OptJobAuthed { authed: _, job_id: token_job_id }: OptJobAuthed,
Extension(db): Extension<DB>,
Path((w_id, job_id)): Path<(String, Uuid)>,
Json(payload): Json<WacInlineCheckpointPayload>,
) -> error::Result<StatusCode> {
// Enforce ephemeral-job-token binding: the presented token must be the
// one issued to *this* specific job. Regular workspace API tokens don't
// have `job_id` populated in their JWT claims, so `token_job_id` is None
// for them — reject unconditionally.
if token_job_id != Some(job_id) {
return Err(error::Error::PermissionDenied(
"wac_inline_checkpoint requires the job's ephemeral token".to_string(),
));
}
// Look up the job's script hash for source-hash validation. We deliberately
// use a minimal query here rather than `fetch_queued(...)` — the latter
// pulls in many extra columns we don't need. Restrict to the job kinds
// that actually run user WAC v2 code (`script` and `preview`) so a
// forged/buggy caller can't write a bogus checkpoint onto a flow,
// dependency, or other job kind that shares the workspace.
let row: Option<(Option<i64>,)> = sqlx::query_as(
"SELECT runnable_id FROM v2_job
WHERE id = $1 AND workspace_id = $2
AND kind IN ('script'::job_kind, 'preview'::job_kind)",
)
.bind(&job_id)
.bind(&w_id)
.fetch_optional(&db)
.await?;
let (runnable_id,) = row.ok_or_else(|| {
error::Error::NotFound(format!("WAC v2 job {job_id} not found in workspace {w_id}"))
})?;
let source_hash = runnable_id.map(|h| h.to_string());
let mut tx = db.begin().await?;
windmill_common::wac::persist_inline_checkpoint_delta(
&mut tx,
&job_id,
source_hash.as_deref(),
&payload.key,
payload.result,
payload.started_at.as_deref(),
payload.duration_ms,
)
.await?;
tx.commit().await?;
Ok(StatusCode::OK)
}
lazy_static::lazy_static! {
static ref JOB_VIEW_AUDIT_LOGS: bool = std::env::var("JOB_VIEW_AUDIT_LOGS")
.ok()

1
backend/windmill-common/src/lib.rs
View File

@@ -108,6 +108,7 @@ pub mod usernames;
pub mod users;
pub mod utils;
pub mod variables;
pub mod wac;
pub mod webhook;
pub mod worker;
pub mod worker_group_job_stats;

300
backend/windmill-common/src/wac.rs Normal file
View File

@@ -0,0 +1,300 @@
//! Workflow-as-Code v2 checkpoint model and persistence primitives.
//!
//! Lives in `windmill-common` (not `windmill-worker`) so the API server can
//! write checkpoint deltas directly from the SDK fast path without pulling in
//! the entire worker crate. The worker still re-exports the same symbols from
//! `windmill_worker::wac_executor` for its own historical call sites.
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sqlx::{Postgres, Transaction};
use uuid::Uuid;
use crate::error::{self, Error};
use crate::DB;
/// Checkpoint state persisted across workflow invocations.
#[derive(Debug, Serialize, Deserialize, Default, Clone)]
pub struct WacCheckpoint {
#[serde(default)]
pub source_hash: String,
#[serde(default)]
pub completed_steps: serde_json::Map<String, Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub pending_steps: Option<WacPendingSteps>,
#[serde(default)]
pub input_args: serde_json::Map<String, Value>,
/// Accumulated map of step_key → child job UUID across all dispatch rounds.
/// Unlike `pending_steps.job_ids` (cleared after completion), this persists
/// so the frontend can always resolve step keys to child job names.
#[serde(default)]
pub job_ids: serde_json::Map<String, Value>,
/// When set on a child job's checkpoint, indicates which step this child
/// should execute directly (instead of dispatching).
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(default)]
pub _executing_key: Option<String>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct WacPendingSteps {
pub mode: String,
pub keys: Vec<String>,
pub job_ids: serde_json::Map<String, Value>,
}
/// Load the WAC checkpoint from `v2_job_status.workflow_as_code_status._checkpoint`.
pub async fn load_checkpoint(db: &DB, job_id: &Uuid) -> error::Result<WacCheckpoint> {
let row: Option<Option<Value>> = sqlx::query_scalar(
"SELECT workflow_as_code_status->'_checkpoint' FROM v2_job_status WHERE id = $1",
)
.bind(job_id)
.fetch_optional(db)
.await?;
match row {
Some(Some(status)) => {
let checkpoint: WacCheckpoint = match serde_json::from_value(status) {
Ok(c) => c,
Err(e) => {
tracing::warn!(
job_id = %job_id,
error = %e,
"Failed to deserialize WAC checkpoint, resetting to empty"
);
WacCheckpoint::default()
}
};
Ok(checkpoint)
}
_ => Ok(WacCheckpoint::default()),
}
}
/// Save the WAC checkpoint to `v2_job_status.workflow_as_code_status._checkpoint`.
/// The top level of workflow_as_code_status is reserved for per-child-job timeline data.
pub async fn save_checkpoint(
db: &DB,
job_id: &Uuid,
checkpoint: &WacCheckpoint,
) -> error::Result<()> {
let status_json = serde_json::to_value(checkpoint)
.map_err(|e| Error::InternalErr(format!("Failed to serialize checkpoint: {e}")))?;
sqlx::query(
"INSERT INTO v2_job_status (id, workflow_as_code_status)
VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb))
ON CONFLICT (id) DO UPDATE SET
workflow_as_code_status = jsonb_set(
COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb),
'{_checkpoint}',
$2::jsonb
)",
)
.bind(job_id)
.bind(&status_json)
.execute(db)
.await
.map_err(|e| Error::InternalErr(format!("Failed to save WAC checkpoint: {e}")))?;
Ok(())
}
/// Process a completed child job result: add to checkpoint's completed_steps.
pub fn add_completed_step(checkpoint: &mut WacCheckpoint, step_key: &str, result: Value) {
checkpoint
.completed_steps
.insert(step_key.to_string(), result);
// If all pending steps are complete, clear pending
if let Some(ref pending) = checkpoint.pending_steps {
let all_done = pending
.keys
.iter()
.all(|k| checkpoint.completed_steps.contains_key(k));
if all_done {
checkpoint.pending_steps = None;
}
}
}
/// Persist a single inline-step checkpoint delta into the given transaction:
/// validate the source hash, add the step to `completed_steps`, save the
/// checkpoint, and write the `_step/<key>` timeline entry.
///
/// The caller owns the transaction and commits it. This lets the worker-side
/// `WacOutput::InlineCheckpoint` fallback arm add its own `UPDATE v2_job_queue
/// SET running = false` in the same transaction — preserving the original
/// all-or-nothing atomicity — while the API fast path simply commits after
/// the helper returns.
///
/// ## Concurrency model
///
/// The helper does a read-modify-write: `SELECT ... FOR UPDATE` → parse
/// `WacCheckpoint` → modify in Rust via `add_completed_step` → write the
/// full serialized `_checkpoint` back via `INSERT ... ON CONFLICT DO UPDATE`
/// plus a separate `UPDATE` for the `_step/<key>` timeline entry. The
/// important property of this pattern: each call **replaces the whole
/// `_checkpoint` object**, not individual `completed_steps[key]` entries.
/// That means distinct step keys do NOT protect concurrent callers from
/// overwriting each other — two writers that start from the same loaded
/// checkpoint will each produce a new serialized object that lacks the
/// other's step.
///
/// **Steady state (row exists)** — `SELECT ... FOR UPDATE` holds the row
/// lock until commit. The second concurrent caller blocks on the lock,
/// then re-reads the post-commit checkpoint (which already contains the
/// first caller's step), applies its own delta, and writes. No loss.
///
/// **First write (row does not yet exist)** — `SELECT ... FOR UPDATE` on a
/// WHERE clause that matches zero rows acquires no lock. Two concurrent
/// callers would both see `None`, both build a fresh `WacCheckpoint` from
/// scratch, and then race on the final `INSERT ... ON CONFLICT DO UPDATE`:
/// the second writer's `DO UPDATE SET workflow_as_code_status = jsonb_set(
/// ..., '{_checkpoint}', $2)` replaces the `_checkpoint` the first writer
/// just inserted, so the first writer's step is lost.
///
/// That race window is closed **on the client side** by the SDKs:
/// `WorkflowCtx._inline_lock` (Python `asyncio.Lock`) and
/// `WorkflowCtx._inlineChain` (TypeScript promise chain) serialize the
/// fast-path POSTs per workflow invocation. The lock wraps only the HTTP
/// call — `fn()` itself still runs in parallel across `asyncio.gather` /
/// `Promise.all` — so the only thing actually ordered is the sequence of
/// API requests, which is exactly what the helper needs to rely on.
///
/// **Future contributors: do not remove the SDK-side lock without also
/// fixing the server-side first-write guarantee (e.g. via a single-statement
/// merge-UPDATE that's cheap enough — see note below — or a pre-created
/// `v2_job_status` row).** The comment used to claim the SDKs could fire in
/// parallel without client-side serialization; that was wrong, because the
/// helper writes the whole `_checkpoint`.
///
/// Cross-process concurrency with the worker-side legacy fallback arm is
/// safe by construction: both paths receive the same `_StepSuspend` payload
/// (same `key`, same `result`, same `started_at`, same `duration_ms`), so
/// even if the fast path's commit and the worker arm's commit land out of
/// order for the same step, the worst case is a redundant idempotent write,
/// not a divergence.
///
/// ## Why not a single-statement merge-UPDATE?
///
/// A pure-SQL single-statement variant (pushing load-modify-save entirely
/// into `jsonb_set` + `jsonb_build_object` so correctness on the first
/// write comes from Postgres row locking rather than a client-side lock)
/// was prototyped and measured at ~80 ms per call in debug mode — the
/// nested `COALESCE(v2_job_status.workflow_as_code_status->'_checkpoint'
/// ->...)` accesses cause Postgres to evaluate the growing JSONB subtree
/// multiple times per call, and the `||` merges re-serialize the whole
/// object. The two-statement Rust-side load-modify-save below is ~10×
/// faster in practice, so we keep it and rely on the SDK-level lock.
pub async fn persist_inline_checkpoint_delta(
tx: &mut Transaction<'_, Postgres>,
job_id: &Uuid,
source_hash_hint: Option<&str>,
key: &str,
result: Value,
started_at: Option<&str>,
duration_ms: Option<u64>,
) -> error::Result<()> {
// Row-lock the existing checkpoint row (if any) for the duration of the
// transaction. NULL if the row doesn't exist yet — see the doc comment
// above for why the first-write race is accepted.
let row: Option<Option<Value>> = sqlx::query_scalar(
"SELECT workflow_as_code_status->'_checkpoint'
FROM v2_job_status WHERE id = $1 FOR UPDATE",
)
.bind(job_id)
.fetch_optional(&mut **tx)
.await?;
let mut checkpoint: WacCheckpoint = match row.flatten() {
Some(status) => serde_json::from_value(status).unwrap_or_else(|e| {
tracing::warn!(
job_id = %job_id,
error = %e,
"Failed to deserialize WAC checkpoint, resetting to empty"
);
WacCheckpoint::default()
}),
None => WacCheckpoint::default(),
};
// Source hash validation: detect if code changed between replays.
match source_hash_hint {
Some(hint) if !hint.is_empty() => {
if checkpoint.source_hash.is_empty() {
checkpoint.source_hash = hint.to_string();
} else if checkpoint.source_hash != hint {
return Err(Error::ExecutionErr(
"Workflow source code changed between replays. \
Cannot safely resume from checkpoint — step keys may have shifted. \
Please restart this workflow."
.to_string(),
));
}
}
_ => {
// Preview / inline jobs have no `runnable_id`, so the caller passes
// None (or Some("")). We can't validate drift for these — log once
// so operators can tell which jobs are running unguarded.
tracing::debug!(
job_id = %job_id,
"WAC v2 inline checkpoint without runnable hash — source-hash drift protection is off for this job"
);
}
}
tracing::info!(
job_id = %job_id,
step_key = %key,
"WAC v2 inline checkpoint — persisting step result"
);
add_completed_step(&mut checkpoint, key, result);
let status_json = serde_json::to_value(&checkpoint)
.map_err(|e| Error::InternalErr(format!("Failed to serialize checkpoint: {e}")))?;
sqlx::query(
"INSERT INTO v2_job_status (id, workflow_as_code_status)
VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb))
ON CONFLICT (id) DO UPDATE SET
workflow_as_code_status = jsonb_set(
COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb),
'{_checkpoint}',
$2::jsonb
)",
)
.bind(job_id)
.bind(&status_json)
.execute(&mut **tx)
.await
.map_err(|e| Error::InternalErr(format!("Failed to save WAC checkpoint: {e}")))?;
// Write the `_step/<key>` timeline entry. Fall back to now() when the
// client doesn't provide started_at (older SDK versions omit it).
let now_str = chrono::Utc::now().to_rfc3339();
let sa = started_at.unwrap_or(&now_str);
let mut timeline_val = serde_json::json!({
"scheduled_for": sa,
"started_at": sa,
"name": key,
});
if let Some(dur) = duration_ms {
timeline_val["duration_ms"] = serde_json::json!(dur);
}
let step_timeline_key = format!("_step/{}", key);
sqlx::query(
"UPDATE v2_job_status SET workflow_as_code_status = jsonb_set(
COALESCE(workflow_as_code_status, '{}'::jsonb),
ARRAY[$2],
$3
) WHERE id = $1",
)
.bind(job_id)
.bind(&step_timeline_key)
.bind(&timeline_val)
.execute(&mut **tx)
.await
.map_err(|e| Error::InternalErr(format!("Failed to write step timeline: {e}")))?;
Ok(())
}

133
backend/windmill-worker/src/bun_executor.rs
View File

@@ -2305,8 +2305,7 @@ pub async fn handle_wac_v2_output(
modules: &Option<std::collections::HashMap<String, windmill_common::scripts::ScriptModule>>,
) -> error::Result<Box<RawValue>> {
use crate::wac_executor::{
add_completed_step, load_checkpoint, parse_wac_output, update_checkpoint_for_dispatch,
WacOutput,
load_checkpoint, parse_wac_output, update_checkpoint_for_dispatch, WacOutput,
};
use serde_json::Value;
use windmill_common::get_latest_flow_version_info_for_path;
@@ -3143,104 +3142,44 @@ pub async fn handle_wac_v2_output(
}
};
let mut checkpoint = load_checkpoint(db, &job.id).await?;
// All-or-nothing: the checkpoint save, the `_step/<key>` timeline
// write, and the `running = false` queue reset must commit
// together. If we split them, a crash or failure in the middle
// would leave the job queued with `running = true` but a
// checkpoint that already contains the current step — any retry
// would then skip the step entirely. Passing the caller's `tx`
// into `persist_inline_checkpoint_delta` preserves the original
// atomicity from before the shared-helper refactor.
let source_hash = job.runnable_id.map(|h| h.0.to_string());
let mut tx = db.begin().await?;
// Source hash validation (same as Dispatch path)
let current_hash = job.runnable_id.map(|h| h.0.to_string()).unwrap_or_default();
if !current_hash.is_empty() {
if checkpoint.source_hash.is_empty() {
checkpoint.source_hash = current_hash.clone();
} else if checkpoint.source_hash != current_hash {
return Err(error::Error::ExecutionErr(
"Workflow source code changed between replays. \
Cannot safely resume from checkpoint — step keys may have shifted. \
Please restart this workflow."
.to_string(),
));
}
}
crate::wac_executor::persist_inline_checkpoint_delta(
&mut tx,
&job.id,
source_hash.as_deref(),
&key,
value,
started_at.as_deref(),
duration_ms,
)
.await?;
tracing::info!(
job_id = %job.id,
step_key = %key,
"WAC v2 inline checkpoint — persisting step result"
);
// Reset running=false so the job is immediately eligible for pickup.
// Unlike dispatch (which sets suspend>0), inline checkpoints don't suspend —
// the job should be re-run right away to continue past the cached step.
sqlx::query!(
"UPDATE v2_job_queue SET running = false, started_at = null WHERE id = $1",
job.id,
)
.execute(&mut *tx)
.await
.map_err(|e| {
error::Error::internal_err(format!(
"Failed to reset running state for inline checkpoint: {e}"
))
})?;
add_completed_step(&mut checkpoint, &key, value);
// Save checkpoint + write step timeline entry + reset running in a single transaction
{
let mut tx = db.begin().await?;
let status_json = serde_json::to_value(&checkpoint).map_err(|e| {
error::Error::internal_err(format!("Failed to serialize checkpoint: {e}"))
})?;
sqlx::query(
"INSERT INTO v2_job_status (id, workflow_as_code_status)
VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb))
ON CONFLICT (id) DO UPDATE SET
workflow_as_code_status = jsonb_set(
COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb),
'{_checkpoint}',
$2::jsonb
)",
)
.bind(&job.id)
.bind(&status_json)
.execute(&mut *tx)
.await
.map_err(|e| {
error::Error::internal_err(format!("Failed to save WAC checkpoint: {e}"))
})?;
// Write timeline entry for the inline step (keyed as _step/<key>).
// Fall back to now() when the client doesn't provide started_at
// (older windmill-client versions omit it).
{
let now_str = chrono::Utc::now().to_rfc3339();
let sa = started_at.as_deref().unwrap_or(&now_str);
let mut timeline_val = serde_json::json!({
"scheduled_for": sa,
"started_at": sa,
"name": key,
});
if let Some(dur) = duration_ms {
timeline_val["duration_ms"] = serde_json::json!(dur);
}
let step_timeline_key = format!("_step/{}", key);
sqlx::query(
"UPDATE v2_job_status SET workflow_as_code_status = jsonb_set(
COALESCE(workflow_as_code_status, '{}'::jsonb),
ARRAY[$2],
$3
) WHERE id = $1",
)
.bind(&job.id)
.bind(&step_timeline_key)
.bind(&timeline_val)
.execute(&mut *tx)
.await
.map_err(|e| {
error::Error::internal_err(format!("Failed to write step timeline: {e}"))
})?;
}
// Reset running=false so the job is immediately eligible for pickup.
// Unlike dispatch (which sets suspend>0), inline checkpoints don't suspend —
// the job should be re-run right away to continue past the cached step.
sqlx::query!(
"UPDATE v2_job_queue SET running = false, started_at = null WHERE id = $1",
job.id,
)
.execute(&mut *tx)
.await
.map_err(|e| {
error::Error::internal_err(format!(
"Failed to reset running state for inline checkpoint: {e}"
))
})?;
tx.commit().await?;
}
tx.commit().await?;
Err(error::Error::WacSuspended(format!(
"WAC v2 job {} inline checkpoint for step {}",

112
backend/windmill-worker/src/wac_executor.rs
View File

@@ -1,4 +1,4 @@
use serde::{Deserialize, Serialize};
use serde::Deserialize;
use serde_json::value::RawValue;
use serde_json::Value;
use uuid::Uuid;
@@ -6,35 +6,13 @@ use uuid::Uuid;
use windmill_common::error::{self, Error};
use windmill_common::DB;
/// Checkpoint state persisted across workflow invocations.
#[derive(Debug, Serialize, Deserialize, Default, Clone)]
pub struct WacCheckpoint {
#[serde(default)]
pub source_hash: String,
#[serde(default)]
pub completed_steps: serde_json::Map<String, Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub pending_steps: Option<WacPendingSteps>,
#[serde(default)]
pub input_args: serde_json::Map<String, Value>,
/// Accumulated map of step_key → child job UUID across all dispatch rounds.
/// Unlike `pending_steps.job_ids` (cleared after completion), this persists
/// so the frontend can always resolve step keys to child job names.
#[serde(default)]
pub job_ids: serde_json::Map<String, Value>,
/// When set on a child job's checkpoint, indicates which step this child
/// should execute directly (instead of dispatching).
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(default)]
pub _executing_key: Option<String>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct WacPendingSteps {
pub mode: String,
pub keys: Vec<String>,
pub job_ids: serde_json::Map<String, Value>,
}
// Checkpoint model + persistence primitives live in windmill-common so the
// API server can use them without pulling in the full worker crate. Re-export
// here for historical call sites inside windmill-worker.
pub use windmill_common::wac::{
load_checkpoint, persist_inline_checkpoint_delta, save_checkpoint, WacCheckpoint,
WacPendingSteps,
};
/// Output from a single WAC invocation (parsed from result.json).
#[derive(Debug, Deserialize)]
@@ -106,63 +84,6 @@ fn default_dispatch_type() -> String {
"inline".to_string()
}
/// Load the WAC checkpoint from `v2_job_status.workflow_as_code_status._checkpoint`.
pub async fn load_checkpoint(db: &DB, job_id: &Uuid) -> error::Result<WacCheckpoint> {
let row: Option<Option<Value>> = sqlx::query_scalar(
"SELECT workflow_as_code_status->'_checkpoint' FROM v2_job_status WHERE id = $1",
)
.bind(job_id)
.fetch_optional(db)
.await?;
match row {
Some(Some(status)) => {
let checkpoint: WacCheckpoint = match serde_json::from_value(status) {
Ok(c) => c,
Err(e) => {
tracing::warn!(
job_id = %job_id,
error = %e,
"Failed to deserialize WAC checkpoint, resetting to empty"
);
WacCheckpoint::default()
}
};
Ok(checkpoint)
}
_ => Ok(WacCheckpoint::default()),
}
}
/// Save the WAC checkpoint to `v2_job_status.workflow_as_code_status._checkpoint`.
/// The top level of workflow_as_code_status is reserved for per-child-job timeline data.
pub async fn save_checkpoint(
db: &DB,
job_id: &Uuid,
checkpoint: &WacCheckpoint,
) -> error::Result<()> {
let status_json = serde_json::to_value(checkpoint)
.map_err(|e| Error::InternalErr(format!("Failed to serialize checkpoint: {e}")))?;
sqlx::query(
"INSERT INTO v2_job_status (id, workflow_as_code_status)
VALUES ($1, jsonb_build_object('_checkpoint', $2::jsonb))
ON CONFLICT (id) DO UPDATE SET
workflow_as_code_status = jsonb_set(
COALESCE(v2_job_status.workflow_as_code_status, '{}'::jsonb),
'{_checkpoint}',
$2::jsonb
)",
)
.bind(job_id)
.bind(&status_json)
.execute(db)
.await
.map_err(|e| Error::InternalErr(format!("Failed to save WAC checkpoint: {e}")))?;
Ok(())
}
/// Parse the WAC result from result.json content.
pub fn parse_wac_output(result: &RawValue) -> error::Result<WacOutput> {
serde_json::from_str(result.get())
@@ -192,23 +113,6 @@ pub fn update_checkpoint_for_dispatch(
checkpoint.pending_steps = Some(pending);
}
/// Process a completed child job result: add to checkpoint's completed_steps.
pub fn add_completed_step(checkpoint: &mut WacCheckpoint, step_key: &str, result: Value) {
checkpoint
.completed_steps
.insert(step_key.to_string(), result);
// If all pending steps are complete, clear pending
if let Some(ref pending) = checkpoint.pending_steps {
let all_done = pending
.keys
.iter()
.all(|k| checkpoint.completed_steps.contains_key(k));
if all_done {
checkpoint.pending_steps = None;
}
}
}
/// Check if all pending parallel steps are complete.
pub fn all_pending_complete(checkpoint: &WacCheckpoint) -> bool {
match &checkpoint.pending_steps {

83
python-client/wmill/wmill/client.py
View File

@@ -2416,6 +2416,23 @@ class WorkflowCtx:
self._counters: dict[str, int] = {}
self._pending: list = []
self._executing_key: str | None = checkpoint.get("_executing_key")
# Reuse a single httpx.AsyncClient across all fast-path step() calls
# in this workflow invocation. Instantiating a fresh client per call
# allocates a new connection pool each time — on localhost this adds
# ~15ms per step, dominating the end-to-end cost. Lazily built so no
# client is created for workflows that never hit the fast path.
self._inline_http_client: "httpx.AsyncClient | None" = None
# Serializes fast-path POSTs across concurrent step() calls within
# one workflow invocation. Wraps only the HTTP call, not fn() — so
# `asyncio.gather(step("a", fn_a), step("b", fn_b))` still runs the
# two fn() bodies in parallel, only the API requests are ordered.
# This closes the first-write race window against `SELECT FOR UPDATE`
# on a not-yet-created `v2_job_status` row: concurrent POSTs would
# both see None and both overwrite each other's checkpoint because
# the helper writes the whole serialized `_checkpoint` object, not
# a single `completed_steps[key]`. Lazily built so the ctx can be
# constructed outside an event loop (tests do this).
self._inline_lock: "_asyncio.Lock | None" = None
def _alloc_key(self, name: str = "step") -> str:
"""Name-based key: ``double`` for first call, ``double_2``, ``double_3`` for subsequent."""
@@ -2543,6 +2560,54 @@ class WorkflowCtx:
result = await result
duration_ms = int((_time_mod.monotonic() - t0) * 1000)
# Fast path: POST the delta to the new per-job API endpoint and return
# the result directly, letting the workflow subprocess continue into
# the next step() without unwinding. On any failure — network, auth,
# timeout, source-hash mismatch, old backend without the endpoint —
# fall through to raising _StepSuspend so the worker takes the legacy
# suspend-and-replay path. Gated by WM_WAC_INLINE_FAST_PATH (default
# on) so the old behavior stays reachable for A/B testing and rollback.
_fast_path_flag = os.environ.get("WM_WAC_INLINE_FAST_PATH", "1").strip().lower()
_fast_path_enabled = _fast_path_flag not in ("0", "false", "off", "no")
_job_id = os.environ.get("WM_JOB_ID")
_workspace = os.environ.get("WM_WORKSPACE")
_base = os.environ.get("BASE_INTERNAL_URL")
_token = os.environ.get("WM_TOKEN")
if _fast_path_enabled and _job_id and _workspace and _base and _token:
try:
if self._inline_lock is None:
self._inline_lock = _asyncio.Lock()
# Lock wraps only the POST, not fn() above — concurrent
# step() calls run fn() in parallel, then serialize on
# the API request.
async with self._inline_lock:
if self._inline_http_client is None:
self._inline_http_client = httpx.AsyncClient(
timeout=httpx.Timeout(10.0),
headers={
"Authorization": f"Bearer {_token}",
"Content-Type": "application/json",
},
)
_resp = await self._inline_http_client.post(
f"{_base}/api/w/{_workspace}/jobs/wac/inline_checkpoint/{_job_id}",
json={
"key": key,
"result": result,
"started_at": started_at,
"duration_ms": duration_ms,
},
)
_resp.raise_for_status()
return result
except Exception as _e:
logger.info(
"WAC v2 inline fast path failed for key %s, falling back to suspend: %s",
key,
_e,
)
# fall through to the legacy suspend path
raise _StepSuspend({
"mode": "inline_checkpoint",
"steps": [],
@@ -2873,7 +2938,23 @@ async def _run_workflow_async(func, checkpoint: dict, input_args: dict):
}
return {"type": "dispatch", **info}
finally:
_workflow_ctx.reset(token)
# Close the lazily-built fast-path httpx client so we don't emit
# asyncio ResourceWarning('unclosed transport') on shutdown and don't
# leak connection pools when this coroutine is driven from a
# long-lived loop (tests, REPL, embedded callers).
#
# Wrapped in its own try/finally so that asyncio.CancelledError
# (which is a BaseException since Python 3.8) during aclose() does
# not skip the _workflow_ctx.reset(token) below.
try:
if ctx._inline_http_client is not None:
try:
await ctx._inline_http_client.aclose()
except Exception:
pass
ctx._inline_http_client = None
finally:
_workflow_ctx.reset(token)
def _run_workflow(func, checkpoint: dict, input_args: dict):

76
typescript-client/client.ts
View File

@@ -1494,6 +1494,16 @@ export class WorkflowCtx {
private _suspended = false;
/** When set, the task matching this key executes its inner function directly */
_executingKey: string | null;
/** Serializes fast-path POSTs across concurrent step() calls within one
* workflow invocation. Wraps only the HTTP call, not fn() — so
* `Promise.all([step("a", fn_a), step("b", fn_b)])` still runs the two
* fn() bodies in parallel, only the API requests are ordered. This
* closes the first-write race window against `SELECT FOR UPDATE` on a
* not-yet-created `v2_job_status` row: two concurrent POSTs would both
* see None and overwrite each other's checkpoint because the helper
* writes the whole serialized `_checkpoint` object, not a single
* `completed_steps[key]`. Initialized to a resolved promise. */
private _inlineChain: Promise<void> = Promise.resolve();
constructor(checkpoint: Record<string, any> = {}) {
this.completed = checkpoint?.completed_steps ?? {};
@@ -1652,6 +1662,72 @@ export class WorkflowCtx {
const t0 = Date.now();
const result = await fn();
const durationMs = Date.now() - t0;
// Fast path: POST the delta to the new per-job API endpoint and return the
// result directly so the workflow subprocess continues into the next step()
// without unwinding. Concurrent step() calls (e.g. inside Promise.all) run
// their fn() bodies in parallel, then serialize the API POSTs via a
// per-ctx promise chain (`this._inlineChain`). Serializing the POSTs is
// required because the backend helper writes the whole serialized
// `_checkpoint` object per call, and two concurrent writes against a
// not-yet-created `v2_job_status` row would both see None under
// `SELECT FOR UPDATE` and overwrite each other.
//
// On any failure — network, auth, timeout, source-hash mismatch, old
// backend without the endpoint — fall through to throwing StepSuspend so
// the worker takes the legacy suspend-and-replay path. Gated by
// WM_WAC_INLINE_FAST_PATH (default on) so the old behavior stays
// reachable for A/B testing and rollback.
const fastPathFlagRaw = (getEnv("WM_WAC_INLINE_FAST_PATH") ?? "1").trim().toLowerCase();
const fastPathEnabled =
fastPathFlagRaw !== "0" &&
fastPathFlagRaw !== "false" &&
fastPathFlagRaw !== "off" &&
fastPathFlagRaw !== "no";
const jobId = getEnv("WM_JOB_ID");
const workspace = getEnv("WM_WORKSPACE");
if (fastPathEnabled && jobId && workspace && OpenAPI.BASE && OpenAPI.TOKEN) {
const chainTail = this._inlineChain.then(async () => {
const ctrl = new AbortController();
const t = setTimeout(() => ctrl.abort(), 10_000);
try {
const resp = await fetch(
`${OpenAPI.BASE}/w/${workspace}/jobs/wac/inline_checkpoint/${jobId}`,
{
method: "POST",
headers: {
"Content-Type": "application/json",
Authorization: `Bearer ${OpenAPI.TOKEN}`,
},
body: JSON.stringify({
key,
result,
started_at: startedAt,
duration_ms: durationMs,
}),
signal: ctrl.signal,
},
);
if (!resp.ok) {
throw new Error(`inline_checkpoint API ${resp.status}`);
}
} finally {
clearTimeout(t);
}
});
// Swallow chain errors so a past failure does not poison future awaits.
this._inlineChain = chainTail.catch(() => {});
try {
await chainTail;
return result as T;
} catch (e) {
console.log(
`WAC v2 inline fast path failed for key ${key}, falling back to suspend: ${e}`,
);
// fall through to the legacy suspend path below
}
}
throw new StepSuspend({ mode: "inline_checkpoint", steps: [], key, result, started_at: startedAt, duration_ms: durationMs });
}
}