From 670ba51d9bd9a9a0b07a2ed064c316234bb5819d Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Fri, 21 Apr 2023 08:49:45 +0200
Subject: [PATCH] feat(backend): workers are instantly ready and sync with
 global cache in background

---
 backend/README.md                             |  9 +++-
 backend/sqlx-data.json                        | 12 +++++
 backend/src/main.rs                           | 19 ++------
 backend/windmill-worker/src/global_cache.rs   | 34 ++++----------
 backend/windmill-worker/src/lib.rs            |  4 --
 .../windmill-worker/src/python_executor.rs    |  3 ++
 backend/windmill-worker/src/worker.rs         | 47 +++++++++++++++++--
 7 files changed, 78 insertions(+), 50 deletions(-)

diff --git a/backend/README.md b/backend/README.md
index 2d3d35b1a1..04d27762a2 100644
--- a/backend/README.md
+++ b/backend/README.md
@@ -1,6 +1,7 @@
 # Windmill Backend
 
-This folder holds all backend components, the [src/](./src/) folder only contains files used to build the "root" binary.
+This folder holds all backend components, the [src/](./src/) folder only
+contains files used to build the "root" binary.
 
 ## Components
 
@@ -13,3 +14,9 @@ This folder holds all backend components, the [src/](./src/) folder only contain
 | [windmill-queue](./windmill-queue/)           | Contains job & flow queuing functionality, commonly written to by the API server and read from by workers |
 | [windmill-worker](./windmill-worker/)         | The worker. Used to process and execute flows & jobs.                                                     |
 | [parsers](./parsers/)                         | Contains code to parse signatures in different langauges.                                                 |
+
+### Compile sqlx for offline ci
+
+```
+cargo sqlx prepare --merged -- --bin windmill --features enterprise
+```
diff --git a/backend/sqlx-data.json b/backend/sqlx-data.json
index 66e32e4f0f..14ca54a791 100644
--- a/backend/sqlx-data.json
+++ b/backend/sqlx-data.json
@@ -238,6 +238,18 @@
     },
     "query": "SELECT script_path FROM queue WHERE id = $1"
   },
