Merge pull request #9161 from Turbo87/diesel-async

Use `diesel-async` database pools

Author: Turbo87

Cargo.lock

@@ -73,6 +73,7 @@ derive_builder = "=0.20.0"
 derive_deref = "=1.1.1"
 dialoguer = "=0.11.0"
 diesel = { version = "=2.2.2", features = ["postgres", "serde_json", "chrono", "numeric"] }
+diesel-async = { version = "=0.5.0", features = ["async-connection-wrapper", "deadpool", "postgres"] }
 diesel_full_text_search = "=2.2.0"
 diesel_migrations = { version = "=2.2.0", features = ["postgres"] }
 dotenvy = "=0.15.7"

Cargo.toml

@@ -11,6 +11,7 @@ workspace = true
 anyhow = "=1.0.86"
 deadpool-diesel = { version = "=0.6.1", features = ["postgres", "tracing"] }
 diesel = { version = "=2.2.2", features = ["postgres", "serde_json"] }
+diesel-async = { version = "=0.5.0", features = ["async-connection-wrapper", "deadpool", "postgres"] }
 futures-util = "=0.3.30"
 sentry-core = { version = "=0.34.0", features = ["client"] }
 serde = { version = "=1.0.204", features = ["derive"] }
@@ -21,4 +22,4 @@ tracing = "=0.1.40"
 
 [dev-dependencies]
 crates_io_test_db = { path = "../crates_io_test_db" }
-tokio = { version = "=1.39.2", features = ["macros", "rt", "sync"]}
+tokio = { version = "=1.39.2", features = ["macros", "rt", "rt-multi-thread", "sync"]}

crates/crates_io_worker/Cargo.toml

@@ -1,7 +1,8 @@
 use crate::errors::EnqueueError;
 use crate::schema::background_jobs;
+use diesel::connection::LoadConnection;
+use diesel::pg::Pg;
 use diesel::prelude::*;
-use diesel::PgConnection;
 use serde::de::DeserializeOwned;
 use serde::Serialize;
 use std::future::Future;
