Add function for loading pending benchmark requests from the DB

Author: Kobzol

database/src/lib.rs

@@ -1,6 +1,6 @@
 use chrono::offset::TimeZone;
 use chrono::{DateTime, Utc};
-use hashbrown::HashMap;
+use hashbrown::{HashMap, HashSet};
 use intern::intern;
 use serde::{Deserialize, Serialize};
 use std::fmt;
@@ -806,18 +806,18 @@ pub enum BenchmarkRequestStatus {
     Completed { completed_at: DateTime<Utc> },
 }
 
-const WAITING_FOR_ARTIFACTS_STR: &str = "waiting_for_artifacts";
-const ARTIFACTS_READY_STR: &str = "artifacts_ready";
-const IN_PROGRESS_STR: &str = "in_progress";
-const COMPLETED_STR: &str = "completed";
+const BENCHMARK_REQUEST_STATUS_WAITING_FOR_ARTIFACTS_STR: &str = "waiting_for_artifacts";
+const BENCHMARK_REQUEST_STATUS_ARTIFACTS_READY_STR: &str = "artifacts_ready";
+const BENCHMARK_REQUEST_STATUS_IN_PROGRESS_STR: &str = "in_progress";
+const BENCHMARK_REQUEST_STATUS_COMPLETED_STR: &str = "completed";
 
 impl BenchmarkRequestStatus {
     pub(crate) fn as_str(&self) -> &str {
         match self {
-            Self::WaitingForArtifacts => WAITING_FOR_ARTIFACTS_STR,
-            Self::ArtifactsReady => ARTIFACTS_READY_STR,
-            Self::InProgress => IN_PROGRESS_STR,
-            Self::Completed { .. } => COMPLETED_STR,
+            Self::WaitingForArtifacts => BENCHMARK_REQUEST_STATUS_WAITING_FOR_ARTIFACTS_STR,
+            Self::ArtifactsReady => BENCHMARK_REQUEST_STATUS_ARTIFACTS_READY_STR,
+            Self::InProgress => BENCHMARK_REQUEST_STATUS_IN_PROGRESS_STR,
+            Self::Completed { .. } => BENCHMARK_REQUEST_STATUS_COMPLETED_STR,
         }
     }
 
@@ -826,10 +826,10 @@ impl BenchmarkRequestStatus {
         completion_date: Option<DateTime<Utc>>,
     ) -> anyhow::Result<Self> {
         match text {
-            WAITING_FOR_ARTIFACTS_STR => Ok(Self::WaitingForArtifacts),
-            ARTIFACTS_READY_STR => Ok(Self::ArtifactsReady),
-            IN_PROGRESS_STR => Ok(Self::InProgress),
-            COMPLETED_STR => Ok(Self::Completed {
+            BENCHMARK_REQUEST_STATUS_WAITING_FOR_ARTIFACTS_STR => Ok(Self::WaitingForArtifacts),
+            BENCHMARK_REQUEST_STATUS_ARTIFACTS_READY_STR => Ok(Self::ArtifactsReady),
+            BENCHMARK_REQUEST_STATUS_IN_PROGRESS_STR => Ok(Self::InProgress),
+            BENCHMARK_REQUEST_STATUS_COMPLETED_STR => Ok(Self::Completed {
                 completed_at: completion_date.ok_or_else(|| {
                     anyhow!("No completion date for a completed BenchmarkRequestStatus")
                 })?,
@@ -930,25 +930,17 @@ impl BenchmarkRequest {
         }
     }
 
-    pub fn create_master(
-        sha: &str,
-        parent_sha: &str,
-        pr: u32,
-        created_at: DateTime<Utc>,
-        status: BenchmarkRequestStatus,
-        backends: &str,
-        profiles: &str,
-    ) -> Self {
+    pub fn create_master(sha: &str, parent_sha: &str, pr: u32, created_at: DateTime<Utc>) -> Self {
         Self {
             commit_type: BenchmarkRequestType::Master {
                 pr,
                 sha: sha.to_string(),
                 parent_sha: parent_sha.to_string(),
             },
             created_at,
-            status,
-            backends: backends.to_string(),
-            profiles: profiles.to_string(),
+            status: BenchmarkRequestStatus::ArtifactsReady,
+            backends: String::new(),
+            profiles: String::new(),
         }
     }
 
@@ -979,3 +971,24 @@ impl BenchmarkRequest {
         }
     }
 }
+
+/// Cached information about benchmark requests in the DB
+/// FIXME: only store non-try requests here
+pub struct BenchmarkRequestIndex {
+    /// Tags (SHA or release name) of all known benchmark requests
+    all: HashSet<String>,
+    /// Tags (SHA or release name) of all benchmark requests in the completed status
+    completed: HashSet<String>,
+}
+
+impl BenchmarkRequestIndex {
+    /// Do we already have a benchmark request for the passed `tag`?
+    pub fn contains_tag(&self, tag: &str) -> bool {
+        self.all.contains(tag)
+    }
+
+    /// Return tags of already completed benchmark requests.
+    pub fn completed_requests(&self) -> &HashSet<String> {
+        &self.completed
+    }
+}

database/src/pool.rs

@@ -1,6 +1,6 @@
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkRequest, BenchmarkRequestStatus,
-    CodegenBackend, CompileBenchmark, Target,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkRequest, BenchmarkRequestIndex,
+    BenchmarkRequestStatus, CodegenBackend, CompileBenchmark, Target,
 };
 use crate::{CollectionId, Index, Profile, QueuedCommit, Scenario, Step};
 use chrono::{DateTime, Utc};
@@ -187,18 +187,18 @@ pub trait Connection: Send + Sync {
         benchmark_request: &BenchmarkRequest,
     ) -> anyhow::Result<()>;
 
-    /// 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>>;
+    /// Load all known benchmark request SHAs and all completed benchmark requests.
+    async fn load_benchmark_request_index(&self) -> anyhow::Result<BenchmarkRequestIndex>;
+
+    /// Load all pending benchmark requests, i.e. those that have artifacts ready, but haven't
+    /// been completed yet. Pending statuses are `ArtifactsReady` and `InProgress`.
+    async fn load_pending_benchmark_requests(&self) -> 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,
+        request: &BenchmarkRequest,
+        status: BenchmarkRequestStatus,
     ) -> anyhow::Result<()>;
 
     /// Update a Try commit to have a `sha` and `parent_sha`. Will update the
@@ -420,7 +420,7 @@ mod tests {
                 "",
             );
 
-            let try_benchmark_request = BenchmarkRequest::create_try(
+            let try_benchmark_request = BenchmarkRequest::create_try_without_artifacts(
                 Some("b-sha-2"),
                 Some("parent-sha-2"),
                 32,
@@ -474,7 +474,7 @@ mod tests {
                 "",
             );
 
-            let try_benchmark_request = BenchmarkRequest::create_try(
+            let try_benchmark_request = BenchmarkRequest::create_try_without_artifacts(
                 Some("b-sha-2"),
                 Some("parent-sha-2"),
                 32,
@@ -502,7 +502,7 @@ mod tests {
                 .unwrap();
 
             let requests = db
-                .get_benchmark_requests_by_status(&[BenchmarkRequestStatus::ArtifactsReady])
+                .load_benchmark_request_index(&[BenchmarkRequestStatus::ArtifactsReady])
                 .await
                 .unwrap();
 
@@ -545,7 +545,7 @@ mod tests {
             .unwrap();
 
             let requests = db
-                .get_benchmark_requests_by_status(&[BenchmarkRequestStatus::InProgress])
+                .load_benchmark_request_index(&[BenchmarkRequestStatus::InProgress])
                 .await
                 .unwrap();
 
@@ -566,7 +566,7 @@ mod tests {
             let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
             let pr = 42;
 
-            let try_benchmark_request = BenchmarkRequest::create_try(
+            let try_benchmark_request = BenchmarkRequest::create_try_without_artifacts(
                 None,
                 None,
                 pr,
@@ -582,7 +582,7 @@ mod tests {
                 .await
                 .unwrap();
             let requests = db
-                .get_benchmark_requests_by_status(&[BenchmarkRequestStatus::ArtifactsReady])
+                .load_benchmark_request_index(&[BenchmarkRequestStatus::ArtifactsReady])
                 .await
                 .unwrap();
 
@@ -604,7 +604,7 @@ mod tests {
             let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
             let pr = 42;
 
-            let completed_try = BenchmarkRequest::create_try(
+            let completed_try = BenchmarkRequest::create_try_without_artifacts(
                 Some("sha-2"),
                 Some("p-sha-1"),
                 pr,
@@ -615,7 +615,7 @@ mod tests {
             );
             db.insert_benchmark_request(&completed_try).await.unwrap();
 
-            let try_benchmark_request = BenchmarkRequest::create_try(
+            let try_benchmark_request = BenchmarkRequest::create_try_without_artifacts(
                 None,
                 None,
                 pr,
@@ -638,7 +638,7 @@ mod tests {
                 .unwrap();
 
             let requests = db
-                .get_benchmark_requests_by_status(&[
+                .load_benchmark_request_index(&[
                     BenchmarkRequestStatus::WaitingForArtifacts,
                     BenchmarkRequestStatus::ArtifactsReady,
                     BenchmarkRequestStatus::InProgress,