Insert try without shas and update try with sha's

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,40 @@ 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;
+    }
 }

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/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,