feat - database schema & enqueue/dequeue logic

database/src/lib.rs

@@ -3,7 +3,7 @@ use chrono::{DateTime, Utc};
 use hashbrown::HashMap;
 use intern::intern;
 use serde::{Deserialize, Serialize};
-use std::fmt;
+use std::fmt::{self, Display, Formatter};
 use std::hash;
 use std::ops::{Add, Sub};
 use std::sync::Arc;
@@ -155,6 +155,15 @@ impl FromStr for CommitType {
     }
 }
 
+impl Display for CommitType {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CommitType::Try => f.write_str("try"),
+            CommitType::Master => f.write_str("master"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
 pub struct Commit {
     pub sha: String,
@@ -794,3 +803,165 @@ pub struct ArtifactCollection {
     pub duration: Duration,
     pub end_time: DateTime<Utc>,
 }
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum CommitJobType {
+    Try { pr: u32 },
+    Master { pr: u32 },
+    Release { tag: String },
+}
+
+impl CommitJobType {
+    /// Get the name of the type as a `str`
+    pub fn name(&self) -> &'static str {
+        match self {
+            CommitJobType::Try { pr: _ } => "try",
+            CommitJobType::Master { pr: _ } => "master",
+            CommitJobType::Release { tag: _ } => "release",
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CommitJob {
+    pub sha: String,
+    pub parent_sha: String,
+    pub commit_time: Date,
+    pub target: Target,
+    pub include: Option<String>,
+    pub exclude: Option<String>,
+    pub runs: Option<i32>,
+    pub backends: Option<String>,
+    pub job_type: CommitJobType,
+    pub state: CommitJobState,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum CommitJobState {
+    Queued,
+    Finished(CommitJobFinished),
+    Failed(CommitJobFailed),
+    InProgress(CommitJobInProgress),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CommitJobInProgress {
+    pub machine_id: String,
+    pub started_at: Date,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CommitJobFinished {
+    pub machine_id: String,
+    pub started_at: Date,
+    pub finished_at: Date,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CommitJobFailed {
+    pub machine_id: String,
+    pub started_at: Date,
+    pub finished_at: Date,
+}
+
+impl CommitJob {
+    /// Get the status as a string
+    pub fn status(&self) -> &'static str {
+        match self.state {
+            CommitJobState::Queued => "queued",
+            CommitJobState::InProgress(_) => "in_progress",
+            CommitJobState::Finished(_) => "finished",
+            CommitJobState::Failed(_) => "failed",
+        }
+    }
+}
+
+/// Maps from the database to a Rust struct
+#[allow(clippy::too_many_arguments)]
+fn commit_job_create(
+    sha: String,
+    parent_sha: String,
+    commit_type: &str,
+    pr: Option<u32>,
+    release_tag: Option<String>,
+    commit_time: Date,
+    target: Target,
+    machine_id: Option<String>,
+    started_at: Option<Date>,
+    finished_at: Option<Date>,
+    status: &str,
+    include: Option<String>,
+    exclude: Option<String>,
+    runs: Option<i32>,
+    backends: Option<String>,
+) -> CommitJob {
+    let job_type = match commit_type {
+        "try" => CommitJobType::Try {
+            pr: pr.expect("`pr` cannot be `None` for a Commit of type `try`"),
+        },
+        "master" => CommitJobType::Master {
+            pr: pr.expect("`pr` cannot be `None` for a Commit of type `master`"),
+        },
+        "release" => CommitJobType::Release {
+            tag: release_tag
+                .expect("`release_tag` cannot be `None` for a Commit of type `release`"),
+        },
+        _ => panic!("Unhandled commit_type {}", commit_type),
+    };
+
+    let state = match status {
+        "queued" => CommitJobState::Queued,
+
+        "in_progress" => {
+            let started_at =
+                started_at.expect("`started_at` must be Some for an `in_progress` job");
+            let machine_id =
+                machine_id.expect("`machine_id` must be Some for an `in_progress` job");
+
+            CommitJobState::InProgress(CommitJobInProgress {
+                started_at,
+                machine_id,
+            })
+        }
+
+        "finished" | "failed" => {
+            let started_at =
+                started_at.expect("`started_at` must be Some for finished or failed job");
+            let finished_at =
+                finished_at.expect("`finished_at` must be Some for finished or failed");
+            let machine_id =
+                machine_id.expect("`machine_id` must be Some for finished or failed a job");
+
+            if status == "finished" {
+                CommitJobState::Finished(CommitJobFinished {
+                    started_at,
+                    finished_at,
+                    machine_id,
+                })
+            } else {
+                CommitJobState::Failed(CommitJobFailed {
+                    started_at,
+                    finished_at,
+                    machine_id,
+                })
+            }
+        }
+
+        other => {
+            panic!("unknown status `{other}` (expected `queued`, `in_progress`, `finished` or `failed`)")
+        }
+    };
+
+    CommitJob {
+        sha,
+        parent_sha,
+        commit_time,
+        target,
+        include,
+        exclude,
+        runs,
+        backends,
+        job_type,
+        state,
+    }
+}

database/src/pool.rs

@@ -1,5 +1,6 @@
 use crate::{
-    ArtifactCollection, ArtifactId, ArtifactIdNumber, CodegenBackend, CompileBenchmark, Target,
+    ArtifactCollection, ArtifactId, ArtifactIdNumber, CodegenBackend, CommitJob, CompileBenchmark,
+    Target,
 };
 use crate::{CollectionId, Index, Profile, QueuedCommit, Scenario, Step};
 use chrono::{DateTime, Utc};
@@ -178,6 +179,16 @@ pub trait Connection: Send + Sync {
 
     /// Removes all data associated with the given artifact.
     async fn purge_artifact(&self, aid: &ArtifactId);
+
+    /// Add a job to the queue
+    async fn enqueue_commit_job(&self, jobs: &CommitJob);
+
+    /// Dequeue jobs, we pass `machine_id` and `target` in case there are jobs
+    /// the machine was previously doing and can pick up again
+    async fn take_commit_job(&self, machine_id: &str, target: Target) -> Option<CommitJob>;
+
+    /// Mark the job as finished
+    async fn finish_commit_job(&self, machine_id: &str, target: Target, sha: String);
 }
 
 #[async_trait::async_trait]
@@ -301,7 +312,7 @@ mod tests {
     use std::str::FromStr;
 
     use super::*;
-    use crate::{tests::run_db_test, Commit, CommitType, Date};
+    use crate::{tests::run_db_test, Commit, CommitJobState, CommitJobType, CommitType, Date};
 
     /// Create a Commit
     fn create_commit(commit_sha: &str, time: chrono::DateTime<Utc>, r#type: CommitType) -> Commit {
@@ -312,6 +323,29 @@ mod tests {
         }
     }
 
+    /// Create a CommitJob
+    fn create_commit_job(
+        sha: &str,
+        parent_sha: &str,
+        commit_time: chrono::DateTime<Utc>,
+        target: Target,
+        job_type: CommitJobType,
+        state: CommitJobState,
+    ) -> CommitJob {
+        CommitJob {
+            sha: sha.to_string(),
+            parent_sha: parent_sha.to_string(),
+            commit_time: Date(commit_time),
+            target,
+            include: None,
+            exclude: None,
+            runs: None,
+            backends: None,
+            job_type,
+            state,
+        }
+    }
+
     #[tokio::test]
     async fn pstat_returns_empty_vector_when_empty() {
         run_db_test(|ctx| async {
@@ -370,4 +404,122 @@ mod tests {
         })
         .await;
     }
+
+    #[tokio::test]
+    async fn take_commit_job() {
+        run_db_test(|ctx| async {
+            // ORDER:
+            // Releases first
+            // Master commits second, order by oldest PR ascending
+            // Try commits last, order by oldest PR ascending
+
+            let db = ctx.db_client().connection().await;
+            let time = chrono::DateTime::from_str("2021-09-01T00:00:00.000Z").unwrap();
+
+            // Try commits
+            let try_job_1 = create_commit_job(
+                "sha1",
+                "p1",
+                time,
+                Target::X86_64UnknownLinuxGnu,
+                CommitJobType::Try { pr: 1 },
+                CommitJobState::Queued,
+            );
+            let try_job_2 = create_commit_job(
+                "sha2",
+                "p2",
+                time,
+                Target::X86_64UnknownLinuxGnu,
+                CommitJobType::Try { pr: 2 },
+                CommitJobState::Queued,
+            );
+
+            // Master commits
+            let master_job_1 = create_commit_job(
+                "sha3",
+                "p3",
+                time,
+                Target::X86_64UnknownLinuxGnu,
+                CommitJobType::Master { pr: 3 },
+                CommitJobState::Queued,
+            );
+            let master_job_2 = create_commit_job(
+                "sha4",
+                "p4",
+                time,
+                Target::X86_64UnknownLinuxGnu,
+                CommitJobType::Master { pr: 4 },
+                CommitJobState::Queued,
+            );
+
+            // Release commits
+            let release_job_1 = create_commit_job(
+                "sha5",
+                "p5",
+                time,
+                Target::X86_64UnknownLinuxGnu,
+                CommitJobType::Release { tag: "tag1".into() },
+                CommitJobState::Queued,
+            );
+            let release_job_2 = create_commit_job(
+                "sha6",
+                "p6",
+                time,
+                Target::X86_64UnknownLinuxGnu,
+                CommitJobType::Release { tag: "tag2".into() },
+                CommitJobState::Queued,
+            );
+
+            // Shuffle the insert order a bit
+            let all_commits = vec![
+                release_job_1,
+                master_job_2,
+                try_job_1,
+                release_job_2,
+                master_job_1,
+                try_job_2,
+            ];
+
+            // queue all the jobs
+            for commit in all_commits {
+                db.enqueue_commit_job(&commit).await;
+            }
+
+            // Now we test the ordering: after each dequeue we immediately mark
+            // the job as finished for the sake of testing so it can't be
+            // returned again in the test.
+            //
+            // The priority should be;
+            //
+            //   1. Release commits           (oldest tag first)
+            //   2. Master  commits           (oldest PR first)
+            //   3. Try     commits           (oldest PR first)
+            //
+            // Given the data we inserted above the expected SHA order is:
+            // sha5, sha6, sha3, sha4, sha1, sha2.
+
+            let machine = "machine-1";
+            let target = Target::X86_64UnknownLinuxGnu;
+            let expected = ["sha5", "sha6", "sha3", "sha4", "sha1", "sha2"];
+
+            for &sha in &expected {
+                let job = db.take_commit_job(machine, target).await;
+                assert!(job.is_some(), "expected a job for sha {sha}");
+                let job = job.unwrap();
+                assert_eq!(job.sha, sha, "jobs dequeued out of priority order");
+
+                // Mark the job finished so it is not returned again.
+                db.finish_commit_job(machine, target, sha.to_string()).await;
+            }
+
+            // After all six jobs have been taken, the queue should be empty.
+            assert!(
+                db.take_commit_job(machine, target).await.is_none(),
+                "queue should be empty after draining all jobs"
+            );
+
+            Ok(ctx)
+        })
+        .await;
+    }
 }