feat: handle worker groups with redis

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
     };

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<u16>,
 ) -> 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(

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<R: rsmq_async::RsmqConnection + Send + Clone>(
                 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<QueuedJob>, QueueTransaction<'c, R>)> {
     let job: Option<QueuedJob> = if let Some(mut rsmq) = rsmq {
         // TODO: REDIS: Race conditions / replace last_ping
-        let msg = rsmq
-            .pop_message::<Vec<u8>>(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::<Vec<u8>>(&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))

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<chrono::DateTime<chrono::Utc>>),
-    DeleteMessage(String),
+    SendMessage(RedisBytes, Option<chrono::DateTime<chrono::Utc>>, String),
+    DeleteMessage(String, String),
 }
 
 unsafe impl Send for RedisOp {}
@@ -14,17 +14,17 @@ unsafe impl Send for RedisOp {}
 impl RedisOp {
     pub async fn apply<R: RsmqConnection>(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<R: RsmqConnection> RedisTransaction<R> {
         &mut self,
         bytes: E,
         delay_until: Option<chrono::DateTime<chrono::Utc>>,
+        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))
     }
 }