Most of the scaffolding for dequeuing a job from a collector

Author: Jamesbarford

collector/src/bin/collector.rs

@@ -35,6 +35,7 @@ use collector::artifact_stats::{
 use collector::codegen::{codegen_diff, CodegenType};
 use collector::compile::benchmark::category::Category;
 use collector::compile::benchmark::codegen_backend::CodegenBackend;
+use collector::compile::benchmark::collector_config::CollectorConfig;
 use collector::compile::benchmark::profile::Profile;
 use collector::compile::benchmark::scenario::Scenario;
 use collector::compile::benchmark::target::Target;
@@ -666,6 +667,19 @@ enum Commands {
         /// The name of the modified artifact to be compared.
         modified: Option<String>,
     },
+
+    /// Polls the job queue for work to benchmark
+    DequeueJob {
+        /// The unique identifier for the collector
+        #[arg(long)]
+        collector_name: String,
+
+        #[arg(long)]
+        target: String,
+
+        #[command(flatten)]
+        db: DbOption,
+    },
 }
 
 #[derive(Debug, clap::Parser)]
@@ -1266,6 +1280,44 @@ Make sure to modify `{dir}/perf-config.json` if the category/artifact don't matc
             rt.block_on(compare_artifacts(conn, metric, base, modified))?;
             Ok(0)
         }
+
+        Commands::DequeueJob {
+            collector_name,
+            db,
+            target,
+        } => {
+            let pool = Pool::open(&db.db);
+            let rt = build_async_runtime();
+            let conn = rt.block_on(pool.connection());
+
+            // Obtain the configuration and validate that it matches the
+            // collector's setup
+            let collector_config: CollectorConfig = rt
+                .block_on(conn.get_collector_config(&collector_name))?
+                .into();
+
+            let collector_target = collector_config.target();
+            if collector_target.as_str() != target {
+                panic!(
+                    "Mismatching target for collector expected `{}` got `{}`",
+                    collector_target, target
+                );
+            }
+
+            // Dequeue a job
+            let benchmark_job = rt.block_on(conn.dequeue_benchmark_job(
+                &collector_name,
+                &collector_config.target().to_db_target(),
+                collector_config.benchmark_set(),
+            ))?;
+
+            if let Some(benchmark_job) = benchmark_job {
+                // TODO; process the job
+                println!("{:?}", benchmark_job);
+            }
+
+            Ok(0)
+        }
     }
 }
 

collector/src/compile/benchmark/collector_config.rs

@@ -0,0 +1,52 @@
+use chrono::{DateTime, Utc};
+use database::BenchmarkSet;
+
+use super::target::Target;
+
+pub struct CollectorConfig {
+    name: String,
+    target: Target,
+    benchmark_set: BenchmarkSet,
+    is_active: bool,
+    last_heartbeat_at: DateTime<Utc>,
+    date_added: DateTime<Utc>,
+}
+
+impl CollectorConfig {
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn target(&self) -> &Target {
+        &self.target
+    }
+
+    pub fn benchmark_set(&self) -> &BenchmarkSet {
+        &self.benchmark_set
+    }
+
+    pub fn is_active(&self) -> bool {
+        self.is_active
+    }
+
+    pub fn last_heartbeat_at(&self) -> DateTime<Utc> {
+        self.last_heartbeat_at
+    }
+
+    pub fn date_added(&self) -> DateTime<Utc> {
+        self.date_added
+    }
+}
+
+impl From<database::CollectorConfig> for CollectorConfig {
+    fn from(value: database::CollectorConfig) -> Self {
+        CollectorConfig {
+            name: value.name().into(),
+            target: Target::from_db_target(value.target()),
+            benchmark_set: value.benchmark_set().clone(),
+            is_active: value.is_active(),
+            last_heartbeat_at: value.last_heartbeat_at(),
+            date_added: value.date_added(),
+        }
+    }
+}

collector/src/compile/benchmark/mod.rs

@@ -19,6 +19,7 @@ use tempfile::TempDir;
 
 pub mod category;
 pub mod codegen_backend;
+pub mod collector_config;
 pub(crate) mod patch;
 pub mod profile;
 pub mod scenario;

collector/src/compile/benchmark/target.rs

@@ -1,3 +1,5 @@
+use std::{fmt, str::FromStr};
+
 /// Target representing an Rust target triple, for a full list of targets and
 /// their support see;
 /// https://doc.rust-lang.org/nightly/rustc/platform-support.html
@@ -15,6 +17,22 @@ impl Default for Target {
     }
 }
 
