Merge pull request #1643 from Kobzol/db-purge-artifact

Kobzol · web-flow · commit a4d589f69579 · 2023-09-28T20:42:49.000+02:00
Add a command and benchmark flags to remove artifact data
diff --git a/collector/src/bin/collector.rs b/collector/src/bin/collector.rs
@@ -413,6 +413,21 @@ struct BenchRustcOption {
     bench_rustc: bool,
 }
 
+#[derive(Clone, Debug, clap::ValueEnum)]
+enum PurgeMode {
+    /// Purge all old data associated with the artifact
+    Old,
+    /// Purge old data of failed benchmarks associated with the artifact
+    Failed,
+}
+
+#[derive(Debug, clap::Args)]
+struct PurgeOption {
+    /// Removes old data for the specified artifact prior to running the benchmarks.
+    #[arg(long = "purge")]
+    purge: Option<PurgeMode>,
+}
+
 // For each subcommand we list the mandatory arguments in the required
 // order, followed by the options in alphabetical order.
 #[derive(Debug, clap::Subcommand)]
@@ -437,6 +452,9 @@ enum Commands {
         /// faster.
         #[arg(long = "no-isolate")]
         no_isolate: bool,
+
+        #[command(flatten)]
+        purge: PurgeOption,
     },
 
     /// Profiles a runtime benchmark.
@@ -493,6 +511,9 @@ enum Commands {
 
         #[command(flatten)]
         self_profile: SelfProfileOption,
+
+        #[command(flatten)]
+        purge: PurgeOption,
     },
 
     /// Benchmarks the next commit or release for perf.rust-lang.org
@@ -552,6 +573,15 @@ enum Commands {
 
     /// Download a crate into collector/benchmarks.
     Download(DownloadCommand),
+
+    /// Removes all data associated with artifact(s) with the given name.
+    PurgeArtifact {
+        /// Name of the artifact.
+        name: String,
+
+        #[command(flatten)]
+        db: DbOption,
+    },
 }
 
 #[derive(Debug, clap::Parser)]
