feat: flow for loop squashing (#7107)

* feat: flow for loop optimization

* fix: make dedicated flows work with flow nodes

* fix: transform dedicated inputs + better error handling + differentiate squash from same worker

* fix builds

* fix build

* fix build

* make dedicated worker private

* update ee ref

* frontend nits

* nit

* add trace to dedicated worker

* update ee ref

* Update ee-repo-ref.txt

---------

Co-authored-by: windmill-internal-app[bot] <windmill-internal-app[bot]@users.noreply.github.com>

Author: hugocasa

backend/ee-repo-ref.txt

@@ -1 +1 @@
-1de62d3f0901aac6ef9ddd7595e1032be56b1938
+f9434efa127dadb0a04c87d6fc368cc289d93501

backend/src/monitor.rs

@@ -2316,13 +2316,14 @@ async fn handle_zombie_jobs(db: &Pool<Postgres>, base_internal_url: &str, node_n
             );
 
             let error_message = format!(
-            "Job timed out after no ping from job since {} (ZOMBIE_JOB_TIMEOUT: {}, reason: {:?}).\nThis likely means that the job died on worker {}, OOM are a common reason for worker crashes.\nCheck the workers around the time of the last ping and the exit code if any.",
-            job.last_ping.unwrap_or_default(),
-            *ZOMBIE_JOB_TIMEOUT,
-            error_kind.to_string(),
-            job.worker.clone().unwrap_or_default(),
-        );
+                "Job timed out after no ping from job since {} (ZOMBIE_JOB_TIMEOUT: {}, reason: {:?}).\nThis likely means that the job died on worker {}, OOM are a common reason for worker crashes.\nCheck the workers around the time of the last ping and the exit code if any.",
+                job.last_ping.unwrap_or_default(),
+                *ZOMBIE_JOB_TIMEOUT,
+                error_kind.to_string(),
+                job.worker.clone().unwrap_or_default(),
+            );
             let memory_peak = job.memory_peak.unwrap_or(0);
