Database groundwork for job queue

Jamesbarford · Jamesbarford · commit 6579d0a56c08 · 2025-07-23T15:25:20.000+01:00
diff --git a/collector/src/compile/benchmark/target.rs b/collector/src/compile/benchmark/target.rs
@@ -14,3 +14,15 @@ impl Default for Target {
         Self::X86_64UnknownLinuxGnu
     }
 }
+
+impl Target {
+    pub fn all() -> Vec<Self> {
+        vec![Self::X86_64UnknownLinuxGnu]
+    }
+
+    pub fn from_db_target(target: &database::Target) -> Target {
+        match target {
+            database::Target::X86_64UnknownLinuxGnu => Self::X86_64UnknownLinuxGnu,
+        }
+    }
+}
diff --git a/database/src/lib.rs b/database/src/lib.rs
@@ -227,6 +227,11 @@ impl Profile {
             Profile::Clippy => "clippy",
         }
     }
+
+    /// 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]
+    }
 }
 
 impl std::str::FromStr for Profile {
@@ -365,6 +370,10 @@ impl Target {
             Target::X86_64UnknownLinuxGnu => "x86_64-unknown-linux-gnu",
         }
     }
+
+    pub fn all() -> Vec<Self> {
+        vec![Self::X86_64UnknownLinuxGnu]
+    }
 }
 
 impl FromStr for Target {
@@ -988,6 +997,35 @@ impl BenchmarkRequest {
     pub fn is_release(&self) -> bool {
         matches!(self.commit_type, BenchmarkRequestType::Release { .. })
     }
+
+    /// Get the codegen backends for the request
+    pub fn backends(&self) -> anyhow::Result<Vec<CodegenBackend>> {
+        // Empty string; default to LLVM.
+        if self.backends.trim().is_empty() {
+            return Ok(vec![CodegenBackend::Llvm]);
+        }
+
+        self.backends
+            .split(',')
+            .map(|s| {
+                CodegenBackend::from_str(s).map_err(|_| anyhow::anyhow!("Invalid backend: {s}"))
+            })
+            .collect()
+    }
+
+    /// Get the profiles for the request
+    pub fn profiles(&self) -> anyhow::Result<Vec<Profile>> {
+        // No profile string; fall back to the library defaults.
+        if self.profiles.trim().is_empty() {
+            return Ok(Profile::default_profiles());
+        }
+
+        self.profiles
+            .split(',')
+            .map(Profile::from_str)
+            .collect::<Result<Vec<_>, _>>()
+            .map_err(|e| anyhow::anyhow!("Invalid backend: {e}"))
+    }
 }
 
 /// Cached information about benchmark requests in the DB
@@ -1010,3 +1048,99 @@ impl BenchmarkRequestIndex {
         &self.completed
     }
 }
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum BenchmarkJobStatus {
+    Queued,
+    InProgress,
+    Success,
+    Failure,
+}
+
+const BENCHMARK_JOB_STATUS_QUEUED_STR: &str = "queued";
+const BENCHMARK_JOB_STATUS_IN_PROGRESS_STR: &str = "in_progress";
+const BENCHMARK_JOB_STATUS_SUCCESS_STR: &str = "success";
+const BENCHMARK_JOB_STATUS_FAILURE_STR: &str = "failure";
+
+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,
+        }
+    }
+}
+
+impl fmt::Display for BenchmarkJobStatus {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.as_str())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct BenchmarkSet(u32);
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct BenchmarkJob {
+    target: Target,
+    backend: CodegenBackend,
+    profile: Profile,
+    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(
+        target: Target,
+        backend: CodegenBackend,
+        profile: Profile,
+        request_tag: &str,
+        benchmark_set: u32,
+        created_at: DateTime<Utc>,
+        status: BenchmarkJobStatus,
+    ) -> Self {
+        BenchmarkJob {
+            target,
+            backend,
+            profile,
+            request_tag: request_tag.to_string(),
+            benchmark_set: BenchmarkSet(benchmark_set),
+            created_at,
+            started_at: None,
+            completed_at: None,
+            status,
+            retry: 0,
+        }
+    }
+
+    pub fn request_tag(&self) -> &str {
+        &self.request_tag
+    }
+
+    pub fn target(&self) -> Target {
+        self.target
+    }
+
+    pub fn backend(&self) -> CodegenBackend {
+        self.backend
+    }
+
+    pub fn profile(&self) -> Profile {
+        self.profile
+    }
+
+    pub fn benchmark_set(&self) -> u32 {
+        self.benchmark_set.0
+    }
+
+    pub fn status(&self) -> &BenchmarkJobStatus {
+        &self.status
+    }
+}
diff --git a/database/src/pool.rs b/database/src/pool.rs
@@ -1,6 +1,6 @@
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkRequest, BenchmarkRequestIndex,
-    BenchmarkRequestStatus, CodegenBackend, CompileBenchmark, Target,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkJob, BenchmarkRequest,
+    BenchmarkRequestIndex, BenchmarkRequestStatus, CodegenBackend, CompileBenchmark, Target,
 };
 use crate::{CollectionId, Index, Profile, QueuedCommit, Scenario, Step};
 use chrono::{DateTime, Utc};
