Feat; mark benchmark requests as complete

Author: Jamesbarford

database/schema.md

@@ -313,3 +313,36 @@ Columns:
   execute.
 * **is_active** (`boolean NOT NULL`): For controlling whether the collector is
   active for use. Useful for adding/removing collectors.
+
+### job_queue
+
+This table stores ephemeral benchmark jobs, which specifically tell the
+collector which benchmarks it should execute. The jobs will be kept in the
+table for ~30 days after being completed, so that we can quickly figure out
+what master parent jobs we need to backfill when handling try builds.
+
+Columns:
+
+- **id** (`bigint` / `serial`): Primary-key identifier for the job row;
+  auto-increments with each new job.
+- **request_id** (`bigint`): References the parent benchmark request that
+  spawned this job.
+- **target** (`text NOT NULL`): Hardware/ISA the benchmarks must run on
+  (e.g. AArch64, x86_64).
+- **backend** (`text NOT NULL`): Code generation backend the collector should
+  test (e.g. llvm, cranelift).
+- **benchmark_set** (`int NOT NULL`): ID of the predefined benchmark suite to
+  execute.
+- **collector_id** (`text`): Id of the collector that claimed the job
+  (populated once the job is started).
+- **created_at** (`timestamptz NOT NULL`): Datetime when the job was queued.
+- **started_at** (`timestamptz`): Datetime when the collector actually began
+  running the benchmarks; NULL until the job is claimed.
+- **completed_at** (`timestampt`): Datetime when the collector finished
+  (successfully or otherwise); used to purge rows after ~30 days.
+- **status** (`text NOT NULL`): Current job state. `queued`, `in_progress`,
+  `success`, or `failure`.
+- **retry** (`int NOT NULL`): Number of times the job has been re-queued after
+  a failure; 0 on the first attempt.
+- **error** (`text`): Optional error message or stack trace from the last
+  failed run; NULL when the job succeeded.

database/src/lib.rs

@@ -1026,6 +1026,10 @@ impl BenchmarkRequest {
             .collect::<Result<Vec<_>, _>>()
             .map_err(|e| anyhow::anyhow!("Invalid backend: {e}"))
     }
+
+    pub fn is_completed(&self) -> bool {
+        matches!(self.status, BenchmarkRequestStatus::Completed { .. })
+    }
 }
 
 /// Cached information about benchmark requests in the DB
@@ -1047,6 +1051,11 @@ impl BenchmarkRequestIndex {
     pub fn completed_requests(&self) -> &HashSet<String> {
         &self.completed
     }
+
+    pub fn add_tag(&mut self, tag: &str) {
+        self.all.insert(tag.to_string());
+        self.completed.insert(tag.to_string());
+    }
 }
 
 #[derive(Debug, Clone, PartialEq)]

database/src/pool.rs

@@ -253,6 +253,13 @@ pub trait Connection: Send + Sync {
         target: &Target,
         benchmark_set: &BenchmarkSet,
     ) -> anyhow::Result<Option<BenchmarkJob>>;
+
+    /// 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
+    async fn mark_benchmark_request_as_completed(
+        &self,
+        benchmark_request: &mut BenchmarkRequest,
+    ) -> anyhow::Result<bool>;
 }
 
 #[async_trait::async_trait]

database/src/pool/postgres.rs

@@ -5,7 +5,8 @@ use crate::{
     BenchmarkRequest, BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkRequestType,
     BenchmarkSet, CodegenBackend, CollectionId, CollectorConfig, Commit, CommitType,
     CompileBenchmark, Date, Index, Profile, QueuedCommit, Scenario, Target,
-    BENCHMARK_JOB_STATUS_IN_PROGRESS_STR, BENCHMARK_JOB_STATUS_QUEUED_STR,
+    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,
     BENCHMARK_REQUEST_STATUS_IN_PROGRESS_STR, BENCHMARK_REQUEST_STATUS_WAITING_FOR_ARTIFACTS_STR,
@@ -1887,6 +1888,82 @@ where
             }
         }
     }
