downloads_counter: persist all counters on shutdown

Author: pietroalbini

src/bin/server.rs

@@ -54,7 +54,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     // Start the background thread periodically persisting download counts to the database.
     downloads_counter_thread(app.clone());
 
-    let handler = cargo_registry::build_handler(app);
+    let handler = cargo_registry::build_handler(app.clone());
 
     // On every server restart, ensure the categories available in the database match
     // the information in *src/categories.toml*.
@@ -167,6 +167,11 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         }
     }
 
+    println!("Persisting remaining downloads counters");
+    if let Err(err) = app.downloads_counter.persist_all_shards(&app) {
+        println!("downloads_counter error: {}", err);
+    }
+
     println!("Server has gracefully shutdown!");
     Ok(())
 }

src/downloads_counter.rs

@@ -53,6 +53,31 @@ impl DownloadsCounter {
         }
     }
 
+    pub fn persist_all_shards(&self, app: &App) -> Result<(), Error> {
+        let conn = app.primary_database.get()?;
+
+        let mut counted_downloads = 0;
+        let mut counted_versions = 0;
+        let mut pending_downloads = 0;
+        for shard in self.inner.shards() {
+            let shard = std::mem::take(&mut *shard.write());
+            let stats = self.persist_shard(&conn, shard)?;
+
+            counted_downloads += stats.counted_downloads;
+            counted_versions += stats.counted_versions;
+            pending_downloads = stats.pending_downloads;
+        }
+
+        println!(
+            "download_counter all_shards counted_versions={} counted_downloads={} pending_downloads={}",
+            counted_versions,
+            counted_downloads,
+            pending_downloads,
+        );
+
+        Ok(())
+    }
+
     pub fn persist_next_shard(&self, app: &App) -> Result<(), Error> {
         let conn = app.primary_database.get()?;