Merge pull request #7640 from Turbo87/multi-queue

Add support for running multiple background worker queues

Author: Turbo87

crates_io_worker/src/background_job.rs

@@ -7,6 +7,8 @@ use serde::de::DeserializeOwned;
 use serde::Serialize;
 use tracing::instrument;
 
+pub const DEFAULT_QUEUE: &str = "default";
+
 #[async_trait]
 pub trait BackgroundJob: Serialize + DeserializeOwned + Send + Sync + 'static {
     /// Unique name of the task.
@@ -19,6 +21,9 @@ pub trait BackgroundJob: Serialize + DeserializeOwned + Send + Sync + 'static {
     /// [Self::enqueue_with_priority] can be used to override the priority value.
     const PRIORITY: i16 = 0;
 
+    /// Job queue where this job will be executed.
+    const QUEUE: &'static str = DEFAULT_QUEUE;
+
     /// The application data provided to this job at runtime.
     type Context: Clone + Send + 'static;
 

crates_io_worker/src/job_registry.rs

@@ -29,6 +29,11 @@ impl<Context: Clone + Send + 'static> JobRegistry<Context> {
     pub fn get(&self, key: &str) -> Option<&Arc<RunTaskFn<Context>>> {
         self.entries.get(key)
     }
+
+    /// Returns a list of all registered job types.
+    pub fn job_types(&self) -> Vec<String> {
+        self.entries.keys().cloned().collect()
+    }
 }
 
 fn runnable<J: BackgroundJob>(ctx: J::Context, payload: serde_json::Value) -> RunTaskFnReturn {
@@ -37,3 +42,30 @@ fn runnable<J: BackgroundJob>(ctx: J::Context, payload: serde_json::Value) -> Ru
         job.run(ctx).await
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::BackgroundJob;
+    use async_trait::async_trait;
+    use serde::{Deserialize, Serialize};
+
+    #[test]
+    fn test_job_types() {
+        #[derive(Serialize, Deserialize)]
+        struct TestJob;
+
+        #[async_trait]
+        impl BackgroundJob for TestJob {
+            const JOB_NAME: &'static str = "test";
+            type Context = ();
+            async fn run(&self, _: Self::Context) -> anyhow::Result<()> {
+                Ok(())
+            }
+        }
+
+        let mut registry = JobRegistry::default();
+        registry.register::<TestJob>();
+        assert_eq!(registry.job_types(), vec!["test"]);
+    }
+}

crates_io_worker/src/runner.rs

@@ -1,10 +1,12 @@
+use crate::background_job::DEFAULT_QUEUE;
 use crate::job_registry::JobRegistry;
 use crate::worker::Worker;
 use crate::{storage, BackgroundJob};
 use anyhow::anyhow;
 use diesel::prelude::*;
 use diesel::r2d2::{ConnectionManager, Pool, PoolError, PooledConnection};
 use futures_util::future::join_all;
+use std::collections::HashMap;
 use std::time::Duration;
 use tokio::runtime::Handle;
 use tokio::task::JoinHandle;
@@ -19,10 +21,8 @@ pub type PooledConn = PooledConnection<ConnectionManager<PgConnection>>;
 pub struct Runner<Context> {
     rt_handle: Handle,
     connection_pool: ConnectionPool,
-    num_workers: usize,
-    job_registry: JobRegistry<Context>,
+    queues: HashMap<String, Queue<Context>>,
     context: Context,
-    poll_interval: Duration,
     shutdown_when_queue_empty: bool,
 }
 
@@ -31,29 +31,34 @@ impl<Context: Clone + Send + 'static> Runner<Context> {
         Self {
             rt_handle: rt_handle.clone(),
             connection_pool,
-            num_workers: 1,
-            job_registry: Default::default(),
+            queues: HashMap::new(),
             context,
-            poll_interval: DEFAULT_POLL_INTERVAL,
             shutdown_when_queue_empty: false,
         }
     }
 
-    /// Set the number of workers to spawn.
-    pub fn num_workers(mut self, num_workers: usize) -> Self {
-        self.num_workers = num_workers;
+    /// Register a new job type for this job runner.
+    pub fn register_job_type<J: BackgroundJob<Context = Context>>(mut self) -> Self {
+        let queue = self.queues.entry(J::QUEUE.into()).or_default();
+        queue.job_registry.register::<J>();
         self
     }
 
-    /// Set the interval after which each worker polls for new jobs.
-    pub fn poll_interval(mut self, poll_interval: Duration) -> Self {
-        self.poll_interval = poll_interval;
-        self
+    /// Adjust the configuration of the [DEFAULT_QUEUE] queue.
+    pub fn configure_default_queue<F>(self, f: F) -> Self
+    where
+        F: FnOnce(&mut Queue<Context>) -> &Queue<Context>,
+    {
+        self.configure_queue(DEFAULT_QUEUE, f)
     }
 
-    /// Register a new job type for this job runner.
-    pub fn register_job_type<J: BackgroundJob<Context = Context>>(mut self) -> Self {
-        self.job_registry.register::<J>();
+    /// Adjust the configuration of a queue. If the queue does not exist,
+    /// it will be created.
+    pub fn configure_queue<F>(mut self, name: &str, f: F) -> Self
+    where
+        F: FnOnce(&mut Queue<Context>) -> &Queue<Context>,
+    {
+        f(self.queues.entry(name.into()).or_default());
         self
     }
 
@@ -67,24 +72,27 @@ impl<Context: Clone + Send + 'static> Runner<Context> {
     ///
     /// This returns a `RunningRunner` which can be used to wait for the workers to shutdown.
     pub fn start(&self) -> RunHandle {
-        let handles = (0..self.num_workers)
-            .map(|i| {
-                let name = format!("background-worker-{i}");
+        let mut handles = Vec::new();
+        for (queue_name, queue) in &self.queues {
+            for i in 1..=queue.num_workers {
+                let name = format!("background-worker-{queue_name}-{i}");
                 info!(worker.name = %name, "Starting worker…");
 
                 let worker = Worker {
                     connection_pool: self.connection_pool.clone(),
                     context: self.context.clone(),
-                    job_registry: self.job_registry.clone(),
+                    job_registry: queue.job_registry.clone(),
                     shutdown_when_queue_empty: self.shutdown_when_queue_empty,
-                    poll_interval: self.poll_interval,
+                    poll_interval: queue.poll_interval,
                 };
 
-                self.rt_handle.spawn_blocking(move || {
+                let handle = self.rt_handle.spawn_blocking(move || {
                     info_span!("worker", worker.name = %name).in_scope(|| worker.run())
-                })
-            })
-            .collect();
+                });
+
+                handles.push(handle);
+            }
+        }
 
         RunHandle { handles }
     }
@@ -121,3 +129,33 @@ impl RunHandle {
         });
     }
 }
