storage: ensure uploaded_files_total metric is recorded

Author: pietroalbini

src/storage/database.rs

@@ -1,16 +1,19 @@
 use super::{Blob, StorageTransaction};
 use crate::db::Pool;
+use crate::Metrics;
 use chrono::{DateTime, NaiveDateTime, Utc};
 use failure::Error;
 use postgres::Transaction;
+use std::sync::Arc;
 
 pub(crate) struct DatabaseBackend {
     pool: Pool,
+    metrics: Arc<Metrics>,
 }
 
 impl DatabaseBackend {
-    pub(crate) fn new(pool: Pool) -> Self {
-        Self { pool }
+    pub(crate) fn new(pool: Pool, metrics: Arc<Metrics>) -> Self {
+        Self { pool, metrics }
     }
 
     pub(super) fn exists(&self, path: &str) -> Result<bool, Error> {
@@ -68,12 +71,14 @@ impl DatabaseBackend {
     pub(super) fn start_connection(&self) -> Result<DatabaseClient, Error> {
         Ok(DatabaseClient {
             conn: self.pool.get()?,
+            metrics: self.metrics.clone(),
         })
     }
 }
 
 pub(super) struct DatabaseClient {
     conn: crate::db::PoolClient,
+    metrics: Arc<Metrics>,
 }
 
 impl DatabaseClient {
@@ -82,12 +87,14 @@ impl DatabaseClient {
     ) -> Result<DatabaseStorageTransaction<'_>, Error> {
         Ok(DatabaseStorageTransaction {
             transaction: self.conn.transaction()?,
+            metrics: &self.metrics,
         })
     }
 }
 
 pub(super) struct DatabaseStorageTransaction<'a> {
     transaction: Transaction<'a>,
+    metrics: &'a Metrics,
 }
 
 impl<'a> StorageTransaction for DatabaseStorageTransaction<'a> {
@@ -101,6 +108,7 @@ impl<'a> StorageTransaction for DatabaseStorageTransaction<'a> {
                     SET mime = EXCLUDED.mime, content = EXCLUDED.content, compression = EXCLUDED.compression",
                 &[&blob.path, &blob.mime, &blob.content, &compression],
             )?;
+            self.metrics.uploaded_files_total.inc();
         }
         Ok(())
     }

src/storage/mod.rs

@@ -81,7 +81,7 @@ impl Storage {
         let backend = if let Some(c) = s3::s3_client() {
             StorageBackend::S3(Box::new(S3Backend::new(c, metrics, config)?))
         } else {
-            StorageBackend::Database(DatabaseBackend::new(pool))
+            StorageBackend::Database(DatabaseBackend::new(pool, metrics))
         };
         Ok(Storage { backend })
     }
@@ -98,9 +98,9 @@ impl Storage {
     }
 
     #[cfg(test)]
-    pub(crate) fn temp_new_db(pool: Pool) -> Result<Self, Error> {
+    pub(crate) fn temp_new_db(pool: Pool, metrics: Arc<Metrics>) -> Result<Self, Error> {
         Ok(Storage {
-            backend: StorageBackend::Database(DatabaseBackend::new(pool)),
+            backend: StorageBackend::Database(DatabaseBackend::new(pool, metrics)),
         })
     }
 
@@ -391,7 +391,7 @@ mod backend_tests {
         Ok(())
     }
 
-    fn test_store_blobs(storage: &Storage) -> Result<(), Error> {
+    fn test_store_blobs(storage: &Storage, metrics: &Metrics) -> Result<(), Error> {
         const NAMES: &[&str] = &[
             "a",
             "b",
@@ -419,10 +419,12 @@ mod backend_tests {
             assert_eq!(blob.mime, actual.mime);
         }
 
+        assert_eq!(NAMES.len(), metrics.uploaded_files_total.get() as usize);
+
         Ok(())
     }
 
-    fn test_store_all(storage: &Storage) -> Result<(), Error> {
+    fn test_store_all(storage: &Storage, metrics: &Metrics) -> Result<(), Error> {
         let dir = tempfile::Builder::new()
             .prefix("docs.rs-upload-test")
             .tempdir()?;
@@ -464,6 +466,8 @@ mod backend_tests {
         expected_algs.insert(CompressionAlgorithm::default());
         assert_eq!(algs, expected_algs);
 
+        assert_eq!(2, metrics.uploaded_files_total.get());
+
         Ok(())
     }
 
@@ -559,7 +563,11 @@ mod backend_tests {
     // Remember to add the test name to the macro below when adding a new one.
 
     macro_rules! backend_tests {
-        (backends($env:ident) { $($backend:ident => $create:expr,)* } tests $tests:tt ) => {
+        (
+            backends($env:ident) { $($backend:ident => $create:expr,)* }
+            tests $tests:tt
+            tests_with_metrics $tests_with_metrics:tt
+        ) => {
             $(
                 mod $backend {
                     use crate::test::TestEnvironment;
@@ -571,6 +579,7 @@ mod backend_tests {
                     }
 
                     backend_tests!(@tests $tests);
+                    backend_tests!(@tests_with_metrics $tests_with_metrics);
                 }
             )*
         };
@@ -584,6 +593,16 @@ mod backend_tests {
                 }
             )*
         };
+        (@tests_with_metrics { $($test:ident,)* }) => {
+            $(
+                #[test]
+                fn $test() {
+                    crate::test::wrapper(|env| {
+                        super::$test(&*get_storage(env), &*env.metrics())
+                    });
+                }
+            )*
+        };
     }
 
     backend_tests! {
@@ -597,10 +616,13 @@ mod backend_tests {
             test_exists,
             test_get_object,
             test_get_too_big,
-            test_store_blobs,
-            test_store_all,
             test_delete_prefix,
             test_delete_percent,
         }
+
+        tests_with_metrics {
+            test_store_blobs,
+            test_store_all,
+        }
     }
 }

src/test/mod.rs

@@ -218,7 +218,7 @@ impl TestEnvironment {
         self.storage_db
             .get_or_init(|| {
                 Arc::new(
-                    Storage::temp_new_db(self.db().pool())
+                    Storage::temp_new_db(self.db().pool(), self.metrics())
                         .expect("failed to initialize the storage"),
                 )
             })