Created AddCollector command line option and created more tests for config and dequeuing a job

Author: Jamesbarford

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};
 
@@ -666,6 +667,24 @@ enum Commands {
         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
@@ -1279,6 +1298,27 @@ Make sure to modify `{dir}/perf-config.json` if the category/artifact don't matc
             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,

collector/src/compile/benchmark/collector_config.rs

@@ -3,6 +3,7 @@ use database::BenchmarkSet;
 
 use super::target::Target;
 
+#[derive(Debug, Clone, PartialEq)]
 pub struct CollectorConfig {
     name: String,
     target: Target,

database/src/lib.rs

@@ -1112,7 +1112,38 @@ pub struct BenchmarkJob {
     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,

database/src/pool.rs

@@ -222,6 +222,15 @@ pub trait Connection: Send + Sync {
         benchmark_set: u32,
     ) -> anyhow::Result<()>;
 
+    /// Add the confiuguration for a collector
+    async fn add_collector_config(
+        &self,
+        collector_name: &str,
+        target: &Target,
+        benchmark_set: u32,
+        is_active: bool,
+    ) -> anyhow::Result<CollectorConfig>;
+
     /// Get the confiuguration for a collector by the name of the collector
     async fn get_collector_config(&self, collector_name: &str) -> anyhow::Result<CollectorConfig>;
 
@@ -652,6 +661,27 @@ mod tests {
         .await;
     }
 
+    #[tokio::test]
+    async fn add_collector_config() {
+        run_postgres_test(|ctx| async {
+            let db = ctx.db_client().connection().await;
+
+            let insert_config_result = db
+                .add_collector_config("collector-1", &Target::X86_64UnknownLinuxGnu, 1, true)
+                .await;
+            assert!(insert_config_result.is_ok());
+
+            let get_config_result = db.get_collector_config("collector-1").await;
+            assert!(get_config_result.is_ok());
+
+            // What we entered into the database should be identical to what is
+            // returned from the database
+            assert_eq!(insert_config_result.unwrap(), get_config_result.unwrap());
+            Ok(ctx)
+        })
+        .await;
+    }
+
     #[tokio::test]
     async fn dequeue_benchmark_job_empty_queue() {
         run_postgres_test(|ctx| async {
@@ -672,4 +702,70 @@ mod tests {
         })
         .await;
     }
+
+    #[tokio::test]
+    async fn dequeue_benchmark_job() {
+        run_postgres_test(|ctx| async {
+            let db = ctx.db_client().connection().await;
+            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
+
+            let insert_result = db
+                .add_collector_config("collector-1", &Target::X86_64UnknownLinuxGnu, 1, true)
+                .await;
+            assert!(insert_result.is_ok());
+
+            let collector_config = insert_result.unwrap();
+
+            let benchmark_request =
+                BenchmarkRequest::create_master("sha-1", "parent-sha-1", 42, time);
+
+            // Insert the request so we don't violate the foreign key
+            db.insert_benchmark_request(&benchmark_request)
+                .await
+                .unwrap();
+
+            // Now we can insert the job
+            let enqueue_result = db
+                .enqueue_benchmark_job(
+                    benchmark_request.tag().unwrap(),
+                    &Target::X86_64UnknownLinuxGnu,
+                    &CodegenBackend::Llvm,
+                    &Profile::Opt,
+                    1u32,
+                )
+                .await;
+            assert!(enqueue_result.is_ok());
+
+            let benchmark_job = db
+                .dequeue_benchmark_job(
+                    collector_config.name(),
+                    collector_config.target(),
+                    collector_config.benchmark_set(),
+                )
+                .await;
+            assert!(benchmark_job.is_ok());
+
+            let benchmark_job = benchmark_job.unwrap();
+            assert!(benchmark_job.is_some());
+
+            // Ensure the properties of the job match both the request and the
+            // collector configuration
+            let benchmark_job = benchmark_job.unwrap();
+            assert_eq!(
+                benchmark_job.request_tag(),
+                benchmark_request.tag().unwrap()
+            );
+            assert_eq!(
+                benchmark_job.benchmark_set(),
+                collector_config.benchmark_set()
+            );
+            assert_eq!(
+                benchmark_job.collector_name().unwrap(),
+                collector_config.name(),
+            );
+
+            Ok(ctx)
+        })
+        .await;
+    }
 }

database/src/pool/postgres.rs

@@ -1695,6 +1695,56 @@ where
         Ok(())
     }
 
+    async fn add_collector_config(
+        &self,
+        collector_name: &str,
+        target: &Target,
+        benchmark_set: u32,
+        is_active: bool,
+    ) -> anyhow::Result<CollectorConfig> {
+        let row = self
+            .conn()
+            .query_one(
+                "INSERT INTO collector_config(
+                     name,
+                     target,
+                     date_added,
+                     last_heartbeat_at,
+                     benchmark_set,
+                     is_active
+                 ) VALUES (
+                     $1,
+                     $2,
+                     NOW(),
+                     NOW(),
+                     $3,
+                     $4
+                 )
+                RETURNING
+                    last_heartbeat_at,
+                    date_added
+                ",
+                &[
+                    &collector_name,
+                    &target,
+                    &(benchmark_set as i32),
+                    &is_active,
+                ],
+            )
+            .await
+            .context("failed to create collector config")?;
+
+        let collector_config = CollectorConfig {
+            name: collector_name.into(),
+            target: *target,
+            benchmark_set: BenchmarkSet(benchmark_set),
+            is_active,
+            last_heartbeat_at: row.get::<_, DateTime<Utc>>(0),
+            date_added: row.get::<_, DateTime<Utc>>(1),
+        };
+        Ok(collector_config)
+    }
+
     async fn get_collector_config(&self, collector_name: &str) -> anyhow::Result<CollectorConfig> {
         let row = self
             .conn()
@@ -1705,11 +1755,10 @@ where
                     is_active,
                     last_heartbeat_at,
                     date_added
-                   
                 FROM
                     collector_config
                 WHERE
-                    collector_name = $1;",
+                    name = $1;",
                 &[&collector_name],
             )
             .await?;

database/src/pool/sqlite.rs

@@ -1320,6 +1320,16 @@ impl Connection for SqliteConnection {
     ) -> anyhow::Result<Option<BenchmarkJob>> {
         no_queue_implementation_abort!()
     }
+
+    async fn add_collector_config(
+        &self,
+        _collector_name: &str,
+        _target: &Target,
+        _benchmark_set: u32,
+        _is_active: bool,
+    ) -> anyhow::Result<CollectorConfig> {
+        no_queue_implementation_abort!()
+    }
 }
 
 fn parse_artifact_id(ty: &str, sha: &str, date: Option<i64>) -> ArtifactId {