Add --purge=[old/failed] flag to bench_local and bench_runtime_local

Author: Kobzol

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
@@ -629,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)?;
@@ -642,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,
@@ -756,6 +780,7 @@ fn main_result() -> anyhow::Result<i32> {
             bench_rustc,
             iterations,
             self_profile,
+            purge,
         } => {
             log_db(&db);
             let profiles = opts.profiles.0;
@@ -784,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,
@@ -1132,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,

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>;

database/src/pool/postgres.rs

@@ -1134,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()

database/src/pool/sqlite.rs

@@ -1074,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