feat: rely on PG time rather than worker time

This commit is contained in:
Ruben Fiszel
2022-08-20 15:16:08 +02:00
parent d76eefe573
commit f679e67de8
8 changed files with 81 additions and 82 deletions
+2 -4
View File
@@ -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,
+39 -30
View File
@@ -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<chrono::DateTime<chrono::Utc>> {
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<Option<chrono::DateTime<chrono::Utc>>> {
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<RunJobQuery>,
) -> 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<Preview>,
Query(sch_query): Query<RunJobQuery>,
) -> 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<PreviewFlow>,
Query(sch_query): Query<RunJobQuery>,
) -> 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<Uuid, Error> {
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<Option<QueuedJob>, crate::Error> {
let now = chrono::Utc::now();
let job: Option<QueuedJob> = 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?;
+4 -5
View File
@@ -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<String> {
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::<TokenResponse>()
.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
+3 -2
View File
@@ -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;
+12 -14
View File
@@ -96,10 +96,9 @@ impl AuthCache {
a @ Some(_) => a,
None => {
let user_o = sqlx::query_as::<_, (Option<String>, Option<String>, 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<String> {
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<Vec<String>, _> = 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<Vec<String>, _> = 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;
+9
View File
@@ -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<chrono::DateTime<chrono::Utc>> {
Ok(sqlx::query_scalar!("SELECT now()")
.fetch_one(db)
.await?
.unwrap())
}
pub fn not_found_if_none<T, U: AsRef<str>>(opt: Option<T>, kind: &str, name: U) -> Result<T> {
if let Some(o) = opt {
Ok(o)
+10 -23
View File
@@ -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
+2 -4
View File
@@ -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?;