+  "062859f1d0e5cfba3115f4241115753b86a4ad239708851c998ff5620ebca5b8": {
+    "describe": {
+      "columns": [],
+      "nullable": [],
+      "parameters": {
+        "Left": [
+          "Uuid"
+        ]
+      }
+    },
+    "query": "UPDATE queue SET last_ping = now() WHERE id = $1"
+  },
   "0715955b7e98cc669a88eca6556cd46b7f9c07fcf32a24f85b69720b54e6e95f": {
     "describe": {
       "columns": [
diff --git a/backend/src/main.rs b/backend/src/main.rs
index 411ccba489..ca0c81c4b0 100644
--- a/backend/src/main.rs
+++ b/backend/src/main.rs
@@ -8,7 +8,7 @@
 
 use std::{
     net::{IpAddr, Ipv4Addr, SocketAddr},
-    sync::{atomic::Ordering, Arc},
+    sync::Arc,
 };
 
 use git_version::git_version;
@@ -16,13 +16,12 @@ use monitor::handle_zombie_jobs_periodically;
 use sqlx::{Pool, Postgres};
 use tokio::{
     fs::{metadata, DirBuilder},
-    join,
     sync::RwLock,
 };
-use windmill_common::{utils::rd_string, IS_READY, METRICS_ADDR};
+use windmill_common::{utils::rd_string, METRICS_ADDR};
 use windmill_worker::{
     DENO_CACHE_DIR, DENO_TMP_CACHE_DIR, GO_CACHE_DIR, GO_TMP_CACHE_DIR, PIP_CACHE_DIR,
-    ROOT_TMP_CACHE_DIR, S3_CACHE_BUCKET, TAR_PIP_TMP_CACHE_DIR,
+    ROOT_TMP_CACHE_DIR, TAR_PIP_TMP_CACHE_DIR,
 };
 
 const GIT_VERSION: &str = git_version!(args = ["--tag", "--always"], fallback = "unknown-version");
@@ -299,18 +298,6 @@ pub async fn run_workers<R: rsmq_async::RsmqConnection + Send + Sync + Clone + '
             .expect("could not create initial worker dir");
     }
 
-    #[cfg(feature = "enterprise")]
-    if let Some(ref s) = S3_CACHE_BUCKET.clone() {
-        join!(
-            windmill_worker::copy_denogo_cache_from_bucket_as_tar(&s),
-            windmill_worker::copy_all_piptars_from_bucket(&s)
-        );
-        if let Err(e) = windmill_worker::untar_all_piptars().await {
-            tracing::info!("Failed to untar piptars. Error: {:?}", e);
-        }
-    }
-
-    IS_READY.store(true, Ordering::Relaxed);
     let sync_barrier = Arc::new(RwLock::new(None));
     for i in 1..(num_workers + 1) {
         let db1 = db.clone();
diff --git a/backend/windmill-worker/src/global_cache.rs b/backend/windmill-worker/src/global_cache.rs
index 313a4a0ab8..704fd0cb8a 100644
--- a/backend/windmill-worker/src/global_cache.rs
+++ b/backend/windmill-worker/src/global_cache.rs
@@ -281,12 +281,12 @@ pub async fn copy_denogo_cache_from_bucket_as_tar(bucket: &str) {
     let start: Instant = Instant::now();
 
     if let Err(e) = execute_command(
-        ROOT_CACHE_DIR,
+        ROOT_TMP_CACHE_DIR,
         "rclone",
         vec![
             "copyto",
             &format!(":s3,env_auth=true:{bucket}/{TAR_CACHE_FILENAME}"),
-            &format!("{ROOT_CACHE_DIR}{TAR_CACHE_FILENAME}"),
+            &format!("{ROOT_TMP_CACHE_DIR}{TAR_CACHE_FILENAME}"),
             "-v",
             "--size-only",
             "--fast-list",
@@ -301,7 +301,10 @@ pub async fn copy_denogo_cache_from_bucket_as_tar(bucket: &str) {
     if let Err(e) = execute_command(
         ROOT_CACHE_DIR,
         "tar",
-        vec!["-xpvf", &format!("{ROOT_CACHE_DIR}{TAR_CACHE_FILENAME}")],
+        vec![
+            "-xpvf",
+            &format!("{ROOT_TMP_CACHE_DIR}{TAR_CACHE_FILENAME}"),
+        ],
     )
     .await
     {
@@ -309,12 +312,14 @@ pub async fn copy_denogo_cache_from_bucket_as_tar(bucket: &str) {
         return;
     }
 
-    let denogen = format!("{ROOT_CACHE_DIR}deno/gen/file/tmp/windmill");
+    let denogen = format!("{ROOT_TMP_CACHE_DIR}deno/gen/file/tmp/windmill");
     if metadata(&denogen).await.is_ok() {
         let _ = tokio::fs::remove_dir_all(denogen).await;
     }
 
-    if let Err(e) = tokio::fs::remove_file(format!("{ROOT_CACHE_DIR}{TAR_CACHE_FILENAME}")).await {
+    if let Err(e) =
+        tokio::fs::remove_file(format!("{ROOT_TMP_CACHE_DIR}{TAR_CACHE_FILENAME}")).await
+    {
         tracing::info!("Failed to remove denotar cache. Error: {:?}", e);
         return;
     };
@@ -323,25 +328,6 @@ pub async fn copy_denogo_cache_from_bucket_as_tar(bucket: &str) {
         "Finished copying denogotar from bucket {bucket} as tar, took: {:?}s",
         start.elapsed().as_secs()
     );
-
-    tracing::info!("Copying denogo cache from bucket {bucket} as tar");
-
-    let start: Instant = Instant::now();
-    for x in ["deno", "go"] {
-        if let Err(e) = execute_command(
-            TMP_DIR,
-            "cp",
-            vec!["-r", &format!("{ROOT_CACHE_DIR}{x}"), &ROOT_TMP_CACHE_DIR],
-        )
-        .await
-        {
-            tracing::info!(error = %e, "Could not copy root dir to tmp root dir");
-        }
-    }
-    tracing::info!(
-        "Finished copying untarred denogo to tmp cache, took: {:?}s",
-        start.elapsed().as_secs()
-    );
 }
 
 #[cfg(feature = "enterprise")]
diff --git a/backend/windmill-worker/src/lib.rs b/backend/windmill-worker/src/lib.rs
index ccaa3fd048..86cb82d4ef 100644
--- a/backend/windmill-worker/src/lib.rs
+++ b/backend/windmill-worker/src/lib.rs
@@ -7,8 +7,4 @@ mod python_executor;
 mod worker;
 mod worker_flow;
 
-#[cfg(feature = "enterprise")]
-pub use global_cache::{
-    copy_all_piptars_from_bucket, copy_denogo_cache_from_bucket_as_tar, untar_all_piptars,
-};
 pub use worker::*;
diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs
index 03485ae81c..092b10e2e6 100644
--- a/backend/windmill-worker/src/python_executor.rs
+++ b/backend/windmill-worker/src/python_executor.rs
@@ -470,6 +470,9 @@ pub async fn handle_python_reqs(
 
         #[cfg(feature = "enterprise")]
         if let Some(ref bucket) = *S3_CACHE_BUCKET {
+            sqlx::query_scalar!("UPDATE queue SET last_ping = now() WHERE id = $1", job_id)
+                .execute(db)
+                .await?;
             if pull_from_tar(bucket, venv_p.clone()).await.is_ok() {
                 req_paths.push(venv_p.clone());
                 continue;
diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs
index a08d68140d..022ebd3bb2 100644
--- a/backend/windmill-worker/src/worker.rs
+++ b/backend/windmill-worker/src/worker.rs
@@ -15,7 +15,7 @@ use windmill_api_client::Client;
 use std::{
     borrow::Borrow, collections::HashMap, io, os::unix::process::ExitStatusExt, panic,
     process::Stdio, time::{Duration},
-    sync::{Arc},
+    sync::{Arc, atomic::Ordering},
 };
 use tracing::{trace_span, Instrument};
 use uuid::Uuid;
@@ -25,7 +25,7 @@ use windmill_common::{
     flows::{FlowModuleValue, FlowValue},
     scripts::{ScriptHash, ScriptLang},
     utils::{rd_string},
-    variables, BASE_URL, users::SUPERADMIN_SECRET_EMAIL, METRICS_ENABLED, jobs::{JobKind, QueuedJob},
+    variables, BASE_URL, users::SUPERADMIN_SECRET_EMAIL, METRICS_ENABLED, jobs::{JobKind, QueuedJob}, IS_READY,
 };
 use windmill_queue::{canceled_job_to_result, get_queued_job, pull, CLOUD_HOSTED};
 
@@ -52,7 +52,7 @@ use async_recursion::async_recursion;
 use rand::Rng;
 
 #[cfg(feature = "enterprise")]
-use crate::global_cache::{copy_cache_to_tmp_cache, cache_global, copy_tmp_cache_to_cache};
+use crate::global_cache::{copy_cache_to_tmp_cache, cache_global, copy_tmp_cache_to_cache, copy_denogo_cache_from_bucket_as_tar, copy_all_piptars_from_bucket};
 
 use crate::{
     jobs::{add_completed_job, add_completed_job_error},
@@ -417,6 +417,34 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
     let mut last_sync =
         Instant::now() + Duration::from_secs(rand::thread_rng().gen_range(0..*GLOBAL_CACHE_INTERVAL));
 
+    #[cfg(feature = "enterprise")]
+    let mut handles = Vec::with_capacity(2);
+
+    #[cfg(feature = "enterprise")]
+    if i_worker == 1 {
+        if let Some(ref s) = S3_CACHE_BUCKET.clone() {
+            let bucket = s.to_string();
+            let copy_to_bucket_tx2 = copy_to_bucket_tx.clone();
+            let worker_name2 = worker_name.clone();
+
+            handles.push(tokio::task::spawn(async move {
+                tracing::info!(worker = %worker_name2, "Started initial denogo tar sync in background");
+                copy_denogo_cache_from_bucket_as_tar(&bucket).await;
+                let _ = copy_to_bucket_tx2.send(()).await;
+            }));
+
+            let bucket = s.to_string();
+            let copy_to_bucket_tx = copy_to_bucket_tx.clone();
+            let worker_name = worker_name.clone();
+
+            handles.push(tokio::task::spawn(async move {
+                tracing::info!(worker = %worker_name, "Started initial piptars cache sync in background");
+                copy_all_piptars_from_bucket(&bucket).await;
+                let _ = copy_to_bucket_tx.send(()).await;
+            }));
+        }
+    }
+
     let (same_worker_tx, mut same_worker_rx) = mpsc::channel::<Uuid>(5);
 
     let (job_completed_tx, mut job_completed_rx) = mpsc::channel::<JobCompleted>(10);
@@ -515,16 +543,23 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
                     tokio::select! {
                         biased;
                         _ = rx.recv() => {
+                            #[cfg(feature = "enterprise")]
                             if let Some(copy_cache_from_bucket_handle) = copy_cache_from_bucket_handle.as_ref() {
                                 if !copy_cache_from_bucket_handle.is_finished() {
                                     copy_cache_from_bucket_handle.abort();
                                 }
                             }
+                            #[cfg(feature = "enterprise")]
+                            for handle in &handles {
+                                if !handle.is_finished() {
+                                    handle.abort();
+                                }
+                            }
                             println!("received killpill for worker {}", i_worker);
                             (true, Ok(None))
                         },
                         _ = copy_to_bucket_rx.recv() => {
-                            tracing::debug!("CAN PULL LOCK START");
+                            tracing::debug!("can_pull lock start");
                             let _lock = CAN_PULL.write().await;
                             if num_workers > 1 {
                                 create_barrier_for_all_workers(num_workers, sync_barrier.clone()).await;
@@ -534,7 +569,7 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
                             if let Err(e) = copy_tmp_cache_to_cache().await {
                                 tracing::error!(worker = %worker_name, "failed to sync tmp cache to cache: {}", e);
                             }
-                            tracing::debug!("CAN PULL LOCK END");
+                            tracing::debug!("can_pull lock end");
                             (false, Ok(None))
                         },
                         Some(job_id) = same_worker_rx.recv() => {
@@ -559,6 +594,8 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
             };
 
             first_run = false;
+            IS_READY.store(true, Ordering::Relaxed);
+
             if do_break {
                 return true;
             }