PR Feedback; split enqueuing job into database into fields, have default_profiles in the database as opposed to the collector

Jamesbarford · Jamesbarford · commit eadb4e243415 · 2025-07-24T22:23:09.000+01:00
diff --git a/collector/src/bin/collector.rs b/collector/src/bin/collector.rs
@@ -1042,7 +1042,12 @@ fn main_result() -> anyhow::Result<i32> {
 
                         let compile_config = CompileBenchmarkConfig {
                             benchmarks,
-                            profiles: Profile::default_profiles(),
+                            profiles: vec![
+                                Profile::Check,
+                                Profile::Debug,
+                                Profile::Doc,
+                                Profile::Opt,
+                            ],
                             scenarios: Scenario::all(),
                             backends,
                             iterations: runs.map(|v| v as usize),
diff --git a/collector/src/compile/benchmark/profile.rs b/collector/src/compile/benchmark/profile.rs
@@ -34,11 +34,6 @@ impl Profile {
         ]
     }
 
-    /// Set of default profiles that should be benchmarked for a master/try artifact.
-    pub fn default_profiles() -> Vec<Self> {
-        vec![Profile::Check, Profile::Debug, Profile::Doc, Profile::Opt]
-    }
-
     pub fn is_doc(&self) -> bool {
         match self {
             Profile::Doc | Profile::DocJson => true,
diff --git a/database/src/lib.rs b/database/src/lib.rs
@@ -1052,9 +1052,14 @@ impl BenchmarkRequestIndex {
 #[derive(Debug, Clone, PartialEq)]
 pub enum BenchmarkJobStatus {
     Queued,
-    InProgress,
-    Success,
-    Failure,
+    InProgress {
+        started_at: DateTime<Utc>,
+    },
+    Completed {
+        started_at: DateTime<Utc>,
+        completed_at: DateTime<Utc>,
+        success: bool,
+    },
 }
 
 const BENCHMARK_JOB_STATUS_QUEUED_STR: &str = "queued";
@@ -1066,9 +1071,14 @@ impl BenchmarkJobStatus {
     pub fn as_str(&self) -> &str {
         match self {
             BenchmarkJobStatus::Queued => BENCHMARK_JOB_STATUS_QUEUED_STR,
-            BenchmarkJobStatus::InProgress => BENCHMARK_JOB_STATUS_IN_PROGRESS_STR,
-            BenchmarkJobStatus::Success => BENCHMARK_JOB_STATUS_SUCCESS_STR,
-            BenchmarkJobStatus::Failure => BENCHMARK_JOB_STATUS_FAILURE_STR,
+            BenchmarkJobStatus::InProgress { .. } => BENCHMARK_JOB_STATUS_IN_PROGRESS_STR,
+            BenchmarkJobStatus::Completed { success, .. } => {
+                if *success {
+                    BENCHMARK_JOB_STATUS_SUCCESS_STR
+                } else {
+                    BENCHMARK_JOB_STATUS_FAILURE_STR
+                }
+            }
         }
     }
 }
@@ -1090,32 +1100,26 @@ pub struct BenchmarkJob {
     request_tag: String,
     benchmark_set: BenchmarkSet,
     created_at: DateTime<Utc>,
-    started_at: Option<DateTime<Utc>>,
-    completed_at: Option<DateTime<Utc>>,
     status: BenchmarkJobStatus,
     retry: u32,
 }
 
 impl BenchmarkJob {
-    pub fn new(
+    pub fn create_queued(
         target: Target,
         backend: CodegenBackend,
         profile: Profile,
         request_tag: &str,
         benchmark_set: u32,
-        created_at: DateTime<Utc>,
-        status: BenchmarkJobStatus,
     ) -> Self {
         BenchmarkJob {
             target,
             backend,
             profile,
+            created_at: Utc::now(),
             request_tag: request_tag.to_string(),
             benchmark_set: BenchmarkSet(benchmark_set),
-            created_at,
-            started_at: None,
-            completed_at: None,
-            status,
+            status: BenchmarkJobStatus::Queued,
             retry: 0,
         }
     }
diff --git a/database/src/pool.rs b/database/src/pool.rs
@@ -1,6 +1,6 @@
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkJob, BenchmarkRequest,
-    BenchmarkRequestIndex, BenchmarkRequestStatus, CodegenBackend, CompileBenchmark, Target,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkRequest, BenchmarkRequestIndex,
+    BenchmarkRequestStatus, CodegenBackend, CompileBenchmark, Target,
 };
 use crate::{CollectionId, Index, Profile, QueuedCommit, Scenario, Step};
 use chrono::{DateTime, Utc};
@@ -212,7 +212,14 @@ pub trait Connection: Send + Sync {
     ) -> anyhow::Result<()>;
 
     /// Add a benchmark job to the job queue.
-    async fn enqueue_benchmark_job(&self, benchmark_job: &BenchmarkJob) -> anyhow::Result<()>;
+    async fn enqueue_benchmark_job(
+        &self,
+        request_tag: &str,
+        target: &Target,
+        backend: &CodegenBackend,
+        profile: &Profile,
+        benchmark_set: u32,
+    ) -> anyhow::Result<()>;
 }
 
 #[async_trait::async_trait]
@@ -338,7 +345,7 @@ mod tests {
     use super::*;
     use crate::{
         tests::{run_db_test, run_postgres_test},
-        BenchmarkRequestStatus, BenchmarkRequestType, Commit, CommitType, Date,
+        BenchmarkJob, BenchmarkRequestStatus, BenchmarkRequestType, Commit, CommitType, Date,
     };
 
     /// Create a Commit
@@ -587,4 +594,43 @@ mod tests {
         })
         .await;
     }
+
+    #[tokio::test]
+    async fn enqueue_benchmark_job() {
+        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 benchmark_request =
+                BenchmarkRequest::create_master("sha-1", "parent-sha-1", 42, time);
+
+            let benchmark_job = BenchmarkJob::create_queued(
+                Target::X86_64UnknownLinuxGnu,
+                CodegenBackend::Llvm,
+                Profile::Opt,
+                benchmark_request.tag().unwrap(),
+                2,
+            );
+
+            // Insert the request so we don't violate the foreign key
+            db.insert_benchmark_request(&benchmark_request)
+                .await
+                .unwrap();
+
+            // Now we can insert the job
+            let result = db
+                .enqueue_benchmark_job(
+                    benchmark_job.request_tag(),
+                    &benchmark_job.target(),
+                    &benchmark_job.backend(),
+                    &benchmark_job.profile(),
+                    0u32,
+                )
+                .await;
+            assert!(result.is_ok());
+
+            Ok(ctx)
+        })
+        .await;
+    }
 }