+impl FromStr for Target {
+    type Err = String;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(match s.to_ascii_lowercase().as_str() {
+            "x86_64-unknown-linux-gnu" => Target::X86_64UnknownLinuxGnu,
+            _ => return Err(format!("{} is not a valid target", s)),
+        })
+    }
+}
+
+impl fmt::Display for Target {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.as_str())
+    }
+}
+
 impl Target {
     pub fn all() -> Vec<Self> {
         vec![Self::X86_64UnknownLinuxGnu]
@@ -29,6 +47,26 @@ impl From<database::Target> for Target {
     }
 }
 
+impl Target {
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Target::X86_64UnknownLinuxGnu => "x86_64-unknown-linux-gnu",
+        }
+    }
+
+    pub fn from_db_target(target: &database::Target) -> Target {
+        match target {
+            database::Target::X86_64UnknownLinuxGnu => Self::X86_64UnknownLinuxGnu,
+        }
+    }
+
+    pub fn to_db_target(&self) -> database::Target {
+        match self {
+            Target::X86_64UnknownLinuxGnu => database::Target::X86_64UnknownLinuxGnu,
+        }
+    }
+}
+
 impl From<Target> for database::Target {
     fn from(value: Target) -> Self {
         match value {

database/src/lib.rs

@@ -1054,10 +1054,12 @@ pub enum BenchmarkJobStatus {
     Queued,
     InProgress {
         started_at: DateTime<Utc>,
+        collector_name: String,
     },
     Completed {
         started_at: DateTime<Utc>,
         completed_at: DateTime<Utc>,
+        collector_name: String,
         success: bool,
     },
 }
@@ -1109,3 +1111,39 @@ pub struct BenchmarkJob {
     status: BenchmarkJobStatus,
     retry: u32,
 }
+
+/// The configuration for a collector
+pub struct CollectorConfig {
+    name: String,
+    target: Target,
+    benchmark_set: BenchmarkSet,
+    is_active: bool,
+    last_heartbeat_at: DateTime<Utc>,
+    date_added: DateTime<Utc>,
+}
+
+impl CollectorConfig {
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn target(&self) -> &Target {
+        &self.target
+    }
+
+    pub fn benchmark_set(&self) -> &BenchmarkSet {
+        &self.benchmark_set
+    }
+
+    pub fn is_active(&self) -> bool {
+        self.is_active
+    }
+
+    pub fn last_heartbeat_at(&self) -> DateTime<Utc> {
+        self.last_heartbeat_at
+    }
+
+    pub fn date_added(&self) -> DateTime<Utc> {
+        self.date_added
+    }
+}

database/src/pool.rs

@@ -1,7 +1,8 @@
 use crate::selector::CompileTestCase;
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkRequest, BenchmarkRequestIndex,
-    BenchmarkRequestStatus, CodegenBackend, CompileBenchmark, Target,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkJob, BenchmarkRequest,
+    BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkSet, CodegenBackend, CollectorConfig,
+    CompileBenchmark, Target,
 };
 use crate::{CollectionId, Index, Profile, QueuedCommit, Scenario, Step};
 use chrono::{DateTime, Utc};
@@ -232,6 +233,17 @@ pub trait Connection: Send + Sync {
         &self,
         artifact_row_id: &ArtifactIdNumber,
     ) -> anyhow::Result<HashSet<CompileTestCase>>;
+
+    /// Get the confiuguration for a collector by the name of the collector
+    async fn get_collector_config(&self, collector_name: &str) -> anyhow::Result<CollectorConfig>;
+
+    /// Get the confiuguration for a collector by the name of the collector
+    async fn dequeue_benchmark_job(
+        &self,
+        collector_name: &str,
+        target: &Target,
+        benchmark_set: &BenchmarkSet,
+    ) -> anyhow::Result<Option<BenchmarkJob>>;
 }
 
 #[async_trait::async_trait]
@@ -690,6 +702,39 @@ mod tests {
                 .await
                 .unwrap()
                 .is_empty());
+            Ok(ctx)
+        })
+        .await;
+    }
+
+    async fn get_collector_config_error_if_not_exist() {
+        run_postgres_test(|ctx| async {
+            let db = ctx.db_client().connection().await;
+
+            let collector_config_result = db.get_collector_config("collector-1").await;
+
+            assert!(collector_config_result.is_err());
+
+            Ok(ctx)
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn dequeue_benchmark_job_empty_queue() {
+        run_postgres_test(|ctx| async {
+            let db = ctx.db_client().connection().await;
+
+            let benchmark_job_result = db
+                .dequeue_benchmark_job(
+                    "collector-1",
+                    &Target::X86_64UnknownLinuxGnu,
+                    &BenchmarkSet(420),
+                )
+                .await;
+
+            assert!(benchmark_job_result.is_ok());
+            assert!(benchmark_job_result.unwrap().is_none());
 
             Ok(ctx)
         })