allow downloads to happen with a broken database

Author: pietroalbini

Cargo.lock

@@ -81,7 +81,7 @@ sha2 = "0.9"
 swirl = { git = "https://github.com/sgrif/swirl.git", rev = "e87cf37" }
 tar = "0.4.16"
 tempfile = "3"
-tokio = { version = "1", features = ["net", "signal", "io-std", "io-util", "rt-multi-thread"]}
+tokio = { version = "1.5.0", features = ["net", "signal", "io-std", "io-util", "rt-multi-thread", "macros"]}
 toml = "0.5"
 tracing = "0.1"
 tracing-subscriber = "0.2"
@@ -92,7 +92,7 @@ claim = "0.5"
 conduit-test = "0.9.0-alpha.4"
 hyper-tls = "0.5"
 lazy_static = "1.0"
-tokio = "1"
+tokio = "1.5.0"
 tower-service = "0.3.0"
 
 [build-dependencies]

Cargo.toml

@@ -2,10 +2,9 @@ use reqwest::blocking::Client;
 use std::panic::AssertUnwindSafe;
 use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
 
-use diesel::r2d2::PoolError;
 use swirl::PerformError;
 
-use crate::db::{DieselPool, DieselPooledConn};
+use crate::db::{DieselPool, DieselPooledConn, PoolError};
 use crate::git::Repository;
 use crate::uploaders::Uploader;
 

src/background_jobs.rs

@@ -2,53 +2,80 @@
 //!
 //! Crate level functionality is located in `krate::downloads`.
 
+use super::{extract_crate_name_and_semver, version_and_crate};
 use crate::controllers::prelude::*;
-
-use chrono::{Duration, NaiveDate, Utc};
-
+use crate::db::PoolError;
 use crate::models::{Crate, VersionDownload};
 use crate::schema::*;
 use crate::views::EncodableVersionDownload;
