Record incremental progress into the database when collecting

Author: Mark-Simulacrum

collector/src/main.rs

@@ -4,7 +4,7 @@
 extern crate clap;
 
 use anyhow::{bail, Context};
-use database::{pool::Connection, ArtifactId, Commit};
+use database::{ArtifactId, Commit};
 use log::debug;
 use std::collections::HashSet;
 use std::fs;
@@ -185,7 +185,7 @@ impl BenchmarkErrors {
 
 fn bench(
     rt: &mut Runtime,
-    mut conn: Box<dyn Connection>,
+    pool: database::Pool,
     cid: &ArtifactId,
     build_kinds: &[BuildKind],
     run_kinds: &[RunKind],
@@ -194,6 +194,8 @@ fn bench(
     iterations: usize,
     self_profile: bool,
 ) -> BenchmarkErrors {
+    let status_conn = rt.block_on(pool.connection());
+    let mut conn = rt.block_on(pool.connection());
     let mut errors = BenchmarkErrors::new();
     eprintln!("Benchmarking {} for triple {}", cid, compiler.triple);
 
@@ -223,7 +225,13 @@ fn bench(
     let interned_cid = rt.block_on(tx.conn().artifact_id(&cid));
 
     let start = Instant::now();
+    let steps = benchmarks
+        .iter()
+        .map(|b| b.name.to_string())
+        .collect::<Vec<_>>();
+    rt.block_on(status_conn.collector_start(interned_cid, &steps));
     for (nth_benchmark, benchmark) in benchmarks.iter().enumerate() {
+        rt.block_on(status_conn.collector_start_step(interned_cid, &benchmark.name.to_string()));
         rt.block_on(
             tx.conn()
                 .record_benchmark(benchmark.name.0.as_str(), benchmark.supports_stable()),
@@ -254,6 +262,7 @@ fn bench(
                 &format!("{:?}", s),
             ));
         };
+        rt.block_on(status_conn.collector_end_step(interned_cid, &benchmark.name.to_string()));
     }
     let end = start.elapsed();
 
@@ -516,15 +525,14 @@ fn main_result() -> anyhow::Result<i32> {
             let self_profile = sub_m.is_present("SELF_PROFILE");
 
             let pool = database::Pool::open(db);
-            let conn = rt.block_on(pool.connection());
 
             let (rustc, rustdoc, cargo) = get_local_toolchain(&build_kinds, rustc, rustdoc, cargo)?;
 
             let benchmarks = get_benchmarks(&benchmark_dir, include, exclude)?;
 
             let res = bench(
                 &mut rt,
-                conn,
+                pool,
                 &ArtifactId::Artifact(id.to_string()),
                 &build_kinds,
                 &run_kinds,
@@ -567,7 +575,6 @@ fn main_result() -> anyhow::Result<i32> {
             let commit = get_commit_or_fake_it(&commit)?;
 
             let pool = database::Pool::open(db);
-            let conn = rt.block_on(pool.connection());
 
             let sysroot = Sysroot::install(commit.sha.to_string(), "x86_64-unknown-linux-gnu")
                 .with_context(|| format!("failed to install sysroot for {:?}", commit))?;
@@ -576,7 +583,7 @@ fn main_result() -> anyhow::Result<i32> {
 
             let res = bench(
                 &mut rt,
-                conn,
+                pool,
                 &ArtifactId::Commit(commit),
                 &BuildKind::all(),
                 &RunKind::all(),
@@ -608,7 +615,6 @@ fn main_result() -> anyhow::Result<i32> {
             }
 
             let pool = database::Pool::open(db);
-            let conn = rt.block_on(pool.connection());
 
             let run_kinds = if collector::version_supports_incremental(toolchain) {
                 RunKind::all()
@@ -647,7 +653,7 @@ fn main_result() -> anyhow::Result<i32> {
 
             let res = bench(
                 &mut rt,
-                conn,
+                pool,
                 &ArtifactId::Artifact(toolchain.to_string()),
                 &build_kinds,
                 &run_kinds,

database/src/pool.rs

@@ -61,6 +61,12 @@ pub trait Connection: Send + Sync {
     async fn pr_attach_commit(&self, pr: u32, sha: &str, parent_sha: &str) -> bool;
     async fn queued_commits(&self) -> Vec<QueuedCommit>;
     async fn mark_complete(&self, sha: &str) -> Option<QueuedCommit>;
+
+    // Collector status API
+
+    async fn collector_start(&self, aid: ArtifactIdNumber, steps: &[String]);
+    async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str);
+    async fn collector_end_step(&self, aid: ArtifactIdNumber, step: &str);
 }
 
 #[async_trait::async_trait]

database/src/pool/postgres.rs

@@ -159,6 +159,14 @@ static MIGRATIONS: &[&str] = &[
         duration integer not null
     );
     "#,
+    r#"
+    create table collector_progress(
+        aid smallint not null references artifact(id) on delete cascade on update cascade,
+        step text not null,
+        start timestamptz,
+        end timestamptz
+    );
+    "#,
 ];
 
 #[async_trait::async_trait]
@@ -771,4 +779,44 @@ where
             .await
             .unwrap();
     }
+
+    async fn collector_start(&self, aid: ArtifactIdNumber, steps: &[String]) {
+        self.conn()
+            .execute(
+                "delete from collector_progress where start is null or end is null;",
+                &[],
+            )
+            .await
+            .unwrap();
+
+        for step in steps {
+            self.conn()
+                .execute(
+                    "insert into collector_progress(aid, step) VALUES ($1, $2)",
+                    &[&(aid.0 as i16), &step],
+                )
+                .await
+                .unwrap();
+        }
+    }
+    async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str) {
+        self.conn()
+            .execute(
+                "update collector_progress set start = CURRENT_TIMESTAMP \
+                where aid = $1 and step = $2 and start is null and end is null;",
+                &[&(aid.0 as i16), &step],
+            )
+            .await
+            .unwrap();
+    }
+    async fn collector_end_step(&self, aid: ArtifactIdNumber, step: &str) {
+        self.conn()
+            .execute(
+                "update collector_progress set end = CURRENT_TIMESTAMP \
+                where aid = $1 and step = $2 and start is not null and end is null;",
+                &[&(aid.0 as i16), &step],
+            )
+            .await
+            .unwrap();
+    }
 }

database/src/pool/sqlite.rs

@@ -140,6 +140,14 @@ static MIGRATIONS: &[&str] = &[
         duration integer not null
     );
     "#,
+    r#"
+    create table collector_progress(
+        aid integer not null references artifact(id) on delete cascade on update cascade,
+        step text not null,
+        start timestamp without time zone,
+        end timestamp without time zone
+    );
+    "#,
 ];
 
 #[async_trait::async_trait]
@@ -626,4 +634,38 @@ impl Connection for SqliteConnection {
             )
             .unwrap();
     }
+    async fn collector_start(&self, aid: ArtifactIdNumber, steps: &[String]) {
+        // Clean out any leftover unterminated steps.
+        self.raw_ref()
+            .execute_batch("delete from collector_progress where start is null or end is null;")
+            .unwrap();
+
+        // Populate unstarted and unfinished steps into collector_progress.
+        for step in steps {
+            self.raw_ref()
+                .execute(
+                    "insert into collector_progress(aid, step) VALUES (?, ?)",
+                    params![&aid.0, step],
+                )
+                .unwrap();
+        }
+    }
+    async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str) {
+        self.raw_ref()
+            .execute(
+                "update collector_progress set start = now \
+            where aid = ? and step = ? and start is null;",
+                params![&aid.0, &step],
+            )
+            .unwrap();
+    }
+    async fn collector_end_step(&self, aid: ArtifactIdNumber, step: &str) {
+        self.raw_ref()
+            .execute(
+                "update collector_progress set end = now \
+            where aid = ? and step = ? and start is not null and end is null;",
+                params![&aid.0, &step],
+            )
+            .unwrap();
+    }
 }