Feat: Logic for finding the next request to be processed and tests

database/Cargo.toml

@@ -26,6 +26,7 @@ csv = "1"
 x509-cert = { version = "0.2.5", features = ["pem"] }
 
 intern = { path = "../intern" }
+uuid = { version = "1.16.0", features = ["v4"] }
 
 [dev-dependencies]
 uuid = { version = "1.16.0", features = ["v4"] }

database/src/lib.rs

@@ -13,9 +13,7 @@ pub mod interpolate;
 pub mod metric;
 pub mod pool;
 pub mod selector;
-
-#[cfg(test)]
-mod tests;
+pub mod tests;
 
 pub use pool::{Connection, Pool};
 
@@ -799,7 +797,7 @@ pub struct ArtifactCollection {
     pub end_time: DateTime<Utc>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum BenchmarkRequestStatus {
     WaitingForArtifacts,
     ArtifactsReady,
@@ -818,12 +816,34 @@ impl fmt::Display for BenchmarkRequestStatus {
     }
 }
 
-#[derive(Debug)]
+impl<'a> tokio_postgres::types::FromSql<'a> for BenchmarkRequestStatus {
+    fn from_sql(
+        ty: &tokio_postgres::types::Type,
+        raw: &'a [u8],
+    ) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
+        // Decode raw bytes into &str with Postgres' own text codec
+        let s: &str = <&str as tokio_postgres::types::FromSql>::from_sql(ty, raw)?;
+
+        match s {
+            "waiting_for_artifacts" => Ok(Self::WaitingForArtifacts),
+            "artifacts_ready" => Ok(Self::ArtifactsReady),
+            "in_progress" => Ok(Self::InProgress),
+            "completed" => Ok(Self::Completed),
+            other => Err(format!("unknown benchmark_request_status '{other}'").into()),
+        }
+    }
+
+    fn accepts(ty: &tokio_postgres::types::Type) -> bool {
+        <&str as tokio_postgres::types::FromSql>::accepts(ty)
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
 pub enum BenchmarkRequestType {
     /// A Try commit
     Try {
         sha: String,
-        parent_sha: String,
+        parent_sha: Option<String>,
         pr: u32,
     },
     /// A Master commit
@@ -864,7 +884,7 @@ impl fmt::Display for BenchmarkRequestType {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct BenchmarkRequest {
     pub commit_type: BenchmarkRequestType,
     pub created_at: DateTime<Utc>,
@@ -896,7 +916,7 @@ impl BenchmarkRequest {
 
     pub fn create_try(
         sha: &str,
-        parent_sha: &str,
+        parent_sha: Option<&str>,
         pr: u32,
         created_at: DateTime<Utc>,
         status: BenchmarkRequestStatus,
@@ -907,7 +927,7 @@ impl BenchmarkRequest {
             commit_type: BenchmarkRequestType::Try {
                 pr,
                 sha: sha.to_string(),
-                parent_sha: parent_sha.to_string(),
+                parent_sha: parent_sha.map(|it| it.to_string()),
             },
             created_at,
             completed_at: None,
@@ -964,8 +984,11 @@ impl BenchmarkRequest {
 
     pub fn parent_sha(&self) -> Option<&str> {
         match &self.commit_type {
-            BenchmarkRequestType::Try { parent_sha, .. }
-            | BenchmarkRequestType::Master { parent_sha, .. } => Some(parent_sha),
+            BenchmarkRequestType::Try { parent_sha, .. } => match parent_sha {
+                Some(parent_sha) => Some(parent_sha),
+                _ => None,
+            },
+            BenchmarkRequestType::Master { parent_sha, .. } => Some(parent_sha),
             BenchmarkRequestType::Release { tag: _ } => None,
         }
     }

database/src/pool.rs

@@ -1,7 +1,7 @@
 use crate::selector::CompileTestCase;
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkRequest, CodegenBackend,
-    CompileBenchmark, Target,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkRequest, BenchmarkRequestStatus,
+    CodegenBackend, CompileBenchmark, Target,
 };
 use crate::{CollectionId, Index, Profile, QueuedCommit, Scenario, Step};
 use chrono::{DateTime, Utc};
@@ -195,6 +195,20 @@ pub trait Connection: Send + Sync {
         &self,
         artifact_row_id: &ArtifactIdNumber,
     ) -> anyhow::Result<HashSet<CompileTestCase>>;
+
+    /// Gets the benchmark requests matching the status. Optionally provide the
+    /// number of days from whence to search from
+    async fn get_benchmark_requests_by_status(
+        &self,
+        statuses: &[BenchmarkRequestStatus],
+    ) -> anyhow::Result<Vec<BenchmarkRequest>>;
+
+    /// Update the status of a `benchmark_request`
+    async fn update_benchmark_request_status(
+        &mut self,
+        benchmark_request: &BenchmarkRequest,
+        benchmark_request_status: BenchmarkRequestStatus,
+    ) -> anyhow::Result<()>;
 }
 
 #[async_trait::async_trait]
@@ -406,7 +420,7 @@ mod tests {
 
             let try_benchmark_request = BenchmarkRequest::create_try(
                 "b-sha-2",
-                "parent-sha-2",
+                Some("parent-sha-2"),
                 32,
                 time,
                 BenchmarkRequestStatus::ArtifactsReady,
@@ -494,4 +508,100 @@ mod tests {
         })
         .await;
     }
+
+    #[tokio::test]
+    async fn get_benchmark_requests_by_status() {
+        // Ensure we get back the requests matching the status with no date
+        // limit
+        run_postgres_test(|ctx| async {
+            let db = ctx.db_client();
+            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
+            let master_benchmark_request = BenchmarkRequest::create_master(
+                "a-sha-1",
+                "parent-sha-1",
+                42,
+                time,
+                BenchmarkRequestStatus::ArtifactsReady,
+                "llvm",
+                "",
+            );
+
+            let try_benchmark_request = BenchmarkRequest::create_try(
+                "b-sha-2",
+                Some("parent-sha-2"),
+                32,
+                time,
+                BenchmarkRequestStatus::Completed,
+                "cranelift",
+                "",
+            );
+
+            let release_benchmark_request = BenchmarkRequest::create_release(
+                "1.8.0",
+                time,
+                BenchmarkRequestStatus::ArtifactsReady,
+                "cranelift,llvm",
+                "",
+            );
+
+            let db = db.connection().await;
+            db.insert_benchmark_request(&master_benchmark_request).await;
+            db.insert_benchmark_request(&try_benchmark_request).await;
+            db.insert_benchmark_request(&release_benchmark_request)
+                .await;
+
+            let requests = db
+                .get_benchmark_requests_by_status(&[BenchmarkRequestStatus::ArtifactsReady])
+                .await
+                .unwrap();
+
+            assert_eq!(requests.len(), 2);
+            assert_eq!(requests[0].status, BenchmarkRequestStatus::ArtifactsReady);
+            assert_eq!(requests[1].status, BenchmarkRequestStatus::ArtifactsReady);
+
+            Ok(ctx)
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn update_benchmark_request_status() {
+        // Insert one item into the database, change the status and then
+        // get the item back out again to ensure it has changed status
+        run_postgres_test(|ctx| async {
+            let db = ctx.db_client();
+            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
+            let master_benchmark_request = BenchmarkRequest::create_master(
+                "a-sha-1",
+                "parent-sha-1",
+                42,
+                time,
+                BenchmarkRequestStatus::ArtifactsReady,
+                "llvm",
+                "",
+            );
+
+            let mut db = db.connection().await;
+            db.insert_benchmark_request(&master_benchmark_request).await;
+
+            db.update_benchmark_request_status(
+                &master_benchmark_request,
+                BenchmarkRequestStatus::InProgress,
+            )
+            .await
+            .unwrap();
+
+            let requests = db
+                .get_benchmark_requests_by_status(&[BenchmarkRequestStatus::InProgress])
+                .await
+                .unwrap();
+
+            assert_eq!(requests.len(), 1);
+            assert_eq!(requests[0].tag(), master_benchmark_request.tag());
+            assert_eq!(requests[0].status, BenchmarkRequestStatus::InProgress);
+
+            Ok(ctx)
+        })
+        .await;
+    }
 }