Change error handling for creating parent jobs

Author: Kobzol

database/src/pool.rs

@@ -15,6 +15,14 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 pub mod postgres;
 pub mod sqlite;
 
+#[derive(Debug)]
+pub enum JobEnqueueResult {
+    JobCreated(u32),
+    JobAlreadyExisted,
+    RequestShaNotFound { error: String },
+    Other(anyhow::Error),
+}
+
 #[async_trait::async_trait]
 pub trait Connection: Send + Sync {
     async fn maybe_create_indices(&mut self);
@@ -238,20 +246,7 @@ pub trait Connection: Send + Sync {
         profile: Profile,
         benchmark_set: u32,
         kind: BenchmarkJobKind,
-    ) -> anyhow::Result<Option<u32>>;
-
-    /// Add a benchmark job which is explicitly using a `parent_sha` we split
-    /// this out to improve our error handling. A `parent_sha` may not have
-    /// an associated request in the `benchmarek`
-    async fn enqueue_parent_benchmark_job(
-        &self,
-        parent_sha: &str,
-        target: Target,
-        backend: CodegenBackend,
-        profile: Profile,
-        benchmark_set: u32,
-        kind: BenchmarkJobKind,
-    ) -> (bool, anyhow::Result<u32>);
+    ) -> JobEnqueueResult;
 
     /// Returns a set of compile-time benchmark test cases that were already computed for the
     /// given artifact.
@@ -460,6 +455,15 @@ mod tests {
         }
     }
 
