Auto merge of #403 - pietroalbini:assign-to, r=pietroalbini

Add a flag to assign an experiment to an agent

cc @Mark-Simulacrum

config.toml

@@ -2452,7 +2452,6 @@ ld = { skip = true } #automatic
 ld_preload = { skip-tests = true } #automatic
 ldscript-parser = { skip-tests = true } #automatic
 leaf = { skip = true } #automatic
-leafers = { skip = true } #automatic
 leanshot = { skip = true } #automatic
 ledcat = { skip = true } #automatic
 ledger = { skip-tests = true } #automatic

docs/bot-usage.md

@@ -117,6 +117,8 @@ beta run you can use:
 * `crates`: the selection of crates to use (default: `full`)
 * `cap-lints`: the lints cap (default: `forbid`, which means no cap)
 * `ignore-blacklist`: whether the blacklist should be ignored (default: `false`)
+* `assign`: assign the experiment to a specific agent (use this only when you
+ know what you're doing)
 * `p`: the priority of the run (default: `0`)
 
 [Go back to the TOC][h-toc]
@@ -143,6 +145,8 @@ priority of the `foo` experiment you can use:
 * `crates`: the selection of crates to use (default: `full`)
 * `cap-lints`: the lints cap (default: `forbid`, which means no cap)
 * `ignore-blacklist`: whether the blacklist should be ignored (default: `false`)
+* `assign`: assign the experiment to a specific agent (use this only when you
+ know what you're doing)
 * `p`: the priority of the run (default: `0`)
 
 [Go back to the TOC][h-toc]

src/actions/experiments/create.rs

@@ -1,6 +1,6 @@
 use crate::actions::{experiments::ExperimentError, Action, ActionsCtx};
 use crate::db::QueryUtils;
-use crate::experiments::{CapLints, CrateSelect, Experiment, GitHubIssue, Mode, Status};
+use crate::experiments::{Assignee, CapLints, CrateSelect, Experiment, GitHubIssue, Mode, Status};
 use crate::prelude::*;
 use crate::toolchain::Toolchain;
 use chrono::Utc;
@@ -14,6 +14,7 @@ pub struct CreateExperiment {
  pub priority: i32,
  pub github_issue: Option<GitHubIssue>,
  pub ignore_blacklist: bool,
+ pub assign: Option<Assignee>,
 }
 
 impl CreateExperiment {
@@ -30,6 +31,7 @@ impl CreateExperiment {
  priority: 0,
  github_issue: None,
  ignore_blacklist: false,
+ assign: None,
  }
  }
 }
@@ -52,8 +54,9 @@ impl Action for CreateExperiment {
  transaction.execute(
  "INSERT INTO experiments \
  (name, mode, cap_lints, toolchain_start, toolchain_end, priority, created_at, \
- status, github_issue, github_issue_url, github_issue_number, ignore_blacklist) \
- VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12);",
+ status, github_issue, github_issue_url, github_issue_number, ignore_blacklist, \
+ assigned_to) \
+ VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13);",
  &[
  &self.name,
  &self.mode.to_str(),
@@ -67,6 +70,7 @@ impl Action for CreateExperiment {
  &self.github_issue.as_ref().map(|i| i.html_url.as_str()),
  &self.github_issue.as_ref().map(|i| i.number),
  &self.ignore_blacklist,
+ &self.assign.map(|a| a.to_string()),
  ],
  )?;
 
@@ -92,7 +96,9 @@ mod tests {
  use crate::config::{Config, CrateConfig};
  use crate::crates::Crate;
  use crate::db::{Database, QueryUtils};
- use crate::experiments::{CapLints, CrateSelect, Experiment, GitHubIssue, Mode, Status};
+ use crate::experiments::{
+ Assignee, CapLints, CrateSelect, Experiment, GitHubIssue, Mode, Status,
+ };
  use crate::toolchain::{MAIN_TOOLCHAIN, TEST_TOOLCHAIN};
 
  #[test]
@@ -119,6 +125,7 @@ mod tests {
  number: 10,
  }),
  ignore_blacklist: true,
+ assign: None,
  }
  .apply(&ctx)
  .unwrap();
@@ -147,6 +154,26 @@ mod tests {
  assert!(ex.ignore_blacklist);
  }
 
