Store commit SHA of collectors in the database

Kobzol · Kobzol · commit 216dcfa49a98 · 2025-08-27T08:46:06.000+02:00
diff --git a/collector/src/bin/collector.rs b/collector/src/bin/collector.rs
@@ -59,7 +59,7 @@ use collector::toolchain::{
 };
 use collector::utils::cachegrind::cachegrind_diff;
 use collector::utils::{is_installed, wait_for_future};
-use collector::{utils, CollectorCtx, CollectorStepBuilder};
+use collector::{command_output, utils, CollectorCtx, CollectorStepBuilder};
 use database::{
     ArtifactId, ArtifactIdNumber, BenchmarkJob, BenchmarkJobConclusion, CollectorConfig, Commit,
     CommitType, Connection, Pool,
@@ -708,6 +708,11 @@ enum Commands {
         #[arg(long)]
         collector_name: String,
 
+        /// Git SHA of the commit that the collector is currently on.
+        /// If not present, the collector will attempt to figure it out from git directly.
+        #[arg(long)]
+        git_sha: Option<String>,
+
         #[command(flatten)]
         db: DbOption,
     },
@@ -1352,17 +1357,33 @@ Make sure to modify `{dir}/perf-config.json` if the category/artifact don't matc
             Ok(0)
         }
 
