worker: Implement async enqueue() fn

Author: Turbo87

crates/crates_io_worker/src/background_job.rs

@@ -3,8 +3,9 @@ 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 diesel::{ExpressionMethods, IntoSql, OptionalExtension, QueryDsl};
+use diesel_async::AsyncPgConnection;
 use serde::de::DeserializeOwned;
 use serde::Serialize;
 use serde_json::Value;
@@ -53,6 +54,24 @@ pub trait BackgroundJob: Serialize + DeserializeOwned + Send + Sync + 'static {
             Ok(Some(enqueue_simple(conn, Self::JOB_NAME, &data, priority)?))
         }
     }
+
+    #[allow(async_fn_in_trait)]
+    #[instrument(name = "swirl.enqueue", skip(self, conn), fields(message = Self::JOB_NAME))]
+    async fn async_enqueue(
+        &self,
+        conn: &mut AsyncPgConnection,
+    ) -> Result<Option<i64>, EnqueueError> {
+        let data = serde_json::to_value(self)?;
+        let priority = Self::PRIORITY;
+
+        if Self::DEDUPLICATED {
+            Ok(async_enqueue_deduplicated(conn, Self::JOB_NAME, &data, priority).await?)
+        } else {
+            Ok(Some(
+                async_enqueue_simple(conn, Self::JOB_NAME, &data, priority).await?,
+            ))
+        }
+    }
 }
 
 fn enqueue_deduplicated(
@@ -61,6 +80,8 @@ fn enqueue_deduplicated(
     data: &Value,
     priority: i16,
 ) -> Result<Option<i64>, EnqueueError> {
+    use diesel::RunQueryDsl;
+
     let similar_jobs = background_jobs::table
         .select(background_jobs::id)
         .filter(background_jobs::job_type.eq(job_type))
@@ -90,12 +111,52 @@ fn enqueue_deduplicated(
     Ok(id)
 }
 
+async fn async_enqueue_deduplicated(
+    conn: &mut AsyncPgConnection,
+    job_type: &str,
+    data: &Value,
+    priority: i16,
+) -> Result<Option<i64>, EnqueueError> {
+    use diesel_async::RunQueryDsl;
+
+    let similar_jobs = background_jobs::table
+        .select(background_jobs::id)
+        .filter(background_jobs::job_type.eq(job_type))
+        .filter(background_jobs::data.eq(data))
+        .filter(background_jobs::priority.eq(priority))
+        .for_update()
+        .skip_locked();
+
+    let deduplicated_select = diesel::select((
+        job_type.into_sql::<Text>(),
+        data.into_sql::<Jsonb>(),
+        priority.into_sql::<Int2>(),
+    ))
+    .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)
+        .await
+        .optional()?;
+
+    Ok(id)
+}
+
 fn enqueue_simple(
     conn: &mut impl LoadConnection<Backend = Pg>,
     job_type: &str,
     data: &Value,
     priority: i16,
 ) -> Result<i64, EnqueueError> {
+    use diesel::RunQueryDsl;
+
     let id = diesel::insert_into(background_jobs::table)
         .values((
             background_jobs::job_type.eq(job_type),
@@ -107,3 +168,24 @@ fn enqueue_simple(
 
     Ok(id)
 }
+
+async fn async_enqueue_simple(
+    conn: &mut AsyncPgConnection,
+    job_type: &str,
+    data: &Value,
+    priority: i16,
+) -> Result<i64, EnqueueError> {
+    use diesel_async::RunQueryDsl;
+
+    let id = diesel::insert_into(background_jobs::table)
+        .values((
+            background_jobs::job_type.eq(job_type),
+            background_jobs::data.eq(data),
+            background_jobs::priority.eq(priority),
+        ))
+        .returning(background_jobs::id)
+        .get_result(conn)
+        .await?;
+
+    Ok(id)
+}