@@ -29,14 +30,14 @@ 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 PgConnection) -> Result<i64, EnqueueError> {
+    fn enqueue(&self, conn: &mut impl LoadConnection<Backend = Pg>) -> Result<i64, EnqueueError> {
         self.enqueue_with_priority(conn, Self::PRIORITY)
     }
 
     #[instrument(name = "swirl.enqueue", skip(self, conn), fields(message = Self::JOB_NAME))]
     fn enqueue_with_priority(
         &self,
-        conn: &mut PgConnection,
+        conn: &mut impl LoadConnection<Backend = Pg>,
         job_priority: i16,
     ) -> Result<i64, EnqueueError> {
         let job_data = serde_json::to_value(self)?;

crates/crates_io_worker/src/background_job.rs

@@ -3,26 +3,28 @@ use crate::job_registry::JobRegistry;
 use crate::worker::Worker;
 use crate::{storage, BackgroundJob};
 use anyhow::anyhow;
-use deadpool_diesel::postgres::Pool;
+use diesel_async::async_connection_wrapper::AsyncConnectionWrapper;
+use diesel_async::pooled_connection::deadpool::Pool;
+use diesel_async::AsyncPgConnection;
 use futures_util::future::join_all;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::task::JoinHandle;
+use tokio::task::{spawn_blocking, JoinHandle};
 use tracing::{info, info_span, warn, Instrument};
 
 const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(1);
 
 /// The core runner responsible for locking and running jobs
 pub struct Runner<Context> {
-    connection_pool: Pool,
+    connection_pool: Pool<AsyncPgConnection>,
     queues: HashMap<String, Queue<Context>>,
     context: Context,
     shutdown_when_queue_empty: bool,
 }
 
 impl<Context: Clone + Send + Sync + 'static> Runner<Context> {
-    pub fn new(connection_pool: Pool, context: Context) -> Self {
+    pub fn new(connection_pool: Pool<AsyncPgConnection>, context: Context) -> Self {
         Self {
             connection_pool,
             queues: HashMap::new(),
@@ -96,7 +98,9 @@ impl<Context: Clone + Send + Sync + 'static> Runner<Context> {
     /// any jobs have failed.
     pub async fn check_for_failed_jobs(&self) -> anyhow::Result<()> {
         let conn = self.connection_pool.get().await?;
-        conn.interact(move |conn| {
+        spawn_blocking(move || {
+            let conn: &mut AsyncConnectionWrapper<_> = &mut conn.into();
+
             let failed_jobs = storage::failed_job_count(conn)?;
             if failed_jobs == 0 {
                 Ok(())

crates/crates_io_worker/src/runner.rs

@@ -1,4 +1,5 @@
 use crate::schema::background_jobs;
+use diesel::connection::LoadConnection;
 use diesel::dsl::now;
 use diesel::pg::Pg;
 use diesel::prelude::*;
@@ -26,7 +27,7 @@ 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,
+    conn: &mut impl LoadConnection<Backend = Pg>,
     job_types: &[String],
 ) -> QueryResult<BackgroundJob> {
     background_jobs::table
@@ -40,15 +41,18 @@ pub(super) fn find_next_unlocked_job(
 }
 
 /// The number of jobs that have failed at least once
-pub(super) fn failed_job_count(conn: &mut PgConnection) -> QueryResult<i64> {
+pub(super) fn failed_job_count(conn: &mut impl LoadConnection<Backend = Pg>) -> QueryResult<i64> {
     background_jobs::table
         .count()
         .filter(background_jobs::retries.gt(0))
         .get_result(conn)
 }
 
 /// Deletes a job that has successfully completed running
-pub(super) fn delete_successful_job(conn: &mut PgConnection, job_id: i64) -> QueryResult<()> {
+pub(super) fn delete_successful_job(
+    conn: &mut impl LoadConnection<Backend = Pg>,
+    job_id: i64,
+) -> QueryResult<()> {
     delete(background_jobs::table.find(job_id)).execute(conn)?;
     Ok(())
 }
@@ -57,7 +61,7 @@ pub(super) fn delete_successful_job(conn: &mut PgConnection, job_id: i64) -> Que
 ///
 /// Ignores any database errors that may have occurred. If the DB has gone away,
 /// we assume that just trying again with a new connection will succeed.
-pub(super) fn update_failed_job(conn: &mut PgConnection, job_id: i64) {
+pub(super) fn update_failed_job(conn: &mut impl LoadConnection<Backend = Pg>, job_id: i64) {
     let _ = update(background_jobs::table.find(job_id))
         .set((
             background_jobs::retries.eq(background_jobs::retries + 1),

crates/crates_io_worker/src/storage.rs

@@ -2,19 +2,22 @@ use crate::job_registry::JobRegistry;
 use crate::storage;
 use crate::util::{try_to_extract_panic_info, with_sentry_transaction};
 use anyhow::anyhow;
-use deadpool_diesel::postgres::Pool;
 use diesel::prelude::*;
+use diesel_async::async_connection_wrapper::AsyncConnectionWrapper;
+use diesel_async::pooled_connection::deadpool::Pool;
+use diesel_async::AsyncPgConnection;
 use futures_util::FutureExt;
 use sentry_core::{Hub, SentryFutureExt};
 use std::panic::AssertUnwindSafe;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::runtime::Handle;
+use tokio::task::spawn_blocking;
 use tokio::time::sleep;
 use tracing::{debug, error, info_span, warn};
 
 pub struct Worker<Context> {
-    pub(crate) connection_pool: Pool,
+    pub(crate) connection_pool: Pool<AsyncPgConnection>,
     pub(crate) context: Context,
     pub(crate) job_registry: Arc<JobRegistry<Context>>,
     pub(crate) shutdown_when_queue_empty: bool,
@@ -58,7 +61,9 @@ impl<Context: Clone + Send + Sync + 'static> Worker<Context> {
         let job_registry = self.job_registry.clone();
         let conn = self.connection_pool.get().await?;
 
-        conn.interact(move |conn| {
+        spawn_blocking(move || {
+            let conn: &mut AsyncConnectionWrapper<_> = &mut conn.into();
+
             let job_types = job_registry.job_types();
             conn.transaction(|conn| {
                 debug!("Looking for next background worker job…");