Add a function to update collector heartbeat

Author: Kobzol

collector/src/bin/collector.rs

@@ -1395,6 +1395,8 @@ async fn run_job_queue_benchmarks(
     all_compile_benchmarks: Vec<Benchmark>,
     host_target_tuple: &str,
 ) -> anyhow::Result<()> {
+    conn.update_collector_heartbeat(collector.name()).await?;
+
     // TODO: reconnect to the DB if there was an error with the previous job
     while let Some((benchmark_job, artifact_id)) = conn
         .dequeue_benchmark_job(
@@ -1478,6 +1480,8 @@ async fn run_job_queue_benchmarks(
                 tokio::time::sleep(Duration::from_secs(30)).await;
             }
         }
+
+        conn.update_collector_heartbeat(collector.name()).await?;
     }
     log::info!("No job found, exiting");
     Ok(())

database/src/pool.rs

@@ -276,6 +276,9 @@ pub trait Connection: Send + Sync {
 
     /// Get all of the configuration for all of the collectors
     async fn get_collector_configs(&self) -> anyhow::Result<Vec<CollectorConfig>>;
+
+    /// Updates the last known heartbeat of a collector to the current time.
+    async fn update_collector_heartbeat(&self, collector_name: &str) -> anyhow::Result<()>;
 }
 
 #[async_trait::async_trait]

database/src/pool/postgres.rs

@@ -2212,6 +2212,20 @@ where
 
         Ok(configs)
     }
+
+    async fn update_collector_heartbeat(&self, collector_name: &str) -> anyhow::Result<()> {
+        self.conn()
+            .query(
+                r#"
+                UPDATE collector_config
+                SET last_heartbeat_at = NOW()
+                WHERE name = $1
+                "#,
+                &[&collector_name],
+            )
+            .await?;
+        Ok(())
+    }
 }
 
 fn row_to_benchmark_request(row: &Row) -> BenchmarkRequest {

database/src/pool/sqlite.rs

@@ -1377,6 +1377,10 @@ impl Connection for SqliteConnection {
     async fn get_collector_configs(&self) -> anyhow::Result<Vec<CollectorConfig>> {
         no_queue_implementation_abort!()
     }
+
+    async fn update_collector_heartbeat(&self, _collector_name: &str) -> anyhow::Result<()> {
+        no_queue_implementation_abort!()
+    }
 }
 
 fn parse_artifact_id(ty: &str, sha: &str, date: Option<i64>) -> ArtifactId {