feat - database schema & enqueue/dequeue logic

Author: Jamesbarford

database/src/lib.rs

@@ -3,9 +3,9 @@ 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::ops::{Add, Deref, Sub};
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -19,6 +19,8 @@ pub use pool::{Connection, Pool};
 intern!(pub struct Metric);
 intern!(pub struct Benchmark);
 
+type PgParam<'a> = &'a (dyn tokio_postgres::types::ToSql + Sync);
+
 #[derive(Debug, PartialEq, Eq)]
 pub struct QueuedCommit {
     pub pr: u32,
@@ -152,6 +154,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,
@@ -791,3 +802,254 @@ pub struct ArtifactCollection {
     pub duration: Duration,
     pub end_time: DateTime<Utc>,
 }
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum CommitJobType {
+    Try(u32),
+    Master(u32),
+    Release(String),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CommitJobEntity {
+    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,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CommitJobInProgress {
+    pub commit_job: CommitJobEntity,
+    pub machine_id: String,
+    pub started_at: Date,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CommitJobFinished {
+    pub commit_job: CommitJobEntity,
+    pub machine_id: String,
+    pub started_at: Date,
+    pub finished_at: Date,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CommitJobFailed {
+    pub commit_job: CommitJobEntity,
+    pub machine_id: String,
+    pub started_at: Date,
+    pub finished_at: Date,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum CommitJob {
+    Queued(CommitJobEntity),
+    InProgress(CommitJobInProgress),
+    Finished(CommitJobFinished),
+    Failed(CommitJobFailed),
+}
+
+impl CommitJob {
+    /// Returns `Some(&CommitJobEntity)` only if the job is still queued.
+    pub fn as_queued(&self) -> Option<&CommitJobEntity> {
+        match self {
+            CommitJob::Queued(e) => Some(e),
+            _ => None,
+        }
+    }
+
+    /// Returns `Some(&CommitJobInProgress)` while the job is running.
+    pub fn as_in_progress(&self) -> Option<&CommitJobInProgress> {
+        match self {
+            CommitJob::InProgress(ip) => Some(ip),
+            _ => None,
+        }
+    }
+
+    /// Returns `Some(&CommitJobFinished)` once the job is done.
+    pub fn as_finished(&self) -> Option<&CommitJobFinished> {
+        match self {
+            CommitJob::Finished(fin) => Some(fin),
+            _ => None,
+        }
+    }
+
+    /// Get the status as a string
+    pub fn status(&self) -> &'static str {
+        match self {
+            CommitJob::Queued(_) => "queued",
+            CommitJob::InProgress(_) => "in_progress",
+            CommitJob::Finished(_) => "finished",
+            CommitJob::Failed(_) => "failed",
+        }
+    }
+
+    /// True when `status == "finished"`.
+    pub fn is_finished(&self) -> bool {
+        matches!(self, CommitJob::Finished(_))
+    }
+
+    /// Will compose the column names for the job type
+    pub fn get_enqueue_column_names(&self) -> Vec<String> {
+        let mut base_columns = vec![
+            String::from("sha"),
+            String::from("parent_sha"),
+            String::from("commit_type"),
+            String::from("commit_time"),
+            String::from("status"),
+            String::from("target"),
+            String::from("include"),
+            String::from("exclude"),
+            String::from("runs"),
+            String::from("backends"),
+        ];
+
+        /* This is the last column */
+        match self.job_type {
+            CommitJobType::Try(_) => base_columns.push("pr".into()),
+            CommitJobType::Master(_) => base_columns.push("pr".into()),
+            CommitJobType::Release(_) => base_columns.push("release_tag".into()),
+        };
+
+        return base_columns;
+    }
+}
+
+impl Deref for CommitJob {
+    type Target = CommitJobEntity;
+    fn deref(&self) -> &Self::Target {
+        match self {
+            CommitJob::Queued(e) => e,
+            CommitJob::InProgress(ip) => &ip.commit_job,
+            CommitJob::Finished(fin) => &fin.commit_job,
+            CommitJob::Failed(fail) => &fail.commit_job,
+        }
+    }
+}
+
+/// Maps from the database to a Rust struct
+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.expect("`pr` cannot be `None` for a Commit of type `try`")),
+        "master" => {
+            CommitJobType::Master(pr.expect("`pr` cannot be `None` for a Commit of type `master`"))
+        }
+        "release" => CommitJobType::Release(
+            release_tag.expect("`release_tag` cannot be `None` for a Commit of type `release`"),
+        ),
+        _ => panic!("Unhandled commit_type {}", commit_type),
+    };
+
+    let commit_job = CommitJobEntity {
+        sha,
+        parent_sha,
+        commit_time,
+        target,
+        include,
+        exclude,
+        runs,
+        backends,
+        job_type,
+    };
+
+    match status {
+        "queued" => CommitJob::Queued(commit_job),
+
+        "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");
+
+            CommitJob::InProgress(CommitJobInProgress {
+                commit_job,
+                started_at,
+                machine_id,
+            })
+        }
+
+        "finished" | "failed" => {
+            let started_at =
+                started_at.expect(&format!("`started_at` must be Some for a `{}` job", status));
+            let finished_at = finished_at.expect(&format!(
+                "`finished_at` must be Some for a `{}` job",
+                status
+            ));
+            let machine_id =
+                machine_id.expect(&format!("`machine_id` must be Some for a `{}` job", status));
+
+            if status == "finished" {
+                CommitJob::Finished(CommitJobFinished {
+                    commit_job,
+                    started_at,
+                    finished_at,
+                    machine_id,
+                })
+            } else {
+                CommitJob::Failed(CommitJobFailed {
+                    commit_job,
+                    started_at,
+                    finished_at,
+                    machine_id,
+                })
+            }
+        }
+
+        other => {
+            panic!("unknown status `{other}` (expected `queued`, `in_progress`, `finished` or `failed`)")
+        }
+    }
+}
+
+/// Given a vector of `CommitJobs` bucket them out into;
+/// `try`, `master` and `release` (in that order)
+pub fn split_queued_commit_jobs(
+    commit_jobs: &[CommitJob],
+) -> (
+    Vec<(&CommitJob, u32)>,
+    Vec<(&CommitJob, u32)>,
+    Vec<(&CommitJob, String)>,
+) {
+    // Split jobs by type as that determines what we enter into the database,
+    // `ToSql` is quite finiky about lifetimes. Moreover the column names
+    // change depending on the commit job type. `master` and `try` have
+    // a `pr` column whereas `release` has a `release_rag` column
+    commit_jobs.iter().fold(
+        (vec![], vec![], vec![]),
+        |(mut try_commits, mut master_commits, mut release_commits), job| {
+            let entity = job
+                .as_queued()
+                .expect("Can only enqueue jobs with a status of `queued`");
+
+            match &entity.job_type {
+                crate::CommitJobType::Try(pr) => try_commits.push((job, *pr)),
+                crate::CommitJobType::Master(pr) => master_commits.push((job, *pr)),
+                crate::CommitJobType::Release(release_tag) => {
+                    release_commits.push((job, release_tag.clone()))
+                }
+            }
+            (try_commits, master_commits, release_commits)
+        },
+    )
+}

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 jobs to the queue
+    async fn enqueue_commit_jobs(&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 dequeue_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) -> bool;
 }
 
 #[async_trait::async_trait]