Skip to content

Commit 8f1763d

Browse files
committed
Database groundwork for job queue
1 parent e31c217 commit 8f1763d

File tree

6 files changed

+314
-12
lines changed

6 files changed

+314
-12
lines changed

collector/src/compile/benchmark/target.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,3 +14,15 @@ impl Default for Target {
1414
Self::X86_64UnknownLinuxGnu
1515
}
1616
}
17+
18+
impl Target {
19+
pub fn all() -> Vec<Self> {
20+
vec![Self::X86_64UnknownLinuxGnu]
21+
}
22+
23+
pub fn from_db_target(target: &database::Target) -> Target {
24+
match target {
25+
database::Target::X86_64UnknownLinuxGnu => Self::X86_64UnknownLinuxGnu,
26+
}
27+
}
28+
}

database/src/lib.rs

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,11 @@ impl Profile {
227227
Profile::Clippy => "clippy",
228228
}
229229
}
230+
231+
/// Set of default profiles that should be benchmarked for a master/try artifact.
232+
pub fn default_profiles() -> Vec<Self> {
233+
vec![Profile::Check, Profile::Debug, Profile::Doc, Profile::Opt]
234+
}
230235
}
231236

232237
impl std::str::FromStr for Profile {
@@ -365,6 +370,10 @@ impl Target {
365370
Target::X86_64UnknownLinuxGnu => "x86_64-unknown-linux-gnu",
366371
}
367372
}
373+
374+
pub fn all() -> Vec<Self> {
375+
vec![Self::X86_64UnknownLinuxGnu]
376+
}
368377
}
369378

370379
impl FromStr for Target {
@@ -988,6 +997,35 @@ impl BenchmarkRequest {
988997
pub fn is_release(&self) -> bool {
989998
matches!(self.commit_type, BenchmarkRequestType::Release { .. })
990999
}
1000+
1001+
/// Get the codegen backends for the request
1002+
pub fn backends(&self) -> anyhow::Result<Vec<CodegenBackend>> {
1003+
// Empty string; default to LLVM.
1004+
if self.backends.trim().is_empty() {
1005+
return Ok(vec![CodegenBackend::Llvm]);
1006+
}
1007+
1008+
self.backends
1009+
.split(',')
1010+
.map(|s| {
1011+
CodegenBackend::from_str(s).map_err(|_| anyhow::anyhow!("Invalid backend: {s}"))
1012+
})
1013+
.collect()
1014+
}
1015+
1016+
/// Get the profiles for the request
1017+
pub fn profiles(&self) -> anyhow::Result<Vec<Profile>> {
1018+
// No profile string; fall back to the library defaults.
1019+
if self.profiles.trim().is_empty() {
1020+
return Ok(Profile::default_profiles());
1021+
}
1022+
1023+
self.profiles
1024+
.split(',')
1025+
.map(Profile::from_str)
1026+
.collect::<Result<Vec<_>, _>>()
1027+
.map_err(|e| anyhow::anyhow!("Invalid backend: {e}"))
1028+
}
9911029
}
9921030

9931031
/// Cached information about benchmark requests in the DB
@@ -1010,3 +1048,99 @@ impl BenchmarkRequestIndex {
10101048
&self.completed
10111049
}
10121050
}
1051+
1052+
#[derive(Debug, Clone, PartialEq)]
1053+
pub enum BenchmarkJobStatus {
1054+
Queued,
1055+
InProgress,
1056+
Success,
1057+
Failure,
1058+
}
1059+
1060+
const BENCHMARK_JOB_STATUS_QUEUED_STR: &str = "queued";
1061+
const BENCHMARK_JOB_STATUS_IN_PROGRESS_STR: &str = "in_progress";
1062+
const BENCHMARK_JOB_STATUS_SUCCESS_STR: &str = "success";
1063+
const BENCHMARK_JOB_STATUS_FAILURE_STR: &str = "failure";
1064+
1065+
impl BenchmarkJobStatus {
1066+
pub fn as_str(&self) -> &str {
1067+
match self {
1068+
BenchmarkJobStatus::Queued => BENCHMARK_JOB_STATUS_QUEUED_STR,
1069+
BenchmarkJobStatus::InProgress => BENCHMARK_JOB_STATUS_IN_PROGRESS_STR,
1070+
BenchmarkJobStatus::Success => BENCHMARK_JOB_STATUS_SUCCESS_STR,
1071+
BenchmarkJobStatus::Failure => BENCHMARK_JOB_STATUS_FAILURE_STR,
1072+
}
1073+
}
1074+
}
1075+
1076+
impl fmt::Display for BenchmarkJobStatus {
1077+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1078+
write!(f, "{}", self.as_str())
1079+
}
1080+
}
1081+
1082+
#[derive(Debug, Clone, PartialEq)]
1083+
pub struct BenchmarkSet(u32);
1084+
1085+
#[derive(Debug, Clone, PartialEq)]
1086+
pub struct BenchmarkJob {
1087+
target: Target,
1088+
backend: CodegenBackend,
1089+
profile: Profile,
1090+
request_tag: String,
1091+
benchmark_set: BenchmarkSet,
1092+
created_at: DateTime<Utc>,
1093+
started_at: Option<DateTime<Utc>>,
1094+
completed_at: Option<DateTime<Utc>>,
1095+
status: BenchmarkJobStatus,
1096+
retry: u32,
1097+
}
1098+
1099+
impl BenchmarkJob {
1100+
pub fn new(
1101+
target: Target,
1102+
backend: CodegenBackend,
1103+
profile: Profile,
1104+
request_tag: &str,
1105+
benchmark_set: u32,
1106+
created_at: DateTime<Utc>,
1107+
status: BenchmarkJobStatus,
1108+
) -> Self {
1109+
BenchmarkJob {
1110+
target,
1111+
backend,
1112+
profile,
1113+
request_tag: request_tag.to_string(),
1114+
benchmark_set: BenchmarkSet(benchmark_set),
1115+
created_at,
1116+
started_at: None,
1117+
completed_at: None,
1118+
status,
1119+
retry: 0,
1120+
}
1121+
}
1122+
1123+
pub fn request_tag(&self) -> &str {
1124+
&self.request_tag
1125+
}
1126+
1127+
pub fn target(&self) -> Target {
1128+
self.target
1129+
}
1130+
1131+
pub fn backend(&self) -> CodegenBackend {
1132+
self.backend
1133+
}
1134+
1135+
pub fn profile(&self) -> Profile {
1136+
self.profile
1137+
}
1138+
1139+
pub fn benchmark_set(&self) -> u32 {
1140+
self.benchmark_set.0
1141+
}
1142+
1143+
pub fn status(&self) -> &BenchmarkJobStatus {
1144+
&self.status
1145+
}
1146+
}

