Skip to content

Commit

Permalink
PR assignment based on work queue availability
Browse files Browse the repository at this point in the history 
Add checks in multiple points when identifying or finding an assignee.

Filter out team members currently not having capacity, return messages
instructing what to do in case the assignment is rejected.
  • Loading branch information
@apiraino
apiraino committed Apr 15, 2024
1 parent bb17381 commit a8b49b1
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 19 deletions.
145 changes: 130 additions & 15 deletions src/handlers/assign.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@
//! * `@rustbot release-assignment`: Removes the commenter's assignment.
//! * `r? @user`: Assigns to the given user (PRs only).
//!
//! Note: this module does not handle review assignments issued from the
//! GitHub "Assignees" dropdown menu
//!
//! This is capable of assigning to any user, even if they do not have write
//! access to the repo. It does this by fake-assigning the bot and adding a
//! "claimed by" section to the top-level comment.
Expand All @@ -20,7 +23,7 @@
use crate::{
config::AssignConfig,
github::{self, Event, FileDiff, Issue, IssuesAction, Selection},
handlers::{Context, GithubClient, IssuesEvent},
handlers::{pr_tracking::has_user_capacity, Context, GithubClient, IssuesEvent},
interactions::EditIssueBody,
};
use anyhow::{bail, Context as _};
Expand All @@ -30,6 +33,7 @@ use rand::seq::IteratorRandom;
use rust_team_data::v1::Teams;
use std::collections::{HashMap, HashSet};
use std::fmt;
use tokio_postgres::Client as DbClient;
use tracing as log;

#[cfg(test)]
Expand Down Expand Up @@ -75,6 +79,27 @@ const NON_DEFAULT_BRANCH: &str =

const SUBMODULE_WARNING_MSG: &str = "These commits modify **submodules**.";