+ #[test]
+ fn test_creation_with_assign() {
+ let db = Database::temp().unwrap();
+ let config = Config::default();
+ let ctx = ActionsCtx::new(&db, &config);
+
+ crate::crates::lists::setup_test_lists(&db, &config).unwrap();
+
+ CreateExperiment {
+ assign: Some(Assignee::CLI),
+ ..CreateExperiment::dummy("foo")
+ }
+ .apply(&ctx)
+ .unwrap();
+
+ let ex = Experiment::get(&db, "foo").unwrap().unwrap();
+ assert_eq!(ex.status, Status::Queued);
+ assert_eq!(ex.assigned_to, Some(Assignee::CLI));
+ }
+
  #[test]
  fn test_ignore_blacklist() {
  fn is_skipped(db: &Database, ex: &str, krate: &str) -> bool {
@@ -223,6 +250,7 @@ mod tests {
  priority: 0,
  github_issue: None,
  ignore_blacklist: false,
+ assign: None,
  }
  .apply(&ctx)
  .unwrap_err();
@@ -251,6 +279,7 @@ mod tests {
  priority: 0,
  github_issue: None,
  ignore_blacklist: false,
+ assign: None,
  }
  .apply(&ctx)
  .unwrap();
@@ -265,6 +294,7 @@ mod tests {
  priority: 0,
  github_issue: None,
  ignore_blacklist: false,
+ assign: None,
  }
  .apply(&ctx)
  .unwrap_err();

src/actions/experiments/edit.rs

@@ -1,6 +1,6 @@
 use crate::actions::{experiments::ExperimentError, Action, ActionsCtx};
 use crate::db::QueryUtils;
-use crate::experiments::{CapLints, CrateSelect, Experiment, Mode, Status};
+use crate::experiments::{Assignee, CapLints, CrateSelect, Experiment, Mode, Status};
 use crate::prelude::*;
 use crate::toolchain::Toolchain;
 
@@ -12,6 +12,7 @@ pub struct EditExperiment {
  pub cap_lints: Option<CapLints>,
  pub priority: Option<i32>,
  pub ignore_blacklist: Option<bool>,
+ pub assign: Option<Assignee>,
 }
 
 impl EditExperiment {
@@ -25,6 +26,7 @@ impl EditExperiment {
  cap_lints: None,
  priority: None,
  ignore_blacklist: None,
+ assign: None,
  }
  }
 }
@@ -134,6 +136,16 @@ impl Action for EditExperiment {
  ex.priority = priority;
  }
 
+ // Try to update the assignee
+ if let Some(assign) = self.assign {
+ let changes = t.execute(
+ "UPDATE experiments SET assigned_to = ?1 WHERE name = ?2;",
+ &[&assign.to_string(), &self.name],
+ )?;
+ assert_eq!(changes, 1);
+ ex.assigned_to = Some(assign);
+ }
+
  Ok(())
  })?;
  Ok(())
