Cache the version_id for download requests

This adds a caching layer in the app backend to reduce the number of
database queries. While this query is extremely fast, a massive number
of concurrent download requests can still result in a request backlog.

In this implementation cached values are never dropped. This means
that memory usage will grow over time, however the memory usage is
expected to be well within our instance size and the application is
automatically rebooted at least once every 24 hours. Additionally, on
the rare occasion that a crate version is deleted, the cache will
continue to serve redirects and attempt to record download counts for
that version.

If a version is deleted and continues to see download requests, it is
expected that the entire shard of downloads containing that version_id
will fail to persist. Even if we expired cached entries after a certain
time period, there would still be a window where a deleted version
could cause download counts for other versions to not be persisted. It
is possible we could make the download counter logic robust against this
scenario, but given that deleting versions is rare we can mitigate this
by restarting the app after deleting crates or versions.

Author: jtgeibel

src/app.rs

@@ -9,6 +9,8 @@ use crate::email::Emails;
 use crate::github::GitHubClient;
 use crate::metrics::{InstanceMetrics, ServiceMetrics};
 use crate::rate_limiter::RateLimiter;
+
+use dashmap::DashMap;
 use diesel::r2d2;
 use oauth2::basic::BasicClient;
 use reqwest::blocking::Client;
@@ -32,6 +34,12 @@ pub struct App {
     /// The server configuration
     pub config: config::Server,
 
+    /// Cache the `version_id` of a `canonical_crate_name:semver` pair
+    ///
+    /// This is used by the download endpoint to reduce the number of database queries. The
+    /// `version_id` is only cached under the canonical spelling of the crate name.
+    pub(crate) version_id_cacher: DashMap<String, i32>,
+
     /// Count downloads and periodically persist them in the database
     pub downloads_counter: DownloadsCounter,
 
@@ -176,6 +184,7 @@ impl App {
             github_oauth,
             rate_limiter: RateLimiter::new(config.rate_limiter.clone()),
             config,
+            version_id_cacher: DashMap::new(),
             downloads_counter: DownloadsCounter::new(),
             emails: Emails::from_environment(),
             service_metrics: ServiceMetrics::new().expect("could not initialize service metrics"),

src/controllers/version/downloads.rs

@@ -8,7 +8,9 @@ use crate::db::PoolError;
 use crate::models::{Crate, VersionDownload};
 use crate::schema::*;
 use crate::views::EncodableVersionDownload;
+
 use chrono::{Duration, NaiveDate, Utc};
+use dashmap::mapref::entry::Entry;
 
 /// Handles the `GET /crates/:crate_id/:version/download` route.
 /// This returns a URL to the location where the crate is stored.
@@ -19,59 +21,82 @@ pub fn download(req: &mut dyn RequestExt) -> EndpointResult {
     let version = req.params()["version"].as_str();
 
     let mut log_metadata = None;
-    match 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) = app
-                .instance_metrics
-                .downloads_select_query_execution_time
-                .observe_closure_duration(|| {
-                    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;
+
+    match app
+        .version_id_cacher
+        .entry(format!("{}:{}", crate_name, version))
+    {
+        // The version_id is cached. This also means that the provided crate_name is canonical
+        // and that no fixup is necessary before redirecting.
+        Entry::Occupied(entry) => {
+            let version_id = *entry.get();
 
             // 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()),
+
+        // The version_id is not cached, so either query the database or fallback to an
+        // unconditional redirect if the database pool is unhealthy.
+        Entry::Vacant(entry) => match 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) = app
+                    .instance_metrics
+                    .downloads_select_query_execution_time
+                    .observe_closure_duration(|| {
+                        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;
+                } else {
+                    // The version_id is only cached if the provided crate name was canonical.
+                    // Non-canonical requests fallback to the "slow" path with a DB query, but
+                    // we typically only get a few hundred non-canonical requests in a day anyway.
+                    entry.insert(version_id);
+                }
+
+                // 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

src/tests/krate/downloads.rs

@@ -1,5 +1,5 @@
 use crate::builders::{CrateBuilder, VersionBuilder};
-use crate::util::{RequestHelper, TestApp};
+use crate::util::{MockAnonymousUser, RequestHelper, TestApp};
 use cargo_registry::views::EncodableVersionDownload;
 use chrono::{Duration, Utc};
 use http::StatusCode;
@@ -9,6 +9,35 @@ struct Downloads {
     version_downloads: Vec<EncodableVersionDownload>,
 }
 
+fn persist_downloads_count(app: &TestApp) {
+    app.as_inner()
+        .downloads_counter
+        .persist_all_shards(app.as_inner())
+        .expect("failed to persist downloads count")
+        .log();
+}
+
+#[track_caller]
+fn assert_dl_count(
+    anon: &MockAnonymousUser,
+    name_and_version: &str,
+    query: Option<&str>,
+    count: i32,
+) {
+    let url = format!("/api/v1/crates/{}/downloads", name_and_version);
+    let downloads: Downloads = if let Some(query) = query {
+        anon.get_with_query(&url, query).good()
+    } else {
+        anon.get(&url).good()
+    };
+    let total_downloads = downloads
+        .version_downloads
+        .iter()
+        .map(|vd| vd.downloads)
+        .sum::<i32>();
+    assert_eq!(total_downloads, count);
+}
+
 #[test]
 fn download() {
     let (app, anon, user) = TestApp::init().with_user();
@@ -20,60 +49,37 @@ fn download() {
             .expect_build(conn);
     });
 
-    let assert_dl_count = |name_and_version: &str, query: Option<&str>, count: i32| {
-        let url = format!("/api/v1/crates/{}/downloads", name_and_version);
-        let downloads: Downloads = if let Some(query) = query {
-            anon.get_with_query(&url, query).good()
-        } else {
-            anon.get(&url).good()
-        };
-        let total_downloads = downloads
-            .version_downloads
-            .iter()
-            .map(|vd| vd.downloads)
-            .sum::<i32>();
-        assert_eq!(total_downloads, count);
-    };
-
     let download = |name_and_version: &str| {
         let url = format!("/api/v1/crates/{}/download", name_and_version);
         let response = anon.get::<()>(&url);
         assert_eq!(response.status(), StatusCode::FOUND);
         // TODO: test the with_json code path
     };
 
-    let persist_downloads_count = || {
-        app.as_inner()
-            .downloads_counter
-            .persist_all_shards(app.as_inner())
-            .expect("failed to persist downloads count")
-            .log();
-    };
-
     download("foo_download/1.0.0");
     // No downloads are counted until the counters are persisted
-    assert_dl_count("foo_download/1.0.0", None, 0);
-    assert_dl_count("foo_download", None, 0);
-    persist_downloads_count();
+    assert_dl_count(&anon, "foo_download/1.0.0", None, 0);
+    assert_dl_count(&anon, "foo_download", None, 0);
+    persist_downloads_count(&app);
     // Now that the counters are persisted the download counts show up.
-    assert_dl_count("foo_download/1.0.0", None, 1);
-    assert_dl_count("foo_download", None, 1);
+    assert_dl_count(&anon, "foo_download/1.0.0", None, 1);
+    assert_dl_count(&anon, "foo_download", None, 1);
 
     download("FOO_DOWNLOAD/1.0.0");
-    persist_downloads_count();
-    assert_dl_count("FOO_DOWNLOAD/1.0.0", None, 2);
-    assert_dl_count("FOO_DOWNLOAD", None, 2);
+    persist_downloads_count(&app);
+    assert_dl_count(&anon, "FOO_DOWNLOAD/1.0.0", None, 2);
+    assert_dl_count(&anon, "FOO_DOWNLOAD", None, 2);
 
     let yesterday = (Utc::today() + Duration::days(-1)).format("%F");
     let query = format!("before_date={}", yesterday);
-    assert_dl_count("FOO_DOWNLOAD/1.0.0", Some(&query), 0);
+    assert_dl_count(&anon, "FOO_DOWNLOAD/1.0.0", Some(&query), 0);
     // crate/downloads always returns the last 90 days and ignores date params
-    assert_dl_count("FOO_DOWNLOAD", Some(&query), 2);
+    assert_dl_count(&anon, "FOO_DOWNLOAD", Some(&query), 2);
 
     let tomorrow = (Utc::today() + Duration::days(1)).format("%F");
     let query = format!("before_date={}", tomorrow);
-    assert_dl_count("FOO_DOWNLOAD/1.0.0", Some(&query), 2);
-    assert_dl_count("FOO_DOWNLOAD", Some(&query), 2);
+    assert_dl_count(&anon, "FOO_DOWNLOAD/1.0.0", Some(&query), 2);
+    assert_dl_count(&anon, "FOO_DOWNLOAD", Some(&query), 2);
 }
 
 #[test]
@@ -105,3 +111,38 @@ fn download_noncanonical_crate_name() {
     anon.get::<()>("/api/v1/crates/foo-download/1.0.0/download")
         .assert_redirect_ends_with("/crates/foo_download/foo_download-1.0.0.crate");
 }
+
+#[test]
+fn download_caches_version_id() {
+    use cargo_registry::schema::crates::dsl::*;
+    use diesel::prelude::*;
+
+    let (app, anon, user) = TestApp::init().with_user();
+    let user = user.as_model();
+
+    app.db(|conn| {
+        CrateBuilder::new("foo_download", user.id)
+            .version(VersionBuilder::new("1.0.0"))
+            .expect_build(conn);
+    });
+
+    anon.get::<()>("/api/v1/crates/foo_download/1.0.0/download")
+        .assert_redirect_ends_with("/crates/foo_download/foo_download-1.0.0.crate");
+
+    // Rename the crate, so that `foo_download` will not be found if its version_id was not cached
+    app.db(|conn| {
+        diesel::update(crates.filter(name.eq("foo_download")))
+            .set(name.eq("other"))
+            .execute(conn)
+            .unwrap();
+    });
+
+    // This would result in a 404 if the endpoint tried to read from the database
+    anon.get::<()>("/api/v1/crates/foo_download/1.0.0/download")
+        .assert_redirect_ends_with("/crates/foo_download/foo_download-1.0.0.crate");
+
+    // Downloads are persisted by version_id, so the rename doesn't matter
+    persist_downloads_count(&app);
+    // Check download count against the new name, rather than rename it back to the original value
+    assert_dl_count(&anon, "other/1.0.0", None, 2);
+}