@@ -210,6 +210,9 @@ pub trait Connection: Send + Sync {
         sha: &str,
         parent_sha: &str,
     ) -> anyhow::Result<()>;
+
+    /// Add a benchmark job to the job queue.
+    async fn enqueue_benchmark_job(&self, benchmark_job: &BenchmarkJob) -> anyhow::Result<()>;
 }
 
 #[async_trait::async_trait]
diff --git a/database/src/pool/postgres.rs b/database/src/pool/postgres.rs
@@ -1,9 +1,9 @@
 use crate::pool::{Connection, ConnectionManager, ManagedConnection, Transaction};
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, Benchmark, BenchmarkRequest,
-    BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkRequestType, CodegenBackend,
-    CollectionId, Commit, CommitType, CompileBenchmark, Date, Index, Profile, QueuedCommit,
-    Scenario, Target, BENCHMARK_REQUEST_MASTER_STR, BENCHMARK_REQUEST_RELEASE_STR,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, Benchmark, BenchmarkJob, BenchmarkJobStatus,
+    BenchmarkRequest, BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkRequestType,
+    CodegenBackend, CollectionId, Commit, CommitType, CompileBenchmark, Date, Index, Profile,
+    QueuedCommit, Scenario, Target, 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_TRY_STR,
 };
@@ -324,6 +324,28 @@ static MIGRATIONS: &[&str] = &[
     CREATE UNIQUE INDEX collector_config_target_bench_active_uniq ON collector_config
         (target, benchmark_set, is_active) WHERE is_active = TRUE;
     "#,
+    r#"
+    CREATE TABLE IF NOT EXISTS job_queue (
+        id            SERIAL PRIMARY KEY,
+        request_tag   TEXT NOT NULL,
+        target        TEXT NOT NULL,
+        backend       TEXT NOT NULL,
+        profile       TEXT NOT NULL,
+        benchmark_set INTEGER NOT NULL,
+        collector_id  TEXT,
+        created_at    TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+        started_at    TIMESTAMPTZ,
+        completed_at  TIMESTAMPTZ,
+        status        TEXT NOT NULL,
+        retry         INTEGER DEFAULT 0,
+
+        CONSTRAINT job_queue_request_fk
+            FOREIGN KEY (request_tag)
+            REFERENCES benchmark_request(tag)
+            ON DELETE CASCADE
+    );
+    CREATE INDEX IF NOT EXISTS job_queue_request_tag_idx ON job_queue (request_tag);
+    "#,
 ];
 
 #[async_trait::async_trait]
@@ -1608,6 +1630,34 @@ where
             .collect();
         Ok(requests)
     }
+
+    async fn enqueue_benchmark_job(&self, benchmark_job: &BenchmarkJob) -> anyhow::Result<()> {
+        self.conn()
+            .execute(
+                r#"
+            INSERT INTO job_queue(
+                request_tag,
+                target,
+                backend,
+                profile,
+                benchmark_set,
+                status
+            )
+            VALUES ($1, $2, $3, $4, $5, $6)
+                "#,
+                &[
+                    &benchmark_job.request_tag(),
+                    &benchmark_job.target(),
+                    &benchmark_job.backend(),
+                    &benchmark_job.profile(),
+                    &(benchmark_job.benchmark_set() as i32),
+                    &benchmark_job.status(),
+                ],
+            )
+            .await
+            .context("failed to insert benchmark_job")?;
+        Ok(())
+    }
 }
 
 fn parse_artifact_id(ty: &str, sha: &str, date: Option<DateTime<Utc>>) -> ArtifactId {
@@ -1653,6 +1703,10 @@ macro_rules! impl_to_postgresql_via_to_string {
 
 impl_to_postgresql_via_to_string!(BenchmarkRequestType);
 impl_to_postgresql_via_to_string!(BenchmarkRequestStatus);
+impl_to_postgresql_via_to_string!(Target);
+impl_to_postgresql_via_to_string!(CodegenBackend);
+impl_to_postgresql_via_to_string!(Profile);
+impl_to_postgresql_via_to_string!(BenchmarkJobStatus);
 
 #[cfg(test)]
 mod tests {
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, BenchmarkRequest, BenchmarkRequestIndex,
-    BenchmarkRequestStatus, CodegenBackend, CollectionId, Commit, CommitType, CompileBenchmark,
-    Date, Profile, Target,
+    ArtifactCollection, ArtifactId, Benchmark, BenchmarkJob, BenchmarkRequest,
+    BenchmarkRequestIndex, BenchmarkRequestStatus, CodegenBackend, CollectionId, Commit,
+    CommitType, CompileBenchmark, Date, Profile, Target,
 };
 use crate::{ArtifactIdNumber, Index, QueuedCommit};
 use chrono::{DateTime, TimeZone, Utc};
@@ -1296,6 +1296,10 @@ impl Connection for SqliteConnection {
     ) -> anyhow::Result<()> {
         no_queue_implementation_abort!()
     }
+
+    async fn enqueue_benchmark_job(&self, _benchmark_job: &BenchmarkJob) -> anyhow::Result<()> {
+        no_queue_implementation_abort!()
+    }
 }
 
 fn parse_artifact_id(ty: &str, sha: &str, date: Option<i64>) -> ArtifactId {
diff --git a/site/src/job_queue/mod.rs b/site/src/job_queue/mod.rs
