Merge pull request #2197 from Jamesbarford/feat/queuing-try-commits

Feat; queuing try commits

Author: Jamesbarford

database/src/pool.rs

@@ -197,6 +197,15 @@ pub trait Connection: Send + Sync {
         benchmark_request: &BenchmarkRequest,
         benchmark_request_status: BenchmarkRequestStatus,
     ) -> anyhow::Result<()>;
+
+    /// Update a Try commit to have a `sha` and `parent_sha`. Will update the
+    /// status of the request too a ready state.
+    async fn attach_shas_to_try_benchmark_request(
+        &self,
+        pr: u32,
+        sha: &str,
+        parent_sha: &str,
+    ) -> anyhow::Result<()>;
 }
 
 #[async_trait::async_trait]
@@ -534,4 +543,108 @@ mod tests {
         })
         .await;
     }
+
+    #[tokio::test]
+    async fn updating_try_commits() {
+        run_postgres_test(|ctx| async {
+            let db = ctx.db_client();
+            let db = db.connection().await;
+            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
+            let pr = 42;
+
+            let try_benchmark_request = BenchmarkRequest::create_try(
+                None,
+                None,
+                pr,
+                time,
+                BenchmarkRequestStatus::WaitingForArtifacts,
+                "cranelift",
+                "",
+            );
+            db.insert_benchmark_request(&try_benchmark_request).await;
+            db.attach_shas_to_try_benchmark_request(pr, "foo", "bar")
+                .await
+                .unwrap();
+            let requests = db
+                .get_benchmark_requests_by_status(&[BenchmarkRequestStatus::ArtifactsReady])
+                .await
+                .unwrap();
+
+            assert_eq!(requests.len(), 1);
+            assert_eq!(requests[0].tag(), Some("foo"));
+            assert_eq!(requests[0].parent_sha(), Some("bar"));
+            assert_eq!(requests[0].status, BenchmarkRequestStatus::ArtifactsReady);
+
+            Ok(ctx)
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn adding_try_commit_to_completed_request() {
+        run_postgres_test(|ctx| async {
+            let db = ctx.db_client();
+            let db = db.connection().await;
+            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
+            let pr = 42;
+
+            let completed_try = BenchmarkRequest::create_try(
+                Some("sha-2"),
+                Some("p-sha-1"),
+                pr,
+                time,
+                BenchmarkRequestStatus::Completed,
+                "cranelift",
+                "",
+            );
+            db.insert_benchmark_request(&completed_try).await;
+
+            let try_benchmark_request = BenchmarkRequest::create_try(
+                None,
+                None,
+                pr,
+                time,
+                BenchmarkRequestStatus::WaitingForArtifacts,
+                "cranelift",
+                "",
+            );
+            // deliberately insert twice
+            db.insert_benchmark_request(&try_benchmark_request).await;
+            // this one should fail
+            db.insert_benchmark_request(&try_benchmark_request).await;
+            db.attach_shas_to_try_benchmark_request(pr, "foo", "bar")
+                .await
+                .unwrap();
+
+            let requests = db
+                .get_benchmark_requests_by_status(&[
+                    BenchmarkRequestStatus::WaitingForArtifacts,
+                    BenchmarkRequestStatus::ArtifactsReady,
+                    BenchmarkRequestStatus::InProgress,
+                    BenchmarkRequestStatus::Completed,
+                ])
+                .await
+                .unwrap();
+
+            assert_eq!(requests.len(), 2);
+            let completed_try = requests
+                .iter()
+                .find(|req| req.status == BenchmarkRequestStatus::Completed);
+            assert!(completed_try.is_some());
+            assert_eq!(completed_try.unwrap().pr(), Some(&pr));
+            assert_eq!(completed_try.unwrap().tag(), Some("sha-2"));
+            assert_eq!(completed_try.unwrap().parent_sha(), Some("p-sha-1"));
+
+            let artifacts_ready_try = requests
+                .iter()
+                .find(|req| req.status == BenchmarkRequestStatus::ArtifactsReady);
+            assert!(artifacts_ready_try.is_some());
+            assert_eq!(artifacts_ready_try.unwrap().pr(), Some(&pr));
+            assert_eq!(artifacts_ready_try.unwrap().tag(), Some("foo"));
+            assert_eq!(artifacts_ready_try.unwrap().parent_sha(), Some("bar"));
+
+            Ok(ctx)
+        })
+        .await;
+    }
 }

database/src/pool/postgres.rs

@@ -1525,6 +1525,39 @@ where
 
         Ok(())
     }
+
+    async fn attach_shas_to_try_benchmark_request(
+        &self,
+        pr: u32,
+        sha: &str,
+        parent_sha: &str,
+    ) -> anyhow::Result<()> {
+        self.conn()
+            .execute(
+                "UPDATE
+                benchmark_request
+            SET
+                tag = $1,
+                parent_sha = $2,
+                status = $3
+            WHERE
+                pr = $4
+                AND commit_type = 'try'
+                AND tag IS NULL
+                AND status = $5;",
+                &[
+                    &sha,
+                    &parent_sha,
+                    &BenchmarkRequestStatus::ArtifactsReady,
+                    &(pr as i32),
+                    &BenchmarkRequestStatus::WaitingForArtifacts,
+                ],
+            )
+            .await
+            .context("failed to execute UPDATE benchmark_request")?;
+
+        Ok(())
+    }
 }
 
 fn parse_artifact_id(ty: &str, sha: &str, date: Option<DateTime<Utc>>) -> ArtifactId {

database/src/pool/sqlite.rs

@@ -1282,6 +1282,15 @@ impl Connection for SqliteConnection {
     ) -> anyhow::Result<()> {
         no_queue_implementation_abort!()
     }
+
+    async fn attach_shas_to_try_benchmark_request(
+        &self,
+        _pr: u32,
+        _sha: &str,
+        _parent_sha: &str,
+    ) -> anyhow::Result<()> {
+        no_queue_implementation_abort!()
+    }
 }
 
 fn parse_artifact_id(ty: &str, sha: &str, date: Option<i64>) -> ArtifactId {

site/src/github.rs

@@ -2,6 +2,7 @@ pub mod client;
 pub mod comparison_summary;
 
 use crate::api::github::Commit;
+use crate::job_queue::run_new_queue;
 use crate::load::{MissingReason, SiteCtxt, TryCommit};
 use std::sync::LazyLock;
 use std::time::Duration;
@@ -233,6 +234,22 @@ pub async fn rollup_pr_number(
         .then_some(issue.number))
 }
 
