starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-10-20 08:03:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

big

Browse Source
This commit is contained in:
dieriba
2025-09-23 00:58:55 +02:00
parent 612d00367c
commit 9741a28b56
30 changed files with 1047 additions and 216 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/windmill-queue/src/jobs.rs
+276 -73
View File
@@ -33,6 +33,7 @@ use windmill_common::add_time;
use windmill_common::auth::JobPerms;
#[cfg(feature = "benchmark")]
use windmill_common::bench::BenchmarkIter;
use windmill_common::db::{get_shard_db_from_job_uuid, get_shard_id_from_job_uuid};
use windmill_common::jobs::{JobTriggerKind, EMAIL_ERROR_HANDLER_USER_EMAIL};
use windmill_common::utils::now_from_db;
use windmill_common::worker::{Connection, SCRIPT_TOKEN_EXPIRY};
@@ -62,6 +63,7 @@ use windmill_common::{
},
DB, METRICS_ENABLED,
};
use windmill_common::{SHARD_DB_INSTANCE, SHARD_ID_TO_DB_INSTANCE, SHARD_MODE};
use backon::ConstantBuilder;
use backon::{BackoffBuilder, Retryable};
@@ -179,6 +181,7 @@ pub async fn cancel_single_job<'c>(
.await;
let add_job = add_completed_job_error(
&db,
&db, //put here for compile will need to pass the proper SHARD DB if shard mode is on
&MiniPulledJob::from(&job_running),
job_running.mem_peak.unwrap_or(0),
Some(CanceledBy { username: Some(username.to_string()), reason: Some(reason) }),
@@ -374,13 +377,25 @@ pub async fn append_logs(
}
match conn {
Connection::Sql(pool) => {
let db = if *SHARD_MODE {
if let Some(worker_db) = SHARD_DB_INSTANCE.get() {
worker_db
} else if let Some(server_shards) = SHARD_ID_TO_DB_INSTANCE.get() {
let shard_id = get_shard_id_from_job_uuid(*job_id);
server_shards.get(&shard_id).unwrap_or(pool)
} else {
pool
}
} else {
pool
};
if let Err(err) = sqlx::query!(
"INSERT INTO job_logs (logs, job_id, workspace_id) VALUES ($1, $2, $3) ON CONFLICT (job_id) DO UPDATE SET logs = concat(job_logs.logs, $1::text)",
logs.as_ref(),
job_id,
workspace.as_ref(),
)
.execute(pool)
.execute(db)
.warn_after_seconds(1)
.await
{
@@ -389,12 +404,17 @@ pub async fn append_logs(
}
Connection::Http(client) => {
if let Err(e) = client
.post::<_, String>(
&format!("/api/w/{}/agent_workers/push_logs/{}", workspace.as_ref(), job_id),
None,
&logs.as_ref(),
)
.await {
.post::<_, String>(
&format!(
"/api/w/{}/agent_workers/push_logs/{}",
workspace.as_ref(),
job_id
),
None,
&logs.as_ref(),
)
.await
{
tracing::error!(%job_id, %e, "error sending logs for job {job_id}: {e}");
};
}
@@ -691,7 +711,8 @@ where
}
pub async fn add_completed_job_error(
db: &Pool<Postgres>,
db: &DB,
job_queue_db: &DB,
queued_job: &MiniPulledJob,
mem_peak: i32,
canceled_by: Option<CanceledBy>,
@@ -728,6 +749,7 @@ pub async fn add_completed_job_error(
);
let _ = add_completed_job(
db,
job_queue_db,
&queued_job,
false,
false,
@@ -748,7 +770,8 @@ lazy_static::lazy_static! {
}
pub async fn add_completed_job<T: Serialize + Send + Sync + ValidableJson>(
db: &Pool<Postgres>,
db: &DB,
job_queue_db: &DB,
queued_job: &MiniPulledJob,
success: bool,
skipped: bool,
@@ -773,6 +796,7 @@ pub async fn add_completed_job<T: Serialize + Send + Sync + ValidableJson>(
let (opt_uuid, duration, _skip_downstream_error_handlers) = (|| {
commit_completed_job(
db,
job_queue_db,
queued_job,
success,
skipped,
@@ -825,6 +849,7 @@ pub async fn add_completed_job<T: Serialize + Send + Sync + ValidableJson>(
async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
db: &Pool<Postgres>,
job_queue_db: &Pool<Postgres>,
queued_job: &MiniPulledJob,
success: bool,
skipped: bool,
@@ -838,7 +863,7 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
// let start = std::time::Instant::now();
let mut tx = db.begin().await?;
let mut job_tx = job_queue_db.begin().await?;
let job_id = queued_job.id;
// tracing::error!("1 {:?}", start.elapsed());
@@ -851,7 +876,7 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
let mem_peak = mem_peak;
// add_time!(bench, "add_completed_job query START");
if let Some(value) = check_result_size(db, queued_job, result).await {
if let Some(value) = check_result_size(db, job_queue_db, queued_job, result).await {
return value;
}
@@ -891,7 +916,7 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
/* $9 */ duration,
/* $10 */ result_columns as Option<&Vec<String>>,
)
.fetch_one(&mut *tx)
.fetch_one(&mut *job_tx)
.await
.map_err(|e| Error::internal_err(format!("Could not add completed job {job_id}: {e:#}")))?;
@@ -904,7 +929,7 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
job_id,
labels as Vec<String>
)
.execute(&mut *tx)
.execute(&mut *job_tx)
.await
.map_err(|e| Error::InternalErr(format!("Could not update job labels: {e:#}")))?;
}
@@ -927,7 +952,7 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
_duration,
parent_job
)
.execute(&mut *tx)
.execute(&mut *job_tx)
.await
.inspect_err(|e| {
tracing::error!(
@@ -960,14 +985,14 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
parent_job,
&queued_job.workspace_id
)
.execute(&mut *tx)
.execute(&mut *job_tx)
.await?;
if flow_is_done {
let r = sqlx::query_scalar!(
"UPDATE parallel_monitor_lock SET last_ping = now() WHERE parent_flow_id = $1 and job_id = $2 RETURNING 1",
parent_job,
&queued_job.id
).fetch_optional(&mut *tx).await?;
).fetch_optional(&mut *job_tx).await?;
if r.is_some() {
tracing::info!(
"parallel flow iteration is done, setting parallel monitor last ping lock for job {}",
@@ -1011,7 +1036,7 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
&queued_job.id,
&queued_job.workspace_id
)
.fetch_optional(&mut *tx)
.fetch_optional(&mut *job_tx)
.await?
.flatten()
.unwrap_or(false);
@@ -1076,25 +1101,35 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
}
}
}
if queued_job.concurrent_limit.is_some() {
let concurrency_key = concurrency_key(db, &queued_job.id).await?;
if *DISABLE_CONCURRENCY_LIMIT || concurrency_key.is_none() {
tracing::warn!("Concurrency limit is disabled, skipping");
} else {
let concurrency_key = concurrency_key.unwrap();
sqlx::query_scalar!(
"UPDATE concurrency_counter SET job_uuids = job_uuids - $2 WHERE concurrency_id = $1",
concurrency_key,
queued_job.id.hyphenated().to_string(),
)
.execute(&mut *tx)
.await
.map_err(|e| {
Error::internal_err(format!(
"Could not decrement concurrency counter for job_id={}: {e:#}",
queued_job.id
))
})?;
match concurrency_key {
Some(concurrency_key) if !*DISABLE_CONCURRENCY_LIMIT => {
sqlx::query_scalar!(
"
UPDATE
concurrency_counter
SET
job_uuids = job_uuids - $2
WHERE
concurrency_id = $1
",
concurrency_key,
queued_job.id.hyphenated().to_string(),
)
.execute(&mut *tx)
.await
.map_err(|e| {
Error::internal_err(format!(
"Could not decrement concurrency counter for job_id={}: {e:#}",
queued_job.id
))
})?;
}
_ => {
tracing::warn!("Concurrency limit is disabled, skipping");
}
}
if let Err(e) = sqlx::query_scalar!(
@@ -1113,10 +1148,49 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
}
sqlx::query!("DELETE FROM job_perms WHERE job_id = $1", job_id)
.execute(&mut *tx)
.execute(&mut *job_tx)
.await?;
tx.commit().await?;
let main_db_result = tx.commit().await;
match main_db_result {
Ok(()) => match job_tx.commit().await {
Ok(()) => {
tracing::debug!(
"Both main DB and shard DB transactions committed successfully for job {}",
job_id
);
}
Err(shard_error) => {
//TODO: looks for solution
tracing::error!(
job_id = %job_id,
workspace_id = %queued_job.workspace_id,
error = %shard_error,
"CRITICAL: Main DB committed but shard DB failed during job completion. Data inconsistency possible."
);
return Err(Error::internal_err(format!(
"Shard DB commit failed after main DB commit for job {}: {}",
job_id, shard_error
)));
}
},
Err(main_error) => {
if let Err(rollback_err) = job_tx.rollback().await {
tracing::error!(
job_id = %job_id,
main_error = %main_error,
rollback_error = %rollback_err,
"Failed to rollback shard DB after main DB commit failure"
);
}
return Err(Error::internal_err(format!(
"Main DB commit failed for job {}: {}",
job_id, main_error
)));
}
}
tracing::info!(
%job_id,
@@ -1142,6 +1216,7 @@ async fn commit_completed_job<T: Serialize + Send + Sync + ValidableJson>(
async fn check_result_size<T: ValidableJson>(
db: &Pool<Postgres>,
job_queue_db: &DB,
queued_job: &MiniPulledJob,
result: Json<&T>,
) -> Option<Result<(Option<Uuid>, i64, bool), Error>> {
@@ -2571,7 +2646,8 @@ pub struct PulledJobResult {
}
pub async fn pull(
db: &Pool<Postgres>,
db: &DB,
job_queue_db: &DB,
suspend_first: bool,
worker_name: &str,
query_o: Option<&(String, String)>,
@@ -2595,7 +2671,7 @@ pub async fn pull(
} else {
sqlx::query_as::<_, PulledJob>(query_suspended)
.bind(worker_name)
.fetch_optional(db)
.fetch_optional(job_queue_db)
.await?
};
if let Some(job) = job {
@@ -2603,7 +2679,7 @@ pub async fn pull(
} else {
let job = sqlx::query_as::<_, PulledJob>(query_no_suspend)
.bind(worker_name)
.fetch_optional(db)
.fetch_optional(job_queue_db)
.await?;
PulledJobResult { job, suspended: false }
}
@@ -2626,7 +2702,7 @@ pub async fn pull(
tag,
job.id
)
.execute(db)
.execute(job_queue_db)
.await?;
continue;
}
@@ -2634,7 +2710,7 @@ pub async fn pull(
return Ok(njob);
};
let (job, suspended) = pull_single_job_and_mark_as_running_no_concurrency_limit(
db,
job_queue_db,
suspend_first,
worker_name,
#[cfg(feature = "benchmark")]
@@ -2718,14 +2794,30 @@ pub async fn pull(
let job_script_path = pulled_job.runnable_path.clone().unwrap_or_default();
let min_started_at = sqlx::query!(
"SELECT COALESCE((SELECT MIN(started_at) as min_started_at
FROM v2_job_queue INNER JOIN v2_job ON v2_job.id = v2_job_queue.id
WHERE v2_job.runnable_path = $1 AND v2_job.kind != 'dependencies' AND v2_job_queue.running = true AND v2_job_queue.workspace_id = $2 AND v2_job_queue.canceled_by IS NULL AND v2_job.concurrent_limit > 0), $3) as min_started_at, now() AS now",
r#"
SELECT
COALESCE(
(
SELECT MIN(started_at) AS min_started_at
FROM v2_job_queue
INNER JOIN v2_job
ON v2_job.id = v2_job_queue.id
WHERE v2_job.runnable_path = $1
AND v2_job.kind != 'dependencies'
AND v2_job_queue.running = true
AND v2_job_queue.workspace_id = $2
AND v2_job_queue.canceled_by IS NULL
AND v2_job.concurrent_limit > 0
),
$3
) AS min_started_at,
NOW() AS now
"#,
job_script_path,
&pulled_job.workspace_id,
max_ended_at
)
.fetch_one(db)
.fetch_one(job_queue_db)
.await
.map_err(|e| {
Error::internal_err(format!(
@@ -2734,15 +2826,27 @@ pub async fn pull(
})?;
let job_uuid: Uuid = pulled_job.id;
let avg_script_duration: Option<i64> = sqlx::query_scalar!(
"SELECT CAST(ROUND(AVG(duration_ms), 0) AS BIGINT) AS avg_duration_s FROM
(SELECT duration_ms FROM concurrency_key LEFT JOIN v2_job_completed ON v2_job_completed.id = concurrency_key.job_id WHERE key = $1 AND ended_at IS NOT NULL
ORDER BY ended_at
DESC LIMIT 10) AS t",
let completed_job_ids = sqlx::query_scalar!(
"SELECT job_id FROM concurrency_key WHERE key = $1 AND ended_at IS NOT NULL
ORDER BY ended_at DESC LIMIT 10",
job_concurrency_key
)
.fetch_one(db)
.fetch_all(db)
.await?;
let avg_script_duration: Option<i64> = if completed_job_ids.is_empty() {
None
} else {
sqlx::query_scalar!(
"SELECT CAST(ROUND(AVG(duration_ms), 0) AS BIGINT) AS avg_duration_s FROM v2_job_completed
WHERE id = ANY($1)",
&completed_job_ids
)
.fetch_optional(job_queue_db)
.await?
.flatten()
};
tracing::debug!(
"avg script duration computed: {}",
avg_script_duration.unwrap_or(0)
@@ -2763,12 +2867,35 @@ pub async fn pull(
(min_started_at_or_now + inc).max(now + Duration::try_seconds(3).unwrap_or_default());
let mut estimated_next_schedule_timestamp = min_started_p_inc;
let all_jobs = sqlx::query_scalar!(
"SELECT scheduled_for FROM v2_job_queue INNER JOIN concurrency_key ON concurrency_key.job_id = v2_job_queue.id
WHERE key = $1 AND running = false AND canceled_by IS NULL AND scheduled_for >= $2",
job_concurrency_key,
estimated_next_schedule_timestamp - inc
).fetch_all(db).await?;
let job_ids_with_key = sqlx::query_scalar!(
"SELECT job_id FROM concurrency_key WHERE key = $1",
job_concurrency_key
)
.fetch_all(db)
.await?;
let all_jobs = if job_ids_with_key.is_empty() {
Vec::new()
} else {
sqlx::query_scalar!(
"
SELECT
scheduled_for
FROM
v2_job_queue
WHERE
id = ANY($1) AND
running = false AND
canceled_by IS NULL AND
scheduled_for >= $2
",
&job_ids_with_key,
estimated_next_schedule_timestamp - inc
)
.fetch_all(job_queue_db)
.await?
};
tracing::debug!(
"all_jobs: {:?}, estimated_next_schedule_timestamp: {:?}, inc: {:?}",
@@ -2826,22 +2953,25 @@ pub async fn pull(
WITH ping AS (
UPDATE v2_job_runtime SET ping = null WHERE id = $2
)
UPDATE v2_job_queue SET
UPDATE
v2_job_queue
SET
running = false,
started_at = null,
scheduled_for = $1
WHERE id = $2",
WHERE
id = $2",
estimated_next_schedule_timestamp,
job_uuid,
)
.execute(db)
.execute(job_queue_db)
.await
.map_err(|e| Error::internal_err(format!("Could not update and re-queue job {job_uuid}. The job will be marked as running but it is not running: {e:#}")))?;
}
}
async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
db: &Pool<Postgres>,
job_queue_db: &DB,
suspend_first: bool,
worker_name: &str,
#[cfg(feature = "benchmark")] bench: &mut BenchmarkIter,
@@ -2865,7 +2995,7 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
// tracing::info!("Pulling job with query: {}", query);
sqlx::query_as::<_, PulledJob>(&query)
.bind(worker_name)
.fetch_optional(db)
.fetch_optional(job_queue_db)
.await?
} else {
None
@@ -2892,7 +3022,7 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit<'c>(
let r = sqlx::query_as::<_, PulledJob>(query)
.bind(worker_name)
.fetch_optional(db)
.fetch_optional(job_queue_db)
.await?;
#[cfg(feature = "benchmark")]
@@ -4559,23 +4689,44 @@ pub async fn push<'c, 'd>(
})
};
let shard_mode = *SHARD_MODE;
let mut existing_job_id = false;
let mut tx = tx.into_tx().await?;
let job_id: Uuid = if let Some(job_id) = job_id {
let conflicting_id = sqlx::query_scalar!("SELECT 1 FROM v2_job WHERE id = $1", job_id)
.fetch_optional(&mut *tx)
.await?;
existing_job_id = true;
job_id
} else {
Ulid::new().into()
};
let mut job_queue_tx = if shard_mode {
if let Some(job_queue_db) = get_shard_db_from_job_uuid(job_id) {
Some(job_queue_db.begin().await?)
} else {
None
}
} else {
None
};
if existing_job_id {
let conflicting_id = if shard_mode {
sqlx::query_scalar!("SELECT 1 FROM v2_job WHERE id = $1", job_id)
.fetch_optional(&mut **job_queue_tx.as_mut().unwrap())
.await?
} else {
sqlx::query_scalar!("SELECT 1 FROM v2_job WHERE id = $1", job_id)
.fetch_optional(&mut *tx)
.await?
};
if conflicting_id.is_some() {
return Err(Error::BadRequest(format!(
"Job with id {job_id} already exists"
)));
}
job_id
} else {
Ulid::new().into()
};
}
if concurrent_limit.is_some() {
insert_concurrency_key(
@@ -4736,7 +4887,9 @@ pub async fn push<'c, 'd>(
trigger_kind as Option<JobTriggerKind>,
running,
)
.execute(&mut *tx)
.execute(if !shard_mode {&mut *tx} else {
job_queue_tx.as_mut().unwrap()
})
.warn_after_seconds(1)
.await?;
@@ -4767,7 +4920,11 @@ pub async fn push<'c, 'd>(
job_id,
Json(flow_status) as Json<FlowStatus>,
)
.execute(&mut *tx)
.execute(if !shard_mode {
&mut *tx
} else {
job_queue_tx.as_mut().unwrap()
})
.warn_after_seconds(1)
.await?;
}
@@ -4842,6 +4999,52 @@ pub async fn push<'c, 'd>(
.await?;
}
if shard_mode {
let job_queue_tx = job_queue_tx.unwrap();
let main_commit_result = tx.commit().await;
let commit_result = match main_commit_result {
Ok(_) => match job_queue_tx.commit().await {
Ok(_) => {
tracing::debug!(
"Both main DB and shard DB committed successfully for job {}",
job_id
);
let new_tx = _db.begin().await?;
Ok((job_id, new_tx))
}
Err(shard_error) => {
tracing::error!(
job_id = %job_id,
workspace_id = %workspace_id,
error = %shard_error,
"CRITICAL: Main DB committed but shard DB failed during job push. Data inconsistency detected."
);
Err(Error::internal_err(format!(
"CRITICAL: Shard DB commit failed after main DB commit for job {} (workspace: {}): {}",
job_id, workspace_id, shard_error
)))
}
},
Err(main_error) => {
if let Err(rollback_err) = job_queue_tx.rollback().await {
tracing::error!(
job_id = %job_id,
main_error = %main_error,
rollback_error = %rollback_err,
"Failed to rollback shard DB after main DB commit failure"
);
}
Err(Error::internal_err(format!(
"Main DB commit failed for job {}: {}",
job_id, main_error
)))
}
};
return commit_result;
}
Ok((job_id, tx))
}