Return the corresponding compiler artifact ID from dequeue_benchmark_job

Author: Kobzol

database/src/lib.rs

@@ -1120,16 +1120,16 @@ impl BenchmarkJob {
         self.id
     }
 
-    pub fn target(&self) -> &Target {
-        &self.target
+    pub fn target(&self) -> Target {
+        self.target
     }
 
-    pub fn backend(&self) -> &CodegenBackend {
-        &self.backend
+    pub fn backend(&self) -> CodegenBackend {
+        self.backend
     }
 
-    pub fn profile(&self) -> &Profile {
-        &self.profile
+    pub fn profile(&self) -> Profile {
+        self.profile
     }
 
     pub fn request_tag(&self) -> &str {

database/src/pool.rs

@@ -251,12 +251,14 @@ pub trait Connection: Send + Sync {
     ) -> anyhow::Result<Option<CollectorConfig>>;
 
     /// Dequeues a single job for the given collector, target and benchmark set.
+    /// Also returns detailed information about the compiler artifact that should be benchmarked
+    /// in the job.
     async fn dequeue_benchmark_job(
         &self,
         collector_name: &str,
         target: Target,
         benchmark_set: BenchmarkSet,
-    ) -> anyhow::Result<Option<BenchmarkJob>>;
+    ) -> anyhow::Result<Option<(BenchmarkJob, ArtifactId)>>;
 
     /// Try and mark the benchmark_request as completed. Will return `true` if
     /// it has been marked as completed else `false` meaning there was no change
@@ -839,12 +841,10 @@ mod tests {
             let db = ctx.db_client().connection().await;
             let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
 
-            let insert_result = db
+            let collector_config = db
                 .add_collector_config("collector-1", Target::X86_64UnknownLinuxGnu, 1, true)
-                .await;
-            assert!(insert_result.is_ok());
-
-            let collector_config = insert_result.unwrap();
+                .await
+                .unwrap();
 
             let benchmark_request =
                 BenchmarkRequest::create_master("sha-1", "parent-sha-1", 42, time);
@@ -855,32 +855,28 @@ mod tests {
                 .unwrap();
 
             // Now we can insert the job
-            let enqueue_result = db
-                .enqueue_benchmark_job(
-                    benchmark_request.tag().unwrap(),
-                    Target::X86_64UnknownLinuxGnu,
-                    CodegenBackend::Llvm,
-                    Profile::Opt,
-                    1u32,
-                )
-                .await;
-            assert!(enqueue_result.is_ok());
+            db.enqueue_benchmark_job(
+                benchmark_request.tag().unwrap(),
+                Target::X86_64UnknownLinuxGnu,
+                CodegenBackend::Llvm,
+                Profile::Opt,
+                1u32,
+            )
+            .await
+            .unwrap();
 
-            let benchmark_job = db
+            let (benchmark_job, artifact_id) = db
                 .dequeue_benchmark_job(
                     collector_config.name(),
                     collector_config.target(),
                     collector_config.benchmark_set(),
                 )
-                .await;
-            assert!(benchmark_job.is_ok());
-
-            let benchmark_job = benchmark_job.unwrap();
-            assert!(benchmark_job.is_some());
+                .await
+                .unwrap()
+                .unwrap();
 
             // Ensure the properties of the job match both the request and the
             // collector configuration
-            let benchmark_job = benchmark_job.unwrap();
             assert_eq!(
                 benchmark_job.request_tag(),
                 benchmark_request.tag().unwrap()
@@ -894,6 +890,15 @@ mod tests {
                 collector_config.name(),
             );
 
+            assert_eq!(
+                artifact_id,
+                ArtifactId::Commit(Commit {
+                    sha: "sha-1".to_string(),
+                    date: Date(time),
+                    r#type: CommitType::Master,
+                })
+            );
+
             Ok(ctx)
         })
         .await;

database/src/pool/postgres.rs

@@ -1837,7 +1837,7 @@ where
         collector_name: &str,
         target: Target,
         benchmark_set: BenchmarkSet,
-    ) -> anyhow::Result<Option<BenchmarkJob>> {
+    ) -> anyhow::Result<Option<(BenchmarkJob, ArtifactId)>> {
         // We take the oldest job from the job_queue matching the benchmark_set,
         // target and status of 'queued'
         let row_opt = self
@@ -1857,7 +1857,7 @@ where
                         BY created_at
                     LIMIT 1
                     FOR UPDATE SKIP LOCKED
-                ), updated_queue AS (
+                ), updated AS (
                     UPDATE
                         job_queue
                     SET
@@ -1868,19 +1868,20 @@ where
                         picked
                     WHERE
                         job_queue.id = picked.id
-                    RETURNING
-                        job_queue.id,
-                        job_queue.backend,
-                        job_queue.profile,
-                        job_queue.request_tag,
-                        job_queue.created_at,
-                        job_queue.started_at,
-                        job_queue.retry
+                    RETURNING *
                 )
                 SELECT
-                    *
-                FROM
-                    updated_queue;",
+                    updated.id,
+                    updated.backend,
+                    updated.profile,
+                    updated.request_tag,
+                    updated.created_at,
+                    updated.started_at,
+                    updated.retry,
+                    br.commit_type,
+                    br.commit_date
+                FROM updated
+                JOIN benchmark_request as br ON br.tag = updated.request_tag;",
                 &[
                     &BENCHMARK_JOB_STATUS_QUEUED_STR,
                     &target,
@@ -1911,7 +1912,31 @@ where
                     },
                     retry: row.get::<_, i32>(6) as u32,
                 };
-                Ok(Some(job))
+                let commit_type = row.get::<_, &str>(7);
+                let commit_date = row.get::<_, Option<DateTime<Utc>>>(8);
+
+                let commit_date = Date(commit_date.ok_or_else(|| {
+                    anyhow::anyhow!("Dequeuing job for a benchmark request without commit date")
+                })?);
+                let artifact_id = match commit_type {
+                    BENCHMARK_REQUEST_TRY_STR => ArtifactId::Commit(Commit {
+                        sha: job.request_tag.clone(),
+                        date: commit_date,
+                        r#type: CommitType::Try,
+                    }),
+                    BENCHMARK_REQUEST_MASTER_STR => ArtifactId::Commit(Commit {
+                        sha: job.request_tag.clone(),
+                        date: commit_date,
+                        r#type: CommitType::Master,
+                    }),
+                    BENCHMARK_REQUEST_RELEASE_STR => ArtifactId::Tag(job.request_tag.clone()),
+                    _ => panic!(
+                        "Invalid commit type {commit_type} for benchmark request {}",
+                        job.request_tag
+                    ),
+                };
+
+                Ok(Some((job, artifact_id)))
             }
         }
     }

database/src/pool/sqlite.rs

@@ -1344,7 +1344,7 @@ impl Connection for SqliteConnection {
         _collector_name: &str,
         _target: Target,
         _benchmark_set: BenchmarkSet,
-    ) -> anyhow::Result<Option<BenchmarkJob>> {
+    ) -> anyhow::Result<Option<(BenchmarkJob, ArtifactId)>> {
         no_queue_implementation_abort!()
     }