worker: Add DEDUPLICATED flag to the BackgroundJob trait

This flag can be set when implementing background jobs to automatically deduplicate jobs when they are enqueued. If an unstarted job already exists in the queue the `enqueue()` fn will return `Ok(None)` instead of the job ID.

Author: Turbo87

Cargo.lock

@@ -21,5 +21,7 @@ tokio = { version = "=1.40.0", features = ["rt", "time"]}
 tracing = "=0.1.40"
 
 [dev-dependencies]
+claims = "=0.7.1"
 crates_io_test_db = { path = "../crates_io_test_db" }
+insta = { version = "=1.40.0", features = ["json"] }
 tokio = { version = "=1.40.0", features = ["macros", "rt", "rt-multi-thread", "sync"]}

crates/crates_io_worker/Cargo.toml

@@ -1,8 +1,10 @@
 use crate::errors::EnqueueError;
 use crate::schema::background_jobs;
 use diesel::connection::LoadConnection;
+use diesel::dsl::{exists, not};
 use diesel::pg::Pg;
 use diesel::prelude::*;
+use diesel::sql_types::{Int2, Jsonb, Text};
 use serde::de::DeserializeOwned;
 use serde::Serialize;
 use std::future::Future;
@@ -21,6 +23,12 @@ pub trait BackgroundJob: Serialize + DeserializeOwned + Send + Sync + 'static {
     /// [Self::enqueue_with_priority] can be used to override the priority value.
     const PRIORITY: i16 = 0;
 
+    /// Whether the job should be deduplicated.
+    ///
+    /// If true, the job will not be enqueued if there is already an unstarted
+    /// job with the same data.
+    const DEDUPLICATED: bool = false;
+
     /// Job queue where this job will be executed.
     const QUEUE: &'static str = DEFAULT_QUEUE;
 
@@ -30,7 +38,10 @@ pub trait BackgroundJob: Serialize + DeserializeOwned + Send + Sync + 'static {
     /// Execute the task. This method should define its logic.
     fn run(&self, ctx: Self::Context) -> impl Future<Output = anyhow::Result<()>> + Send;
 