+            let (_, killpill_rx_never_used) = KillpillSender::new(1);
             let _ = handle_job_error(
                 db,
                 &client,
@@ -2335,6 +2336,7 @@ async fn handle_zombie_jobs(db: &Pool<Postgres>, base_internal_url: &str, node_n
                 "",
                 node_name,
                 send_result_never_used,
+                &killpill_rx_never_used,
                 #[cfg(feature = "benchmark")]
                 &mut windmill_common::bench::BenchmarkIter::new(),
             )

backend/tests/worker.rs

@@ -217,6 +217,7 @@ async fn test_deno_flow(db: Pool<Postgres>) -> anyhow::Result<()> {
                         iterator: InputTransform::Javascript { expr: "result".to_string() },
                         skip_failures: false,
                         parallel: false,
+                        squash: None,
                         parallelism: None,
                         modules: vec![FlowModule {
                             id: "c".to_string(),
@@ -396,6 +397,7 @@ async fn test_deno_flow_same_worker(db: Pool<Postgres>) -> anyhow::Result<()> {
                         iterator: InputTransform::Static { value: windmill_common::worker::to_raw_value(&[1, 2, 3]) },
                         skip_failures: false,
                         parallel: false,
+                        squash: None,
                         parallelism: None,
                         modules: vec![
                             FlowModule {

backend/windmill-api/src/flows.rs

@@ -1583,6 +1583,7 @@ mod tests {
                         skip_failures: true,
                         parallel: false,
                         parallelism: None,
+                        squash: None,
                     }),
                     stop_after_if: Some(StopAfterIf {
                         expr: "previous.isEmpty()".to_string(),

backend/windmill-common/src/flows.rs

@@ -869,6 +869,8 @@ pub enum FlowModuleValue {
         parallel: bool,
         #[serde(skip_serializing_if = "Option::is_none")]
         parallelism: Option<InputTransform>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        squash: Option<bool>,
     },
 
     /// While loop node
@@ -878,6 +880,8 @@ pub enum FlowModuleValue {
         modules_node: Option<FlowNodeId>,
         #[serde(default = "default_false")]
         skip_failures: bool,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        squash: Option<bool>,
     },
 
     /// Branch-one node
@@ -989,6 +993,7 @@ struct UntaggedFlowModuleValue {
     assets: Option<Vec<AssetWithAltAccessType>>,
     tools: Option<Vec<AgentTool>>,
     pass_flow_input_directly: Option<bool>,
+    squash: Option<bool>,
 }
 
 impl<'de> Deserialize<'de> for FlowModuleValue {
@@ -1027,13 +1032,15 @@ impl<'de> Deserialize<'de> for FlowModuleValue {
                 skip_failures: untagged.skip_failures.unwrap_or(true),
                 parallel: untagged.parallel.unwrap_or(false),
                 parallelism: untagged.parallelism,
+                squash: untagged.squash,
             }),
             "whileloopflow" => Ok(FlowModuleValue::WhileloopFlow {
                 modules: untagged
                     .modules
                     .ok_or_else(|| serde::de::Error::missing_field("modules"))?,
                 modules_node: untagged.modules_node,
                 skip_failures: untagged.skip_failures.unwrap_or(false),
+                squash: untagged.squash,
             }),
             "branchone" => Ok(FlowModuleValue::BranchOne {
                 branches: untagged

backend/windmill-queue/src/jobs.rs

@@ -24,6 +24,9 @@ use serde::{ser::SerializeMap, Serialize};
 use serde_json::{json, value::RawValue};
 use sqlx::{types::Json, Pool, Postgres, Transaction};
 use sqlx::{Encode, PgExecutor};
+use tokio::sync::mpsc::Sender;
+use tokio::sync::oneshot;
+use tokio::task::JoinHandle;
 use tokio::{sync::RwLock, time::sleep};
 use ulid::Ulid;
 use uuid::Uuid;
@@ -137,7 +140,7 @@ pub struct CanceledBy {
     pub reason: Option<String>,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct JobCompleted {
     pub job: MiniCompletedJob,
     pub preprocessed_args: Option<HashMap<String, Box<RawValue>>>,
@@ -151,6 +154,10 @@ pub struct JobCompleted {
     pub duration: Option<i64>,
     pub has_stream: Option<bool>,
     pub from_cache: Option<bool>,
+    #[serde(skip)]
+    pub flow_runners: Option<Arc<FlowRunners>>,
+    #[serde(skip)]
+    pub done_tx: Option<oneshot::Sender<()>>,
 }
 
 pub async fn cancel_single_job<'c>(
@@ -2271,6 +2278,8 @@ pub struct JobAndPerms {
     pub parent_runnable_path: Option<String>,
     pub token: String,
     pub precomputed_agent_info: Option<PrecomputedAgentInfo>,
+    #[serde(skip)]
+    pub flow_runners: Option<Arc<FlowRunners>>,
 }
 impl PulledJob {
     pub async fn get_job_and_perms(self, db: &DB) -> JobAndPerms {
@@ -2302,6 +2311,7 @@ impl PulledJob {
             parent_runnable_path: self.parent_runnable_path,
             token,
             precomputed_agent_info: None,
+            flow_runners: None,
         }
     }
 }
@@ -2492,6 +2502,8 @@ impl PulledJobResult {
                     duration: None,
                     has_stream: Some(false),
                     from_cache: None,
+                    flow_runners: None,
+                    done_tx: None,
                 }),
             ),
             PulledJobResult { job: Some(job), error_while_preprocessing: Some(e), .. } => Err(
@@ -2511,6 +2523,8 @@ impl PulledJobResult {
                     duration: None,
                     has_stream: Some(false),
                     from_cache: None,
+                    flow_runners: None,
+                    done_tx: None,
                 }),
             ),
             PulledJobResult { job, .. } => Ok(job),
@@ -5715,10 +5729,74 @@ async fn restarted_flows_resolution(
     ))
 }
 
+
+
+// Wrapper struct to send both job and optional flow_runners to dedicated workers
+pub struct DedicatedWorkerJob {
+    pub job: Arc<MiniPulledJob>,
+    pub flow_runners: Option<Arc<FlowRunners>>,
+    pub done_tx: Option<oneshot::Sender<()>>,
+}
+
+#[derive(Debug)]
+pub struct FlowRunners {
+    pub runners: HashMap<String, Sender<DedicatedWorkerJob>>,
+    pub handles: Vec<JoinHandle<()>>,
+    pub job_id: Uuid,
+}
+
+impl Drop for FlowRunners {
+    fn drop(&mut self) {
+        let total_runners = self.handles.len();
+        tracing::info!("dropping {} flow runners for job {}", total_runners, self.job_id);
+
+        // First, drop all senders to signal workers to stop gracefully
+        self.runners.clear();
+
+        // Spawn a background task to wait with timeout and abort if needed
+        let handles = std::mem::take(&mut self.handles);
+        let job_id = self.job_id;
+
+        tokio::spawn(async move {
+            // Extract abort handles before consuming the join handles
+            let abort_handles: Vec<_> = handles.iter().map(|h| h.abort_handle()).collect();
+
+            // Wait up to 5 seconds for natural termination
+            let timeout_result = tokio::time::timeout(
+                tokio::time::Duration::from_secs(5),
+                futures::future::join_all(handles)
+            ).await;
+
+            match timeout_result {
+                Ok(_) => {
+                    tracing::info!("all {} flow runners for job {} terminated gracefully", total_runners, job_id);
+                }
+                Err(_) => {
+                    // Timeout reached, abort only the handles that haven't finished
+                    let mut aborted = 0;
+                    for abort_handle in abort_handles {
+                        if !abort_handle.is_finished() {
+                            abort_handle.abort();
+                            aborted += 1;
+                        }
+                    }
+                    let graceful = total_runners - aborted;
+                    tracing::warn!(
+                        "flow runners for job {}: {} terminated gracefully, {} aborted after 5s timeout",
+                        job_id, graceful, aborted
+                    );
+                }
+            }
+        });
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct SameWorkerPayload {
     pub job_id: Uuid,
     pub recoverable: bool,
+    #[serde(skip)]
+    pub flow_runners: Option<Arc<FlowRunners>>,
 }
 
 pub async fn get_same_worker_job(

backend/windmill-queue/src/lib.rs

@@ -6,7 +6,7 @@
  * LICENSE-AGPL for a copy of the license.
  */
 
-mod jobs;
+pub mod jobs;
 #[cfg(feature = "private")]
 pub mod jobs_ee;
 pub mod jobs_oss;

backend/windmill-worker/src/ai/tools.rs

@@ -529,6 +529,7 @@ async fn execute_windmill_tool(
             &mut occupancy_metrics_spawn,
             &mut killpill_rx_spawn,
             None,
+            None,
             #[cfg(feature = "benchmark")]
             &mut bench_spawn,
         )
@@ -656,8 +657,9 @@ async fn handle_tool_execution_success(
         ..
     }) = send_result.as_ref()
     {
+        let result = result.clone();
         ctx.job_completed_tx
-            .send(send_result.as_ref().unwrap().result.clone(), true)
+            .send(send_result.unwrap().result, true)
             .await
             .map_err(to_anyhow)?;
         result

backend/windmill-worker/src/bun_executor.rs

@@ -9,10 +9,8 @@ use serde_json::value::RawValue;
 
 use uuid::Uuid;
 use windmill_parser_ts::remove_pinned_imports;
-use windmill_queue::{append_logs, CanceledBy, MiniPulledJob, PrecomputedAgentInfo};
 
-#[cfg(feature = "enterprise")]
-use crate::common::build_envs_map;
+use windmill_queue::{append_logs, CanceledBy, MiniPulledJob, PrecomputedAgentInfo};
 
 use crate::{
     common::{
@@ -38,12 +36,6 @@ use tokio::{fs::File, process::Command};
 
 use tokio::io::AsyncReadExt;
 
-#[cfg(feature = "enterprise")]
-use tokio::sync::mpsc::Receiver;
-
-#[cfg(feature = "enterprise")]
-use windmill_common::variables;
-
 use windmill_common::{
     error::{self, Result},
     get_latest_hash_for_path,
@@ -1567,10 +1559,18 @@ pub async fn get_common_bun_proc_envs(base_internal_url: Option<&str>) -> HashMa
     return bun_envs;
 }
 
-#[cfg(feature = "enterprise")]
-use crate::{dedicated_worker::handle_dedicated_process, JobCompletedSender};
+#[cfg(feature = "private")]
+use crate::{
+    common::build_envs_map, dedicated_worker_oss::handle_dedicated_process, JobCompletedSender,
+};
+#[cfg(feature = "private")]
+use tokio::sync::mpsc::Receiver;
+#[cfg(feature = "private")]
+use windmill_common::variables;
+#[cfg(feature = "private")]
+use windmill_queue::DedicatedWorkerJob;
 
-#[cfg(feature = "enterprise")]
+#[cfg(feature = "private")]
 pub async fn start_worker(
     requirements_o: Option<String>,
     codebase: Option<String>,
@@ -1584,8 +1584,9 @@ pub async fn start_worker(
     script_path: &str,
     token: &str,
     job_completed_tx: JobCompletedSender,
-    jobs_rx: Receiver<std::sync::Arc<MiniPulledJob>>,
+    jobs_rx: Receiver<DedicatedWorkerJob>,
     killpill_rx: tokio::sync::broadcast::Receiver<()>,
+    client: windmill_common::client::AuthedClient,
 ) -> Result<()> {
     let mut logs = "".to_string();
     let mut mem_peak: i32 = 0;
@@ -1737,16 +1738,21 @@ BigInt.prototype.toJSON = function () {{
 
 console.log('start');
 
+function getArgs(line) {{
+    let {{ {spread} }} = JSON.parse(line)
+    {dates}
+    return [ {spread} ];
+}}
+
 for await (const line of Readline.createInterface({{ input: process.stdin }})) {{
     {print_lines}
 
     if (line === "end") {{
         process.exit(0);
     }}
     try {{
-        let {{ {spread} }} = JSON.parse(line)
-        {dates}
-        let res = await Main.main(...[ {spread} ]);
+        const args = getArgs(line);
+        const res = await Main.main(...args);
         console.log("wm_res[success]:" + JSON.stringify(res ?? null, (key, value) => typeof value === 'undefined' ? null : value));
     }} catch (e) {{
         console.log("wm_res[error]:" + JSON.stringify({{ message: e.message, name: e.name, stack: e.stack, line: line }}));
@@ -1812,6 +1818,7 @@ for await (const line of Readline.createInterface({{ input: process.stdin }})) {
             db,
             &script_path,
             "nodejs",
+            client,
         )
         .await
     } else {
@@ -1838,6 +1845,7 @@ for await (const line of Readline.createInterface({{ input: process.stdin }})) {
             db,
             script_path,
             "bun",
+            client,
         )
         .await
     }