-        Commands::BenchmarkJobQueue { collector_name, db } => {
+        Commands::BenchmarkJobQueue {
+            collector_name,
+            git_sha,
+            db,
+        } => {
             log_db(&db);
 
+            let git_sha = match git_sha {
+                Some(sha) => sha,
+                None => {
+                    let mut cmd = Command::new("git");
+                    cmd.args(&["rev-parse", "HEAD"]);
+                    let stdout = command_output(&mut cmd)
+                        .context("Cannot determine current commit SHA")?
+                        .stdout;
+                    String::from_utf8(stdout).unwrap().trim().to_string()
+                }
+            };
+
             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 host target
             let collector_config = rt
-                .block_on(conn.get_collector_config(&collector_name))?
+                .block_on(conn.start_collector(&collector_name, &git_sha))?
                 .ok_or_else(|| {
                     anyhow::anyhow!("Collector with name `{collector_name}` not found")
                 })?;
@@ -1398,8 +1419,6 @@ async fn run_job_queue_benchmarks(
     collector: &CollectorConfig,
     all_compile_benchmarks: Vec<Benchmark>,
 ) -> anyhow::Result<()> {
-    conn.update_collector_heartbeat(collector.name()).await?;
-
     // TODO: check collector SHA vs site SHA
     while let Some((benchmark_job, artifact_id)) = conn
         .dequeue_benchmark_job(
diff --git a/database/src/lib.rs b/database/src/lib.rs
@@ -1204,6 +1204,8 @@ pub struct CollectorConfig {
     is_active: bool,
     last_heartbeat_at: DateTime<Utc>,
     date_added: DateTime<Utc>,
+    /// The commit SHA of `rustc-perf` that the collector currently has checked out.
+    commit_sha: Option<String>,
 }
 
 impl CollectorConfig {
@@ -1230,6 +1232,10 @@ impl CollectorConfig {
     pub fn date_added(&self) -> DateTime<Utc> {
         self.date_added
     }
+
+    pub fn commit_sha(&self) -> Option<&str> {
+        self.commit_sha.as_deref()
+    }
 }
 
 /// The data that can be retrived from the database directly to populate the
diff --git a/database/src/pool.rs b/database/src/pool.rs
@@ -244,10 +244,13 @@ pub trait Connection: Send + Sync {
         is_active: bool,
     ) -> anyhow::Result<CollectorConfig>;
 
-    /// Get the configuration for a collector by its name.
-    async fn get_collector_config(
+    /// Call this function when a job queue collector starts.
+    /// It ensures that a collector with the given name exists, updates its commit SHA and heartbeat
+    /// and returns its collector config.
+    async fn start_collector(
         &self,
         collector_name: &str,
+        commit_sha: &str,
     ) -> anyhow::Result<Option<CollectorConfig>>;
 
     /// Dequeues a single job for the given collector, target and benchmark set.
@@ -788,7 +791,7 @@ mod tests {
         run_postgres_test(|ctx| async {
             let db = ctx.db_client().connection().await;
 
-            let collector_config_result = db.get_collector_config("collector-1").await.unwrap();
+            let collector_config_result = db.start_collector("collector-1", "foo").await.unwrap();
 
             assert!(collector_config_result.is_none());
 
@@ -802,19 +805,20 @@ mod tests {
         run_postgres_test(|ctx| async {
             let db = ctx.db_client().connection().await;
 
-            let inserted_config = db
+            let mut inserted_config = db
                 .add_collector_config("collector-1", Target::X86_64UnknownLinuxGnu, 1, true)
                 .await
                 .unwrap();
 
             let config = db
-                .get_collector_config("collector-1")
+                .start_collector("collector-1", "foo")
                 .await
                 .unwrap()
                 .expect("collector config not found");
 
-            // What we entered into the database should be identical to what is
-            // returned from the database
+            inserted_config.commit_sha = Some("foo".to_string());
+            inserted_config.last_heartbeat_at = config.last_heartbeat_at;
+
             assert_eq!(inserted_config, config);
             Ok(ctx)
         })
diff --git a/database/src/pool/postgres.rs b/database/src/pool/postgres.rs
@@ -384,6 +384,9 @@ static MIGRATIONS: &[&str] = &[
     r#"
     ALTER TABLE benchmark_request ADD COLUMN duration_ms INTEGER NULL;
     "#,
+    r#"
+    ALTER TABLE collector_config ADD COLUMN commit_sha TEXT NULL;
+    "#,
 ];
 
 #[async_trait::async_trait]
@@ -1732,28 +1735,33 @@ where
             is_active,
             last_heartbeat_at: row.get::<_, DateTime<Utc>>(0),
             date_added: row.get::<_, DateTime<Utc>>(1),
+            commit_sha: None,
         };
         Ok(collector_config)
     }
 
-    async fn get_collector_config(
+    async fn start_collector(
         &self,
         collector_name: &str,
+        commit_sha: &str,
     ) -> anyhow::Result<Option<CollectorConfig>> {
         let row = self
             .conn()
             .query_opt(
-                "SELECT
+                "
+                UPDATE collector_config
+                SET
+                    last_heartbeat_at = NOW(),
+                    commit_sha = $2
+                WHERE
+                    name = $1
+                RETURNING
                     target,
                     benchmark_set,
                     is_active,
                     last_heartbeat_at,
-                    date_added
-                FROM
-                    collector_config
-                WHERE
-                    name = $1;",
-                &[&collector_name],
+                    date_added",
+                &[&collector_name, &commit_sha],
             )
             .await?;
 
@@ -1767,6 +1775,7 @@ where
                     is_active: row.get::<_, bool>(2),
                     last_heartbeat_at: row.get::<_, DateTime<Utc>>(3),
                     date_added: row.get::<_, DateTime<Utc>>(4),
+                    commit_sha: Some(commit_sha.to_string()),
                 })
             })
             .transpose()?)
@@ -2180,25 +2189,29 @@ where
                     benchmark_set,
                     is_active,
                     last_heartbeat_at,
-                    date_added
+                    date_added,
+                    commit_sha
                 FROM
                     collector_config;",
                 &[],
             )
             .await?;
 
-        let mut configs = vec![];
-        for row in rows {
-            let config = CollectorConfig {
-                name: row.get::<_, String>(0),
-                target: Target::from_str(row.get::<_, &str>(1)).map_err(|e| anyhow::anyhow!(e))?,
-                benchmark_set: BenchmarkSet(row.get::<_, i32>(2) as u32),
-                is_active: row.get::<_, bool>(3),
-                last_heartbeat_at: row.get::<_, DateTime<Utc>>(4),
-                date_added: row.get::<_, DateTime<Utc>>(5),
-            };
-            configs.push(config);
-        }
+        let configs = rows
+            .into_iter()
+            .map(|row| {
+                Ok(CollectorConfig {
+                    name: row.get::<_, String>(0),
+                    target: Target::from_str(row.get::<_, &str>(1))
+                        .map_err(|e| anyhow::anyhow!(e))?,
+                    benchmark_set: BenchmarkSet(row.get::<_, i32>(2) as u32),
+                    is_active: row.get::<_, bool>(3),
+                    last_heartbeat_at: row.get::<_, DateTime<Utc>>(4),
+                    date_added: row.get::<_, DateTime<Utc>>(5),
+                    commit_sha: row.get::<_, Option<String>>(6),
+                })
+            })
+            .collect::<anyhow::Result<Vec<_>>>()?;
 
         Ok(configs)
     }
diff --git a/database/src/pool/sqlite.rs b/database/src/pool/sqlite.rs
@@ -1332,9 +1332,10 @@ impl Connection for SqliteConnection {
             .collect::<Result<_, _>>()?)
     }
 
-    async fn get_collector_config(
+    async fn start_collector(
         &self,
         _collector_name: &str,
+        _commit_sha: &str,
     ) -> anyhow::Result<Option<CollectorConfig>> {
         no_queue_implementation_abort!()
     }

Original file line number	Diff line number	Diff line change
`@@ -1204,6 +1204,8 @@ pub struct CollectorConfig {`
`1204`	`1204`	`is_active: bool,`
`1205`	`1205`	`last_heartbeat_at: DateTime<Utc>,`
`1206`	`1206`	`date_added: DateTime<Utc>,`
	`1207`	+ /// The commit SHA of `rustc-perf` that the collector currently has checked out.
	`1208`	`+ commit_sha: Option<String>,`
`1207`	`1209`	`}`
`1208`	`1210`
`1209`	`1211`	`impl CollectorConfig {`
`@@ -1230,6 +1232,10 @@ impl CollectorConfig {`
`1230`	`1232`	`pub fn date_added(&self) -> DateTime<Utc> {`
`1231`	`1233`	`self.date_added`
`1232`	`1234`	`}`
	`1235`	`+`
	`1236`	`+ pub fn commit_sha(&self) -> Option<&str> {`
	`1237`	`+ self.commit_sha.as_deref()`
	`1238`	`+ }`
`1233`	`1239`	`}`
`1234`	`1240`
`1235`	`1241`	`/// The data that can be retrived from the database directly to populate the`
Original file line number	Diff line number	Diff line change
`@@ -1332,9 +1332,10 @@ impl Connection for SqliteConnection {`
`1332`	`1332`	`.collect::<Result<_, _>>()?)`
`1333`	`1333`	`}`
`1334`	`1334`
`1335`		`- async fn get_collector_config(`
	`1335`	`+ async fn start_collector(`
`1336`	`1336`	`&self,`
`1337`	`1337`	`_collector_name: &str,`
	`1338`	`+ _commit_sha: &str,`
`1338`	`1339`	`) -> anyhow::Result<Option<CollectorConfig>> {`
`1339`	`1340`	`no_queue_implementation_abort!()`
`1340`	`1341`	`}`