From 08a3f7964272abe4ab5ade36ff36f35ae38c85a0 Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Fri, 12 Jan 2024 00:21:52 +0100
Subject: [PATCH] feat: make dedicated workers for flows able to share runtime
 for the same scripts

---
 backend/windmill-worker/src/worker.rs | 61 ++++++++++++++++++---------
 1 file changed, 41 insertions(+), 20 deletions(-)

diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs
index 74ee933a09..b7d3162b02 100644
--- a/backend/windmill-worker/src/worker.rs
+++ b/backend/windmill-worker/src/worker.rs
@@ -1183,8 +1183,8 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
                         if let Ok(flow) = value {
                             let workers = spawn_dedicated_workers_for_flow(
                                 &flow.modules,
-                                &_wp.path,
                                 &_wp.workspace_id,
+                                &_wp.path,
                                 killpill_tx.clone(),
                                 &killpill_rx,
                                 db,
@@ -1195,7 +1195,13 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
                             )
                             .await;
                             workers.into_iter().for_each(|(path, sender, handle)| {
-                                dedicated_handles.push(handle);
+                                tracing::info!(
+                                    "spawned dedicated worker for flow: {}",
+                                    path.as_str()
+                                );
+                                if let Some(h) = handle {
+                                    dedicated_handles.push(h);
+                                }
                                 hm.insert(path, sender);
                             });
                         }
@@ -1224,7 +1230,9 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
                 )
                 .await
                 {
-                    dedicated_handles.push(handle);
+                    if let Some(h) = handle {
+                        dedicated_handles.push(h);
+                    }
                     hm.insert(path, sender);
                 } else {
                     tracing::error!(
@@ -1754,7 +1762,7 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
     println!("worker {} exited", i_worker);
 }
 
-type DedicatedWorker = (String, Sender<Arc<QueuedJob>>, JoinHandle<()>);
+type DedicatedWorker = (String, Sender<Arc<QueuedJob>>, Option<JoinHandle<()>>);
 
 // spawn one dedicated worker per compatible steps of the flow, associating the node id to the dedicated worker channel send
 #[async_recursion]
@@ -1771,24 +1779,37 @@ async fn spawn_dedicated_workers_for_flow(
     job_completed_tx: &JobCompletedSender,
 ) -> Vec<DedicatedWorker> {
     let mut workers = vec![];
+    let mut script_path_to_worker: HashMap<String, Sender<Arc<QueuedJob>>> = HashMap::new();
     for module in modules.iter() {
         match &module.value {
             FlowModuleValue::Script { path, hash, .. } => {
-                if let Some(dedi_w) = spawn_dedicated_worker(
-                    SpawnWorker::Script { path: path.to_string(), hash: hash.clone() },
-                    w_id,
-                    killpill_tx.clone(),
-                    killpill_rx,
-                    db,
-                    worker_dir,
-                    base_internal_url,
-                    worker_name,
-                    job_completed_tx,
-                    Some(module.id.clone()),
-                )
-                .await
-                {
-                    workers.push(dedi_w);
+                let key = format!(
+                    "{}:{}",
+                    path,
+                    hash.clone()
+                        .map(|x| x.to_string())
+                        .unwrap_or_else(|| "".to_string())
+                );
+                if let Some(sender) = script_path_to_worker.get(&key) {
+                    workers.push((module.id.clone(), sender.clone(), None));
+                } else {
+                    if let Some(dedi_w) = spawn_dedicated_worker(
+                        SpawnWorker::Script { path: path.to_string(), hash: hash.clone() },
+                        w_id,
+                        killpill_tx.clone(),
+                        killpill_rx,
+                        db,
+                        worker_dir,
+                        base_internal_url,
+                        worker_name,
+                        job_completed_tx,
+                        Some(module.id.clone()),
+                    )
+                    .await
+                    {
+                        script_path_to_worker.insert(key, dedi_w.1.clone());
+                        workers.push(dedi_w);
+                    }
                 }
             }
             FlowModuleValue::ForloopFlow { modules, .. } => {
@@ -2051,7 +2072,7 @@ async fn spawn_dedicated_worker(
                 tracing::error!("error in dedicated worker: {:?}", e);
             };
         });
-        return Some((node_id.unwrap_or(path2), dedicated_worker_tx, handle));
+        return Some((node_id.unwrap_or(path2), dedicated_worker_tx, Some(handle)));
         // (Some(dedi_path), Some(dedicated_worker_tx), Some(handle))
     }
 }