improve transparency for slow queries

backend/Cargo.lock

@@ -10644,6 +10644,7 @@ dependencies = [
  "magic-crypt",
  "mail-send",
  "object_store",
+ "pin-project-lite",
  "prometheus",
  "quick_cache",
  "rand 0.8.5",

backend/Cargo.toml

@@ -284,6 +284,7 @@ tikv-jemalloc-sys  = { version = "^0.5" }
 tikv-jemalloc-ctl = { version = "^0.5" }
 
 triomphe = "^0"
+pin-project-lite = "^0"
 
 tantivy = "0.22.0"
 

backend/migrations/20241204154025_grant_all_concurrency_key.down.sql

@@ -0,0 +1 @@
+-- Add down migration script here

backend/migrations/20241204154025_grant_all_concurrency_key.up.sql

@@ -0,0 +1,3 @@
+-- Add up migration script here
+GRANT ALL ON concurrency_key TO windmill_admin;
+GRANT ALL ON concurrency_key TO windmill_user;

backend/src/monitor.rs

@@ -85,12 +85,12 @@ lazy_static::lazy_static! {
     static ref ZOMBIE_JOB_TIMEOUT: String = std::env::var("ZOMBIE_JOB_TIMEOUT")
     .ok()
     .and_then(|x| x.parse::<String>().ok())
-    .unwrap_or_else(|| "30".to_string());
+    .unwrap_or_else(|| "60".to_string());
 
     static ref FLOW_ZOMBIE_TRANSITION_TIMEOUT: String = std::env::var("FLOW_ZOMBIE_TRANSITION_TIMEOUT")
     .ok()
     .and_then(|x| x.parse::<String>().ok())
-    .unwrap_or_else(|| "30".to_string());
+    .unwrap_or_else(|| "60".to_string());
 
 
     pub static ref RESTART_ZOMBIE_JOBS: bool = std::env::var("RESTART_ZOMBIE_JOBS")

backend/windmill-common/Cargo.toml

@@ -62,6 +62,7 @@ windmill-macros.workspace = true
 semver.workspace = true
 croner = "2.0.6"
 quick_cache.workspace = true
+pin-project-lite.workspace = true
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 tikv-jemalloc-ctl = { optional = true, workspace = true }

backend/windmill-common/src/utils.rs

@@ -34,8 +34,8 @@ pub const GIT_VERSION: &str =
     git_version!(args = ["--tag", "--always"], fallback = "unknown-version");
 
 use crate::CRITICAL_ALERT_MUTE_UI_ENABLED;
-use std::sync::atomic::Ordering;
 use std::panic::{self, AssertUnwindSafe};
+use std::sync::atomic::Ordering;
 
 use crate::worker::CLOUD_HOSTED;
 
@@ -78,14 +78,18 @@ pub fn require_admin(is_admin: bool, username: &str) -> Result<()> {
     }
 }
 
-pub async fn require_admin_or_devops(is_admin: bool, username: &str, email: &str, db: &DB) -> Result<()> {
+pub async fn require_admin_or_devops(
+    is_admin: bool,
+    username: &str,
+    email: &str,
+    db: &DB,
+) -> Result<()> {
     if !is_admin {
         if !is_devops_email(db, email).await? {
             return Err(Error::RequireAdmin(username.to_string()));
         }
     }
     Ok(())
-
 }
 
 pub fn hostname() -> String {
@@ -94,7 +98,7 @@ pub fn hostname() -> String {
             .to_str()
             .map(|x| x.to_string())
             .unwrap_or_else(|| rd_string(5))
-    }) 
+    })
 }
 
 pub fn paginate(pagination: Pagination) -> (usize, usize) {
@@ -440,9 +444,7 @@ impl ScheduleType {
             Some("v2") | Some(_) => {
                 // Use Croner for v2
                 let schedule_type_result = panic::catch_unwind(AssertUnwindSafe(|| {
-                    Cron::new(schedule_str)
-                        .with_seconds_optional()
-                        .parse()
+                    Cron::new(schedule_str).with_seconds_optional().parse()
                 }))
                 .map_err(|_| {
                     tracing::error!(
@@ -478,8 +480,14 @@ impl ScheduleType {
                     }
 
                     if let Err(e) = result {
-                        tracing::error!("An error occurred while finding the next occurrence: {:?}", e);
-                        return Err(Error::BadRequest(format!("cron: error during find_next_occurrence: {:?}", e)));
+                        tracing::error!(
+                            "An error occurred while finding the next occurrence: {:?}",
+                            e
+                        );
+                        return Err(Error::BadRequest(format!(
+                            "cron: error during find_next_occurrence: {:?}",
+                            e
+                        )));
                     }
                 }
 
@@ -526,3 +534,75 @@ impl ScheduleType {
         Ok(events)
     }
 }
+
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context as TContext, Poll};
+use tokio::time::{self, Duration, Sleep};
+
+use pin_project_lite::pin_project;
+
+pub trait WarnAfterExt: Future + Sized {
+    /// Warns if the future takes longer than the specified number of seconds to complete.
+    fn warn_after_seconds(self, seconds: u8, location: &'static str) -> WarnAfterFuture<Self> {
+        WarnAfterFuture {
+            future: self,
+            timeout: time::sleep(Duration::from_secs(seconds as u64)),
+            warned: false,
+            start_time: std::time::Instant::now(),
+            location,
+            seconds,
+        }
+    }
+}
+
+// Blanket implementation for all futures.
+impl<F: Future> WarnAfterExt for F {}
+
+pin_project! {
+    /// A future that wraps another future and prints a warning if it takes too long.
+    pub struct WarnAfterFuture<F> {
+        #[pin]
+        future: F,
+        #[pin]
+        timeout: Sleep,
+        warned: bool,
+        location: &'static str,
+        start_time: std::time::Instant,
+        seconds: u8,
+    }
+}
+
+impl<F: Future> Future for WarnAfterFuture<F> {
+    type Output = F::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut TContext<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+
+        // Poll the timeout future to check if it has elapsed.
+        if !*this.warned {
+            if this.timeout.poll(cx).is_ready() {
+                tracing::warn!(location = this.location, "SLOW_QUERY: query to db taking longer than expected (> {} seconds). This is a sign the database is under heavy load, query is too heavy or database is undersized",
+                    this.seconds,
+                );
+                *this.warned = true;
+            }
+        }
+
+        // Poll the wrapped future.
+        match this.future.poll(cx) {
+            Poll::Ready(output) => {
+                if *this.warned {
+                    let elapsed = this.start_time.elapsed();
+                    tracing::warn!(
+                        location = this.location,
+                        "SLOW QUERY: completed with total duration: {:.2?}",
+                        elapsed
+                    );
+                }
+                Poll::Ready(output)
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}

backend/windmill-worker/src/result_processor.rs

@@ -14,6 +14,7 @@ use windmill_common::{
     add_time,
     error::{self, Error},
     jobs::{JobKind, QueuedJob},
+    utils::WarnAfterExt,
     worker::{to_raw_value, WORKER_GROUP},
     DB,
 };
@@ -431,6 +432,7 @@ pub async fn process_completed_job(
                     #[cfg(feature = "benchmark")]
                     bench,
                 )
+                .warn_after_seconds(10, "update_flow_status_after_job_completion success")
                 .await?;
             }
         }
@@ -469,6 +471,7 @@ pub async fn process_completed_job(
                     #[cfg(feature = "benchmark")]
                     bench,
                 )
+                .warn_after_seconds(10, "update_flow_status_after_job_completion error")
                 .await?;
             }
         }

backend/windmill-worker/src/worker.rs

@@ -9,6 +9,7 @@
 use windmill_common::{
     auth::{fetch_authed_from_permissioned_as, JWTAuthClaims, JobPerms, JWT_SECRET},
     scripts::PREVIEW_IS_TAR_CODEBASE_HASH,
+    utils::WarnAfterExt,
     worker::{
         get_memory, get_vcpus, get_windmill_memory_usage, get_worker_memory_usage, write_file,
         ROOT_CACHE_DIR, TMP_DIR,
@@ -159,6 +160,7 @@ pub async fn create_token_for_owner_in_bg(
                 &email,
                 &job_id,
             )
+            .warn_after_seconds(5, "creating token for owner")
             .await
             .expect("could not create job token");
             *locked = token;
@@ -1823,7 +1825,9 @@ async fn handle_queued_job(
     if job.parent_job.is_none() && job.created_by.starts_with("email-") {
         let daily_count = sqlx::query!(
             "SELECT value FROM metrics WHERE id = 'email_trigger_usage' AND created_at > NOW() - INTERVAL '1 day' ORDER BY created_at DESC LIMIT 1"
-        ).fetch_optional(db).await?.map(|x| serde_json::from_value::<i64>(x.value).unwrap_or(1));
+        ).fetch_optional(db)
+        .warn_after_seconds(5, "getting email_trigger_usage")
+        .await?.map(|x| serde_json::from_value::<i64>(x.value).unwrap_or(1));
 
         if let Some(count) = daily_count {
             if count >= 100 {
@@ -1836,6 +1840,7 @@ async fn handle_queued_job(
                     serde_json::json!(count + 1)
                 )
                 .execute(db)
+                .warn_after_seconds(5, "updating email_trigger_usage")
                 .await?;
             }
         } else {
@@ -1843,6 +1848,7 @@ async fn handle_queued_job(
                 "INSERT INTO metrics (id, value) VALUES ('email_trigger_usage', to_jsonb(1))"
             )
             .execute(db)
+            .warn_after_seconds(5, "inserting email_trigger_usage")
             .await?;
         }
     }
@@ -1855,6 +1861,7 @@ async fn handle_queued_job(
                 .ok_or_else(|| Error::InternalErr(format!("expected parent job")))?,
             job.id,
         )
+        .warn_after_seconds(5, "updating flow status in progress")
         .await?;
 
         Some(r)
@@ -1867,6 +1874,7 @@ async fn handle_queued_job(
                 &job.workspace_id
             )
             .execute(db)
+            .warn_after_seconds(5, "updating parent job started_at flow_status")
             .await {
                 tracing::error!("Could not update parent job started_at flow_status: {}", e);
             }
@@ -1883,6 +1891,7 @@ async fn handle_queued_job(
             job.workspace_id
         )
         .fetch_one(db)
+        .warn_after_seconds(5, "getting job raw values")
         .await
         .map(|record| (record.raw_code, record.raw_lock, record.raw_flow))
         .unwrap_or_default(),
@@ -1923,6 +1932,7 @@ async fn handle_queued_job(
                 &job.parent_job.unwrap()
             )
             .fetch_one(db)
+            .warn_after_seconds(5, "getting script path from queue for caching purposes")
             .await
             .map_err(|e| {
                 Error::InternalErr(format!(
@@ -1957,6 +1967,7 @@ async fn handle_queued_job(
             &job.workspace_id,
             &cached_res_path,
         )
+        .warn_after_seconds(5, "getting cached resource value")
         .await;
         if let Some(cached_resource_value) = cached_resource_value_maybe {
             {
@@ -1995,6 +2006,7 @@ async fn handle_queued_job(
             worker_dir,
             job_completed_tx.0.clone(),
         )
+        .warn_after_seconds(10, "handling flow")
         .await?;
         Ok(true)
     } else {
@@ -2285,9 +2297,11 @@ async fn handle_code_execution_job(
             .await?
         }
         JobKind::FlowScript => {
-            let (lockfile, content) = cache::flow::fetch_script(db, FlowNodeId(
-                job.script_hash.unwrap_or(ScriptHash(0)).0
-            )).await?;
+            let (lockfile, content) = cache::flow::fetch_script(
+                db,
+                FlowNodeId(job.script_hash.unwrap_or(ScriptHash(0)).0),
+            )
+            .await?;
             ContentReqLangEnvs {
                 content,
                 lockfile,

backend/windmill-worker/src/worker_flow.rs

@@ -41,6 +41,7 @@ use windmill_common::jobs::{
     script_hash_to_tag_and_limits, script_path_to_payload, BranchResults, JobPayload, QueuedJob,
     RawCode, ENTRYPOINT_OVERRIDE,
 };
+use windmill_common::utils::WarnAfterExt;
 use windmill_common::worker::to_raw_value;
 use windmill_common::{
     error::{self, to_anyhow, Error},
@@ -1109,6 +1110,7 @@ pub async fn update_flow_status_after_job_completion_internal(
             worker_dir,
                 job_completed_tx,
         )
+        .warn_after_seconds(10, "handle_flow in update_flow_status")
         .await
         {
             Err(err) => {
@@ -1510,7 +1512,7 @@ pub async fn handle_flow(
         let schedule_path = flow_job.schedule_path.as_ref().unwrap();
 
         let schedule =
-            get_schedule_opt(&mut tx, &flow_job.workspace_id, schedule_path).await?;
+            get_schedule_opt(&mut tx, &flow_job.workspace_id, schedule_path).warn_after_seconds(5, "get schedule_opt in handle_flow").await?;
 
         tx.commit().await?;
 
@@ -1522,6 +1524,7 @@ pub async fn handle_flow(
                 flow_job.script_path.as_ref().unwrap(),
                 &flow_job.workspace_id,
             )
+            .warn_after_seconds(5, "handle_maybe_scheduled_job in handle_flow")
             .await
             {
                 match err {
@@ -1549,6 +1552,7 @@ pub async fn handle_flow(
         worker_dir,
         job_completed_tx,
     )
+    .warn_after_seconds(10, "push next flow job")
     .await?;
     Ok(())
 }