Merge pull request #11167 from Turbo87/create-database-pool

Use same database pool creation code for API server and background worker

Author: Turbo87

src/app.rs

@@ -94,51 +94,8 @@ impl<S: app_builder::State> AppBuilder<S> {
         S::PrimaryDatabase: app_builder::IsUnset,
         S::ReplicaDatabase: app_builder::IsUnset,
     {
-        let primary_database = {
-            use secrecy::ExposeSecret;
-
-            let primary_db_connection_config = ConnectionConfig {
-                statement_timeout: config.statement_timeout,
-                read_only: config.primary.read_only_mode,
-            };
-
-            let url = connection_url(config, config.primary.url.expose_secret());
-            let manager_config = make_manager_config(config.enforce_tls);
-            let manager = AsyncDieselConnectionManager::new_with_config(url, manager_config);
-
-            DeadpoolPool::builder(manager)
-                .runtime(Runtime::Tokio1)
-                .max_size(config.primary.pool_size)
-                .wait_timeout(Some(config.connection_timeout))
-                .post_create(primary_db_connection_config)
-                .build()
-                .unwrap()
-        };
-
-        let replica_database = if let Some(pool_config) = config.replica.as_ref() {
-            use secrecy::ExposeSecret;
-
-            let replica_db_connection_config = ConnectionConfig {
-                statement_timeout: config.statement_timeout,
-                read_only: pool_config.read_only_mode,
-            };
-
-            let url = connection_url(config, pool_config.url.expose_secret());
-            let manager_config = make_manager_config(config.enforce_tls);
-            let manager = AsyncDieselConnectionManager::new_with_config(url, manager_config);
-
-            let pool = DeadpoolPool::builder(manager)
-                .runtime(Runtime::Tokio1)
-                .max_size(pool_config.pool_size)
-                .wait_timeout(Some(config.connection_timeout))
-                .post_create(replica_db_connection_config)
-                .build()
-                .unwrap();
-
-            Some(pool)
-        } else {
-            None
-        };
+        let primary_database = create_database_pool(&config.primary);
+        let replica_database = config.replica.as_ref().map(create_database_pool);
 
         self.primary_database(primary_database)
             .maybe_replica_database(replica_database)
@@ -165,6 +122,25 @@ impl<S: app_builder::State> AppBuilder<S> {
     }
 }
 
+pub fn create_database_pool(config: &config::DbPoolConfig) -> DeadpoolPool<AsyncPgConnection> {
+    let connection_config = ConnectionConfig {
+        statement_timeout: config.statement_timeout,
+        read_only: config.read_only_mode,
+    };
+
+    let url = connection_url(config);
+    let manager_config = make_manager_config(config.enforce_tls);
+    let manager = AsyncDieselConnectionManager::new_with_config(url, manager_config);
+
+    DeadpoolPool::builder(manager)
+        .runtime(Runtime::Tokio1)
+        .max_size(config.pool_size)
+        .wait_timeout(Some(config.connection_timeout))
+        .post_create(connection_config)
+        .build()
+        .unwrap()
+}
+
 impl App {
     /// A unique key to generate signed cookies
     pub fn session_key(&self) -> &cookie::Key {

src/bin/background-worker.rs

@@ -14,22 +14,19 @@
 extern crate tracing;
 
 use anyhow::Context;
+use crates_io::app::create_database_pool;
 use crates_io::cloudfront::CloudFront;
-use crates_io::db::make_manager_config;
 use crates_io::fastly::Fastly;
+use crates_io::ssh;
 use crates_io::storage::Storage;
 use crates_io::worker::{Environment, RunnerExt};
 use crates_io::{Emails, config};
-use crates_io::{db, ssh};
 use crates_io_env_vars::var;
 use crates_io_index::RepositoryConfig;
 use crates_io_team_repo::TeamRepoImpl;
 use crates_io_worker::Runner;
-use diesel_async::pooled_connection::AsyncDieselConnectionManager;
-use diesel_async::pooled_connection::deadpool::Pool;
 use object_store::prefix::PrefixStore;
 use reqwest::Client;
-use secrecy::ExposeSecret;
 use std::sync::Arc;
 use std::thread::sleep;
 use std::time::Duration;
@@ -44,7 +41,10 @@ fn main() -> anyhow::Result<()> {
 
     info!("Booting runner");
 
-    let config = config::Server::from_environment()?;
+    let mut config = config::Server::from_environment()?;
+
+    // Override the pool size to 10 for the background worker
+    config.db.primary.pool_size = 10;
 
     let runtime = tokio::runtime::Builder::new_multi_thread()
         .enable_all()
@@ -61,8 +61,6 @@ fn main() -> anyhow::Result<()> {
         }
     }
 
-    let db_url = db::connection_url(&config.db, config.db.primary.url.expose_secret());
-
     if var("HEROKU")?.is_some() {
         ssh::write_known_hosts_file()?;
     }
@@ -84,9 +82,7 @@ fn main() -> anyhow::Result<()> {
     let fastly = Fastly::from_environment(client.clone());
     let team_repo = TeamRepoImpl::default();
 
-    let manager_config = make_manager_config(config.db.enforce_tls);
-    let manager = AsyncDieselConnectionManager::new_with_config(db_url, manager_config);
-    let deadpool = Pool::builder(manager).max_size(10).build()?;
+    let deadpool = create_database_pool(&config.db.primary);
 
     let environment = Environment::builder()
         .config(Arc::new(config))

src/config/database_pools.rs

@@ -24,6 +24,14 @@ pub struct DatabasePools {
     pub primary: DbPoolConfig,
     /// An optional follower database. Always read-only.
     pub replica: Option<DbPoolConfig>,
+}
+
+#[derive(Debug)]
+pub struct DbPoolConfig {
+    pub url: SecretString,
+    pub read_only_mode: bool,
+    pub pool_size: usize,
+    pub min_idle: Option<u32>,
     /// Number of seconds to wait for unacknowledged TCP packets before treating the connection as
     /// broken. This value will determine how long crates.io stays unavailable in case of full
     /// packet loss between the application and the database: setting it too high will result in an
@@ -43,14 +51,6 @@ pub struct DatabasePools {
     pub enforce_tls: bool,
 }
 
-#[derive(Debug)]
-pub struct DbPoolConfig {
-    pub url: SecretString,
-    pub read_only_mode: bool,
-    pub pool_size: usize,
-    pub min_idle: Option<u32>,
-}
-
 impl DatabasePools {
     pub fn are_all_read_only(&self) -> bool {
         self.primary.read_only_mode
@@ -102,13 +102,13 @@ impl DatabasePools {
                     read_only_mode: true,
                     pool_size: primary_async_pool_size,
                     min_idle: primary_min_idle,
+                    tcp_timeout_ms,
+                    connection_timeout,
+                    statement_timeout,
+                    helper_threads,
+                    enforce_tls,
                 },
                 replica: None,
-                tcp_timeout_ms,
-                connection_timeout,
-                statement_timeout,
-                helper_threads,
-                enforce_tls,
             },
             // The follower is down, don't configure the replica.
             Some("follower") => Self {
@@ -117,20 +117,25 @@ impl DatabasePools {
                     read_only_mode,
                     pool_size: primary_async_pool_size,
                     min_idle: primary_min_idle,
+                    tcp_timeout_ms,
+                    connection_timeout,
+                    statement_timeout,
+                    helper_threads,
+                    enforce_tls,
                 },
                 replica: None,
-                tcp_timeout_ms,
-                connection_timeout,
-                statement_timeout,
-                helper_threads,
-                enforce_tls,
             },
             _ => Self {
                 primary: DbPoolConfig {
                     url: leader_url,
                     read_only_mode,
                     pool_size: primary_async_pool_size,
                     min_idle: primary_min_idle,
+                    tcp_timeout_ms,
+                    connection_timeout,
+                    statement_timeout,
+                    helper_threads,
+                    enforce_tls,
                 },
                 replica: follower_url.map(|url| DbPoolConfig {
                     url,
@@ -140,12 +145,12 @@ impl DatabasePools {
                     read_only_mode: true,
                     pool_size: replica_async_pool_size,
                     min_idle: replica_min_idle,
+                    tcp_timeout_ms,
+                    connection_timeout,
+                    statement_timeout,
+                    helper_threads,
+                    enforce_tls,
                 }),
-                tcp_timeout_ms,
-                connection_timeout,
-                statement_timeout,
-                helper_threads,
-                enforce_tls,
             },
         })
     }

src/db.rs

@@ -14,17 +14,17 @@ use crate::config;
 pub async fn oneoff_connection_with_config(
     config: &config::DatabasePools,
 ) -> ConnectionResult<AsyncPgConnection> {
-    let url = connection_url(config, config.primary.url.expose_secret());
-    establish_async_connection(&url, config.enforce_tls).await
+    let url = connection_url(&config.primary);
+    establish_async_connection(&url, config.primary.enforce_tls).await
 }
 
 pub async fn oneoff_connection() -> anyhow::Result<AsyncPgConnection> {
     let config = config::DatabasePools::full_from_environment(&config::Base::from_environment()?)?;
     Ok(oneoff_connection_with_config(&config).await?)
 }
 
-pub fn connection_url(config: &config::DatabasePools, url: &str) -> String {
-    let mut url = Url::parse(url).expect("Invalid database URL");
+pub fn connection_url(config: &config::DbPoolConfig) -> String {
+    let mut url = Url::parse(config.url.expose_secret()).expect("Invalid database URL");
 
     if config.enforce_tls {
         maybe_append_url_param(&mut url, "sslmode", "require");

src/lib.rs

@@ -25,7 +25,7 @@ use tikv_jemallocator::Jemalloc;
 #[global_allocator]
 static ALLOC: Jemalloc = Jemalloc;
 
-mod app;
+pub mod app;
 pub mod auth;
 pub mod boot;
 pub mod certs;

src/tests/util/test_app.rs

@@ -401,6 +401,11 @@ impl TestAppBuilder {
             read_only_mode: true,
             pool_size: primary.pool_size,
             min_idle: primary.min_idle,
+            tcp_timeout_ms: primary.tcp_timeout_ms,
+            connection_timeout: primary.connection_timeout,
+            statement_timeout: primary.statement_timeout,
+            helper_threads: primary.helper_threads,
+            enforce_tls: primary.enforce_tls,
         });
 
         self
@@ -419,13 +424,13 @@ fn simple_config() -> config::Server {
             read_only_mode: false,
             pool_size: 5,
             min_idle: None,
+            tcp_timeout_ms: 1000, // 1 second
+            connection_timeout: Duration::from_secs(1),
+            statement_timeout: Duration::from_secs(1),
+            helper_threads: 1,
+            enforce_tls: false,
         },
         replica: None,
-        tcp_timeout_ms: 1000, // 1 second
-        connection_timeout: Duration::from_secs(1),
-        statement_timeout: Duration::from_secs(1),
-        helper_threads: 1,
-        enforce_tls: false,
     };
 
     let mut storage = StorageConfig::in_memory();