-    fn enqueue(&self, conn: &mut impl LoadConnection<Backend = Pg>) -> Result<i64, EnqueueError> {
+    fn enqueue(
+        &self,
+        conn: &mut impl LoadConnection<Backend = Pg>,
+    ) -> Result<Option<i64>, EnqueueError> {
         self.enqueue_with_priority(conn, Self::PRIORITY)
     }
 
@@ -39,16 +50,48 @@ pub trait BackgroundJob: Serialize + DeserializeOwned + Send + Sync + 'static {
         &self,
         conn: &mut impl LoadConnection<Backend = Pg>,
         job_priority: i16,
-    ) -> Result<i64, EnqueueError> {
+    ) -> Result<Option<i64>, EnqueueError> {
         let job_data = serde_json::to_value(self)?;
-        let id = diesel::insert_into(background_jobs::table)
-            .values((
-                background_jobs::job_type.eq(Self::JOB_NAME),
-                background_jobs::data.eq(job_data),
-                background_jobs::priority.eq(job_priority),
+
+        if Self::DEDUPLICATED {
+            let similar_jobs = background_jobs::table
+                .select(background_jobs::id)
+                .filter(background_jobs::job_type.eq(Self::JOB_NAME))
+                .filter(background_jobs::data.eq(&job_data))
+                .filter(background_jobs::priority.eq(job_priority))
+                .for_update()
+                .skip_locked();
+
+            let deduplicated_select = diesel::select((
+                Self::JOB_NAME.into_sql::<Text>(),
+                (&job_data).into_sql::<Jsonb>(),
+                job_priority.into_sql::<Int2>(),
             ))
-            .returning(background_jobs::id)
-            .get_result(conn)?;
-        Ok(id)
+            .filter(not(exists(similar_jobs)));
+
+            let id = diesel::insert_into(background_jobs::table)
+                .values(deduplicated_select)
+                .into_columns((
+                    background_jobs::job_type,
+                    background_jobs::data,
+                    background_jobs::priority,
+                ))
+                .returning(background_jobs::id)
+                .get_result::<i64>(conn)
+                .optional()?;
+
+            Ok(id)
+        } else {
+            let id = diesel::insert_into(background_jobs::table)
+                .values((
+                    background_jobs::job_type.eq(Self::JOB_NAME),
+                    background_jobs::data.eq(job_data),
+                    background_jobs::priority.eq(job_priority),
+                ))
+                .returning(background_jobs::id)
+                .get_result(conn)?;
+
+            Ok(Some(id))
+        }
     }
 }

crates/crates_io_worker/src/background_job.rs

@@ -1,14 +1,25 @@
+use claims::{assert_none, assert_some};
 use crates_io_test_db::TestDatabase;
 use crates_io_worker::schema::background_jobs;
 use crates_io_worker::{BackgroundJob, Runner};
 use diesel::prelude::*;
 use diesel_async::pooled_connection::deadpool::Pool;
 use diesel_async::pooled_connection::AsyncDieselConnectionManager;
 use diesel_async::AsyncPgConnection;
+use insta::assert_compact_json_snapshot;
 use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::sync::atomic::{AtomicU8, Ordering};
 use std::sync::Arc;
 use tokio::sync::Barrier;
 
+fn all_jobs(conn: &mut PgConnection) -> Vec<(String, Value)> {
+    background_jobs::table
+        .select((background_jobs::job_type, background_jobs::data))
+        .get_results(conn)
+        .unwrap()
+}
+
 fn job_exists(id: i64, conn: &mut PgConnection) -> bool {
     background_jobs::table
         .find(id)
@@ -63,7 +74,7 @@ async fn jobs_are_locked_when_fetched() {
     let runner = runner(test_database.url(), test_context.clone()).register_job_type::<TestJob>();
 
     let mut conn = test_database.connect();
-    let job_id = TestJob.enqueue(&mut conn).unwrap();
+    let job_id = TestJob.enqueue(&mut conn).unwrap().unwrap();
 
     assert!(job_exists(job_id, &mut conn));
     assert!(!job_is_locked(job_id, &mut conn));
@@ -193,7 +204,7 @@ async fn panicking_in_jobs_updates_retry_counter() {
 
     let mut conn = test_database.connect();
 
-    let job_id = TestJob.enqueue(&mut conn).unwrap();
+    let job_id = TestJob.enqueue(&mut conn).unwrap().unwrap();
 
     let runner = runner.start();
     runner.wait_for_shutdown().await;
@@ -207,6 +218,85 @@ async fn panicking_in_jobs_updates_retry_counter() {
     assert_eq!(tries, 1);
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn jobs_can_be_deduplicated() {
+    #[derive(Clone)]
+    struct TestContext {
+        runs: Arc<AtomicU8>,
+        job_started_barrier: Arc<Barrier>,
+        assertions_finished_barrier: Arc<Barrier>,
+    }
+
+    #[derive(Serialize, Deserialize)]
+    struct TestJob {
+        value: String,
+    }
+
+    impl TestJob {
+        fn new(value: impl Into<String>) -> Self {
+            let value = value.into();
+            Self { value }
+        }
+    }
+
+    impl BackgroundJob for TestJob {
+        const JOB_NAME: &'static str = "test";
+        const DEDUPLICATED: bool = true;
+        type Context = TestContext;
+
+        async fn run(&self, ctx: Self::Context) -> anyhow::Result<()> {
+            let runs = ctx.runs.fetch_add(1, Ordering::SeqCst);
+            if runs == 0 {
+                ctx.job_started_barrier.wait().await;
+                ctx.assertions_finished_barrier.wait().await;
+            }
+            Ok(())
+        }
+    }
+
+    let test_database = TestDatabase::new();
+
+    let test_context = TestContext {
+        runs: Arc::new(AtomicU8::new(0)),
+        job_started_barrier: Arc::new(Barrier::new(2)),
+        assertions_finished_barrier: Arc::new(Barrier::new(2)),
+    };
+
+    let runner = runner(test_database.url(), test_context.clone()).register_job_type::<TestJob>();
+
+    let mut conn = test_database.connect();
+
+    // Enqueue first job
+    assert_some!(TestJob::new("foo").enqueue(&mut conn).unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}]]"#);
+
+    // Try to enqueue the same job again, which should be deduplicated
+    assert_none!(TestJob::new("foo").enqueue(&mut conn).unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}]]"#);
+
+    // Start processing the first job
+    let runner = runner.start();
+    test_context.job_started_barrier.wait().await;
+
+    // Enqueue the same job again, which should NOT be deduplicated,
+    // since the first job already still running
+    assert_some!(TestJob::new("foo").enqueue(&mut conn).unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}], ["test", {"value": "foo"}]]"#);
+
+    // Try to enqueue the same job again, which should be deduplicated again
+    assert_none!(TestJob::new("foo").enqueue(&mut conn).unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}], ["test", {"value": "foo"}]]"#);
+
+    // Enqueue the same job but with different data, which should
+    // NOT be deduplicated
+    assert_some!(TestJob::new("bar").enqueue(&mut conn).unwrap());
+    assert_compact_json_snapshot!(all_jobs(&mut conn), @r#"[["test", {"value": "foo"}], ["test", {"value": "foo"}], ["test", {"value": "bar"}]]"#);
+
+    // Resolve the final barrier to finish the test
+    test_context.assertions_finished_barrier.wait().await;
+    runner.wait_for_shutdown().await;
+}
+
 fn runner<Context: Clone + Send + Sync + 'static>(
     database_url: &str,
     context: Context,