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

Author: Kobzol

collector/src/bin/collector.rs

@@ -469,6 +469,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)]
@@ -490,6 +505,9 @@ enum Commands {
         /// faster.
         #[arg(long = "no-isolate")]
         no_isolate: bool,
+
+        #[command(flatten)]
+        purge: PurgeOption,
     },
     /// Benchmarks a local rustc
     BenchLocal {
@@ -511,6 +529,9 @@ enum Commands {
 
         #[command(flatten)]
         self_profile: SelfProfileOption,
+
+        #[command(flatten)]
+        purge: PurgeOption,
     },
 
     /// Benchmarks the next commit or release for perf.rust-lang.org
@@ -642,6 +663,7 @@ fn main_result() -> anyhow::Result<i32> {
             iterations,
             db,
             no_isolate,
+            purge,
         } => {
             log_db(&db);
             let toolchain = get_local_toolchain(
@@ -666,6 +688,8 @@ fn main_result() -> anyhow::Result<i32> {
                 isolation_mode,
             )?;
             let artifact_id = ArtifactId::Tag(toolchain.id.clone());
+            let mut conn = rt.block_on(pool.connection());
+            rt.block_on(purge_old_data(conn.as_mut(), &artifact_id, purge.purge));
 
             let shared = SharedBenchmarkConfig {
                 artifact_id,
@@ -676,7 +700,7 @@ fn main_result() -> anyhow::Result<i32> {
                 filter: BenchmarkFilter::new(local.exclude, local.include),
                 iterations,
             };
-            let conn = rt.block_on(pool.connection());
+
             run_benchmarks(&mut rt, conn, shared, None, Some(config))?;
             Ok(0)
         }
@@ -687,6 +711,7 @@ fn main_result() -> anyhow::Result<i32> {
             bench_rustc,
             iterations,
             self_profile,
+            purge,
         } => {
             log_db(&db);
             let profiles = opts.profiles.0;
@@ -713,7 +738,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,
@@ -977,6 +1004,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

@@ -149,6 +149,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

@@ -1117,6 +1117,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

@@ -1061,6 +1061,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