Merge pull request #1610 from Kobzol/runtime-site-backend

Generalize site `pstat` querying for runtime benchmarks

Author: Mark-Simulacrum

collector/src/runtime/mod.rs

@@ -26,10 +26,6 @@ pub async fn bench_runtime(
     filter: BenchmarkFilter,
     iterations: u32,
 ) -> anyhow::Result<()> {
-    for benchmark in suite.benchmark_names() {
-        conn.record_runtime_benchmark(benchmark).await;
-    }
-
     let total_benchmark_count = suite.total_benchmark_count();
     let filtered = suite.filtered_benchmark_count(&filter);
     println!(

database/src/bin/import-sqlite.rs

@@ -1,4 +1,4 @@
-use database::{CompileBenchmark, Lookup, Pool};
+use database::{CompileBenchmark, Pool};
 use hashbrown::HashMap;
 use std::collections::HashSet;
 
@@ -45,7 +45,9 @@ async fn main() {
         let sqlite_aid = sqlite_conn.artifact_id(&aid).await;
         let postgres_aid = postgres_conn.artifact_id(&aid).await;
 
-        for &(benchmark, profile, scenario, metric) in sqlite_idx.all_statistic_descriptions() {
+        for (&(benchmark, profile, scenario, metric), id) in
+            sqlite_idx.compile_statistic_descriptions()
+        {
             if benchmarks.insert(benchmark) {
                 postgres_conn
                     .record_compile_benchmark(
@@ -56,15 +58,6 @@ async fn main() {
                     .await;
             }
 
-            let id = database::DbLabel::StatisticDescription {
-                benchmark,
-                profile,
-                scenario,
-                metric,
-            }
-            .lookup(&sqlite_idx)
-            .unwrap();
-
             let stat = sqlite_conn
                 .get_pstats(&[id], &[Some(sqlite_aid)])
                 .await
@@ -86,5 +79,20 @@ async fn main() {
                     .await;
             }
         }
+
+        for (&(benchmark, metric), id) in sqlite_idx.runtime_statistic_descriptions() {
+            let stat = sqlite_conn
+                .get_runtime_pstats(&[id], &[Some(sqlite_aid)])
+                .await
+                .pop()
+                .unwrap()
+                .pop()
+                .unwrap();
+            if let Some(stat) = stat {
+                postgres_conn
+                    .record_runtime_statistic(cid, postgres_aid, &benchmark, metric.as_str(), stat)
+                    .await;
+            }
+        }
     }
 }

database/src/lib.rs

@@ -406,11 +406,11 @@ pub struct Index {
     commits: Indexed<Commit>,
     /// Id lookup of published release artifacts
     artifacts: Indexed<Box<str>>,
-    /// Id lookup of the errors for a crate
-    errors: Indexed<Benchmark>,
-    /// Id lookup of stat description ids
+    /// Id lookup of compile stat description ids
     /// For legacy reasons called `pstat_series` in the database, and so the name is kept here.
     pstat_series: Indexed<(Benchmark, Profile, Scenario, Metric)>,
+    /// Id lookup of runtime stat description ids
+    runtime_pstat_series: Indexed<(Benchmark, Metric)>,
 }
 
 /// An index lookup
@@ -524,9 +524,6 @@ mod index_serde {
 
 #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
 pub enum DbLabel {
-    Errors {
-        benchmark: Benchmark,
-    },
     StatisticDescription {
         benchmark: Benchmark,
         profile: Profile,
@@ -544,7 +541,6 @@ impl Lookup for DbLabel {
     type Id = u32;
     fn lookup(&self, index: &Index) -> Option<Self::Id> {
         match self {
-            DbLabel::Errors { benchmark } => index.errors.get(benchmark),
             DbLabel::StatisticDescription {
                 benchmark,
                 profile,
@@ -567,6 +563,8 @@ impl Lookup for ArtifactId {
     }
 }
 
+pub type StatisticalDescriptionId = u32;
+
 impl Index {
     pub async fn load(conn: &mut dyn pool::Connection) -> Index {
         conn.load_index().await
@@ -607,18 +605,31 @@ impl Index {
             .collect()
     }
 
-    pub fn all_errors(&self) -> impl Iterator<Item = Benchmark> + '_ {
-        self.errors.map.keys().copied()
-    }
-
     // FIXME: in theory this won't scale indefinitely as there's potentially
     // millions of queries and labels and iterating all of them is eventually
     // going to be impractical. But for now it performs quite well, so we'll go
     // for it as keeping indices around would be annoying.
-    pub fn all_statistic_descriptions(
+    pub fn compile_statistic_descriptions(
         &self,
-    ) -> impl Iterator<Item = &'_ (Benchmark, Profile, Scenario, Metric)> + '_ {
-        self.pstat_series.map.keys()
+    ) -> impl Iterator<
+        Item = (
+            &(Benchmark, Profile, Scenario, Metric),
+            StatisticalDescriptionId,
+        ),
+    > + '_ {
+        self.pstat_series
+            .map
+            .iter()
+            .map(|(test_case, &row_id)| (test_case, row_id))
+    }
+
+    pub fn runtime_statistic_descriptions(
+        &self,
+    ) -> impl Iterator<Item = (&(Benchmark, Metric), StatisticalDescriptionId)> + '_ {
+        self.runtime_pstat_series
+            .map
+            .iter()
+            .map(|(test_case, &row_id)| (test_case, row_id))
     }
 
     pub fn artifact_id_for_commit(&self, commit: &str) -> Option<ArtifactId> {

database/src/pool.rs

@@ -26,8 +26,6 @@ pub trait Connection: Send + Sync {
     );
     async fn get_compile_benchmarks(&self) -> Vec<CompileBenchmark>;
 
-    async fn record_runtime_benchmark(&self, name: &str);
-
     async fn artifact_by_name(&self, artifact: &str) -> Option<ArtifactId>;
 
     /// This records the duration of a collection run, i.e., collecting all of
@@ -112,6 +110,11 @@ pub trait Connection: Send + Sync {
         pstat_series_row_ids: &[u32],
         artifact_row_id: &[Option<ArtifactIdNumber>],
     ) -> Vec<Vec<Option<f64>>>;
+    async fn get_runtime_pstats(
+        &self,
+        runtime_pstat_series_row_ids: &[u32],
+        artifact_row_id: &[Option<ArtifactIdNumber>],
+    ) -> Vec<Vec<Option<f64>>>;
     async fn get_error(&self, artifact_row_id: ArtifactIdNumber) -> HashMap<String, String>;
 
     async fn queue_pr(

database/src/pool/postgres.rs

@@ -215,6 +215,21 @@ static MIGRATIONS: &[&str] = &[
     r#"
     alter table pull_request_build add column commit_date timestamptz;
     "#,
+    r#"
+    create table runtime_pstat_series(
+        id integer primary key generated always as identity,
+        benchmark text not null,
+        metric text not null,
+        UNIQUE(benchmark, metric)
+    );
+    create table runtime_pstat(
+        series integer references runtime_pstat_series(id) on delete cascade on update cascade,
+        aid integer references artifact(id) on delete cascade on update cascade,
+        cid integer references collection(id) on delete cascade on update cascade,
+        value double precision not null,
+        PRIMARY KEY(series, aid, cid)
+    );
+"#,
 ];
 
 #[async_trait::async_trait]
@@ -293,6 +308,10 @@ pub struct CachedStatements {
     record_duration: Statement,
     in_progress_steps: Statement,
     get_benchmarks: Statement,
+    insert_runtime_pstat_series: Statement,
+    select_runtime_pstat_series: Statement,
+    insert_runtime_pstat: Statement,
+    get_runtime_pstat: Statement,
 }
 
 pub struct PostgresTransaction<'a> {
@@ -449,7 +468,36 @@ impl PostgresConnection {
                 get_benchmarks: conn.prepare("
                     select name, category
                     from benchmark
-                ").await.unwrap()
+                ").await.unwrap(),
+                insert_runtime_pstat_series: conn.prepare("insert into runtime_pstat_series (benchmark, metric) VALUES ($1, $2) ON CONFLICT DO NOTHING RETURNING id").await.unwrap(),
+                select_runtime_pstat_series: conn.prepare("select id from runtime_pstat_series where benchmark = $1 and metric = $2").await.unwrap(),
+                insert_runtime_pstat: conn
+                    .prepare("insert into runtime_pstat (series, aid, cid, value) VALUES ($1, $2, $3, $4)")
+                    .await
+                    .unwrap(),
+                get_runtime_pstat: conn
+                    .prepare("
+                         WITH aids AS (
+                             select aid, num from unnest($2::int[]) with ordinality aids(aid, num)
+                         ),
+                         sids AS (
+                             select sid, idx from unnest($1::int[]) with ordinality sids(sid, idx)
+                         )
+                         select ARRAY(
+                             (
+                                 select min(runtime_pstat.value) from aids
+                                     left outer join runtime_pstat
+                                     on (aids.aid = runtime_pstat.aid and runtime_pstat.series = sids.sid)
+                                     group by aids.num
+                                     order by aids.num
+                             )
+                         ) from
+                         sids
+                         group by (sids.idx, sids.sid)
+                         order by sids.idx
+                     ")
+                    .await
+                    .unwrap()
             }),
             conn,
         }
@@ -522,23 +570,31 @@ where
                     )
                 })
                 .collect(),
-            errors: self
+            pstat_series: self
                 .conn()
-                .query("select id, crate from error_series", &[])
+                .query(
+                    "select id, crate, profile, cache, statistic from pstat_series;",
+                    &[],
+                )
                 .await
                 .unwrap()
                 .into_iter()
                 .map(|row| {
                     (
                         row.get::<_, i32>(0) as u32,
-                        row.get::<_, String>(1).as_str().into(),
+                        (
+                            Benchmark::from(row.get::<_, String>(1).as_str()),
+                            Profile::from_str(row.get::<_, String>(2).as_str()).unwrap(),
+                            row.get::<_, String>(3).as_str().parse().unwrap(),
+                            row.get::<_, String>(4).as_str().into(),
+                        ),
                     )
                 })
                 .collect(),
-            pstat_series: self
+            runtime_pstat_series: self
                 .conn()
                 .query(
-                    "select id, crate, profile, cache, statistic from pstat_series;",
+                    "select id, benchmark, metric from runtime_pstat_series;",
                     &[],
                 )
                 .await
@@ -548,10 +604,8 @@ where
                     (
                         row.get::<_, i32>(0) as u32,
                         (
-                            Benchmark::from(row.get::<_, String>(1).as_str()),
-                            Profile::from_str(row.get::<_, String>(2).as_str()).unwrap(),
-                            row.get::<_, String>(3).as_str().parse().unwrap(),
-                            row.get::<_, String>(4).as_str().into(),
+                            row.get::<_, String>(1).as_str().into(),
+                            row.get::<_, String>(2).as_str().into(),
                         ),
                     )
                 })
@@ -597,6 +651,31 @@ where
             .map(|row| row.get::<_, Vec<Option<f64>>>(0))
             .collect()
     }
+    async fn get_runtime_pstats(
+        &self,
+        runtime_pstat_series_row_ids: &[u32],
+        artifact_row_ids: &[Option<crate::ArtifactIdNumber>],
+    ) -> Vec<Vec<Option<f64>>> {
+        let runtime_pstat_series_row_ids = runtime_pstat_series_row_ids
+            .iter()
+            .map(|sid| *sid as i32)
+            .collect::<Vec<_>>();
+        let artifact_row_ids = artifact_row_ids
+            .iter()
+            .map(|id| id.map(|id| id.0 as i32))
+            .collect::<Vec<_>>();
+        let rows = self
+            .conn()
+            .query(
+                &self.statements().get_runtime_pstat,
+                &[&runtime_pstat_series_row_ids, &artifact_row_ids],
+            )
+            .await
+            .unwrap();
+        rows.into_iter()
+            .map(|row| row.get::<_, Vec<Option<f64>>>(0))
+            .collect()
+    }
     async fn get_error(&self, artifact_row_id: crate::ArtifactIdNumber) -> HashMap<String, String> {
         let rows = self
             .conn()
@@ -744,13 +823,47 @@ where
     }
     async fn record_runtime_statistic(
         &self,
-        _collection: CollectionId,
-        _artifact: ArtifactIdNumber,
-        _benchmark: &str,
-        _metric: &str,
-        _value: f64,
+        collection: CollectionId,
+        artifact: ArtifactIdNumber,
+        benchmark: &str,
+        metric: &str,
+        value: f64,
     ) {
-        unimplemented!()
+        let sid = self
+            .conn()
+            .query_opt(
+                &self.statements().select_runtime_pstat_series,
+                &[&benchmark, &metric],
+            )
+            .await
+            .unwrap();
+        let sid: i32 = match sid {
+            Some(id) => id.get(0),
+            None => {
+                self.conn()
+                    .query_opt(
+                        &self.statements().insert_runtime_pstat_series,
+                        &[&benchmark, &metric],
+                    )
+                    .await
+                    .unwrap();
+                self.conn()
+                    .query_one(
+                        &self.statements().select_runtime_pstat_series,
+                        &[&benchmark, &metric],
+                    )
+                    .await
+                    .unwrap()
+                    .get(0)
+            }
+        };
+        self.conn()
+            .execute(
+                &self.statements().insert_runtime_pstat,
+                &[&sid, &(artifact.0 as i32), &{ collection.0 }, &value],
+            )
+            .await
+            .unwrap();
     }
 
     async fn record_rustc_crate(
@@ -935,9 +1048,6 @@ where
                 .unwrap();
         }
     }
-    async fn record_runtime_benchmark(&self, _name: &str) {
-        unimplemented!()
-    }
 
     async fn collector_start(&self, aid: ArtifactIdNumber, steps: &[String]) {
         // Clean up -- we'll re-insert any missing things in the loop below.