From ea8a456accacfe5048926cb0bee4e1514a5b313c Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Wed, 26 Jul 2023 15:26:54 +0200 Subject: [PATCH] feat: handle worker groups with redis --- backend/src/main.rs | 6 +---- backend/windmill-api/src/lib.rs | 10 +++++++++ backend/windmill-queue/src/jobs.rs | 22 ++++++++++++------- .../windmill-queue/src/queue_transaction.rs | 19 ++++++++-------- 4 files changed, 35 insertions(+), 22 deletions(-) diff --git a/backend/src/main.rs b/backend/src/main.rs index 8111a1ac05..e0812edac2 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -91,11 +91,7 @@ async fn main() -> anyhow::Result<()> { tracing::info!("Database connected"); let rsmq = if let Some(config) = rsmq_config { - let mut rsmq = rsmq_async::MultiplexedRsmq::new(config).await.unwrap(); - - let _ = rsmq_async::RsmqConnection::create_queue(&mut rsmq, "main_queue", None, None, None) - .await; - Some(rsmq) + Some(rsmq_async::MultiplexedRsmq::new(config).await.unwrap()) } else { None }; diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index dff743fcf6..ab7aa041e5 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -10,6 +10,7 @@ use crate::oauth2::AllClients; use crate::saml::{SamlSsoLogin, ServiceProviderExt}; use crate::scim::has_scim_token; use crate::tracing_init::MyOnFailure; +use crate::workers::ALL_TAGS; use crate::{ db::UserDB, oauth2::{build_oauth_clients, SlackVerifier}, @@ -149,6 +150,15 @@ pub async fn run_server( mut rx: tokio::sync::broadcast::Receiver<()>, port_tx: tokio::sync::oneshot::Sender, ) -> anyhow::Result<()> { + if let Some(mut rsmq) = rsmq.clone() { + for tag in ALL_TAGS.clone() { + let r = + rsmq_async::RsmqConnection::create_queue(&mut rsmq, &tag, None, None, None).await; + if r.is_ok() { + tracing::info!("Redis queue {tag} created"); + } + } + } let user_db = UserDB::new(db.clone()); let auth_cache = Arc::new(users::AuthCache::new( diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 144c44e00f..398b149b04 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -121,8 +121,6 @@ const MAX_FREE_EXECS: i32 = 1000; #[cfg(feature = "enterprise")] const MAX_FREE_CONCURRENT_RUNS: i32 = 15; -const RSMQ_MAIN_QUEUE: &'static str = "main_queue"; - #[async_recursion] pub async fn cancel_job<'c: 'async_recursion>( username: &str, @@ -169,7 +167,7 @@ pub async fn cancel_job<'c: 'async_recursion>( } } if let Some(mut rsmq) = rsmq.clone() { - rsmq.change_message_visibility(RSMQ_MAIN_QUEUE, &id.to_string(), 0) + rsmq.change_message_visibility(&job_running.tag, &id.to_string(), 0) .await .map_err(|e| anyhow::anyhow!(e))?; } @@ -806,6 +804,7 @@ pub async fn pull( rsmq.send_message( job_uuid.to_bytes_le().to_vec(), Option::Some(estimated_next_schedule_timestamp), + _requeued_job.tag, ); } tx.commit().await?; @@ -837,10 +836,17 @@ async fn pull_single_job_and_mark_as_running_no_concurrency_limit< ) -> windmill_common::error::Result<(Option, QueueTransaction<'c, R>)> { let job: Option = if let Some(mut rsmq) = rsmq { // TODO: REDIS: Race conditions / replace last_ping - let msg = rsmq - .pop_message::>(RSMQ_MAIN_QUEUE) - .await - .map_err(|e| anyhow::anyhow!(e))?; + + // TODO: shuffle this list to have fairness + let mut all_tags = ACCEPTED_TAGS.clone(); + + let mut msg: Option<_> = None; + while msg.is_none() && !all_tags.is_empty() { + msg = rsmq + .pop_message::>(&all_tags.pop().unwrap()) + .await + .map_err(|e| anyhow::anyhow!(e))?; + } // println!("3.1: {:?} {rs}", instant.elapsed()); if let Some(msg) = msg { @@ -1412,7 +1418,7 @@ pub async fn push<'c, R: rsmq_async::RsmqConnection + Send + 'c>( .await?; } if let Some(ref mut rsmq) = tx.rsmq { - rsmq.send_message(job_id.to_bytes_le().to_vec(), scheduled_for_o); + rsmq.send_message(job_id.to_bytes_le().to_vec(), scheduled_for_o, tag); } Ok((uuid, tx)) diff --git a/backend/windmill-queue/src/queue_transaction.rs b/backend/windmill-queue/src/queue_transaction.rs index 5dc137b0dd..04a4b380a3 100644 --- a/backend/windmill-queue/src/queue_transaction.rs +++ b/backend/windmill-queue/src/queue_transaction.rs @@ -5,8 +5,8 @@ use rsmq_async::{RedisBytes, RsmqConnection}; use sqlx::{Postgres, Transaction}; pub enum RedisOp { - SendMessage(RedisBytes, Option>), - DeleteMessage(String), + SendMessage(RedisBytes, Option>, String), + DeleteMessage(String, String), } unsafe impl Send for RedisOp {} @@ -14,17 +14,17 @@ unsafe impl Send for RedisOp {} impl RedisOp { pub async fn apply(self, rsmq: &mut R) -> Result<(), rsmq_async::RsmqError> { match self { - RedisOp::SendMessage(bytes, time) => { + RedisOp::SendMessage(bytes, time, queue) => { rsmq.send_message( - "main_queue", + &queue, bytes, time.map(|t| (t - chrono::Utc::now()).num_seconds()) .and_then(|e| e.try_into().ok()), ) .await?; } - RedisOp::DeleteMessage(id) => { - rsmq.delete_message("main_queue", &id).await?; + RedisOp::DeleteMessage(id, queue) => { + rsmq.delete_message(&queue, &id).await?; } }; @@ -56,13 +56,14 @@ impl RedisTransaction { &mut self, bytes: E, delay_until: Option>, + queue: String, ) { self.queued_ops - .push(RedisOp::SendMessage(bytes.into(), delay_until)) + .push(RedisOp::SendMessage(bytes.into(), delay_until, queue)) } - pub fn delete_message(&mut self, id: String) { - self.queued_ops.push(RedisOp::DeleteMessage(id)) + pub fn delete_message(&mut self, id: String, queue: String) { + self.queued_ops.push(RedisOp::DeleteMessage(id, queue)) } }