Feat; collector dequeue job

collector/src/bin/collector.rs

@@ -12,6 +12,7 @@ use std::marker::PhantomData;
 use std::path::{Path, PathBuf};
 use std::process;
 use std::process::Command;
+use std::str::FromStr;
 use std::time::Duration;
 use std::{str, time::Instant};
 
@@ -35,6 +36,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 +668,37 @@ enum Commands {
         /// The name of the modified artifact to be compared.
         modified: Option<String>,
     },
+
+    /// Registers a collector in the database
+    AddCollector {
+        #[command(flatten)]
+        db: DbOption,
+
+        #[arg(long)]
+        collector_name: String,
+
+        #[arg(long)]
+        target: String,
+
+        #[arg(long)]
+        is_active: bool,
+
+        #[arg(long)]
+        benchmark_set: u32,
+    },
+
+    /// 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 +1299,65 @@ 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::AddCollector {
+            db,
+            collector_name,
+            target,
+            is_active,
+            benchmark_set,
+        } => {
+            let pool = Pool::open(&db.db);
+            let rt = build_async_runtime();
+            let conn = rt.block_on(pool.connection());
+
+            let target = database::Target::from_str(&target).map_err(|e| anyhow::anyhow!(e))?;
+            rt.block_on(conn.add_collector_config(
+                &collector_name,
+                &target,
+                benchmark_set,
+                is_active,
+            ))?;
+            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,53 @@
+use chrono::{DateTime, Utc};
+use database::BenchmarkSet;
+
+use super::target::Target;
+
+#[derive(Debug, Clone, PartialEq)]
+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,70 @@ pub struct BenchmarkJob {
     status: BenchmarkJobStatus,
     retry: u32,
 }
+
+impl BenchmarkJob {
+    pub fn target(&self) -> &Target {
+        &self.target
+    }
+
+    pub fn backend(&self) -> &CodegenBackend {
+        &self.backend
+    }
+
+    pub fn profile(&self) -> &Profile {
+        &self.profile
+    }
+
+    pub fn request_tag(&self) -> &str {
+        &self.request_tag
+    }
+
+    pub fn benchmark_set(&self) -> &BenchmarkSet {
+        &self.benchmark_set
+    }
+
+    pub fn collector_name(&self) -> Option<&str> {
+        match &self.status {
+            BenchmarkJobStatus::Queued => None,
+            BenchmarkJobStatus::InProgress { collector_name, .. }
+            | BenchmarkJobStatus::Completed { collector_name, .. } => Some(collector_name),
+        }
+    }
+}
+
+/// The configuration for a collector
+#[derive(Debug, PartialEq)]
+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
+    }
+}