diff --git a/database/src/pool/postgres.rs b/database/src/pool/postgres.rs
@@ -1,6 +1,6 @@
 use crate::pool::{Connection, ConnectionManager, ManagedConnection, Transaction};
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, Benchmark, BenchmarkJob, BenchmarkJobStatus,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, Benchmark, BenchmarkJobStatus,
     BenchmarkRequest, BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkRequestType,
     CodegenBackend, CollectionId, Commit, CommitType, CompileBenchmark, Date, Index, Profile,
     QueuedCommit, Scenario, Target, BENCHMARK_REQUEST_MASTER_STR, BENCHMARK_REQUEST_RELEASE_STR,
@@ -332,7 +332,7 @@ static MIGRATIONS: &[&str] = &[
         backend       TEXT NOT NULL,
         profile       TEXT NOT NULL,
         benchmark_set INTEGER NOT NULL,
-        collector_id  TEXT,
+        collector_id  INTEGER,
         created_at    TIMESTAMPTZ NOT NULL DEFAULT NOW(),
         started_at    TIMESTAMPTZ,
         completed_at  TIMESTAMPTZ,
@@ -344,7 +344,14 @@ static MIGRATIONS: &[&str] = &[
             REFERENCES benchmark_request(tag)
             ON DELETE CASCADE,
 
-       CONSTRAINT job_queue_unique UNIQUE(request_tag, target, backend, profile)
+       CONSTRAINT job_queue_unique
+       UNIQUE (
+           request_tag,
+           target,
+           backend,
+           profile,
+           benchmark_set
+        )
     );
     CREATE INDEX IF NOT EXISTS job_queue_request_tag_idx ON job_queue (request_tag);
     "#,
@@ -1633,7 +1640,14 @@ where
         Ok(requests)
     }
 