+async fn attach_shas_to_try_benchmark_request(
+    conn: &dyn database::pool::Connection,
+    pr: u32,
+    sha: &str,
+    parent_sha: &str,
+) {
+    if run_new_queue() {
+        if let Err(e) = conn
+            .attach_shas_to_try_benchmark_request(pr, sha, parent_sha)
+            .await
+        {
+            log::error!("Failed to add shas to try commit {}", e);
+        }
+    }
+}
+
 pub async fn enqueue_shas(
     ctxt: &SiteCtxt,
     gh_client: &client::Client,
@@ -258,6 +275,15 @@ pub async fn enqueue_shas(
             parent_sha: commit_response.parents.remove(0).sha,
         };
         let conn = ctxt.conn().await;
+
+        attach_shas_to_try_benchmark_request(
+            &*conn,
+            pr_number,
+            &try_commit.sha,
+            &try_commit.parent_sha,
+        )
+        .await;
+
         let queued = conn
             .pr_attach_commit(
                 pr_number,

site/src/job_queue.rs

@@ -12,6 +12,13 @@ use parking_lot::RwLock;
 use regex::Regex;
 use tokio::time::{self, Duration};
 
+pub fn run_new_queue() -> bool {
+    std::env::var("RUN_CRON")
+        .ok()
+        .and_then(|x| x.parse().ok())
+        .unwrap_or(false)
+}
+
 /// Store the latest master commits or do nothing if all of them are
 /// already in the database
 async fn create_benchmark_request_master_commits(

site/src/main.rs

@@ -1,5 +1,6 @@
 use futures::future::FutureExt;
 use parking_lot::RwLock;
+use site::job_queue::{cron_main, run_new_queue};
 use site::load;
 use std::env;
 use std::sync::Arc;
@@ -33,10 +34,6 @@ async fn main() {
         .ok()
         .and_then(|x| x.parse().ok())
         .unwrap_or(30);
-    let run_cron_job = env::var("RUN_CRON")
-        .ok()
-        .and_then(|x| x.parse().ok())
-        .unwrap_or(false);
 
     let fut = tokio::task::spawn_blocking(move || {
         tokio::task::spawn(async move {
@@ -62,9 +59,9 @@ async fn main() {
 
     let server = site::server::start(ctxt.clone(), port).fuse();
 
-    if run_cron_job {
+    if run_new_queue() {
         task::spawn(async move {
-            site::job_queue::cron_main(ctxt.clone(), queue_update_interval_seconds).await;
+            cron_main(ctxt.clone(), queue_update_interval_seconds).await;
         });
     }
 

site/src/request_handlers/github.rs

@@ -3,8 +3,10 @@ use crate::github::{
     client, enqueue_shas, parse_homu_comment, rollup_pr_number, unroll_rollup,
     COMMENT_MARK_TEMPORARY, RUST_REPO_GITHUB_API_URL,
 };
+use crate::job_queue::run_new_queue;
 use crate::load::SiteCtxt;
 
+use database::{BenchmarkRequest, BenchmarkRequestStatus};
 use hashbrown::HashMap;
 use std::sync::Arc;
 
@@ -72,6 +74,29 @@ async fn handle_issue(
     Ok(github::Response)
 }
 
+/// The try does not have a `sha` or a `parent_sha` but we need to keep a record
+/// of this commit existing. We make sure there can only be one `pr` with a
+/// status of `WaitingForArtifacts` to ensure we don't have duplicates.
+async fn queue_partial_try_benchmark_request(
+    conn: &dyn database::pool::Connection,
+    pr: u32,
+    backends: &str,
+) {
+    // We only want to run this if the new system is running
+    if run_new_queue() {
+        let try_request = BenchmarkRequest::create_try(
+            None,
+            None,
+            pr,
+            chrono::Utc::now(),
+            BenchmarkRequestStatus::WaitingForArtifacts,
+            backends,
+            "",
+        );
+        conn.insert_benchmark_request(&try_request).await;
+    }
+}
+
 async fn handle_rust_timer(
     ctxt: Arc<SiteCtxt>,
     main_client: &client::Client,
@@ -97,6 +122,13 @@ async fn handle_rust_timer(
         let msg = match queue {
             Ok(cmd) => {
                 let conn = ctxt.conn().await;
+
+                queue_partial_try_benchmark_request(
+                    &*conn,
+                    issue.number,
+                    cmd.params.backends.unwrap_or(""),
+                )
+                .await;
                 conn.queue_pr(
                     issue.number,
                     cmd.params.include,
@@ -137,6 +169,12 @@ async fn handle_rust_timer(
     {
         let conn = ctxt.conn().await;
         for command in &valid_build_cmds {
+            queue_partial_try_benchmark_request(
+                &*conn,
+                issue.number,
+                command.params.backends.unwrap_or(""),
+            )
+            .await;
             conn.queue_pr(
                 issue.number,
                 command.params.include,