From 4661dd2cea343572099ca66459d3db5f3c7879c8 Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Fri, 10 Nov 2023 20:41:43 +0100
Subject: [PATCH] feat: exit dedicated worker upon new deployment

---
 backend/src/monitor.rs                         |  6 ++++--
 backend/windmill-worker/src/python_executor.rs |  2 ++
 backend/windmill-worker/src/worker.rs          | 18 ++++++++++--------
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/backend/src/monitor.rs b/backend/src/monitor.rs
index 8f2428c00a..da0d52069e 100644
--- a/backend/src/monitor.rs
+++ b/backend/src/monitor.rs
@@ -531,9 +531,11 @@ pub async fn reload_worker_config(
     } else {
         let wc = WORKER_CONFIG.read().await;
         let config = config.unwrap();
-        if *wc != config {
+        if *wc != config || config.dedicated_worker.is_some() {
             if kill_if_change {
-                if (*wc).dedicated_worker != config.dedicated_worker {
+                if config.dedicated_worker.is_some()
+                    || (*wc).dedicated_worker != config.dedicated_worker
+                {
                     tracing::info!("Dedicated worker config changed, sending killpill. Expecting to be restarted by supervisor.");
                     let _ = tx.send(());
                 }
diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs
index 263b105527..4113a0c8fe 100644
--- a/backend/windmill-worker/src/python_executor.rs
+++ b/backend/windmill-worker/src/python_executor.rs
@@ -872,6 +872,8 @@ replace_nan = re.compile(r'\bNaN\b')
 sys.stdout.write('start\n')
 
 for line in sys.stdin:
+    if line == 'end\n':
+        break
     kwargs = json.loads(line, strict=False)
     args = {{}}
 {indented_transforms}
diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs
index 1d85afdb35..cb952443e7 100644
--- a/backend/windmill-worker/src/worker.rs
+++ b/backend/windmill-worker/src/worker.rs
@@ -992,10 +992,11 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
                     #[cfg(feature = "benchmark")]
                     let process_start = Instant::now();
 
+                    let is_dependency_job = matches!(jc.job.job_kind, JobKind::Dependencies);
                     handle_receive_completed_job(
                         jc,
                         base_internal_url2,
-                        db2,
+                        db2.clone(),
                         worker_dir2,
                         same_worker_tx2,
                         rsmq2,
@@ -1040,6 +1041,13 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
                     }
 
                     thread_count.fetch_sub(1, Ordering::SeqCst);
+                    if is_dependency_job && is_dedicated_worker {
+                        tracing::error!("Dedicated worker executed a dependency job, a new script has been deployed. Exiting expecting to be restarted.");
+                        sqlx::query!("UPDATE config SET config = config WHERE name = $1", format!("worker__{}", *WORKER_GROUP))
+                            .execute(&db2)
+                            .await
+                            .expect("update config to trigger restart of all dedicated workers at that config");
+                    }
                 });
             } else {
                 handle_receive_completed_job(
@@ -1062,12 +1070,6 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
             tokio::time::sleep(Duration::from_millis(50)).await;
         }
         tracing::info!("finished processing all completed jobs");
-
-        // if let Err(e) =
-        //     add_completed_job(&db2, &job, success, false, result, logs, rsmq2.clone()).await
-        // {
-        //     tracing::error!(worker = %worker_name2, "failed to add completed job: {}", e);
-        // }
     });
 
     let mut last_executed_job: Option<Instant> = None;
@@ -1125,7 +1127,7 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
                 .await
                 {
                     tracing::error!("failed to create token for dedicated worker: {:?}", e);
-                    killpill_tx.send(()).expect("send");
+                    killpill_tx.clone().send(()).expect("send");
                 };
 
                 let (content, lock, language, envs) = {