Pr Feedback; expose BenchmarkJob id and create BenchmakrJobConclusion, refactored usages and tests

Jamesbarford · Jamesbarford · commit 3074e3614fcd · 2025-08-04T16:22:43.000+01:00
diff --git a/database/src/lib.rs b/database/src/lib.rs
@@ -1104,6 +1104,7 @@ pub struct BenchmarkSet(pub u32);
 /// they are responsible for.
 #[derive(Debug, Clone, PartialEq)]
 pub struct BenchmarkJob {
+    id: u32,
     target: Target,
     backend: CodegenBackend,
     profile: Profile,
@@ -1115,6 +1116,10 @@ pub struct BenchmarkJob {
 }
 
 impl BenchmarkJob {
+    pub fn id(&self) -> u32 {
+        self.id
+    }
+
     pub fn target(&self) -> &Target {
         &self.target
     }
@@ -1144,6 +1149,22 @@ impl BenchmarkJob {
     }
 }
 
+/// Describes the final state of a job
+#[derive(Debug, Clone, PartialEq)]
+pub enum BenchmarkJobConclusion {
+    Failure,
+    Success,
+}
+
+impl BenchmarkJobConclusion {
+    pub fn as_str(&self) -> &str {
+        match self {
+            BenchmarkJobConclusion::Failure => BENCHMARK_JOB_STATUS_FAILURE_STR,
+            BenchmarkJobConclusion::Success => BENCHMARK_JOB_STATUS_SUCCESS_STR,
+        }
+    }
+}
+
 /// The configuration for a collector
 #[derive(Debug, PartialEq)]
 pub struct CollectorConfig {
diff --git a/database/src/pool.rs b/database/src/pool.rs
@@ -1,6 +1,6 @@
 use crate::selector::CompileTestCase;
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkJob, BenchmarkJobStatus,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkJob, BenchmarkJobConclusion,
     BenchmarkRequest, BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkSet, CodegenBackend,
     CollectorConfig, CompileBenchmark, Target,
 };
@@ -262,10 +262,8 @@ pub trait Connection: Send + Sync {
     /// depending on the enum's completed state being a success
     async fn mark_benchmark_job_as_completed(
         &self,
-        request_tag: &str,
-        benchmark_set: u32,
-        target: &Target,
-        status: &BenchmarkJobStatus,
+        id: u32,
+        benchmark_job_conculsion: &BenchmarkJobConclusion,
     ) -> anyhow::Result<()>;
 }
 
@@ -409,7 +407,6 @@ mod tests {
         benchmark_set: u32,
         target: &Target,
     ) {
-        let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
         /* Create job for the request */
         db.enqueue_benchmark_job(
             request_tag,
@@ -430,19 +427,9 @@ mod tests {
         assert_eq!(job.request_tag(), request_tag);
 
         /* Mark the job as complete */
-        db.mark_benchmark_job_as_completed(
-            request_tag,
-            benchmark_set,
-            &target,
-            &BenchmarkJobStatus::Completed {
-                started_at: time,
-                completed_at: time,
-                collector_name: collector_name.into(),
-                success: true,
-            },
-        )
-        .await
-        .unwrap();
+        db.mark_benchmark_job_as_completed(job.id(), &BenchmarkJobConclusion::Success)
+            .await
+            .unwrap();
 
         db.mark_benchmark_request_as_completed(request_tag)
             .await
@@ -968,19 +955,9 @@ mod tests {
             assert_eq!(job.request_tag(), benchmark_request.tag().unwrap());
 
             /* Mark the job as complete */
-            db.mark_benchmark_job_as_completed(
-                tag,
-                benchmark_set.0,
-                &target,
-                &BenchmarkJobStatus::Completed {
-                    started_at: time,
-                    completed_at: time,
-                    collector_name: collector_name.into(),
-                    success: false,
-                },
-            )
-            .await
-            .unwrap();
+            db.mark_benchmark_job_as_completed(job.id(), &BenchmarkJobConclusion::Success)
+                .await
+                .unwrap();
 
             db.mark_benchmark_request_as_completed(tag).await.unwrap();
 
diff --git a/database/src/pool/postgres.rs b/database/src/pool/postgres.rs
@@ -1,11 +1,11 @@
 use crate::pool::{Connection, ConnectionManager, ManagedConnection, Transaction};
 use crate::selector::CompileTestCase;
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, Benchmark, BenchmarkJob, BenchmarkJobStatus,
-    BenchmarkRequest, BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkRequestType,
-    BenchmarkSet, CodegenBackend, CollectionId, CollectorConfig, Commit, CommitType,
-    CompileBenchmark, Date, Index, Profile, QueuedCommit, Scenario, Target,
-    BENCHMARK_JOB_STATUS_FAILURE_STR, BENCHMARK_JOB_STATUS_IN_PROGRESS_STR,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, Benchmark, BenchmarkJob,
+    BenchmarkJobConclusion, BenchmarkJobStatus, BenchmarkRequest, BenchmarkRequestIndex,
+    BenchmarkRequestStatus, BenchmarkRequestType, BenchmarkSet, CodegenBackend, CollectionId,
+    CollectorConfig, Commit, CommitType, CompileBenchmark, Date, Index, Profile, QueuedCommit,
+    Scenario, Target, BENCHMARK_JOB_STATUS_FAILURE_STR, BENCHMARK_JOB_STATUS_IN_PROGRESS_STR,
     BENCHMARK_JOB_STATUS_QUEUED_STR, BENCHMARK_JOB_STATUS_SUCCESS_STR,
     BENCHMARK_REQUEST_MASTER_STR, BENCHMARK_REQUEST_RELEASE_STR,
     BENCHMARK_REQUEST_STATUS_ARTIFACTS_READY_STR, BENCHMARK_REQUEST_STATUS_COMPLETED_STR,
@@ -1849,6 +1849,7 @@ where
                     WHERE
                         job_queue.id = picked.id
                     RETURNING
+                        job_queue.id,
                         job_queue.backend,
                         job_queue.profile,
                         job_queue.request_tag,
@@ -1874,20 +1875,21 @@ where
             None => Ok(None),
             Some(row) => {
                 let job = BenchmarkJob {
+                    id: row.get::<_, i32>(0) as u32,
                     target: *target,
-                    backend: CodegenBackend::from_str(&row.get::<_, String>(0))
+                    backend: CodegenBackend::from_str(&row.get::<_, String>(1))
                         .map_err(|e| anyhow::anyhow!(e))?,
-                    profile: Profile::from_str(&row.get::<_, String>(1))
+                    profile: Profile::from_str(&row.get::<_, String>(2))
                         .map_err(|e| anyhow::anyhow!(e))?,
-                    request_tag: row.get::<_, String>(2),
+                    request_tag: row.get::<_, String>(3),
                     benchmark_set: benchmark_set.clone(),
-                    created_at: row.get::<_, DateTime<Utc>>(3),
+                    created_at: row.get::<_, DateTime<Utc>>(4),
                     // The job is now in an in_progress state
                     status: BenchmarkJobStatus::InProgress {
-                        started_at: row.get::<_, DateTime<Utc>>(4),
+                        started_at: row.get::<_, DateTime<Utc>>(5),
                         collector_name: collector_name.into(),
                     },
-                    retry: row.get::<_, i32>(5) as u32,
+                    retry: row.get::<_, i32>(6) as u32,
                 };
                 Ok(Some(job))
             }
@@ -1949,40 +1951,24 @@ where
 
     async fn mark_benchmark_job_as_completed(
         &self,
-        request_tag: &str,
-        benchmark_set: u32,
-        target: &Target,
-        status: &BenchmarkJobStatus,
+        id: u32,
+        benchmark_job_conclusion: &BenchmarkJobConclusion,
     ) -> anyhow::Result<()> {
-        match status {
-            BenchmarkJobStatus::Queued | BenchmarkJobStatus::InProgress { .. } => {
-                panic!("Can only mark a job as complete")
-            }
-            BenchmarkJobStatus::Completed { success, .. } => {
-                let status = if *success {
-                    BENCHMARK_JOB_STATUS_SUCCESS_STR
-                } else {
-                    BENCHMARK_JOB_STATUS_FAILURE_STR
-                };
-                self.conn()
-                    .execute(
-                        "
-                        UPDATE
-                            job_queue 
-                        SET
-                            status = $1,
-                            completed_at = NOW()
-                        WHERE
-                            request_tag = $2
-                            AND benchmark_set = $3
-                            AND target = $4",
-                        &[&status, &request_tag, &(benchmark_set as i32), &target],
-                    )
-                    .await
-                    .context("Failed to mark benchmark job as completed")?;
-                Ok(())
-            }
-        }
+        self.conn()
+            .execute(
+                "
+                UPDATE
+                    job_queue 
+                SET
+                    status = $1,
+                    completed_at = NOW()
+                WHERE
+                    id = $2",
+                &[&benchmark_job_conclusion.as_str(), &(id as i32)],
+            )
+            .await
+            .context("Failed to mark benchmark job as completed")?;
+        Ok(())
     }
 }
 
diff --git a/database/src/pool/sqlite.rs b/database/src/pool/sqlite.rs
@@ -1,9 +1,9 @@
 use crate::pool::{Connection, ConnectionManager, ManagedConnection, Transaction};
 use crate::selector::CompileTestCase;
 use crate::{
-    ArtifactCollection, ArtifactId, Benchmark, BenchmarkJob, BenchmarkJobStatus, BenchmarkRequest,
-    BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkSet, CodegenBackend, CollectionId,
-    CollectorConfig, Commit, CommitType, CompileBenchmark, Date, Profile, Target,
+    ArtifactCollection, ArtifactId, Benchmark, BenchmarkJob, BenchmarkJobConclusion,
+    BenchmarkRequest, BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkSet, CodegenBackend,
+    CollectionId, CollectorConfig, Commit, CommitType, CompileBenchmark, Date, Profile, Target,
 };
 use crate::{ArtifactIdNumber, Index, QueuedCommit};
 use chrono::{DateTime, TimeZone, Utc};
@@ -1360,10 +1360,8 @@ impl Connection for SqliteConnection {
 
     async fn mark_benchmark_job_as_completed(
         &self,
-        _request_tag: &str,
-        _benchmark_set: u32,
-        _target: &Target,
-        _status: &BenchmarkJobStatus,
+        _id: u32,
+        _benchmark_job_conculsion: &BenchmarkJobConclusion,
     ) -> anyhow::Result<()> {
         no_queue_implementation_abort!()
     }
diff --git a/site/src/job_queue/mod.rs b/site/src/job_queue/mod.rs
@@ -308,7 +308,7 @@ mod tests {
     use super::*;
     use chrono::{Datelike, Duration, TimeZone, Utc};
     use database::{
-        tests::run_postgres_test, BenchmarkJobStatus, BenchmarkSet, CodegenBackend, Profile,
+        tests::run_postgres_test, BenchmarkJobConclusion, BenchmarkSet, CodegenBackend, Profile,
     };
 
     fn days_ago(day_str: &str) -> chrono::DateTime<Utc> {
@@ -360,7 +360,6 @@ mod tests {
         benchmark_set: u32,
         target: &Target,
     ) {
-        let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
         /* Create job for the request */
         db.enqueue_benchmark_job(
             request_tag,
@@ -381,19 +380,9 @@ mod tests {
         assert_eq!(job.request_tag(), request_tag);
 
         /* Mark the job as complete */
-        db.mark_benchmark_job_as_completed(
-            request_tag,
-            benchmark_set,
-            &target,
-            &BenchmarkJobStatus::Completed {
-                started_at: time,
-                completed_at: time,
-                collector_name: collector_name.into(),
-                success: true,
-            },
-        )
-        .await
-        .unwrap();
+        db.mark_benchmark_job_as_completed(job.id(), &BenchmarkJobConclusion::Success)
+            .await
+            .unwrap();
 
         db.mark_benchmark_request_as_completed(request_tag)
             .await