-
-use super::{extract_crate_name_and_semver, version_and_crate};
+use chrono::{Duration, NaiveDate, Utc};
 
 /// Handles the `GET /crates/:crate_id/:version/download` route.
 /// This returns a URL to the location where the crate is stored.
 pub fn download(req: &mut dyn RequestExt) -> EndpointResult {
+    let app = req.app().clone();
     let recorder = req.timing_recorder();
 
-    let crate_name = &req.params()["crate_id"];
-    let version = &req.params()["version"];
-
-    let (version_id, canonical_crate_name): (_, String) = {
-        use self::versions::dsl::*;
-
-        let conn = recorder.record("get_conn", || req.db_conn())?;
-
-        // Returns the crate name as stored in the database, or an error if we could
-        // not load the version ID from the database.
-        recorder.record("get_version", || {
-            versions
-                .inner_join(crates::table)
-                .select((id, crates::name))
-                .filter(Crate::with_name(crate_name))
-                .filter(num.eq(version))
-                .first(&*conn)
-        })?
-    };
-
-    // The increment does not happen instantly, but it's deferred to be executed in a batch
-    // along with other downloads. See crate::downloads_counter for the implementation.
-    req.app().downloads_counter.increment(version_id);
+    let mut crate_name = req.params()["crate_id"].clone();
+    let version = req.params()["version"].as_str();
+
+    let mut log_metadata = None;
+    match recorder.record("get_conn", || req.db_conn()) {
+        Ok(conn) => {
+            use self::versions::dsl::*;
+
+            // Returns the crate name as stored in the database, or an error if we could
+            // not load the version ID from the database.
+            let (version_id, canonical_crate_name) = recorder.record("get_version", || {
+                versions
+                    .inner_join(crates::table)
+                    .select((id, crates::name))
+                    .filter(Crate::with_name(&crate_name))
+                    .filter(num.eq(version))
+                    .first::<(i32, String)>(&*conn)
+            })?;
+
+            if canonical_crate_name != crate_name {
+                app.instance_metrics.downloads_non_canonical_crate_name_total.inc();
+                log_metadata = Some(("bot", "dl"));
+            }
+            crate_name = canonical_crate_name;
+
+            // The increment does not happen instantly, but it's deferred to be executed in a batch
+            // along with other downloads. See crate::downloads_counter for the implementation.
+            app.downloads_counter.increment(version_id);
+        }
+        Err(PoolError::UnhealthyPool) => {
+            // The download endpoint is the most critical route in the whole crates.io application,
+            // as it's relied upon by users and automations to download crates. Keeping it working
+            // is the most important thing for us.
+            //
+            // The endpoint relies on the database to fetch the canonical crate name (with the
+            // right capitalization and hyphenation), but that's only needed to serve clients who
+            // don't call the endpoint with the crate's canonical name.
+            //
+            // Thankfully Cargo always uses the right name when calling the endpoint, and we can
+            // keep it working during a full database outage by unconditionally redirecting without
+            // checking whether the crate exists or the rigth name is used. Non-Cargo clients might
+            // get a 404 response instead of a 500, but that's worth it.
+            //
+            // Without a working database we also can't count downloads, but that's also less
+            // critical than keeping Cargo downloads operational.
+
+            app.instance_metrics.downloads_unconditional_redirects_total.inc();
+            log_metadata = Some(("unconditional_redirect", "true"));
+        }
+        Err(err) => return Err(err.into()),
+    }
 
     let redirect_url = req
         .app()
         .config
         .uploader
-        .crate_location(&canonical_crate_name, version);
+        .crate_location(&crate_name, &*version);
 
-    if &canonical_crate_name != crate_name {
-        req.log_metadata("bot", "dl");
+    if let Some((key, value)) = log_metadata {
+        req.log_metadata(key, value);
     }
 
     if req.wants_json() {

src/controllers/version/downloads.rs

@@ -2,8 +2,8 @@ use conduit::RequestExt;
 use diesel::prelude::*;
 use diesel::r2d2::{self, ConnectionManager, CustomizeConnection};
 use parking_lot::{ReentrantMutex, ReentrantMutexGuard};
-use std::ops::Deref;
 use std::sync::Arc;
+use std::{ops::Deref, time::Duration};
 use url::Url;
 
 use crate::middleware::app::RequestApp;
@@ -32,9 +32,17 @@ impl DieselPool {
         DieselPool::Test(Arc::new(ReentrantMutex::new(conn)))
     }
 
-    pub fn get(&self) -> Result<DieselPooledConn<'_>, r2d2::PoolError> {
+    pub fn get(&self) -> Result<DieselPooledConn<'_>, PoolError> {
         match self {
-            DieselPool::Pool(pool) => Ok(DieselPooledConn::Pool(pool.get()?)),
+            DieselPool::Pool(pool) => {
+                if let Some(conn) = pool.try_get() {
+                    Ok(DieselPooledConn::Pool(conn))
+                } else if !self.is_healthy() {
+                    Err(PoolError::UnhealthyPool)
+                } else {
+                    Ok(DieselPooledConn::Pool(pool.get().map_err(PoolError::R2D2)?))
+                }
+            }
             DieselPool::Test(conn) => Ok(DieselPooledConn::Test(conn.lock())),
         }
     }
@@ -54,6 +62,21 @@ impl DieselPool {
             },
         }
     }
+
+    pub fn wait_until_healthy(&self, timeout: Duration) -> Result<(), PoolError> {
+        match self {
+            DieselPool::Pool(pool) => match pool.get_timeout(timeout) {
+                Ok(_) => Ok(()),
+                Err(_) if !self.is_healthy() => Err(PoolError::UnhealthyPool),
+                Err(err) => Err(PoolError::R2D2(err)),
+            },
+            DieselPool::Test(_) => Ok(()),
+        }
+    }
+
+    fn is_healthy(&self) -> bool {
+        self.state().connections > 0
+    }
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -96,23 +119,23 @@ pub fn connection_url(url: &str) -> String {
 
 pub trait RequestTransaction {
     /// Obtain a read/write database connection from the primary pool
-    fn db_conn(&self) -> Result<DieselPooledConn<'_>, r2d2::PoolError>;
+    fn db_conn(&self) -> Result<DieselPooledConn<'_>, PoolError>;
 
     /// 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) -> Result<DieselPooledConn<'_>, r2d2::PoolError>;
+    fn db_read_only(&self) -> Result<DieselPooledConn<'_>, PoolError>;
 }
 
 impl<T: RequestExt + ?Sized> RequestTransaction for T {
-    fn db_conn(&self) -> Result<DieselPooledConn<'_>, r2d2::PoolError> {
-        self.app().primary_database.get().map_err(Into::into)
+    fn db_conn(&self) -> Result<DieselPooledConn<'_>, PoolError> {
+        self.app().primary_database.get()
     }
 
-    fn db_read_only(&self) -> Result<DieselPooledConn<'_>, r2d2::PoolError> {
+    fn db_read_only(&self) -> Result<DieselPooledConn<'_>, PoolError> {
         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),
+            Some(pool) => pool.get(),
+            None => self.app().primary_database.get(),
         }
     }
 }
