From f679e67de8f40bc8e6b585b3af7e599dc8dabf4e Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Sat, 20 Aug 2022 15:16:08 +0200 Subject: [PATCH] feat: rely on PG time rather than worker time --- backend/src/flows.rs | 6 ++-- backend/src/jobs.rs | 69 +++++++++++++++++++++----------------- backend/src/oauth2.rs | 9 +++-- backend/src/schedule.rs | 5 +-- backend/src/users.rs | 26 +++++++------- backend/src/utils.rs | 9 +++++ backend/src/worker.rs | 33 ++++++------------ backend/src/worker_flow.rs | 6 ++-- 8 files changed, 81 insertions(+), 82 deletions(-) diff --git a/backend/src/flows.rs b/backend/src/flows.rs index d85b719de1..904db755b0 100644 --- a/backend/src/flows.rs +++ b/backend/src/flows.rs @@ -229,14 +229,13 @@ async fn create_flow( sqlx::query!( "INSERT INTO flow (workspace_id, path, summary, description, value, edited_by, edited_at, \ - schema) VALUES ($1, $2, $3, $4, $5, $6, $7, $8::text::json)", + schema) VALUES ($1, $2, $3, $4, $5, $6, now(), $7::text::json)", w_id, nf.path, nf.summary, nf.description, nf.value, &authed.username, - &chrono::Utc::now(), nf.schema.and_then(|x| serde_json::to_string(&x.0).ok()), ) .execute(&mut tx) @@ -299,13 +298,12 @@ async fn update_flow( let schema = nf.schema.map(|x| x.0); let flow = sqlx::query_scalar!( "UPDATE flow SET path = $1, summary = $2, description = $3, value = $4, edited_by = $5, \ - edited_at = $6, schema = $7 WHERE path = $8 AND workspace_id = $9 RETURNING path", + edited_at = now(), schema = $6 WHERE path = $7 AND workspace_id = $8 RETURNING path", nf.path, nf.summary, nf.description, nf.value, &authed.username, - &chrono::Utc::now(), schema, flow_path, w_id, diff --git a/backend/src/jobs.rs b/backend/src/jobs.rs index ce0553f3a1..c68b5ab345 100644 --- a/backend/src/jobs.rs +++ b/backend/src/jobs.rs @@ -7,7 +7,6 @@ */ use axum::extract::Host; -use chrono::Duration; use sql_builder::prelude::*; use sqlx::{query_scalar, Postgres, Transaction}; @@ -23,7 +22,7 @@ use crate::{ schedule::get_schedule_opt, scripts::{get_hub_script_by_path, ScriptHash, ScriptLang}, users::{owner_to_token_owner, Authed}, - utils::{require_admin, Pagination, StripPath}, + utils::{require_admin, Pagination, StripPath, now_from_db}, worker, worker_flow::init_flow_status, }; @@ -140,11 +139,18 @@ pub struct RunJobQuery { } impl RunJobQuery { - fn get_scheduled_for(self) -> Option> { - self.scheduled_for.or_else(|| { - self.scheduled_in_secs - .map(|s| chrono::Utc::now() + Duration::seconds(s)) - }) + async fn get_scheduled_for<'c>( + self, + db: &mut Transaction<'c, Postgres>, + ) -> error::Result>> { + if let Some(scheduled_for) = self.scheduled_for { + Ok(Some(scheduled_for)) + } else if let Some(scheduled_in_secs) = self.scheduled_in_secs { + let now = now_from_db(db).await?; + Ok(Some(now + chrono::Duration::seconds(scheduled_in_secs))) + } else { + Ok(None) + } } } @@ -156,7 +162,8 @@ pub async fn run_flow_by_path( Query(run_query): Query, ) -> error::Result<(StatusCode, String)> { let flow_path = flow_path.to_path(); - let tx = user_db.begin(&authed).await?; + let mut tx = user_db.begin(&authed).await?; + let scheduled_for = run_query.get_scheduled_for(&mut tx).await?; let (uuid, tx) = push( tx, &w_id, @@ -164,7 +171,7 @@ pub async fn run_flow_by_path( args, &authed.username, owner_to_token_owner(&authed.username, false), - run_query.get_scheduled_for(), + scheduled_for, None, run_query.parent_job, false, @@ -184,6 +191,8 @@ pub async fn run_job_by_path( let script_path = script_path.to_path(); let mut tx = user_db.begin(&authed).await?; let job_payload = script_path_to_payload(script_path, &mut tx, &w_id).await?; + let scheduled_for = run_query.get_scheduled_for(&mut tx).await?; + let (uuid, tx) = push( tx, &w_id, @@ -191,7 +200,7 @@ pub async fn run_job_by_path( args, &authed.username, owner_to_token_owner(&authed.username, false), - run_query.get_scheduled_for(), + scheduled_for, None, run_query.parent_job, false, @@ -211,6 +220,8 @@ pub async fn run_wait_result_job_by_path( let script_path = script_path.to_path(); let mut tx = user_db.clone().begin(&authed).await?; let job_payload = script_path_to_payload(script_path, &mut tx, &w_id).await?; + let scheduled_for = run_query.get_scheduled_for(&mut tx).await?; + let (uuid, tx) = push( tx, &w_id, @@ -218,7 +229,7 @@ pub async fn run_wait_result_job_by_path( args, &authed.username, owner_to_token_owner(&authed.username, false), - run_query.get_scheduled_for(), + scheduled_for, None, run_query.parent_job, false, @@ -299,6 +310,8 @@ pub async fn run_job_by_hash( let hash = script_hash.0; let mut tx = user_db.begin(&authed).await?; let path = get_path_for_hash(&mut tx, &w_id, hash).await?; + let scheduled_for = run_query.get_scheduled_for(&mut tx).await?; + let (uuid, tx) = push( tx, &w_id, @@ -306,7 +319,7 @@ pub async fn run_job_by_hash( args, &authed.username, owner_to_token_owner(&authed.username, false), - run_query.get_scheduled_for(), + scheduled_for, None, run_query.parent_job, false, @@ -344,7 +357,9 @@ async fn run_preview_job( Json(preview): Json, Query(sch_query): Query, ) -> error::Result<(StatusCode, String)> { - let tx = user_db.begin(&authed).await?; + let mut tx = user_db.begin(&authed).await?; + let scheduled_for = sch_query.get_scheduled_for(&mut tx).await?; + let (uuid, tx) = push( tx, &w_id, @@ -356,7 +371,7 @@ async fn run_preview_job( preview.args, &authed.username, owner_to_token_owner(&authed.username, false), - sch_query.get_scheduled_for(), + scheduled_for, None, None, false, @@ -373,7 +388,8 @@ async fn run_preview_flow_job( Json(raw_flow): Json, Query(sch_query): Query, ) -> error::Result<(StatusCode, String)> { - let tx = user_db.begin(&authed).await?; + let mut tx = user_db.begin(&authed).await?; + let scheduled_for = sch_query.get_scheduled_for(&mut tx).await?; let (uuid, tx) = push( tx, &w_id, @@ -381,7 +397,7 @@ async fn run_preview_flow_job( raw_flow.args, &authed.username, owner_to_token_owner(&authed.username, false), - sch_query.get_scheduled_for(), + scheduled_for, None, None, false, @@ -1294,24 +1310,21 @@ pub async fn add_completed_job( logs: String, ) -> Result { let job_id = queued_job.id.clone(); - let duration_ms = (chrono::Utc::now() - queued_job.started_at.unwrap_or(queued_job.created_at)) - .num_milliseconds() as i32; - let _ = sqlx::query!( + sqlx::query!( "INSERT INTO completed_job as cj (workspace_id, id, parent_job, created_by, created_at, duration_ms, success, \ script_hash, script_path, args, result, logs, \ raw_code, canceled, canceled_by, canceled_reason, job_kind, schedule_path, \ permissioned_as, flow_status, raw_flow, is_flow_step, is_skipped) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, \ + VALUES ($1, $2, $3, $4, $5, EXTRACT(milliseconds FROM (now() - $6)), $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, \ $18, $19, $20, $21, $22, $23) - ON CONFLICT (id) DO UPDATE SET success = $7, result = $11, logs = concat(cj.logs, $12) - RETURNING id", + ON CONFLICT (id) DO UPDATE SET success = $7, result = $11, logs = concat(cj.logs, $12)", queued_job.workspace_id, queued_job.id, queued_job.parent_job, queued_job.created_by, queued_job.created_at, - duration_ms, + queued_job.started_at, success, queued_job.script_hash.map(|x| x.0), queued_job.script_path, @@ -1330,7 +1343,7 @@ pub async fn add_completed_job( queued_job.is_flow_step, skipped ) - .fetch_one(db) + .execute(db) .await .map_err(|e| Error::InternalErr(format!("Could not add completed job {job_id}: {e}")))?; tracing::debug!("Added completed job {}", queued_job.id); @@ -1380,23 +1393,19 @@ pub async fn schedule_again_if_scheduled( } pub async fn pull(db: &DB) -> Result, crate::Error> { - let now = chrono::Utc::now(); - let job: Option = sqlx::query_as::<_, QueuedJob>( "UPDATE queue - SET running = true, started_at = $1 + SET running = true, started_at = now() WHERE id IN ( SELECT id FROM queue - WHERE running = false AND scheduled_for <= $2 + WHERE running = false AND scheduled_for <= now() ORDER BY scheduled_for FOR UPDATE SKIP LOCKED LIMIT 1 ) RETURNING *", ) - .bind(now) - .bind(now) .fetch_optional(db) .await?; diff --git a/backend/src/oauth2.rs b/backend/src/oauth2.rs index 34f9223e24..6a808c5495 100644 --- a/backend/src/oauth2.rs +++ b/backend/src/oauth2.rs @@ -10,7 +10,6 @@ use axum::{ routing::{get, post}, Json, Router, }; -use chrono::{Duration, Utc}; use hyper::StatusCode; use itertools::Itertools; @@ -22,6 +21,7 @@ use sqlx::{Postgres, Transaction}; use tokio::{fs::File, io::AsyncReadExt}; use tower_cookies::{Cookie, Cookies}; +use crate::utils::now_from_db; use crate::IsSecure; use crate::{ audit::{audit_log, ActionKind}, @@ -274,14 +274,13 @@ async fn create_account( ) -> error::Result { let mut tx = user_db.begin(&authed).await?; - let expires_at = chrono::Utc::now() + Duration::seconds(payload.expires_in); let id = sqlx::query_scalar!( "INSERT INTO account (workspace_id, client, owner, expires_at, refresh_token) VALUES ($1, \ - $2, $3, $4, $5) RETURNING id", + $2, $3, now() + ($4 || ' seconds')::interval, $5) RETURNING id", w_id, payload.client, payload.owner, - expires_at, + payload.expires_in.to_string(), payload.refresh_token ) .fetch_one(&mut tx) @@ -473,7 +472,7 @@ pub async fn _refresh_token<'c>( .execute::() .await .map_err(to_anyhow)?; - let expires_at = Utc::now() + let expires_at = now_from_db(&mut tx).await? + chrono::Duration::seconds( token .expires_in diff --git a/backend/src/schedule.rs b/backend/src/schedule.rs index e2265b516e..b44c25c74b 100644 --- a/backend/src/schedule.rs +++ b/backend/src/schedule.rs @@ -14,7 +14,7 @@ use crate::{ error::{self, Error, JsonResult, Result}, jobs::{self, push, JobPayload}, users::Authed, - utils::{get_owner_from_path, Pagination, StripPath}, + utils::{get_owner_from_path, now_from_db, Pagination, StripPath}, }; use axum::{ extract::{Extension, Path, Query}, @@ -76,8 +76,9 @@ pub async fn push_scheduled_job<'c>( .map_err(|e| error::Error::BadRequest(e.to_string()))?; let offset = Duration::minutes(schedule.offset_.into()); + let now = now_from_db(&mut tx).await?; let next = sched - .after(&(chrono::Utc::now() - offset + Duration::seconds(1))) + .after(&(now - offset + Duration::seconds(1))) .next() .expect("a schedule should have a next event") + offset; diff --git a/backend/src/users.rs b/backend/src/users.rs index 5e24e8c782..6f350d9a83 100644 --- a/backend/src/users.rs +++ b/backend/src/users.rs @@ -96,10 +96,9 @@ impl AuthCache { a @ Some(_) => a, None => { let user_o = sqlx::query_as::<_, (Option, Option, bool)>( - "UPDATE token SET last_used_at = $1 WHERE token = $2 AND (expiration > NOW() \ + "UPDATE token SET last_used_at = now() WHERE token = $1 AND (expiration > NOW() \ OR expiration IS NULL) RETURNING owner, email, super_admin", ) - .bind(chrono::Utc::now()) .bind(token) .fetch_optional(&self.db) .await @@ -1272,11 +1271,11 @@ pub async fn create_session_token<'c>( sqlx::query!( "INSERT INTO token (token, email, label, expiration, super_admin) - VALUES ($1, $2, $3, $4, $5)", + VALUES ($1, $2, $3, now() + ($4 || ' hours')::interval, $5)", token, email, "session", - chrono::Utc::now() + chrono::Duration::hours(TTL_TOKEN_DB_H as i64), + TTL_TOKEN_DB_H.to_string(), super_admin ) .execute(tx) @@ -1305,7 +1304,8 @@ pub async fn create_token_for_owner( db: &DB, w_id: &str, owner: &str, - NewToken { label, expiration }: NewToken, + label: &str, + expires_in: i32, username: &str, ) -> Result { use rand::prelude::*; @@ -1324,18 +1324,18 @@ pub async fn create_token_for_owner( .await? .unwrap_or(false); - sqlx::query!( + let expiration = sqlx::query_scalar!( "INSERT INTO token (workspace_id, token, owner, label, expiration, super_admin) - VALUES ($1, $2, $3, $4, $5, $6)", + VALUES ($1, $2, $3, $4, now() + ($5 || ' seconds')::interval, $6) RETURNING expiration", &w_id, token, owner, label, - expiration, + expires_in.to_string(), is_super_admin ) - .execute(&mut tx) + .fetch_one(&mut tx) .await?; audit_log( &mut tx, @@ -1346,7 +1346,7 @@ pub async fn create_token_for_owner( Some(&truncate_token(&token)), Some( [ - label.as_ref().map(|label| ("label", &label[..])), + Some(("label", label)), expiration .map(|x| x.to_string()) .as_ref() @@ -1494,10 +1494,9 @@ pub async fn delete_expired_items_perdiodically( ) -> () { loop { let tokens_deleted_r: std::result::Result, _> = sqlx::query_scalar( - "DELETE FROM token WHERE expiration <= $1 + "DELETE FROM token WHERE expiration <= now() RETURNING concat(substring(token for 10), '*****')", ) - .bind(chrono::Utc::now()) .fetch_all(db) .await; @@ -1507,10 +1506,9 @@ pub async fn delete_expired_items_perdiodically( } let magic_links_deleted_r: std::result::Result, _> = sqlx::query_scalar( - "DELETE FROM magic_link WHERE expiration <= $1 + "DELETE FROM magic_link WHERE expiration <= now() RETURNING concat(substring(token for 10), '*****')", ) - .bind(chrono::Utc::now()) .fetch_all(db) .await; diff --git a/backend/src/utils.rs b/backend/src/utils.rs index a059f9cb25..be2e6dce9e 100644 --- a/backend/src/utils.rs +++ b/backend/src/utils.rs @@ -83,6 +83,15 @@ pub fn paginate(pagination: Pagination) -> (usize, usize) { (per_page, offset) } +pub async fn now_from_db<'c>( + db: &mut Transaction<'c, Postgres>, +) -> Result> { + Ok(sqlx::query_scalar!("SELECT now()") + .fetch_one(db) + .await? + .unwrap()) +} + pub fn not_found_if_none>(opt: Option, kind: &str, name: U) -> Result { if let Some(o) = opt { Ok(o) diff --git a/backend/src/worker.rs b/backend/src/worker.rs index f28d731684..1f8731c0eb 100644 --- a/backend/src/worker.rs +++ b/backend/src/worker.rs @@ -135,8 +135,7 @@ pub async fn run_worker( loop { if last_ping.elapsed().as_secs() > NUM_SECS_ENV_CHECK { sqlx::query!( - "UPDATE worker_ping SET ping_at = $1, jobs_executed = $2 WHERE worker = $3", - chrono::Utc::now(), + "UPDATE worker_ping SET ping_at = now(), jobs_executed = $1 WHERE worker = $2", jobs_executed, &worker_name ) @@ -651,12 +650,8 @@ async fn handle_nondep_job( &db, &job.workspace_id, &job.permissioned_as, - crate::users::NewToken { - label: Some("ephemeral-script".to_string()), - expiration: Some( - chrono::Utc::now() + chrono::Duration::seconds((timeout * 2).into()), - ), - }, + "ephemeral-script", + timeout * 2, &job.created_by, ) .await?; @@ -778,12 +773,8 @@ print(res_json) &db, &job.workspace_id, &job.permissioned_as, - crate::users::NewToken { - label: Some("ephemeral-script".to_string()), - expiration: Some( - chrono::Utc::now() + chrono::Duration::seconds((timeout * 2).into()), - ), - }, + "ephemeral-script", + timeout * 2, &job.created_by, ) .await?; @@ -1078,13 +1069,9 @@ async fn handle_child( tokio::spawn(async move { while !&done3.load(Ordering::Relaxed) { - let q = sqlx::query!( - "UPDATE queue SET last_ping = $1 WHERE id = $2", - chrono::Utc::now(), - id - ) - .execute(&db2) - .await; + let q = sqlx::query!("UPDATE queue SET last_ping = now() WHERE id = $1", id) + .execute(&db2) + .await; if q.is_err() { tracing::error!("error setting last ping for id {}", id); @@ -1211,8 +1198,8 @@ pub async fn restart_zombie_jobs_periodically( ) { loop { let restarted = sqlx::query!( - "UPDATE queue SET running = false WHERE last_ping < $1 and running = true RETURNING id, workspace_id", - chrono::Utc::now() - chrono::Duration::seconds(timeout as i64 * 5) + "UPDATE queue SET running = false WHERE last_ping < now() + ($1 || ' seconds')::interval and running = true RETURNING id, workspace_id", + (timeout * 5).to_string(), ) .fetch_all(db) .await diff --git a/backend/src/worker_flow.rs b/backend/src/worker_flow.rs index 9f65b42869..4ea780a32c 100644 --- a/backend/src/worker_flow.rs +++ b/backend/src/worker_flow.rs @@ -496,10 +496,8 @@ async fn push_next_flow_job( &db, &flow_job.workspace_id, &flow_job.permissioned_as, - crate::users::NewToken { - label: Some("transform-input".to_string()), - expiration: Some(chrono::Utc::now() + chrono::Duration::seconds(10)), - }, + "transform-input", + 10, &flow_job.created_by, ) .await?;