Merge pull request #706 from Mark-Simulacrum/better-status

Progress reporting data collection

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,15 @@ fn bench(
     iterations: usize,
     self_profile: bool,
 ) -> BenchmarkErrors {
+    let mut conn = rt.block_on(pool.connection());
+    let status_conn;
+    let status_conn: Option<&dyn database::Connection> =
+        if conn.separate_transaction_for_collector() {
+            status_conn = rt.block_on(pool.connection());
+            Some(&*status_conn)
+        } else {
+            None
+        };
     let mut errors = BenchmarkErrors::new();
     eprintln!("Benchmarking {} for triple {}", cid, compiler.triple);
 
@@ -223,7 +232,21 @@ 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
+            .unwrap_or_else(|| &*tx.conn())
+            .collector_start(interned_cid, &steps),
+    );
     for (nth_benchmark, benchmark) in benchmarks.iter().enumerate() {
+        rt.block_on(
+            status_conn
+                .unwrap_or_else(|| &*tx.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 +277,11 @@ fn bench(
                 &format!("{:?}", s),
             ));
         };
+        rt.block_on(
+            status_conn
+                .unwrap_or_else(|| &*tx.conn())
+                .collector_end_step(interned_cid, &benchmark.name.to_string()),
+        );
     }
     let end = start.elapsed();
 
@@ -516,15 +544,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 +594,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 +602,7 @@ fn main_result() -> anyhow::Result<i32> {
 
             let res = bench(
                 &mut rt,
-                conn,
+                pool,
                 &ArtifactId::Commit(commit),
                 &BuildKind::all(),
                 &RunKind::all(),
@@ -608,7 +634,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 +672,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/bin/ingest-json.rs

@@ -600,13 +600,6 @@ async fn main() {
         Pool::Postgres(mut p) => {
             postgres = Some(p.raw().open().await.into());
         }
-        Pool::Both {
-            sqlite: mut s,
-            postgres: mut p,
-        } => {
-            sqlite = Some(s.raw().open().await.into_inner().unwrap());
-            postgres = Some(p.raw().open().await.into());
-        }
     }
 
     if let Some(s) = &mut sqlite {

database/src/pool.rs

@@ -5,7 +5,6 @@ use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 
-pub mod both;
 pub mod postgres;
 pub mod sqlite;
 
@@ -62,6 +61,24 @@ 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);
+
+    // This returns `true` if the collector commands can be placed in a separate
+    // transaction.
+    //
+    // Currently, the sqlite backend does not support "regular" usage where they
+    // are used for genuine progress reporting. sqlite does not support
+    // concurrent writers -- it will return an error (or wait, if a busy timeout
+    // is configured).
+    //
+    // For now we don't care much as sqlite is not used in production and in
+    // local usage you can just look at the logs.
+    fn separate_transaction_for_collector(&self) -> bool;
 }
 
 #[async_trait::async_trait]
@@ -160,43 +177,19 @@ where
 pub enum Pool {
     Sqlite(ConnectionPool<sqlite::Sqlite>),
     Postgres(ConnectionPool<postgres::Postgres>),
-    Both {
-        sqlite: ConnectionPool<sqlite::Sqlite>,
-        postgres: ConnectionPool<postgres::Postgres>,
-    },
 }
 
 impl Pool {
     pub async fn connection(&self) -> Box<dyn Connection> {
         match self {
             Pool::Sqlite(p) => Box::new(sqlite::SqliteConnection::new(p.get().await)),
             Pool::Postgres(p) => Box::new(p.get().await),
-            Pool::Both { sqlite, postgres } => Box::new(both::BothConnection::new(
-                sqlite::SqliteConnection::new(sqlite.get().await),
-                postgres.get().await,
-            )),
         }
     }
 
     pub fn open(uri: &str) -> Pool {
         if uri.starts_with("postgres") {
             Pool::Postgres(ConnectionPool::new(postgres::Postgres::new(uri.into())))
-        } else if uri.starts_with("both://") {
-            let mut parts = uri["both://".len()..].rsplitn(2, ';');
-            let p1 = parts.next().unwrap();
-            let p2 = parts.next().unwrap();
-            match (Pool::open(p1), Pool::open(p2)) {
-                (Pool::Sqlite(s), Pool::Postgres(p)) | (Pool::Postgres(p), Pool::Sqlite(s)) => {
-                    Pool::Both {
-                        sqlite: s,
-                        postgres: p,
-                    }
-                }
-                _ => panic!(
-                    "unsupported inputs, must be sqlite and postgres: {} and {}",
-                    p1, p2
-                ),
-            }
         } else {
             Pool::Sqlite(ConnectionPool::new(sqlite::Sqlite::new(uri.into())))
         }