pub const SELF_ASSIGN_HAS_NO_CAPACITY: &str = "
You have insufficient capacity to be assigned the pull request at this time. PR assignment has been reverted.
Please choose another assignee or increase your assignment limit.
(see [documentation](https://forge.rust-lang.org/triagebot/pr-assignment-tracking.html))";

pub const REVIEWER_HAS_NO_CAPACITY: &str = "
`{username}` has insufficient capacity to be assigned the pull request at this time. PR assignment has been reverted.
Please choose another assignee.
(see [documentation](https://forge.rust-lang.org/triagebot/pr-assignment-tracking.html))";

const NO_REVIEWER_HAS_CAPACITY: &str = "
Could not find a reviewer with enough capacity to be assigned at this time. This is a problem.
Please contact us on [#t-infra](https://rust-lang.zulipchat.com/#narrow/stream/242791-t-infra) on Zulip.
cc: @**Jack Huey** @**apiraino**";

fn on_vacation_msg(user: &str) -> String {
ON_VACATION_WARNING.replace("{username}", user)
}
Expand Down Expand Up @@ -272,6 +297,8 @@ async fn set_assignee(issue: &Issue, github: &GithubClient, username: &str) {
/// Determines who to assign the PR to based on either an `r?` command, or
/// based on which files were modified.
///
/// Will also check if candidates have capacity in their work queue.
///
/// Returns `(assignee, from_comment)` where `assignee` is who to assign to
/// (or None if no assignee could be found). `from_comment` is a boolean
/// indicating if the assignee came from an `r?` command (it is false if
Expand All @@ -282,13 +309,14 @@ async fn determine_assignee(
config: &AssignConfig,
diff: &[FileDiff],
) -> anyhow::Result<(Option<String>, bool)> {
let db_client = ctx.db.get().await;
let teams = crate::team_data::teams(&ctx.github).await?;
if let Some(name) = find_assign_command(ctx, event) {
if is_self_assign(&name, &event.issue.user.login) {
return Ok((Some(name.to_string()), true));
}
// User included `r?` in the opening PR body.
match find_reviewer_from_names(&teams, config, &event.issue, &[name]) {
match find_reviewer_from_names(&db_client, &teams, config, &event.issue, &[name]).await {
Ok(assignee) => return Ok((Some(assignee), true)),
Err(e) => {
event
Expand All @@ -302,7 +330,9 @@ async fn determine_assignee(
// Errors fall-through to try fallback group.
match find_reviewers_from_diff(config, diff) {
Ok(candidates) if !candidates.is_empty() => {
match find_reviewer_from_names(&teams, config, &event.issue, &candidates) {
match find_reviewer_from_names(&db_client, &teams, config, &event.issue, &candidates)
.await
{
Ok(assignee) => return Ok((Some(assignee), false)),
Err(FindReviewerError::TeamNotFound(team)) => log::warn!(
"team {team} not found via diff from PR {}, \
Expand All @@ -312,7 +342,9 @@ async fn determine_assignee(
// TODO: post a comment on the PR if the reviewers were filtered due to being on vacation
Err(
e @ FindReviewerError::NoReviewer { .. }
| e @ FindReviewerError::AllReviewersFiltered { .. },
| e @ FindReviewerError::AllReviewersFiltered { .. }
| e @ FindReviewerError::NoReviewerHasCapacity
| e @ FindReviewerError::ReviewerHasNoCapacity { .. },
) => log::trace!(
"no reviewer could be determined for PR {}: {e}",
event.issue.global_id()
Expand All @@ -330,7 +362,7 @@ async fn determine_assignee(
}

if let Some(fallback) = config.adhoc_groups.get("fallback") {
match find_reviewer_from_names(&teams, config, &event.issue, fallback) {
match find_reviewer_from_names(&db_client, &teams, config, &event.issue, fallback).await {
Ok(assignee) => return Ok((Some(assignee), false)),
Err(e) => {
log::trace!(
Expand Down Expand Up @@ -492,7 +524,20 @@ pub(super) async fn handle_command(
// welcome message).
return Ok(());
}
let db_client = ctx.db.get().await;
if is_self_assign(&name, &event.user().login) {
match has_user_capacity(&db_client, &name).await {
Ok(work_queue) => work_queue.username,
Err(_) => {
issue
.post_comment(
&ctx.github,
&REVIEWER_HAS_NO_CAPACITY.replace("{username}", &name),
)
.await?;
return Ok(());
}
};
name.to_string()
} else {
let teams = crate::team_data::teams(&ctx.github).await?;
Expand All @@ -512,8 +557,14 @@ pub(super) async fn handle_command(
}
}
}

match find_reviewer_from_names(&teams, config, issue, &[team_name.to_string()])
match find_reviewer_from_names(
&db_client,
&teams,
config,
issue,
&[team_name.to_string()],
)
.await
{
Ok(assignee) => assignee,
Err(e) => {
Expand All @@ -524,7 +575,11 @@ pub(super) async fn handle_command(
}
}
};

// This user is validated and can accept the PR
set_assignee(issue, &ctx.github, &username).await;
// This PR will now be registered in the reviewer's work queue
// by the `pr_tracking` handler
return Ok(());
}

Expand Down Expand Up @@ -582,6 +637,7 @@ pub(super) async fn handle_command(

e.apply(&ctx.github, String::new(), &data).await?;

// Assign the PR: user's work queue has been checked and can accept this PR
match issue.set_assignee(&ctx.github, &to_assign).await {
Ok(()) => return Ok(()), // we are done
Err(github::AssignmentError::InvalidAssignee) => {
Expand All @@ -603,7 +659,7 @@ pub(super) async fn handle_command(
}

#[derive(PartialEq, Debug)]
enum FindReviewerError {
pub enum FindReviewerError {
/// User specified something like `r? foo/bar` where that team name could
/// not be found.
TeamNotFound(String),
Expand All @@ -621,6 +677,11 @@ enum FindReviewerError {
initial: Vec<String>,
filtered: Vec<String>,
},
/// No reviewer has capacity to accept a pull request assignment at this time
NoReviewerHasCapacity,
/// The requested reviewer has no capacity to accept a pull request
/// assignment at this time
ReviewerHasNoCapacity { username: String },
}

impl std::error::Error for FindReviewerError {}
Expand Down Expand Up @@ -650,13 +711,23 @@ impl fmt::Display for FindReviewerError {
write!(
f,
"Could not assign reviewer from: `{}`.\n\
User(s) `{}` are either the PR author, already assigned, or on vacation, \
and there are no other candidates.\n\
Use `r?` to specify someone else to assign.",
User(s) `{}` are either the PR author, already assigned, or on vacation. \
If it's a team, we could not find any candidates.\n\
Please use `r?` to specify someone else to assign.",
initial.join(","),
filtered.join(","),
)
}
FindReviewerError::ReviewerHasNoCapacity { username } => {
write!(
f,
"{}",
REVIEWER_HAS_NO_CAPACITY.replace("{username}", username)
)
}
FindReviewerError::NoReviewerHasCapacity => {
write!(f, "{}", NO_REVIEWER_HAS_CAPACITY)
}
}
}
}
Expand All @@ -667,7 +738,8 @@ impl fmt::Display for FindReviewerError {
/// `@octocat`, or names from the owners map. It can contain GitHub usernames,
/// auto-assign groups, or rust-lang team names. It must have at least one
/// entry.
fn find_reviewer_from_names(
async fn find_reviewer_from_names(
db: &DbClient,
teams: &Teams,
config: &AssignConfig,
issue: &Issue,
Expand All @@ -693,14 +765,57 @@ fn find_reviewer_from_names(
//
// These are all ideas for improving the selection here. However, I'm not
// sure they are really worth the effort.
Ok(candidates

// filter out team members without capacity
let filtered_candidates = filter_by_capacity(db, &candidates)
.await
.expect("Error while filtering out team members");

if filtered_candidates.is_empty() {
return Err(FindReviewerError::AllReviewersFiltered {
initial: names.to_vec(),
filtered: names.to_vec(),
});
}

log::debug!("Filtered list of candidates: {:?}", filtered_candidates);

Ok(filtered_candidates
.into_iter()
.choose(&mut rand::thread_rng())
.expect("candidate_reviewers_from_names always returns at least one entry")
.expect("candidate_reviewers_from_names should return at least one entry")
.to_string())
}

/// Returns a list of candidate usernames to choose as a reviewer.
/// Filter out candidates not having review capacity
async fn filter_by_capacity(
db: &DbClient,
candidates: &HashSet<&str>,
) -> anyhow::Result<HashSet<String>> {
let usernames = candidates
.iter()
.map(|c| *c)
.collect::<Vec<&str>>()
.join(",");

let q = format!(
"
SELECT username
FROM review_prefs r
JOIN users on users.user_id=r.user_id
AND username = ANY('{{ {} }}')
AND ARRAY_LENGTH(r.assigned_prs, 1) < max_assigned_prs",
usernames
);
let result = db.query(&q, &[]).await.context("Select DB error")?;
let mut candidates: HashSet<String> = HashSet::new();
for row in result {
candidates.insert(row.get("username"));
}
Ok(candidates)
}

/// returns a list of candidate usernames (from relevant teams) to choose as a reviewer.
fn candidate_reviewers_from_names<'a>(
teams: &'a Teams,
config: &'a AssignConfig,
Expand Down
7 changes: 3 additions & 4 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,10 +137,9 @@ pub struct ReviewPrefs {

impl ReviewPrefs {
fn to_string(&self) -> String {
let capacity = if self.max_assigned_prs.is_some() {
format!("{}", self.max_assigned_prs.expect("This can't fail"))
} else {
String::from("Not set (i.e. unlimited)")
let capacity = match self.max_assigned_prs {
Some(max) => format!("{}", max),
None => String::from("Not set (i.e. unlimited)"),
};
let prs = self
.assigned_prs
Expand Down

0 comments on commit a8b49b1

Please sign in to comment.