+
+pub struct Queue<Context> {
+    job_registry: JobRegistry<Context>,
+    num_workers: usize,
+    poll_interval: Duration,
+}
+
+impl<Context> Default for Queue<Context> {
+    fn default() -> Self {
+        Self {
+            job_registry: JobRegistry::default(),
+            num_workers: 1,
+            poll_interval: DEFAULT_POLL_INTERVAL,
+        }
+    }
+}
+
+impl<Context> Queue<Context> {
+    /// Set the number of workers to spawn for this queue.
+    pub fn num_workers(&mut self, num_workers: usize) -> &mut Self {
+        self.num_workers = num_workers;
+        self
+    }
+
+    /// Set the interval after which each worker of this queue polls for new jobs.
+    pub fn poll_interval(&mut self, poll_interval: Duration) -> &mut Self {
+        self.poll_interval = poll_interval;
+        self
+    }
+}

crates_io_worker/src/storage.rs

@@ -25,9 +25,13 @@ fn retriable() -> Box<dyn BoxableExpression<background_jobs::table, Pg, SqlType
 
 /// Finds the next job that is unlocked, and ready to be retried. If a row is
 /// found, it will be locked.
-pub(super) fn find_next_unlocked_job(conn: &mut PgConnection) -> QueryResult<BackgroundJob> {
+pub(super) fn find_next_unlocked_job(
+    conn: &mut PgConnection,
+    job_types: &[String],
+) -> QueryResult<BackgroundJob> {
     background_jobs::table
         .select(BackgroundJob::as_select())
+        .filter(background_jobs::job_type.eq_any(job_types))
         .filter(retriable())
         .order((background_jobs::priority.desc(), background_jobs::id))
         .for_update()

crates_io_worker/src/worker.rs

@@ -51,11 +51,13 @@ impl<Context: Clone + Send + 'static> Worker<Context> {
     /// - `Ok(None)` if no jobs were waiting
     /// - `Err(...)` if there was an error retrieving the job
     fn run_next_job(&self) -> anyhow::Result<Option<i64>> {
+        let job_types = &self.job_registry.job_types();
+
         let conn = &mut *self.connection_pool.get()?;
 
         conn.transaction(|conn| {
             debug!("Looking for next background worker job…");
-            let Some(job) = storage::find_next_unlocked_job(conn).optional()? else {
+            let Some(job) = storage::find_next_unlocked_job(conn, job_types).optional()? else {
                 return Ok(None);
             };
 

crates_io_worker/tests/runner.rs

@@ -221,6 +221,6 @@ fn runner<Context: Clone + Send + 'static>(
         .build_unchecked(ConnectionManager::new(database_url));
 
     Runner::new(&Handle::current(), connection_pool, context)
-        .num_workers(2)
+        .configure_default_queue(|queue| queue.num_workers(2))
         .shutdown_when_queue_empty()
 }

src/bin/background-worker.rs

@@ -104,7 +104,8 @@ fn main() -> anyhow::Result<()> {
     });
 
     let runner = Runner::new(runtime.handle(), connection_pool, environment.clone())
-        .num_workers(5)
+        .configure_default_queue(|queue| queue.num_workers(5))
+        .configure_queue("repository", |queue| queue.num_workers(1))
         .register_crates_io_job_types()
         .start();
 

src/tests/util/test_app.rs

@@ -272,7 +272,6 @@ impl TestAppBuilder {
                 (*app.primary_database).clone(),
                 Arc::new(environment),
             )
-            .num_workers(1)
             .shutdown_when_queue_empty()
             .register_crates_io_job_types();
 

src/worker/jobs/git.rs

@@ -29,6 +29,7 @@ impl SyncToGitIndex {
 impl BackgroundJob for SyncToGitIndex {
     const JOB_NAME: &'static str = "sync_to_git_index";
     const PRIORITY: i16 = 100;
+    const QUEUE: &'static str = "repository";
 
     type Context = Arc<Environment>;
 
@@ -165,6 +166,7 @@ pub struct SquashIndex;
 #[async_trait]
 impl BackgroundJob for SquashIndex {
     const JOB_NAME: &'static str = "squash_index";
+    const QUEUE: &'static str = "repository";
 
     type Context = Arc<Environment>;
 
@@ -223,6 +225,7 @@ impl NormalizeIndex {
 #[async_trait]
 impl BackgroundJob for NormalizeIndex {
     const JOB_NAME: &'static str = "normalize_index";
+    const QUEUE: &'static str = "repository";
 
     type Context = Arc<Environment>;