-    async fn enqueue_benchmark_job(&self, benchmark_job: &BenchmarkJob) -> anyhow::Result<()> {
+    async fn enqueue_benchmark_job(
+        &self,
+        request_tag: &str,
+        target: &Target,
+        backend: &CodegenBackend,
+        profile: &Profile,
+        benchmark_set: u32,
+    ) -> anyhow::Result<()> {
         self.conn()
             .execute(
                 r#"
@@ -1646,14 +1660,15 @@ where
                 status
             )
             VALUES ($1, $2, $3, $4, $5, $6)
+            ON CONFLICT DO NOTHING
                 "#,
                 &[
-                    &benchmark_job.request_tag(),
-                    &benchmark_job.target(),
-                    &benchmark_job.backend(),
-                    &benchmark_job.profile(),
-                    &(benchmark_job.benchmark_set() as i32),
-                    &benchmark_job.status(),
+                    &request_tag,
+                    &target,
+                    &backend,
+                    &profile,
+                    &(benchmark_set as i32),
+                    &BenchmarkJobStatus::Queued,
                 ],
             )
             .await
diff --git a/database/src/pool/sqlite.rs b/database/src/pool/sqlite.rs
@@ -1,8 +1,8 @@
 use crate::pool::{Connection, ConnectionManager, ManagedConnection, Transaction};
 use crate::{
-    ArtifactCollection, ArtifactId, Benchmark, BenchmarkJob, BenchmarkRequest,
-    BenchmarkRequestIndex, BenchmarkRequestStatus, CodegenBackend, CollectionId, Commit,
-    CommitType, CompileBenchmark, Date, Profile, Target,
+    ArtifactCollection, ArtifactId, Benchmark, BenchmarkRequest, BenchmarkRequestIndex,
+    BenchmarkRequestStatus, CodegenBackend, CollectionId, Commit, CommitType, CompileBenchmark,
+    Date, Profile, Target,
 };
 use crate::{ArtifactIdNumber, Index, QueuedCommit};
 use chrono::{DateTime, TimeZone, Utc};
@@ -1297,7 +1297,14 @@ impl Connection for SqliteConnection {
         no_queue_implementation_abort!()
     }
 
-    async fn enqueue_benchmark_job(&self, _benchmark_job: &BenchmarkJob) -> anyhow::Result<()> {
+    async fn enqueue_benchmark_job(
+        &self,
+        _request_tag: &str,
+        _target: &Target,
+        _backend: &CodegenBackend,
+        _profile: &Profile,
+        _benchmark_set: u32,
+    ) -> anyhow::Result<()> {
         no_queue_implementation_abort!()
     }
 }
diff --git a/site/src/job_queue/mod.rs b/site/src/job_queue/mod.rs
@@ -7,8 +7,7 @@ use crate::load::{partition_in_place, SiteCtxt};
 use chrono::Utc;
 use collector::benchmark_set::benchmark_set_count;
 use database::{
-    BenchmarkJob, BenchmarkJobStatus, BenchmarkRequest, BenchmarkRequestIndex,
-    BenchmarkRequestStatus, Target,
+    BenchmarkJob, BenchmarkRequest, BenchmarkRequestIndex, BenchmarkRequestStatus, Target,
 };
 use hashbrown::HashSet;
 use parking_lot::RwLock;
@@ -198,14 +197,12 @@ pub fn create_benchmark_jobs(
         ) {
             for backend in backends.iter() {
                 for profile in profiles.iter() {
-                    let job = BenchmarkJob::new(
+                    let job = BenchmarkJob::create_queued(
                         target,
                         *backend,
                         *profile,
                         benchmark_request.tag().unwrap(),
                         benchmark_set as u32,
-                        Utc::now(),
-                        BenchmarkJobStatus::Queued,
                     );
                     jobs.push(job);
                 }
@@ -218,35 +215,46 @@ pub fn create_benchmark_jobs(
 
 /// Enqueue the job into the job_queue
 async fn enqueue_next_job(
-    conn: &dyn database::pool::Connection,
+    conn: &mut dyn database::pool::Connection,
     index: &mut BenchmarkRequestIndex,
 ) -> anyhow::Result<()> {
     let queue = build_queue(conn, index).await?;
+    let mut tx = conn.transaction().await;
     for request in queue {
         if request.status() != BenchmarkRequestStatus::InProgress {
-            for job in create_benchmark_jobs(&request)? {
-                conn.enqueue_benchmark_job(&job).await?;
+            for benchmark_job in create_benchmark_jobs(&request)? {
+                tx.conn()
+                    .enqueue_benchmark_job(
+                        benchmark_job.request_tag(),
+                        &benchmark_job.target(),
+                        &benchmark_job.backend(),
+                        &benchmark_job.profile(),
+                        benchmark_job.benchmark_set(),
+                    )
+                    .await?;
             }
-            conn.update_benchmark_request_status(
-                request.tag().unwrap(),
-                BenchmarkRequestStatus::InProgress,
-            )
-            .await?;
+            tx.conn()
+                .update_benchmark_request_status(
+                    request.tag().unwrap(),
+                    BenchmarkRequestStatus::InProgress,
+                )
+                .await?;
             break;
         }
     }
+    tx.commit().await?;
     Ok(())
 }
 
 /// For queueing jobs, add the jobs you want to queue to this function
 async fn cron_enqueue_jobs(site_ctxt: &Arc<SiteCtxt>) -> anyhow::Result<()> {
-    let conn = site_ctxt.conn().await;
+    let mut conn = site_ctxt.conn().await;
     let mut index = conn.load_benchmark_request_index().await?;
     // Put the master commits into the `benchmark_requests` queue
     create_benchmark_request_master_commits(site_ctxt, &*conn, &index).await?;
     // Put the releases into the `benchmark_requests` queue
     create_benchmark_request_releases(&*conn, &index).await?;
-    enqueue_next_job(&*conn, &mut index).await?;
+    enqueue_next_job(&mut *conn, &mut index).await?;
     Ok(())
 }
 
@@ -455,14 +463,12 @@ mod tests {
         let jobs = create_benchmark_jobs(&request).unwrap();
 
         let create_job = |profile: Profile| -> BenchmarkJob {
-            BenchmarkJob::new(
+            BenchmarkJob::create_queued(
                 Target::X86_64UnknownLinuxGnu,
                 CodegenBackend::Llvm,
                 profile,
                 request.tag().unwrap(),
                 0u32,
-                Utc::now(),
-                BenchmarkJobStatus::Queued,
             )
         };
 

Original file line number	Diff line number	Diff line change
`@@ -34,11 +34,6 @@ impl Profile {`
`34`	`34`	`]`
`35`	`35`	`}`
`36`	`36`
`37`		`- /// Set of default profiles that should be benchmarked for a master/try artifact.`
`38`		`- pub fn default_profiles() -> Vec<Self> {`
`39`		`- vec![Profile::Check, Profile::Debug, Profile::Doc, Profile::Opt]`
`40`		`- }`
`41`		`-`
`42`	`37`	`pub fn is_doc(&self) -> bool {`
`43`	`38`	`match self {`
`44`	`39`	`Profile::Doc \| Profile::DocJson => true,`