Merge pull request #9632 from Turbo87/async-enqueue

Turbo87 · web-flow · commit c553294a700b · 2024-10-13T11:35:28.000+02:00
worker: Implement async `enqueue()` fn
diff --git a/crates/crates_io_worker/src/background_job.rs b/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)
+}
diff --git a/src/admin/delete_crate.rs b/src/admin/delete_crate.rs
@@ -1,13 +1,12 @@
 use crate::schema::{crate_owners, teams, users};
-use crate::tasks::spawn_blocking;
 use crate::worker::jobs;
 use crate::{admin::dialoguer, db, schema::crates};
 use anyhow::Context;
 use crates_io_worker::BackgroundJob;
 use diesel::dsl::sql;
 use diesel::sql_types::Text;
 use diesel::{ExpressionMethods, JoinOnDsl, QueryDsl};
-use diesel_async::async_connection_wrapper::AsyncConnectionWrapper;
+use diesel_async::RunQueryDsl;
 use std::collections::HashMap;
 
 #[derive(clap::Parser, Debug)]
@@ -34,24 +33,21 @@ pub async fn run(opts: Opts) -> anyhow::Result<()> {
     let mut crate_names = opts.crate_names;
     crate_names.sort();
 
-    let query_result = {
-        use diesel_async::RunQueryDsl;
-
-        crates::table
-            .select((
-                crates::name,
-                crates::id,
-                sql::<Text>(
-                    "CASE WHEN crate_owners.owner_kind = 1 THEN teams.login ELSE users.gh_login END",
-                ),
-            ))
-            .left_join(crate_owners::table.on(crate_owners::crate_id.eq(crates::id)))
-            .left_join(teams::table.on(teams::id.eq(crate_owners::owner_id)))
-            .left_join(users::table.on(users::id.eq(crate_owners::owner_id)))
-            .filter(crates::name.eq_any(&crate_names))
-            .load::<(String, i32, String)>(&mut conn).await
-            .context("Failed to look up crate name from the database")
-    }?;
+    let query_result = crates::table
+        .select((
+            crates::name,
+            crates::id,
+            sql::<Text>(
+                "CASE WHEN crate_owners.owner_kind = 1 THEN teams.login ELSE users.gh_login END",
+            ),
+        ))
+        .left_join(crate_owners::table.on(crate_owners::crate_id.eq(crates::id)))
+        .left_join(teams::table.on(teams::id.eq(crate_owners::owner_id)))
+        .left_join(users::table.on(users::id.eq(crate_owners::owner_id)))
+        .filter(crates::name.eq_any(&crate_names))
+        .load::<(String, i32, String)>(&mut conn)
+        .await
+        .context("Failed to look up crate name from the database")?;
 
     let mut existing_crates: HashMap<String, (i32, Vec<String>)> = HashMap::new();
     for (name, id, login) in query_result {
@@ -81,37 +77,36 @@ pub async fn run(opts: Opts) -> anyhow::Result<()> {
         return Ok(());
     }
 
-    spawn_blocking(move || {
-        use diesel::RunQueryDsl;
-
-        let conn: &mut AsyncConnectionWrapper<_> = &mut conn.into();
-
-        for name in &crate_names {
-            if let Some((id, _)) = existing_crates.get(name) {
-                info!("{name}: Deleting crate from the database…");
-                if let Err(error) = diesel::delete(crates::table.find(id)).execute(conn) {
-                    warn!(%id, "{name}: Failed to delete crate from the database: {error}");
-                }
-            } else {
-                info!("{name}: Skipped missing crate");
-            };
-
-            info!("{name}: Enqueuing index sync jobs…");
-            if let Err(error) = jobs::SyncToGitIndex::new(name).enqueue(conn) {
-                warn!("{name}: Failed to enqueue SyncToGitIndex job: {error}");
-            }
-            if let Err(error) = jobs::SyncToSparseIndex::new(name).enqueue(conn) {
-                warn!("{name}: Failed to enqueue SyncToSparseIndex job: {error}");
+    for name in &crate_names {
+        if let Some((id, _)) = existing_crates.get(name) {
+            info!("{name}: Deleting crate from the database…");
+            if let Err(error) = diesel::delete(crates::table.find(id))
+                .execute(&mut conn)
+                .await
+            {
+                warn!(%id, "{name}: Failed to delete crate from the database: {error}");
             }
+        } else {
+            info!("{name}: Skipped missing crate");
+        };
+
+        info!("{name}: Enqueuing index sync jobs…");
+        let job = jobs::SyncToGitIndex::new(name);
+        if let Err(error) = job.async_enqueue(&mut conn).await {
+            warn!("{name}: Failed to enqueue SyncToGitIndex job: {error}");
+        }
 
-            info!("{name}: Enqueuing DeleteCrateFromStorage job…");
-            let job = jobs::DeleteCrateFromStorage::new(name.into());
-            if let Err(error) = job.enqueue(conn) {
-                warn!("{name}: Failed to enqueue DeleteCrateFromStorage job: {error}");
-            }
+        let job = jobs::SyncToSparseIndex::new(name);
+        if let Err(error) = job.async_enqueue(&mut conn).await {
+            warn!("{name}: Failed to enqueue SyncToSparseIndex job: {error}");
         }
 
-        Ok(())
-    })
-    .await
+        info!("{name}: Enqueuing DeleteCrateFromStorage job…");
+        let job = jobs::DeleteCrateFromStorage::new(name.into());
+        if let Err(error) = job.async_enqueue(&mut conn).await {
+            warn!("{name}: Failed to enqueue DeleteCrateFromStorage job: {error}");
+        }
+    }
+
+    Ok(())
 }
diff --git a/src/admin/enqueue_job.rs b/src/admin/enqueue_job.rs