+
+    async fn mark_benchmark_request_as_completed(
+        &self,
+        benchmark_request: &mut BenchmarkRequest,
+    ) -> anyhow::Result<bool> {
+        anyhow::ensure!(
+            benchmark_request.tag().is_some(),
+            "Benchmark request has no tag"
+        );
+        anyhow::ensure!(
+            benchmark_request.status == BenchmarkRequestStatus::InProgress,
+            "Can only mark benchmark request whos status is in_progress as complete"
+        );
+
+        // Find if the benchmark is completed and update it's status to completed
+        // in one SQL block
+        let row = self
+            .conn()
+            .query_opt(
+                "
+                UPDATE
+                    benchmark_request
+                SET
+                    status = $1,
+                    completed_at = NOW()
+                WHERE
+                    banchmark_request.tag = $2
+                    AND benchmark_request.status != $1
+                    AND NOT EXISTS (
+                        SELECT
+                            1
+                        FROM
+                            job_queue
+                        WHERE
+                            request_tag = benchmark_request.tag
+                            AND status NOT IN ($3, $4)
+                    )
+                    AND (
+                        benchmark_request.parent_sha IS NULL
+                        OR NOT EXISTS (
+                            SELECT
+                                1
+                            FROM
+                                job_queue
+                            JOIN
+                                benchmark_request AS parent_request 
+                            ON
+                                parent_request.tag = benchmark_request.parent_tag
+                            WHERE
+                                job_queue.request_tag = parent_request.tag
+                                AND job_queue.status NOT IN ($3, $4)
+                       )
+                )
+                RETURNING
+                    benchmark_request.completed_at;
+                ",
+                &[
+                    &BENCHMARK_REQUEST_STATUS_COMPLETED_STR,
+                    &benchmark_request.tag(),
+                    &BENCHMARK_JOB_STATUS_SUCCESS_STR,
+                    &BENCHMARK_JOB_STATUS_FAILURE_STR,
+                ],
+            )
+            .await
+            .context("Failed to mark benchmark_request as completed")?;
+        // The affected id is returned by the query thus we can use the row's
+        // presence to determine if the request was marked as completed
+        if let Some(row) = row {
+            let completed_at = row.get::<_, DateTime<Utc>>(0);
+            // Also mutate our object
+            benchmark_request.status = BenchmarkRequestStatus::Completed { completed_at };
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
 }
 
 fn parse_artifact_id(ty: &str, sha: &str, date: Option<DateTime<Utc>>) -> ArtifactId {

database/src/pool/sqlite.rs

@@ -1353,6 +1353,13 @@ impl Connection for SqliteConnection {
     ) -> anyhow::Result<CollectorConfig> {
         no_queue_implementation_abort!()
     }
+
+    async fn mark_benchmark_request_as_completed(
+        &self,
+        _benchmark_request: &mut BenchmarkRequest,
+    ) -> anyhow::Result<bool> {
+        no_queue_implementation_abort!()
+    }
 }
 
 fn parse_artifact_id(ty: &str, sha: &str, date: Option<i64>) -> ArtifactId {

site/src/job_queue/mod.rs

@@ -246,14 +246,20 @@ async fn try_enqueue_next_benchmark_request(
     let queue = build_queue(conn, index).await?;
 
     #[allow(clippy::never_loop)]
-    for request in queue {
+    for mut request in queue {
         match request.status() {
             BenchmarkRequestStatus::ArtifactsReady => {
                 enqueue_benchmark_request(conn, &request).await?;
                 break;
             }
             BenchmarkRequestStatus::InProgress => {
-                // TODO: Try to mark as completed
+                if conn
+                    .mark_benchmark_request_as_completed(&mut request)
+                    .await?
+                {
+                    index.add_tag(request.tag().unwrap());
+                    continue;
+                }
                 break;
             }
             BenchmarkRequestStatus::WaitingForArtifacts