+    impl JobEnqueueResult {
+        pub fn unwrap(self) -> u32 {
+            match self {
+                JobEnqueueResult::JobCreated(id) => id,
+                error => panic!("Unexpected job enqueue result: {error:?}"),
+            }
+        }
+    }
+
     #[tokio::test]
     async fn pstat_returns_empty_vector_when_empty() {
         run_db_test(|ctx| async {
@@ -715,7 +719,10 @@ mod tests {
                     BenchmarkJobKind::Runtime,
                 )
                 .await;
-            assert!(result.is_ok());
+            match result {
+                JobEnqueueResult::JobCreated(_) => {}
+                error => panic!("Invalid result: {error:?}"),
+            }
 
             Ok(ctx)
         })
@@ -861,16 +868,20 @@ mod tests {
                 .unwrap();
 
             // Now we can insert the job
-            db.enqueue_benchmark_job(
-                benchmark_request.tag().unwrap(),
-                Target::X86_64UnknownLinuxGnu,
-                CodegenBackend::Llvm,
-                Profile::Opt,
-                1u32,
-                BenchmarkJobKind::Runtime,
-            )
-            .await
-            .unwrap();
+            match db
+                .enqueue_benchmark_job(
+                    benchmark_request.tag().unwrap(),
+                    Target::X86_64UnknownLinuxGnu,
+                    CodegenBackend::Llvm,
+                    Profile::Opt,
+                    1u32,
+                    BenchmarkJobKind::Runtime,
+                )
+                .await
+            {
+                JobEnqueueResult::JobCreated(_) => {}
+                error => panic!("Invalid result: {error:?}"),
+            };
 
             let (benchmark_job, artifact_id) = db
                 .dequeue_benchmark_job(
@@ -1205,29 +1216,6 @@ mod tests {
         .await;
     }
 
-    #[tokio::test]
-    async fn enqueue_parent_benchmark_job() {
-        run_postgres_test(|ctx| async {
-            let db = ctx.db();
-
-            let (violates_foreign_key, _) = db
-                .enqueue_parent_benchmark_job(
-                    "sha-0",
-                    Target::X86_64UnknownLinuxGnu,
-                    CodegenBackend::Llvm,
-                    Profile::Debug,
-                    0,
-                    BenchmarkJobKind::Runtime,
-                )
-                .await;
-
-            assert!(violates_foreign_key);
-
-            Ok(ctx)
-        })
-        .await;
-    }
-
     #[tokio::test]
     async fn purge_artifact() {
         run_postgres_test(|ctx| async {
@@ -1244,7 +1232,6 @@ mod tests {
                 BenchmarkJobKind::Compiletime,
             )
             .await
-            .unwrap()
             .unwrap();
             db.purge_artifact(&ArtifactId::Tag("foo".to_string())).await;
 

database/src/pool/postgres.rs

@@ -1,4 +1,6 @@
-use crate::pool::{Connection, ConnectionManager, ManagedConnection, Transaction};
+use crate::pool::{
+    Connection, ConnectionManager, JobEnqueueResult, ManagedConnection, Transaction,
+};
 use crate::selector::CompileTestCase;
 use crate::{
     ArtifactCollection, ArtifactId, ArtifactIdNumber, Benchmark, BenchmarkJob,
@@ -1773,18 +1775,19 @@ where
         })
     }
 
-    async fn enqueue_parent_benchmark_job(
+    async fn enqueue_benchmark_job(
         &self,
-        parent_sha: &str,
+        request_tag: &str,
         target: Target,
         backend: CodegenBackend,
         profile: Profile,
         benchmark_set: u32,
         kind: BenchmarkJobKind,
-    ) -> (bool, anyhow::Result<u32>) {
-        let row_result = self
+    ) -> JobEnqueueResult {
+        // This will return zero rows if the job already exists
+        let result = self
             .conn()
-            .query_one(
+            .query(
                 r#"
             INSERT INTO job_queue(
                 request_tag,
@@ -1800,7 +1803,7 @@ where
             RETURNING job_queue.id
                 "#,
                 &[
-                    &parent_sha,
+                    &request_tag,
                     &target,
                     &backend,
                     &profile,
@@ -1811,75 +1814,28 @@ where
             )
             .await;
 
-        match row_result {
-            Ok(row) => (false, Ok(row.get::<_, i32>(0) as u32)),
+        match result {
+            Ok(rows) => {
+                let Some(row) = rows.into_iter().next() else {
+                    return JobEnqueueResult::JobAlreadyExisted;
+                };
+                JobEnqueueResult::JobCreated(row.get::<_, i32>(0) as u32)
+            }
             Err(e) => {
                 if let Some(db_err) = e.as_db_error() {
                     if db_err.code() == &SqlState::FOREIGN_KEY_VIOLATION {
                         let constraint = db_err.constraint().unwrap_or("benchmark_tag constraint");
                         let detail = db_err.detail().unwrap_or("");
-                        return (
-                            true,
-                            Err(anyhow::anyhow!(
-                                "Foreign key violation on {} for request_tag='{}'. {}",
-                                constraint,
-                                parent_sha,
-                                detail
-                            )),
-                        );
+                        return JobEnqueueResult::RequestShaNotFound {
+                            error: format!("Foreign key violation on `{constraint}`: {detail}"),
+                        };
                     }
                 }
-                (false, Err(e.into()))
+                JobEnqueueResult::Other(e.into())
             }
         }
     }
 
-    async fn enqueue_benchmark_job(
-        &self,
-        request_tag: &str,
-        target: Target,
-        backend: CodegenBackend,
-        profile: Profile,
-        benchmark_set: u32,
-        kind: BenchmarkJobKind,
-    ) -> anyhow::Result<Option<u32>> {
-        // This will return zero rows if the job already exists
-        let rows = self
-            .conn()
-            .query(
-                r#"
-            INSERT INTO job_queue(
-                request_tag,
-                target,
-                backend,
-                profile,
-                benchmark_set,
-                status,
-                kind
-            )
-            VALUES ($1, $2, $3, $4, $5, $6, $7)
-            ON CONFLICT DO NOTHING
-            RETURNING job_queue.id
-                "#,
-                &[
-                    &request_tag,
-                    &target,
-                    &backend,
-                    &profile,
-                    &(benchmark_set as i32),
-                    &BENCHMARK_JOB_STATUS_QUEUED_STR,
-                    &kind,
-                ],
-            )
-            .await
-            .context("failed to insert benchmark_job")?;
-        if let Some(row) = rows.first() {
-            Ok(Some(row.get::<_, i32>(0) as u32))
-        } else {
-            Ok(None)
-        }
-    }
-
     async fn get_compile_test_cases_with_measurements(
         &self,
         artifact_row_id: &ArtifactIdNumber,

database/src/pool/sqlite.rs

@@ -1,4 +1,6 @@
-use crate::pool::{Connection, ConnectionManager, ManagedConnection, Transaction};
+use crate::pool::{
+    Connection, ConnectionManager, JobEnqueueResult, ManagedConnection, Transaction,
+};
 use crate::selector::CompileTestCase;
 use crate::{
     ArtifactCollection, ArtifactId, Benchmark, BenchmarkJob, BenchmarkJobConclusion,
@@ -1356,19 +1358,7 @@ impl Connection for SqliteConnection {
         _profile: Profile,
         _benchmark_set: u32,
         _kind: BenchmarkJobKind,
-    ) -> anyhow::Result<Option<u32>> {
-        no_queue_implementation_abort!()
-    }
-
-    async fn enqueue_parent_benchmark_job(
-        &self,
-        _parent_sha: &str,
-        _target: Target,
-        _backend: CodegenBackend,
-        _profile: Profile,
-        _benchmark_set: u32,
-        _kind: BenchmarkJobKind,
-    ) -> (bool, anyhow::Result<u32>) {
+    ) -> JobEnqueueResult {
         no_queue_implementation_abort!()
     }
 

database/src/tests/builder.rs

@@ -1,3 +1,4 @@
+use crate::pool::JobEnqueueResult;
 use crate::{
     BenchmarkJob, BenchmarkJobConclusion, BenchmarkJobKind, BenchmarkRequest,
     BenchmarkRequestStatus, BenchmarkSet, CodegenBackend, CollectorConfig, Connection, Profile,
@@ -40,7 +41,7 @@ impl RequestBuilder {
 
     pub async fn add_jobs(mut self, db: &dyn Connection, jobs: &[JobBuilder]) -> Self {
         for job in jobs {
-            let id = db
+            let id = match db
                 .enqueue_benchmark_job(
                     self.tag(),
                     job.target,
@@ -50,8 +51,11 @@ impl RequestBuilder {
                     job.kind,
                 )
                 .await
-                .unwrap();
-            self.jobs.push((job.clone(), id.unwrap()));
+            {
+                JobEnqueueResult::JobCreated(id) => id,
+                error => panic!("Unexpected job enqueue result: {error:?}"),
+            };
+            self.jobs.push((job.clone(), id));
         }
         self
     }