database/src/pool.rs

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::{
2-
ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkRequest, BenchmarkRequestIndex,
3-
BenchmarkRequestStatus, CodegenBackend, CompileBenchmark, Target,
2+
ArtifactCollection, ArtifactId, ArtifactIdNumber, BenchmarkJob, BenchmarkRequest,
3+
BenchmarkRequestIndex, BenchmarkRequestStatus, CodegenBackend, CompileBenchmark, Target,
44
};
55
use crate::{CollectionId, Index, Profile, QueuedCommit, Scenario, Step};
66
use chrono::{DateTime, Utc};
@@ -210,6 +210,9 @@ pub trait Connection: Send + Sync {
210210
sha: &str,
211211
parent_sha: &str,
212212
) -> anyhow::Result<()>;
213+
214+
/// Add a benchmark job to the job queue.
215+
async fn enqueue_benchmark_job(&self, benchmark_job: &BenchmarkJob) -> anyhow::Result<()>;
213216
}
214217

215218
#[async_trait::async_trait]

database/src/pool/postgres.rs

Lines changed: 58 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
use crate::pool::{Connection, ConnectionManager, ManagedConnection, Transaction};
22
use crate::{
3-
ArtifactCollection, ArtifactId, ArtifactIdNumber, Benchmark, BenchmarkRequest,
4-
BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkRequestType, CodegenBackend,
5-
CollectionId, Commit, CommitType, CompileBenchmark, Date, Index, Profile, QueuedCommit,
6-
Scenario, Target, BENCHMARK_REQUEST_MASTER_STR, BENCHMARK_REQUEST_RELEASE_STR,
3+
ArtifactCollection, ArtifactId, ArtifactIdNumber, Benchmark, BenchmarkJob, BenchmarkJobStatus,
4+
BenchmarkRequest, BenchmarkRequestIndex, BenchmarkRequestStatus, BenchmarkRequestType,
5+
CodegenBackend, CollectionId, Commit, CommitType, CompileBenchmark, Date, Index, Profile,
6+
QueuedCommit, Scenario, Target, BENCHMARK_REQUEST_MASTER_STR, BENCHMARK_REQUEST_RELEASE_STR,
77
BENCHMARK_REQUEST_STATUS_ARTIFACTS_READY_STR, BENCHMARK_REQUEST_STATUS_COMPLETED_STR,
88
BENCHMARK_REQUEST_STATUS_IN_PROGRESS_STR, BENCHMARK_REQUEST_TRY_STR,
99
};
@@ -324,6 +324,28 @@ static MIGRATIONS: &[&str] = &[
324324
CREATE UNIQUE INDEX collector_config_target_bench_active_uniq ON collector_config
325325
(target, benchmark_set, is_active) WHERE is_active = TRUE;
326326
"#,
327+
r#"
328+
CREATE TABLE IF NOT EXISTS job_queue (
329+
id SERIAL PRIMARY KEY,
330+
request_tag TEXT NOT NULL,
331+
target TEXT NOT NULL,
332+
backend TEXT NOT NULL,
333+
profile TEXT NOT NULL,
334+
benchmark_set INTEGER NOT NULL,
335+
collector_id TEXT,
336+
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
337+
started_at TIMESTAMPTZ,
338+
completed_at TIMESTAMPTZ,
339+
status TEXT NOT NULL,
340+
retry INTEGER DEFAULT 0,
341+
342+
CONSTRAINT job_queue_request_fk
343+
FOREIGN KEY (request_tag)
344+
REFERENCES benchmark_request(tag)
345+
ON DELETE CASCADE
346+
);
347+
CREATE INDEX IF NOT EXISTS job_queue_request_tag_idx ON job_queue (request_tag);
348+
"#,
327349
];
328350