@@ -147,7 +159,7 @@ mod tests {
  use crate::config::{Config, CrateConfig};
  use crate::crates::Crate;
  use crate::db::{Database, QueryUtils};
- use crate::experiments::{CapLints, CrateSelect, Experiment, Mode, Status};
+ use crate::experiments::{Assignee, CapLints, CrateSelect, Experiment, Mode, Status};
  use crate::toolchain::{MAIN_TOOLCHAIN, TEST_TOOLCHAIN};
 
  #[test]
@@ -180,6 +192,7 @@ mod tests {
  priority: 0,
  github_issue: None,
  ignore_blacklist: false,
+ assign: None,
  }
  .apply(&ctx)
  .unwrap();
@@ -196,6 +209,7 @@ mod tests {
  cap_lints: Some(CapLints::Warn),
  priority: Some(10),
  ignore_blacklist: Some(true),
+ assign: Some(Assignee::CLI),
  }
  .apply(&ctx)
  .unwrap();
@@ -209,6 +223,7 @@ mod tests {
  assert_eq!(ex.cap_lints, CapLints::Warn);
  assert_eq!(ex.priority, 10);
  assert_eq!(ex.ignore_blacklist, true);
+ assert_eq!(ex.assigned_to, Some(Assignee::CLI));
 
  assert_eq!(
  ex.get_crates(&ctx.db).unwrap(),

src/cli.rs

@@ -123,6 +123,8 @@ pub enum Crater {
  priority: i32,
  #[structopt(name = "ignore-blacklist", long = "ignore-blacklist")]
  ignore_blacklist: bool,
+ #[structopt(name = "assign", long = "assign")]
+ assign: Option<Assignee>,
  },
 
  #[structopt(name = "edit", about = "edit an experiment configuration")]
@@ -165,6 +167,8 @@ pub enum Crater {
  conflicts_with = "ignore-blacklist"
  )]
  no_ignore_blacklist: bool,
+ #[structopt(name = "assign", long = "assign")]
+ assign: Option<Assignee>,
  },
 
  #[structopt(name = "delete-ex", about = "delete shared data for experiment")]
@@ -307,6 +311,7 @@ impl Crater {
  ref cap_lints,
  ref priority,
  ref ignore_blacklist,
+ ref assign,
  } => {
  let config = Config::load()?;
  let db = Database::open()?;
@@ -321,6 +326,7 @@ impl Crater {
  priority: *priority,
  github_issue: None,
  ignore_blacklist: *ignore_blacklist,
+ assign: assign.clone(),
  }
  .apply(&ctx)?;
  }
@@ -334,6 +340,7 @@ impl Crater {
  ref priority,
  ref ignore_blacklist,
  ref no_ignore_blacklist,
+ ref assign,
  } => {
  let config = Config::load()?;
  let db = Database::open()?;
@@ -355,6 +362,7 @@ impl Crater {
  cap_lints: *cap_lints,
  priority: *priority,
  ignore_blacklist,
+ assign: assign.clone(),
  }
  .apply(&ctx)?;
  }

src/experiments.rs

@@ -179,22 +179,36 @@ impl Experiment {
  return Ok(Some((false, experiment)));
  }
 
- let record = db.get_row(
- "SELECT * FROM experiments \
- WHERE status = \"queued\" \
- ORDER BY priority DESC, created_at;",
- &[],
- |r| ExperimentDBRecord::from_row(r),
- )?;
-
- if let Some(record) = record {
- let mut experiment = record.into_experiment()?;
- experiment.set_status(&db, Status::Running)?;
- experiment.set_assigned_to(&db, Some(assignee))?;
- Ok(Some((true, experiment)))
- } else {
- Ok(None)
+ // First look at queued experiments explicitly assigned to this agent, and then queued
+ // experiments without an assigned agent.
+ for assigned_to in &[Some(assignee), None] {
+ let record = if let Some(assigned_to) = assigned_to {
+ db.get_row(
+ "SELECT * FROM experiments \
+ WHERE status = \"queued\" AND assigned_to = ?1 \
+ ORDER BY priority DESC, created_at;",
+ &[&assigned_to.to_string()],
+ |r| ExperimentDBRecord::from_row(r),
+ )?
+ } else {
+ db.get_row(
+ "SELECT * FROM experiments \
+ WHERE status = \"queued\" AND assigned_to IS NULL \
+ ORDER BY priority DESC, created_at;",
+ &[],
+ |r| ExperimentDBRecord::from_row(r),
+ )?
+ };
+
+ if let Some(record) = record {
+ let mut experiment = record.into_experiment()?;
+ experiment.set_status(&db, Status::Running)?;
+ experiment.set_assigned_to(&db, Some(assignee))?;
+ return Ok(Some((true, experiment)));
+ }
  }
+
+ Ok(None)
  }
 
  pub fn get(db: &Database, name: &str) -> Fallible<Option<Experiment>> {
@@ -488,6 +502,48 @@ mod tests {
  assert!(Experiment::next(&db, &agent3).unwrap().is_none());
  }
 
+ #[test]
+ fn test_assigning_experiment_with_preassigned_agent() {
+ let db = Database::temp().unwrap();
+ let config = Config::load().unwrap();
+
+ crate::crates::lists::setup_test_lists(&db, &config).unwrap();
+
+ let mut tokens = Tokens::default();
+ tokens.agents.insert("token1".into(), "agent-1".into());
+ tokens.agents.insert("token2".into(), "agent-2".into());
+
+ let agent1 = Assignee::Agent("agent-1".to_string());
+ let agent2 = Assignee::Agent("agent-2".to_string());
+
+ // Populate the `agents` table
+ let _ = Agents::new(db.clone(), &tokens).unwrap();
+
+ let config = Config::default();
+ let ctx = ActionsCtx::new(&db, &config);
+
+ let mut create_assigned = CreateExperiment::dummy("assigned");
+ create_assigned.assign = Some(agent1.clone());
+ create_assigned.apply(&ctx).unwrap();
+
+ let mut create_important = CreateExperiment::dummy("important");
+ create_important.priority = 10;
+ create_important.apply(&ctx).unwrap();
+
+ // Try to get an experiment for agent 1, it should pick 'assigned' even if 'important' has
+ // an higher priority.
+ let (new, ex) = Experiment::next(&db, &agent1).unwrap().unwrap();
+ assert!(new);
+ assert_eq!(ex.assigned_to.unwrap(), agent1);
+ assert_eq!(ex.name.as_str(), "assigned");
+
+ // Then the 'important' experiment will be picked by agent 2
+ let (new, ex) = Experiment::next(&db, &agent2).unwrap().unwrap();
+ assert!(new);
+ assert_eq!(ex.assigned_to.unwrap(), agent2);
+ assert_eq!(ex.name.as_str(), "important");
+ }
+
  #[test]
  fn test_completed_crates() {
  use crate::prelude::*;

src/server/routes/webhooks/args.rs

@@ -1,4 +1,4 @@
-use crate::experiments::{CapLints, CrateSelect, Mode};
+use crate::experiments::{Assignee, CapLints, CrateSelect, Mode};
 use crate::toolchain::Toolchain;
 
 #[derive(Debug, Fail)]
@@ -110,6 +110,7 @@ generate_parser!(pub enum Command {
  cap_lints: Option<CapLints> = "cap-lints",
  priority: Option<i32> = "p",
  ignore_blacklist: Option<bool> = "ignore-blacklist",
+ assign: Option<Assignee> = "assign",
  })
 
  "abort" => Abort(AbortArgs {
@@ -137,6 +138,7 @@ generate_parser!(pub enum Command {
  cap_lints: Option<CapLints> = "cap-lints",
  priority: Option<i32> = "p",
  ignore_blacklist: Option<bool> = "ignore-blacklist",
+ assign: Option<Assignee> = "assign",
  })
 });