From d8897d80545530ba58f6e47d4de02afd66ea39d1 Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Fri, 20 Jan 2023 19:39:25 +0100
Subject: [PATCH] feat: add QUEUE_LIMIT_WAIT_RESULT

---
 README.md                        |  3 +++
 backend/windmill-api/src/jobs.rs |  5 +++--
 backend/windmill-api/src/lib.rs  | 13 +++++++++----
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index a8f21b3442..d4d5cbb0ce 100644
--- a/README.md
+++ b/README.md
@@ -305,6 +305,9 @@ upcoming CLI tool.
 | DATABASE_CONNECTIONS      | 50 (Server)/3 (Worker) | The max number of connections in the database connection pool                                                                                                                                      | All                   |
 | SUPERADMIN_SECRET         | None                   | A token that would let the caller act as a virtual superadmin superadmin@windmill.dev                                                                                                              | Server                |
 | TIMEOUT_WAIT_RESULT       | 20                     | The number of seconds to wait before timeout on the 'run_wait_result' endpoint                                                                                                                     | Worker                |
+| QUEUE_LIMIT_WAIT_RESULT   | None                   | The number of max jobs in the queue before rejecting immediately the request in 'run_wait_result' endpoint. Takes precedence on the query arg. If none is specified, there are no limit.           | Worker                |
+
+
 
 
 ## Run a local dev setup
diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs
index 300eb4e786..78892088af 100644
--- a/backend/windmill-api/src/jobs.rs
+++ b/backend/windmill-api/src/jobs.rs
@@ -38,7 +38,7 @@ use crate::{
     db::{UserDB, DB},
     users::{require_owner_of_path, Authed},
     variables::get_workspace_key,
-    BaseUrl, TimeoutWaitResult,
+    BaseUrl, QueueLimitWaitResult, TimeoutWaitResult,
 };
 
 pub fn workspaced_service() -> Router {
@@ -1298,12 +1298,13 @@ pub async fn run_wait_result_job_by_path(
     Extension(user_db): Extension<UserDB>,
     Extension(db): Extension<DB>,
     Extension(timeout): Extension<Arc<TimeoutWaitResult>>,
+    Extension(queue_limit): Extension<Arc<QueueLimitWaitResult>>,
     Path((w_id, script_path)): Path<(String, StripPath)>,
     Query(run_query): Query<RunJobQuery>,
     headers: HeaderMap,
     Json(args): Json<Option<serde_json::Map<String, serde_json::Value>>>,
 ) -> error::JsonResult<serde_json::Value> {
-    check_queue_too_long(db, run_query.queue_limit).await?;
+    check_queue_too_long(db, queue_limit.0.or(run_query.queue_limit)).await?;
     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?;
diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs
index c245754cf8..8db6297e78 100644
--- a/backend/windmill-api/src/lib.rs
+++ b/backend/windmill-api/src/lib.rs
@@ -54,6 +54,7 @@ pub struct CookieDomain(Option<String>);
 pub struct CloudHosted(bool);
 pub struct ContentSecurityPolicy(String);
 pub struct TimeoutWaitResult(i32);
+pub struct QueueLimitWaitResult(Option<i64>);
 
 pub use users::delete_expired_items_perdiodically;
 
@@ -117,14 +118,18 @@ pub async fn run_server(
                         .nest("/scripts", scripts::workspaced_service())
                         .nest(
                             "/jobs",
-                            jobs::workspaced_service().layer(Extension(Arc::new(
-                                TimeoutWaitResult(
+                            jobs::workspaced_service()
+                                .layer(Extension(Arc::new(TimeoutWaitResult(
                                     std::env::var("TIMEOUT_WAIT_RESULT")
                                         .ok()
                                         .and_then(|x| x.parse().ok())
                                         .unwrap_or(20),
-                                ),
-                            ))),
+                                ))))
+                                .layer(Extension(Arc::new(QueueLimitWaitResult(
+                                    std::env::var("QUEUE_LIMIT_WAIT_RESULT")
+                                        .ok()
+                                        .and_then(|x| x.parse().ok()),
+                                )))),
                         )
                         .nest(
                             "/users",