@@ -620,6 +650,7 @@ fn main_result() -> anyhow::Result<i32> {
             iterations,
             db,
             no_isolate,
+            purge,
         } => {
             log_db(&db);
             let toolchain = get_local_toolchain_for_runtime_benchmarks(&local, &target_triple)?;
@@ -633,6 +664,8 @@ fn main_result() -> anyhow::Result<i32> {
 
             let mut conn = rt.block_on(pool.connection());
             let artifact_id = ArtifactId::Tag(toolchain.id.clone());
+            rt.block_on(purge_old_data(conn.as_mut(), &artifact_id, purge.purge));
+
             let runtime_suite = rt.block_on(load_runtime_benchmarks(
                 conn.as_mut(),
                 &runtime_benchmark_dir,
@@ -747,6 +780,7 @@ fn main_result() -> anyhow::Result<i32> {
             bench_rustc,
             iterations,
             self_profile,
+            purge,
         } => {
             log_db(&db);
             let profiles = opts.profiles.0;
@@ -775,7 +809,9 @@ fn main_result() -> anyhow::Result<i32> {
             benchmarks.retain(|b| b.category().is_primary_or_secondary());
 
             let artifact_id = ArtifactId::Tag(toolchain.id.clone());
-            let conn = rt.block_on(pool.connection());
+            let mut conn = rt.block_on(pool.connection());
+            rt.block_on(purge_old_data(conn.as_mut(), &artifact_id, purge.purge));
+
             let shared = SharedBenchmarkConfig {
                 toolchain,
                 artifact_id,
@@ -1057,6 +1093,14 @@ Make sure to modify `{dir}/perf-config.json` if the category/artifact don't matc
             );
             Ok(0)
         }
+        Commands::PurgeArtifact { name, db } => {
+            let pool = Pool::open(&db.db);
+            let conn = rt.block_on(pool.connection());
+            rt.block_on(conn.purge_artifact(&ArtifactId::Tag(name.clone())));
+
+            println!("Data of artifact {name} were removed");
+            Ok(0)
+        }
     }
 }
 
@@ -1115,6 +1159,28 @@ fn log_db(db_option: &DbOption) {
     println!("Using database `{}`", db_option.db);
 }
 
+async fn purge_old_data(
+    conn: &mut dyn Connection,
+    artifact_id: &ArtifactId,
+    purge_mode: Option<PurgeMode>,
+) {
+    match purge_mode {
+        Some(PurgeMode::Old) => {
+            // Delete everything associated with the artifact
+            conn.purge_artifact(artifact_id).await;
+        }
+        Some(PurgeMode::Failed) => {
+            // Delete all benchmarks that have an error for the given artifact
+            let artifact_row_id = conn.artifact_id(artifact_id).await;
+            let errors = conn.get_error(artifact_row_id).await;
+            for krate in errors.keys() {
+                conn.collector_remove_step(artifact_row_id, krate).await;
+            }
+        }
+        None => {}
+    }
+}
+
 /// Record a collection entry into the database, specifying which benchmark steps will be executed.
 async fn init_collection(
     connection: &mut dyn Connection,
diff --git a/database/manual-modifications.md b/database/manual-modifications.md
@@ -1,5 +1,6 @@
-# Useful queries
-This document contains useful queries that should be performed manually in exceptional situations.
+# Useful queries and commands
+This document contains useful queries and commands that should be performed manually in exceptional
+situations.
 
 ## Remove data for a stable artifact from the DB
 This is important for situations where there is some compilation error for a stable benchmark,
@@ -8,9 +9,16 @@ of future incompatibility lints turning into errors.
 
 The benchmark should be fixed first, and then the DB should be altered (see below).
 
-The easiest solution is to simply completely remove the artifact from the DB. There are
-`ON DELETE CASCADE` clauses for `aid` (artifact ID) on tables that reference it, so it should be
-enough to just delete the artifact from the `artifact` table.
+The easiest solution is to simply completely remove the artifact from the DB.
+You can do that either using the following command:
+
+```bash
+$ cargo run --bin collector purge_artifact <artifact-name>
+# $ cargo run --bin collector purge_artifact 1.70.0 # Remove stable artifact 1.70.0
+```
+
+Or using SQL queries. There are `ON DELETE CASCADE` clauses for `aid` (artifact ID) on tables that
+reference it, so it should be enough to just delete the artifact from the `artifact` table.
 The set of queries below show an example of removing the measured data and errors for Rust `1.69`
 and `1.70`:
 ```sql
diff --git a/database/src/lib.rs b/database/src/lib.rs
@@ -384,6 +384,30 @@ impl From<Commit> for ArtifactId {
     }
 }
 
+struct ArtifactInfo<'a> {
+    name: &'a str,
+    date: Option<DateTime<Utc>>,
+    kind: &'static str,
+}
+
+impl ArtifactId {
+    fn info(&self) -> ArtifactInfo {
+        let (name, date, ty) = match self {
+            Self::Commit(commit) => (
+                commit.sha.as_str(),
+                Some(commit.date.0),
+                if commit.is_try() { "try" } else { "master" },
+            ),
+            Self::Tag(a) => (a.as_str(), None, "release"),
+        };
+        ArtifactInfo {
+            name,
+            date,
+            kind: ty,
+        }
+    }
+}
+
 intern!(pub struct QueryLabel);
 
 #[derive(PartialEq, Eq, Clone, Debug)]
diff --git a/database/src/pool.rs b/database/src/pool.rs
@@ -152,6 +152,8 @@ pub trait Connection: Send + Sync {
     async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str) -> bool;
     async fn collector_end_step(&self, aid: ArtifactIdNumber, step: &str);
 
+    async fn collector_remove_step(&self, aid: ArtifactIdNumber, step: &str);
+
     async fn in_progress_artifacts(&self) -> Vec<ArtifactId>;
 
     async fn in_progress_steps(&self, aid: &ArtifactId) -> Vec<Step>;
@@ -179,6 +181,9 @@ pub trait Connection: Send + Sync {
         profile: &str,
         cache: &str,
     ) -> Vec<(ArtifactIdNumber, i32)>;
+
+    /// Removes all data associated with the given artifact.
+    async fn purge_artifact(&self, aid: &ArtifactId);
 }
 
 #[async_trait::async_trait]
diff --git a/database/src/pool/postgres.rs b/database/src/pool/postgres.rs
@@ -948,33 +948,26 @@ where
     }
 
     async fn artifact_id(&self, artifact: &ArtifactId) -> ArtifactIdNumber {
-        let (name, date, ty) = match artifact {
-            ArtifactId::Commit(commit) => (
-                commit.sha.to_string(),
-                Some(commit.date.0),
-                if commit.is_try() { "try" } else { "master" },
-            ),
-            ArtifactId::Tag(a) => (a.clone(), None, "release"),
-        };
+        let info = artifact.info();
         let aid = self
             .conn()
-            .query_opt("select id from artifact where name = $1", &[&name])
+            .query_opt("select id from artifact where name = $1", &[&info.name])
             .await
             .unwrap();
 
         let aid = match aid {
             Some(aid) => aid.get::<_, i32>(0) as u32,
             None => {
                 self.conn()
-                    .query_opt("insert into artifact (name, date, type) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING RETURNING id", &[
-                        &name,
-                        &date,
-                        &ty,
-                    ])
-                    .await
-                    .unwrap();
+            .query_opt("insert into artifact (name, date, type) VALUES ($1, $2, $3) ON CONFLICT DO NOTHING RETURNING id", &[
+                &info.name,
+                &info.date,
+                &info.kind,
+            ])
+            .await
+            .unwrap();
                 self.conn()
-                    .query_one("select id from artifact where name = $1", &[&name])
+                    .query_one("select id from artifact where name = $1", &[&info.name])
                     .await
                     .unwrap()
                     .get::<_, i32>(0) as u32
@@ -1141,6 +1134,16 @@ where
             log::error!("did not end {} for {:?}", step, aid);
         }
     }
+    async fn collector_remove_step(&self, aid: ArtifactIdNumber, step: &str) {
+        self.conn()
+            .execute(
+                "delete from collector_progress \
+                where aid = $1 and step = $2;",
+                &[&(aid.0 as i32), &step],
+            )
+            .await
+            .unwrap();
+    }
     async fn in_progress_artifacts(&self) -> Vec<ArtifactId> {
         let rows = self
             .conn()
@@ -1387,4 +1390,14 @@ where
             _ => panic!("unknown artifact type: {:?}", ty),
         }
     }
+
+    async fn purge_artifact(&self, aid: &ArtifactId) {
+        // Once we delete the artifact, all data associated with it should also be deleted
+        // thanks to ON DELETE CASCADE.
+        let info = aid.info();
+        self.conn()
+            .execute("delete from artifact where name = $1", &[&info.name])
+            .await
+            .unwrap();
+    }
 }
diff --git a/database/src/pool/sqlite.rs b/database/src/pool/sqlite.rs
@@ -624,26 +624,19 @@ impl Connection for SqliteConnection {
         )
     }
     async fn artifact_id(&self, artifact: &crate::ArtifactId) -> ArtifactIdNumber {
-        let (name, date, ty) = match artifact {
-            crate::ArtifactId::Commit(commit) => (
-                commit.sha.to_string(),
-                Some(commit.date.0),
-                if commit.is_try() { "try" } else { "master" },
-            ),
-            crate::ArtifactId::Tag(a) => (a.clone(), None, "release"),
-        };
+        let info = artifact.info();
 
         self.raw_ref()
             .execute(
                 "insert or ignore into artifact (name, date, type) VALUES (?, ?, ?)",
-                params![&name, &date.map(|d| d.timestamp()), &ty,],
+                params![&info.name, &info.date.map(|d| d.timestamp()), &info.kind,],
             )
             .unwrap();
         ArtifactIdNumber(
             self.raw_ref()
                 .query_row(
                     "select id from artifact where name = $1",
-                    params![&name],
+                    params![&info.name],
                     |r| r.get::<_, i32>(0),
                 )
                 .unwrap() as u32,
@@ -1081,6 +1074,17 @@ impl Connection for SqliteConnection {
             log::error!("did not end {} for {:?}", step, aid);
         }
     }
+
+    async fn collector_remove_step(&self, aid: ArtifactIdNumber, step: &str) {
+        self.raw_ref()
+            .execute(
+                "delete from collector_progress \
+                where aid = ? and step = ?",
+                params![&aid.0, &step],
+            )
+            .unwrap();
+    }
+
     async fn in_progress_artifacts(&self) -> Vec<ArtifactId> {
         let conn = self.raw_ref();
         let mut aids = conn
@@ -1238,4 +1242,13 @@ impl Connection for SqliteConnection {
             .collect::<Result<_, _>>()
             .unwrap()
     }
+
+    async fn purge_artifact(&self, aid: &ArtifactId) {
+        // Once we delete the artifact, all data associated with it should also be deleted
+        // thanks to ON DELETE CASCADE.
+        let info = aid.info();
+        self.raw_ref()
+            .execute("delete from artifact where name = ?1", [info.name])
+            .unwrap();
+    }
 }
