Don't lock more than 1 row at a time in update-downloads

Prior to this commit, `bin/update-downloads` will hold a lock on every
row it updates for the entire batch of 1000. Each row gets locked when
it is `updated`, but the lock isn't released until the whole transaction
is comitted. This means that the first row in the batch is locked for
the whole transaction, but the last row in the batch isn't locked longer
than it should be.

Processing 1000 rows is routinely taking more than 30 seconds to
complete. Whenever it does, this causes requests to time out. Even when
it doesn't, this locking behavior is the reason that our 99th percentile
response time is getting to be so bad.

The endpoint affected by this is `/crates/{name}/{version}/download`. In
other words, this bug is causing `cargo install` to fail. Even when it
doesn't, this bug is causing it to take seconds when it should take
milliseconds.

By establishing a transaction per row instead of per batch, we retain
atomicity, without holding a lock for the whole batch. To keep
everything atomic, we need to update the global counter per row instead
of per batch.

This relies on the invariant that no rows in `version_downloads` get
touched once `date != CURRENT_DATE` other than by this script.

I believe this is also the root cause of rust-lang#1304, and that we can maybe
revert rust-lang#1312

Author: sgrif

src/bin/update-downloads.rs

@@ -30,19 +30,17 @@ fn main() {
 
 fn update(conn: &PgConnection) -> QueryResult<()> {
     use version_downloads::dsl::*;
+
     let mut max = Some(0);
     while let Some(m) = max {
-        conn.transaction::<_, diesel::result::Error, _>(|| {
-            let rows = version_downloads
-                .filter(processed.eq(false))
-                .filter(id.gt(m))
-                .order(id)
-                .limit(LIMIT)
-                .load(conn)?;
-            collect(conn, &rows)?;
-            max = rows.last().map(|d| d.id);
-            Ok(())
-        })?;
+        let rows = version_downloads
+            .filter(processed.eq(false))
+            .filter(id.gt(m))
+            .order(id)
+            .limit(LIMIT)
+            .load(conn)?;
+        collect(conn, &rows)?;
+        max = rows.last().map(|d| d.id);
     }
     Ok(())
 }
@@ -61,50 +59,52 @@ fn collect(conn: &PgConnection, rows: &[VersionDownload]) -> QueryResult<()> {
         now.to_rfc2822()
     );
 
-    let mut total = 0;
     for download in rows {
         let amt = download.downloads - download.counted;
-        total += i64::from(amt);
-
-        // Flag this row as having been processed if we're passed the cutoff,
-        // and unconditionally increment the number of counted downloads.
-        update(version_downloads::table.find(download.id))
-            .set((
-                version_downloads::processed.eq(download.date < cutoff),
-                version_downloads::counted.eq(version_downloads::counted + amt),
-            ))
-            .execute(conn)?;
-
-        // Update the total number of version downloads
-        let crate_id = update(versions::table.find(download.version_id))
-            .set(versions::downloads.eq(versions::downloads + amt))
-            .returning(versions::crate_id)
-            .get_result::<i32>(conn)?;
 
-        // Update the total number of crate downloads
-        update(crates::table.find(crate_id))
-            .set(crates::downloads.eq(crates::downloads + amt))
-            .execute(conn)?;
+        conn.transaction::<_, diesel::result::Error, _>(|| {
+            // Flag this row as having been processed if we're passed the cutoff,
+            // and unconditionally increment the number of counted downloads.
+            update(version_downloads::table.find(download.id))
+                .set((
+                    version_downloads::processed.eq(download.date < cutoff),
+                    version_downloads::counted.eq(version_downloads::counted + amt),
+                ))
+                .execute(conn)?;
+
+            // Update the total number of version downloads
+            let crate_id = update(versions::table.find(download.version_id))
+                .set(versions::downloads.eq(versions::downloads + amt))
+                .returning(versions::crate_id)
+                .get_result::<i32>(conn)?;
+
+            // Update the total number of crate downloads
+            update(crates::table.find(crate_id))
+                .set(crates::downloads.eq(crates::downloads + amt))
+                .execute(conn)?;
+
+            // Update the total number of crate downloads for today
+            insert_into(crate_downloads::table)
+                .values((
+                    crate_downloads::crate_id.eq(crate_id),
+                    crate_downloads::downloads.eq(amt),
+                    crate_downloads::date.eq(download.date),
+                ))
+                .on_conflict(crate_downloads::table.primary_key())
+                .do_update()
+                .set(crate_downloads::downloads.eq(crate_downloads::downloads + amt))
+                .execute(conn)?;
+
+            // Now that everything else for this crate is done, update the global counter of total
+            // downloads
+            update(metadata::table)
+                .set(metadata::total_downloads.eq(metadata::total_downloads + (amt as i64)))
+                .execute(conn)?;
 
-        // Update the total number of crate downloads for today
-        insert_into(crate_downloads::table)
-            .values((
-                crate_downloads::crate_id.eq(crate_id),
-                crate_downloads::downloads.eq(amt),
-                crate_downloads::date.eq(download.date),
-            ))
-            .on_conflict(crate_downloads::table.primary_key())
-            .do_update()
-            .set(crate_downloads::downloads.eq(crate_downloads::downloads + amt))
-            .execute(conn)?;
+            Ok(())
+        })?;
     }
 
-    // After everything else is done, update the global counter of total
-    // downloads.
-    update(metadata::table)
-        .set(metadata::total_downloads.eq(metadata::total_downloads + total))
-        .execute(conn)?;
-
     Ok(())
 }