publish: Save default_versions.total value

Author: eth3lbert

src/controllers/krate/publish.rs

@@ -15,8 +15,7 @@ use diesel::dsl::{exists, select};
 use diesel::prelude::*;
 use diesel_async::scoped_futures::ScopedFutureExt;
 use diesel_async::{AsyncConnection, AsyncPgConnection, RunQueryDsl};
-use futures_util::TryFutureExt;
-use futures_util::TryStreamExt;
+use futures_util::{FutureExt, TryFutureExt, TryStreamExt};
 use hex::ToHex;
 use http::request::Parts;
 use http::StatusCode;
@@ -433,10 +432,10 @@ pub async fn publish(app: AppState, req: Parts, body: Body) -> AppResult<Json<Go
         // Link this new version to all dependencies
         add_dependencies(conn, &deps, version.id).await?;
 
-        let existing_default_version = default_versions::table
+        let existing_default_version: Option<(DefaultVersion, i32)> = default_versions::table
             .inner_join(versions::table)
             .filter(default_versions::crate_id.eq(krate.id))
-            .select(DefaultVersion::as_select())
+            .select((DefaultVersion::as_select(), default_versions::total))
             .first(conn)
             .await
             .optional()?;
@@ -450,22 +449,45 @@ pub async fn publish(app: AppState, req: Parts, body: Body) -> AppResult<Json<Go
         // Compared to only using a background job, this prevents us from getting into a
         // situation where a crate exists in the `crates` table but doesn't have a default
         // version in the `default_versions` table.
-        if let Some(existing_default_version) = existing_default_version {
+        if let Some((existing_default_version, existing_total)) = existing_default_version {
             let published_default_version = DefaultVersion {
                 id: version.id,
                 num: semver,
                 yanked: false,
             };
-
-            if existing_default_version < published_default_version {
-                diesel::update(default_versions::table)
-                    .filter(default_versions::crate_id.eq(krate.id))
-                    .set(default_versions::version_id.eq(version.id))
+            let update_base = diesel::update(default_versions::table)
+                    .filter(default_versions::crate_id.eq(krate.id));
+
+            // If `total` isn't set (it defaults to 0), this only occurs when a crate already
+            // exists and releases a new version before the admin sets the total.
+            // In this scenario, the `total` value is updated in a background job. This case should
+            // be safe to remove once all `total` values have been set.
+            let is_new = existing_default_version < published_default_version;
+            let default_total = existing_total == 0;
+            let fut = match (is_new, default_total) {
+                (true, true) => {
+                    update_base.set(default_versions::version_id.eq(version.id)).execute(conn)
+                },
+                (false, true) => {
+                    default_version = Some(existing_default_version.num.to_string());
+                    async { Ok(0) }.boxed()
+                },
+                (true, false) => {
+                    update_base.set((
+                        default_versions::version_id.eq(version.id),
+                        default_versions::total.eq(default_versions::total + 1)
+                    ))
                     .execute(conn)
-                    .await?;
-            } else {
-                default_version = Some(existing_default_version.num.to_string());
-            }
+                },
+                (false, false) => {
+                    default_version = Some(existing_default_version.num.to_string());
+                    update_base.set(
+                        default_versions::total.eq(default_versions::total + 1)
+                    )
+                    .execute(conn)
+                },
+            };
+            fut.await?;
 
             // Update the default version asynchronously in a background job
             // to ensure correctness and eventual consistency.
@@ -475,6 +497,7 @@ pub async fn publish(app: AppState, req: Parts, body: Body) -> AppResult<Json<Go
                 .values((
                     default_versions::crate_id.eq(krate.id),
                     default_versions::version_id.eq(version.id),
+                    default_versions::total.eq(1),
                 ))
                 .execute(conn)
                 .await?;

src/models/default_versions.rs

@@ -1,5 +1,5 @@
 use crate::schema::{default_versions, versions};
-use crates_io_diesel_helpers::SemverVersion;
+use crates_io_diesel_helpers::{greatest, SemverVersion};
 use diesel::prelude::*;
 use diesel_async::{AsyncPgConnection, RunQueryDsl};
 
@@ -60,7 +60,7 @@ pub async fn update_default_version(
     crate_id: i32,
     conn: &mut AsyncPgConnection,
 ) -> QueryResult<()> {
-    let default_version = calculate_default_version(crate_id, conn).await?;
+    let (default_version, total) = calculate_default_version(crate_id, conn).await?;
 
     debug!(
         "Updating default version to {} (id: {})…",
@@ -74,7 +74,10 @@ pub async fn update_default_version(
         ))
         .on_conflict(default_versions::crate_id)
         .do_update()
-        .set(default_versions::version_id.eq(default_version.id))
+        .set((
+            default_versions::version_id.eq(default_version.id),
+            default_versions::total.eq(greatest(total as i32, default_versions::total)),
+        ))
         .execute(conn)
         .await?;
 
@@ -87,22 +90,22 @@ pub async fn verify_default_version(
     crate_id: i32,
     conn: &mut AsyncPgConnection,
 ) -> QueryResult<()> {
-    let calculated = calculate_default_version(crate_id, conn).await?;
+    let (calculated, total) = calculate_default_version(crate_id, conn).await?;
 
     let saved = default_versions::table
-        .select(default_versions::version_id)
+        .select((default_versions::version_id, default_versions::total))
         .filter(default_versions::crate_id.eq(crate_id))
-        .first::<i32>(conn)
+        .first::<(i32, i32)>(conn)
         .await
         .optional()?;
 
-    if let Some(saved) = saved {
-        if saved == calculated.id {
+    if let Some((saved_id, saved_total)) = saved {
+        if saved_id == calculated.id && saved_total == (total as i32) {
             debug!("Default version for crate {crate_id} is up to date");
         } else {
             warn!(
-                "Default version for crate {crate_id} is outdated (expected: {saved}, actual: {})",
-                calculated.id,
+                "Default version for crate {crate_id} is outdated (expected: {saved_id} with total: {saved_total}, actual: {} with total: {})",
+                calculated.id, total
             );
         }
     } else {
@@ -118,7 +121,7 @@ pub async fn verify_default_version(
 async fn calculate_default_version(
     crate_id: i32,
     conn: &mut AsyncPgConnection,
-) -> QueryResult<Version> {
+) -> QueryResult<(Version, usize)> {
     use diesel::result::Error::NotFound;
 
     debug!("Loading all versions for the crate…");
@@ -128,9 +131,14 @@ async fn calculate_default_version(
         .load::<Version>(conn)
         .await?;
 
-    debug!("Found {} versions", versions.len());
+    let total = versions.len();
+    debug!("Found {} versions", total);
 
-    versions.into_iter().max().ok_or(NotFound)
+    versions
+        .into_iter()
+        .max()
+        .map(|version| (version, total))
+        .ok_or(NotFound)
 }
 
 #[cfg(test)]