diff --git a/backend/tests/worker.rs b/backend/tests/worker.rs index a00dbda8ff..31f6a1d120 100644 --- a/backend/tests/worker.rs +++ b/backend/tests/worker.rs @@ -262,7 +262,7 @@ mod suspend_resume { // print_job(second, &db).await; let tx = db.begin().await.unwrap(); - let (tx, token) = windmill_worker::create_token_for_owner(tx, "test-workspace", "u/test-user", "", 100).await.unwrap(); + let (tx, token) = windmill_worker::create_token_for_owner(tx, "test-workspace", "u/test-user", "", 100, "").await.unwrap(); tx.commit().await.unwrap(); let secret = reqwest::get(format!( "http://localhost:{port}/api/w/test-workspace/jobs/job_signature/{second}/0?token={token}&approver=ruben" @@ -365,7 +365,7 @@ mod suspend_resume { let second = completed.next().await.unwrap(); let tx = db.begin().await.unwrap(); - let (tx, token) = windmill_worker::create_token_for_owner(tx, "test-workspace", "u/test-user", "", 100).await.unwrap(); + let (tx, token) = windmill_worker::create_token_for_owner(tx, "test-workspace", "u/test-user", "", 100, "").await.unwrap(); tx.commit().await.unwrap(); let secret = reqwest::get(format!( "http://localhost:{port}/api/w/test-workspace/jobs/job_signature/{second}/0?token={token}" diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 7152cac233..a2581a91bd 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -45,6 +45,7 @@ lazy_static::lazy_static! { } const MAX_FREE_EXECS: i32 = 1000; +const MAX_FREE_CONCURRENT_RUNS: i32 = 3; pub async fn cancel_job<'c>( username: &str, @@ -318,10 +319,26 @@ pub async fn push<'c>( .fetch_optional(&mut tx) .await? .unwrap_or(false); - if !is_super_admin && usage > MAX_FREE_EXECS { - return Err(error::Error::BadRequest(format!( + + if !is_super_admin { + if usage > MAX_FREE_EXECS { + return Err(error::Error::BadRequest(format!( "User {email} has exceeded the free usage limit of {MAX_FREE_EXECS} that applies outside of premium workspaces." ))); + } + let concurrent_runs = sqlx::query_scalar!( + "SELECT COUNT(id) FROM queue WHERE running = true AND email = $1", + email + ) + .fetch_one(&mut tx) + .await? + .unwrap_or(0); + + if concurrent_runs > MAX_FREE_CONCURRENT_RUNS.into() { + return Err(error::Error::BadRequest(format!( + "User {email} has exceeded the concurrent runs limit of {MAX_FREE_CONCURRENT_RUNS} that applies outside of premium workspaces." + ))); + } } } } diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index e81e4462ca..5272ba6afd 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -234,6 +234,7 @@ pub async fn create_token_for_owner<'c>( owner: &str, label: &str, expires_in: i32, + email: &str, ) -> error::Result<(Transaction<'c, Postgres>, String)> { // TODO: Bad implementation. We should not have access to this DB here. let token: String = rd_string(30); @@ -245,14 +246,15 @@ pub async fn create_token_for_owner<'c>( sqlx::query_scalar!( "INSERT INTO token - (workspace_id, token, owner, label, expiration, super_admin) - VALUES ($1, $2, $3, $4, now() + ($5 || ' seconds')::interval, $6)", + (workspace_id, token, owner, label, expiration, super_admin, email) + VALUES ($1, $2, $3, $4, now() + ($5 || ' seconds')::interval, $6, $7)", &w_id, token, owner, label, expires_in.to_string(), - is_super_admin + is_super_admin, + email ) .execute(&mut tx) .await?; @@ -579,6 +581,7 @@ pub async fn run_worker( &job.permissioned_as, "ephemeral-script", timeout * 2, + &job.email, ) .await.expect("could not create job token"); tx.commit().await.expect("could not commit job token"); @@ -2638,6 +2641,7 @@ async fn handle_zombie_jobs(db: &Pool, timeout: i32, base_url: &str) { &job.permissioned_as, "ephemeral-zombie-jobs", timeout * 2, + &job.email, ) .await .expect("could not create job token"); diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index 3b6155b846..4df7866f5d 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -1883,6 +1883,7 @@ async fn get_transform_context( &flow_job.permissioned_as, "transform-input", 10, + &flow_job.email, ) .await?; //we need to commit asap otherwise the token won't be valid for auth to check outside of this transaction