329351
#[async_trait::async_trait]
@@ -1608,6 +1630,34 @@ where
16081630
.collect();
16091631
Ok(requests)
16101632
}
1633+
1634+
async fn enqueue_benchmark_job(&self, benchmark_job: &BenchmarkJob) -> anyhow::Result<()> {
1635+
self.conn()
1636+
.execute(
1637+
r#"
1638+
INSERT INTO job_queue(
1639+
request_tag,
1640+
target,
1641+
backend,
1642+
profile,
1643+
benchmark_set,
1644+
status
1645+
)
1646+
VALUES ($1, $2, $3, $4, $5, $6)
1647+
"#,
1648+
&[
1649+
&benchmark_job.request_tag(),
1650+
&benchmark_job.target(),
1651+
&benchmark_job.backend(),
1652+
&benchmark_job.profile(),
1653+
&(benchmark_job.benchmark_set() as i32),
1654+
&benchmark_job.status(),
1655+
],
1656+
)
1657+
.await
1658+
.context("failed to insert benchmark_job")?;
1659+
Ok(())
1660+
}
16111661
}
16121662

16131663
fn parse_artifact_id(ty: &str, sha: &str, date: Option<DateTime<Utc>>) -> ArtifactId {
@@ -1653,6 +1703,10 @@ macro_rules! impl_to_postgresql_via_to_string {
16531703

16541704
impl_to_postgresql_via_to_string!(BenchmarkRequestType);
16551705
impl_to_postgresql_via_to_string!(BenchmarkRequestStatus);
1706+
impl_to_postgresql_via_to_string!(Target);
1707+
impl_to_postgresql_via_to_string!(CodegenBackend);
1708+
impl_to_postgresql_via_to_string!(Profile);
1709+
impl_to_postgresql_via_to_string!(BenchmarkJobStatus);
16561710

16571711
#[cfg(test)]
16581712
mod tests {

database/src/pool/sqlite.rs

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
use crate::pool::{Connection, ConnectionManager, ManagedConnection, Transaction};
22
use crate::{
3-
ArtifactCollection, ArtifactId, Benchmark, BenchmarkRequest, BenchmarkRequestIndex,
4-
BenchmarkRequestStatus, CodegenBackend, CollectionId, Commit, CommitType, CompileBenchmark,
5-
Date, Profile, Target,
3+
ArtifactCollection, ArtifactId, Benchmark, BenchmarkJob, BenchmarkRequest,
4+
BenchmarkRequestIndex, BenchmarkRequestStatus, CodegenBackend, CollectionId, Commit,
5+
CommitType, CompileBenchmark, Date, Profile, Target,
66
};
77
use crate::{ArtifactIdNumber, Index, QueuedCommit};
88
use chrono::{DateTime, TimeZone, Utc};
@@ -1296,6 +1296,10 @@ impl Connection for SqliteConnection {
12961296
) -> anyhow::Result<()> {
12971297
no_queue_implementation_abort!()
12981298
}
1299+
1300+
async fn enqueue_benchmark_job(&self, _benchmark_job: &BenchmarkJob) -> anyhow::Result<()> {
1301+
no_queue_implementation_abort!()
1302+
}
12991303
}
13001304

13011305
fn parse_artifact_id(ty: &str, sha: &str, date: Option<i64>) -> ArtifactId {

0 commit comments

Comments
 (0)