@@ -148,3 +171,20 @@ pub(crate) fn test_conn() -> PgConnection {
     conn.begin_test_transaction().unwrap();
     conn
 }
+
+#[derive(Debug)]
+pub enum PoolError {
+    R2D2(r2d2::PoolError),
+    UnhealthyPool,
+}
+
+impl std::error::Error for PoolError {}
+
+impl std::fmt::Display for PoolError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PoolError::R2D2(err) => write!(f, "{}", err),
+            PoolError::UnhealthyPool => write!(f, "unhealthy database pool"),
+        }
+    }
+}

src/db.rs

@@ -32,6 +32,11 @@ metrics! {
         pub requests_total: IntCounter,
         /// Number of requests currently being processed
         pub requests_in_flight: IntGauge,
+
+        /// Number of download requests that were served with an unconditional redirect.
+        pub downloads_unconditional_redirects_total: IntCounter,
+        /// Number of download requests with a non-canonical crate name.
+        pub downloads_non_canonical_crate_name_total: IntCounter,
     }
 
     // All instance metrics will be prefixed with this namespace.

src/metrics/instance.rs

@@ -37,6 +37,7 @@ mod categories;
 mod category;
 mod dump_db;
 mod git;
+mod unhealthy_database;
 mod keyword;
 mod krate;
 mod metrics;

src/tests/all.rs

@@ -0,0 +1,66 @@
+use crate::{
+    builders::CrateBuilder,
+    util::{MockAnonymousUser, RequestHelper, TestApp},
+};
+use std::time::Duration;
+
+#[test]
+fn download_crate_with_broken_networking_primary_database() {
+    let (app, anon, _, owner) = TestApp::init().with_slow_real_db_pool().with_token();
+    app.db(|conn| {
+        CrateBuilder::new("crate_name", owner.as_model().user_id)
+            .version("1.0.0")
+            .expect_build(conn)
+    });
+
+    // When the database connection is healthy downloads are redirected with the proper
+    // capitalization, and missing crates or versions return a 404.
+
+    assert_checked_redirects(&anon);
+
+    // After networking breaks, preventing new database connections, the download endpoint should
+    // do an unconditional redirect to the CDN, without checking whether the crate exists or what
+    // the exact capitalization of crate name is.
+
+    app.db_chaosproxy().break_networking();
+    assert_unconditional_redirects(&anon);
+
+    // After restoring the network and waiting for the database pool to get healthy again redirects
+    // should be checked again.
+
+    app.db_chaosproxy().restore_networking();
+    app.as_inner()
+        .primary_database
+        .wait_until_healthy(Duration::from_millis(500))
+        .expect("the database did not return healthy");
+
+    assert_checked_redirects(&anon);
+}
+
+fn assert_checked_redirects(anon: &MockAnonymousUser) {
+    anon.get::<()>("/api/v1/crates/crate_name/1.0.0/download")
+        .assert_redirect_ends_with("/crate_name/crate_name-1.0.0.crate");
+
+    anon.get::<()>("/api/v1/crates/Crate-Name/1.0.0/download")
+        .assert_redirect_ends_with("/crate_name/crate_name-1.0.0.crate");
+
+    anon.get::<()>("/api/v1/crates/crate_name/2.0.0/download")
+        .assert_not_found();
+
+    anon.get::<()>("/api/v1/crates/awesome-project/1.0.0/download")
+        .assert_not_found();
+}
+
+fn assert_unconditional_redirects(anon: &MockAnonymousUser) {
+    anon.get::<()>("/api/v1/crates/crate_name/1.0.0/download")
+        .assert_redirect_ends_with("/crate_name/crate_name-1.0.0.crate");
+
+    anon.get::<()>("/api/v1/crates/Crate-Name/1.0.0/download")
+        .assert_redirect_ends_with("/Crate-Name/Crate-Name-1.0.0.crate");
+
+    anon.get::<()>("/api/v1/crates/crate_name/2.0.0/download")
+        .assert_redirect_ends_with("/crate_name/crate_name-2.0.0.crate");
+
+    anon.get::<()>("/api/v1/crates/awesome-project/1.0.0/download")
+        .assert_redirect_ends_with("/awesome-project/awesome-project-1.0.0.crate");
+}

src/tests/unhealthy_database.rs

@@ -36,6 +36,7 @@ use std::collections::HashMap;
 mod response;
 mod test_app;
 mod fresh_schema;
+mod chaosproxy;
 
 pub use response::Response;
 pub use test_app::TestApp;