From 8e215326848ca445e889511eef5201f880ce7536 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 19 Feb 2024 22:10:34 +0100 Subject: [PATCH] refactor future wrapper for queries --- backend/ee-repo-ref.txt | 2 +- backend/windmill-worker/src/common.rs | 48 +++++++++++++++ backend/windmill-worker/src/mysql_executor.rs | 60 +++++++++++++------ backend/windmill-worker/src/pg_executor.rs | 32 ++-------- backend/windmill-worker/src/worker.rs | 11 +++- 5 files changed, 108 insertions(+), 45 deletions(-) diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 036bd91159..b30f447bea 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -3de554fc89db8f9c6cc544c6a2ba004e65a7c133 \ No newline at end of file +fd0c4dab2ae54bed1c7cd7041669fc3463cd163e \ No newline at end of file diff --git a/backend/windmill-worker/src/common.rs b/backend/windmill-worker/src/common.rs index 55c71d51d2..2bcb8a4427 100644 --- a/backend/windmill-worker/src/common.rs +++ b/backend/windmill-worker/src/common.rs @@ -14,6 +14,7 @@ use serde_json::{json, Value}; use sqlx::{Pool, Postgres}; use tokio::process::Command; use tokio::{fs::File, io::AsyncReadExt}; +use windmill_common::error::to_anyhow; use windmill_common::s3_helpers::{ get_etag_or_empty, AzureBlobResource, LargeFileStorage, ObjectStoreResource, S3Object, S3Resource, @@ -447,6 +448,53 @@ async fn get_mem_peak(pid: Option, nsjail: bool) -> i32 { } } +pub async fn run_future_with_polling_update_job_poller( + job_id: Uuid, + timeout: Option, + db: &DB, + mem_peak: &mut i32, + canceled_by_ref: &mut Option, + result_f: Fut, + worker_name: &str, + w_id: &str, +) -> anyhow::Result +where + Fut: Future>, +{ + let (tx, rx) = broadcast::channel::<()>(3); + + let update_job = update_job_poller( + job_id, + db, + mem_peak, + canceled_by_ref, + || async { 0 }, + worker_name, + w_id, + rx, + ); + + let timeout_ms = u64::try_from( + resolve_job_timeout(&db, &w_id, job_id, timeout) + .await + .0 + .as_millis(), + ) + .unwrap_or(200000); + + let rows = tokio::select! { + biased; + result = tokio::time::timeout(std::time::Duration::from_millis(timeout_ms), result_f) => result + .map_err(|e| { + tracing::error!("Query timeout: {}", e); + Error::ExecutionErr("Query timeout".to_string()) + })?, + _ = update_job, if job_id != Uuid::nil() => Err(Error::ExecutionErr("Job cancelled".to_string())).map_err(to_anyhow)?, + }?; + drop(tx); + Ok(rows) +} + pub async fn update_job_poller( job_id: Uuid, db: &DB, diff --git a/backend/windmill-worker/src/mysql_executor.rs b/backend/windmill-worker/src/mysql_executor.rs index e144a5ffc1..7c4689e95a 100644 --- a/backend/windmill-worker/src/mysql_executor.rs +++ b/backend/windmill-worker/src/mysql_executor.rs @@ -12,8 +12,12 @@ use windmill_common::{ jobs::QueuedJob, }; use windmill_parser_sql::{parse_mysql_sig, RE_ARG_MYSQL_NAMED}; +use windmill_queue::CanceledBy; -use crate::{common::build_args_map, AuthedClientBackgroundTask}; +use crate::{ + common::{build_args_map, run_future_with_polling_update_job_poller}, + AuthedClientBackgroundTask, +}; #[derive(Deserialize)] struct MysqlDatabase { @@ -30,6 +34,9 @@ pub async fn do_mysql( client: &AuthedClientBackgroundTask, query: &str, db: &sqlx::Pool, + mem_peak: &mut i32, + canceled_by: &mut Option, + worker_name: &str, ) -> windmill_common::error::Result> { let args = build_args_map(job, client, db).await?.map(Json); let job_args = if args.is_some() { @@ -129,28 +136,47 @@ pub async fn do_mysql( _ => {} } } - let rows: Vec = conn - .exec( - query, - match statement_values { - Params::Positional(v) => Params::Positional(v), - Params::Named(m) => Params::Named(m), - _ => Params::Empty, - }, - ) - .await - .map_err(to_anyhow)?; - let rows = rows - .into_iter() - .map(|x| convert_row_to_value(x)) - .collect::>(); + + let result_f = async { + let rows: Vec = conn + .exec( + query, + match statement_values { + Params::Positional(v) => Params::Positional(v), + Params::Named(m) => Params::Named(m), + _ => Params::Empty, + }, + ) + .await + .map_err(to_anyhow)?; + Ok(rows + .into_iter() + .map(|x| convert_row_to_value(x)) + .collect::>()) + as Result, anyhow::Error> + }; + + let result = run_future_with_polling_update_job_poller( + job.id, + job.timeout, + db, + mem_peak, + canceled_by, + result_f, + worker_name, + &job.workspace_id, + ) + .await?; drop(conn); pool.disconnect().await.map_err(to_anyhow)?; + let raw_result = windmill_common::worker::to_raw_value(&json!(result)); + *mem_peak = (raw_result.get().len() / 1000) as i32; + // And then check that we got back the same string we sent over. - return Ok(windmill_common::worker::to_raw_value(&json!(rows))); + return Ok(raw_result); } fn string_date_to_mysql_date(s: &str) -> mysql_async::Value { diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index dd177c83ae..ae40689985 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -13,7 +13,7 @@ use serde::Deserialize; use serde_json::value::RawValue; use serde_json::Map; use serde_json::Value; -use tokio::sync::{broadcast, Mutex}; +use tokio::sync::Mutex; use tokio_postgres::types::IsNull; use tokio_postgres::{ types::{to_sql_checked, ToSql}, @@ -31,7 +31,7 @@ use windmill_parser::Typ; use windmill_parser_sql::parse_pgsql_sig; use windmill_queue::CanceledBy; -use crate::common::{build_args_values, resolve_job_timeout, update_job_poller}; +use crate::common::{build_args_values, run_future_with_polling_update_job_poller}; use crate::AuthedClientBackgroundTask; use bytes::{Buf, BytesMut}; use lazy_static::lazy_static; @@ -206,37 +206,17 @@ pub async fn do_postgresql( .collect::, _>>()?) as anyhow::Result> }; - let (tx, rx) = broadcast::channel::<()>(3); - - let update_job = update_job_poller( + let result = run_future_with_polling_update_job_poller( job.id, + job.timeout, db, mem_peak, canceled_by, - || async { 0 }, + result_f, worker_name, &job.workspace_id, - rx, - ); - - let timeout_ms = u64::try_from( - resolve_job_timeout(&db, &job.workspace_id, job.id, job.timeout) - .await - .0 - .as_millis(), ) - .unwrap_or(200000); - tracing::info!("Query timeout: {}ms", timeout_ms); - let result = tokio::select! { - biased; - result = tokio::time::timeout(std::time::Duration::from_millis(timeout_ms), result_f) => result - .map_err(|e| { - tracing::error!("Query timeout: {}", e); - Error::ExecutionErr("Query timeout".to_string()) - })?, - _ = update_job, if job.id != Uuid::nil() => Err(Error::ExecutionErr("Job cancelled".to_string())).map_err(to_anyhow)?, - }?; - drop(tx); + .await?; RUNNING.store(false, std::sync::atomic::Ordering::Relaxed); diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 8f1b1023ba..cad874a3e0 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -2893,7 +2893,16 @@ async fn handle_code_execution_job( ) .await; } else if language == Some(ScriptLang::Mysql) { - return do_mysql(job, &client, &inner_content, db).await; + return do_mysql( + job, + &client, + &inner_content, + db, + mem_peak, + canceled_by, + worker_name, + ) + .await; } else if language == Some(ScriptLang::Bigquery) { #[cfg(not(feature = "enterprise"))] {