Setup read-only replica pool if READ_ONLY_REPLICA_URL is set

Author: jtgeibel

src/app.rs

@@ -13,8 +13,11 @@ use scheduled_thread_pool::ScheduledThreadPool;
 // The db, oauth, and git2 types don't implement debug.
 #[allow(missing_debug_implementations)]
 pub struct App {
-    /// The database connection pool
-    pub diesel_database: db::DieselPool,
+    /// The primary database connection pool
+    pub primary_database: db::DieselPool,
+
+    /// The read-only replica database connection pool
+    pub read_only_replica_database: Option<db::DieselPool>,
 
     /// The GitHub OAuth2 configuration
     pub github: BasicClient,
@@ -78,29 +81,53 @@ impl App {
             _ => 1,
         };
 
+        // Used as the connection and statement timeout value for the database pool(s)
         let db_connection_timeout = match (dotenv::var("DB_TIMEOUT"), config.env) {
             (Ok(num), _) => num.parse().expect("couldn't parse DB_TIMEOUT"),
             (_, Env::Production) => 10,
             (_, Env::Test) => 1,
             _ => 30,
         };
+
+        // Determine if the primary pool is also read-only
         let read_only_mode = dotenv::var("READ_ONLY_MODE").is_ok();
-        let connection_config = db::ConnectionConfig {
+        let primary_db_connection_config = db::ConnectionConfig {
             statement_timeout: db_connection_timeout,
             read_only: read_only_mode,
         };
 
         let thread_pool = Arc::new(ScheduledThreadPool::new(db_helper_threads));
 
-        let diesel_db_config = r2d2::Pool::builder()
+        let primary_db_config = r2d2::Pool::builder()
             .max_size(db_pool_size)
             .min_idle(db_min_idle)
             .connection_timeout(Duration::from_secs(db_connection_timeout))
-            .connection_customizer(Box::new(connection_config))
-            .thread_pool(thread_pool);
+            .connection_customizer(Box::new(primary_db_connection_config))
+            .thread_pool(thread_pool.clone());
+
+        let primary_database = db::diesel_pool(&config.db_url, config.env, primary_db_config);
+
+        let read_only_replica_database = if let Some(url) = &config.replica_db_url {
+            let replica_db_connection_config = db::ConnectionConfig {
+                statement_timeout: db_connection_timeout,
+                read_only: true,
+            };
+
+            let replica_db_config = r2d2::Pool::builder()
+                .max_size(db_pool_size)
+                .min_idle(db_min_idle)
+                .connection_timeout(Duration::from_secs(db_connection_timeout))
+                .connection_customizer(Box::new(replica_db_connection_config))
+                .thread_pool(thread_pool);
+
+            Some(db::diesel_pool(&url, config.env, replica_db_config))
+        } else {
+            None
+        };
 
         App {
-            diesel_database: db::diesel_pool(&config.db_url, config.env, diesel_db_config),
+            primary_database,
+            read_only_replica_database,
             github,
             session_key: config.session_key.clone(),
             git_repo_checkout: config.git_repo_checkout.clone(),

src/config.rs

@@ -10,6 +10,7 @@ pub struct Config {
     pub gh_client_id: String,
     pub gh_client_secret: String,
     pub db_url: String,
+    pub replica_db_url: Option<String>,
     pub env: Env,
     pub max_upload_size: u64,
     pub max_unpack_size: u64,
@@ -33,14 +34,15 @@ impl Default for Config {
     /// - `MIRROR`: Is this instance of cargo_registry a mirror of crates.io.
     /// - `HEROKU`: Is this instance of cargo_registry currently running on Heroku.
     /// - `S3_BUCKET`: The S3 bucket used to store crate files. If not present during development,
-    /// cargo_registry will fall back to a local uploader.
+    ///    cargo_registry will fall back to a local uploader.
     /// - `S3_REGION`: The region in which the bucket was created. Optional if US standard.
     /// - `S3_ACCESS_KEY`: The access key to interact with S3. Optional if running a mirror.
     /// - `S3_SECRET_KEY`: The secret key to interact with S3. Optional if running a mirror.
     /// - `SESSION_KEY`: The key used to sign and encrypt session cookies.
     /// - `GH_CLIENT_ID`: The client ID of the associated GitHub application.
     /// - `GH_CLIENT_SECRET`: The client secret of the associated GitHub application.
     /// - `DATABASE_URL`: The URL of the postgres database to use.
+    /// - `READ_ONLY_REPLICA_URL`: The URL of an optional postgres read-only replica database.
     /// - `BLOCKED_TRAFFIC`: A list of headers and environment variables to use for blocking
     ///.  traffic. See the `block_traffic` module for more documentation.
     fn default() -> Config {
@@ -129,6 +131,7 @@ impl Default for Config {
             gh_client_id: env("GH_CLIENT_ID"),
             gh_client_secret: env("GH_CLIENT_SECRET"),
             db_url: env("DATABASE_URL"),
+            replica_db_url: dotenv::var("READ_ONLY_REPLICA_URL").ok(),
             env: cargo_env,
             max_upload_size: 10 * 1024 * 1024, // 10 MB default file upload size limit
             max_unpack_size: 512 * 1024 * 1024, // 512 MB max when decompressed

src/controllers/krate/publish.rs

@@ -42,7 +42,7 @@ pub fn publish(req: &mut dyn Request) -> AppResult<Response> {
 
     let (new_crate, user) = parse_new_headers(req)?;
 
-    let conn = app.diesel_database.get()?;
+    let conn = app.primary_database.get()?;
 
     let verified_email_address = user.verified_email(&conn)?;
     let verified_email_address = verified_email_address.ok_or_else(|| {

src/db.rs

@@ -85,16 +85,25 @@ pub fn diesel_pool(
 }
 
 pub trait RequestTransaction {
-    /// Return the lazily initialized postgres connection for this request.
-    ///
-    /// The connection will live for the lifetime of the request.
-    // FIXME: This description does not match the implementation below.
+    /// Obtain a read/write database connection from the primary pool
     fn db_conn(&self) -> AppResult<DieselPooledConn<'_>>;
+
+    /// Obtain a readonly database connection from the replica pool
+    ///
+    /// If there is no replica pool, the primary pool is used instead.
+    fn db_read_only(&self) -> AppResult<DieselPooledConn<'_>>;
 }
 
 impl<T: Request + ?Sized> RequestTransaction for T {
     fn db_conn(&self) -> AppResult<DieselPooledConn<'_>> {
-        self.app().diesel_database.get().map_err(Into::into)
+        self.app().primary_database.get().map_err(Into::into)
+    }
+
+    fn db_read_only(&self) -> AppResult<DieselPooledConn<'_>> {
+        match &self.app().read_only_replica_database {
+            Some(pool) => pool.get().map_err(Into::into),
+            None => self.app().primary_database.get().map_err(Into::into),
+        }
     }
 }
 

src/middleware/log_connection_pool_status.rs

@@ -37,8 +37,12 @@ impl Middleware for LogConnectionPoolStatus {
         if last_log_time.elapsed() >= Duration::from_secs(60) {
             *last_log_time = Instant::now();
             println!(
-                "connection_pool_status=\"{:?}\" in_flight_requests={}",
-                self.app.diesel_database.state(),
+                "primary_pool_status=\"{:?}\" read_only_pool_status=\"{:?}\" in_flight_requests={}",
+                self.app.primary_database.state(),
+                self.app
+                    .read_only_replica_database
+                    .as_ref()
+                    .map(|pool| pool.state()),
                 in_flight_requests
             );
         }

src/tests/all.rs

@@ -133,6 +133,7 @@ fn simple_config() -> Config {
         gh_client_id: dotenv::var("GH_CLIENT_ID").unwrap_or_default(),
         gh_client_secret: dotenv::var("GH_CLIENT_SECRET").unwrap_or_default(),
         db_url: env("TEST_DATABASE_URL"),
+        replica_db_url: None,
         env: Env::Test,
         max_upload_size: 3000,
         max_unpack_size: 2000,
@@ -159,7 +160,7 @@ fn build_app(
     };
 
     let app = App::new(&config, client);
-    t!(t!(app.diesel_database.get()).begin_test_transaction());
+    t!(t!(app.primary_database.get()).begin_test_transaction());
     let app = Arc::new(app);
     let handler = cargo_registry::build_handler(Arc::clone(&app));
     (app, handler)
@@ -278,8 +279,8 @@ fn multiple_live_references_to_the_same_connection_can_be_checked_out() {
     use std::ptr;
 
     let (app, _) = app();
-    let conn1 = app.diesel_database.get().unwrap();
-    let conn2 = app.diesel_database.get().unwrap();
+    let conn1 = app.primary_database.get().unwrap();
+    let conn2 = app.primary_database.get().unwrap();
     let conn1_ref: &PgConnection = &conn1;
     let conn2_ref: &PgConnection = &conn2;
 

src/tests/util.rs

@@ -79,7 +79,7 @@ impl Drop for TestAppInner {
 
         // Manually verify that all jobs have completed successfully
         // This will catch any tests that enqueued a job but forgot to initialize the runner
-        let conn = self.app.diesel_database.get().unwrap();
+        let conn = self.app.primary_database.get().unwrap();
         let job_count: i64 = background_jobs.count().get_result(&*conn).unwrap();
         assert_eq!(
             0, job_count,
@@ -122,7 +122,7 @@ impl TestApp {
     /// connection before making any API calls.  Once the closure returns, the connection is
     /// dropped, ensuring it is returned to the pool and available for any future API calls.
     pub fn db<T, F: FnOnce(&PgConnection) -> T>(&self, f: F) -> T {
-        let conn = self.0.app.diesel_database.get().unwrap();
+        let conn = self.0.app.primary_database.get().unwrap();
         f(&conn)
     }
 
@@ -215,7 +215,7 @@ impl TestAppBuilder {
         let (app, middle) = crate::build_app(self.config, self.proxy);
 
         let runner = if self.build_job_runner {
-            let connection_pool = app.diesel_database.clone();
+            let connection_pool = app.primary_database.clone();
             let repository_config = RepositoryConfig {
                 index_location: Url::from_file_path(&git::bare()).unwrap(),
                 credentials: Credentials::Missing,