From d310b05e8c77f2d6eb9b5b848dc664ed1a79122c Mon Sep 17 00:00:00 2001 From: apiraino Date: Wed, 3 Apr 2024 12:30:02 +0200 Subject: [PATCH 1/3] Add work queue tracking in triagebot DB --- src/db.rs | 1 + src/lib.rs | 12 +++++++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/src/db.rs b/src/db.rs index 1da96372..0bb3f8c7 100644 --- a/src/db.rs +++ b/src/db.rs @@ -335,4 +335,5 @@ CREATE table review_prefs ( CREATE EXTENSION intarray; CREATE UNIQUE INDEX review_prefs_user_id ON review_prefs(user_id); ", + "ALTER TABLE review_prefs ADD COLUMN max_assigned_prs INTEGER DEFAULT NULL;", ]; diff --git a/src/lib.rs b/src/lib.rs index bf3df35a..8aedf117 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -132,17 +132,26 @@ pub struct ReviewPrefs { pub username: String, pub user_id: i64, pub assigned_prs: Vec, + pub max_assigned_prs: Option, } 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 prs = self .assigned_prs .iter() .map(|pr| format!("#{}", pr)) .collect::>() .join(", "); - format!("Username: {}\nAssigned PRs: {}", self.username, prs) + format!( + "Username: {}\nAssigned PRs: {}\nReview capacity: {}", + self.username, prs, capacity + ) } } @@ -153,6 +162,7 @@ impl From for ReviewPrefs { username: row.get("username"), user_id: row.get("user_id"), assigned_prs: row.get("assigned_prs"), + max_assigned_prs: row.get("max_assigned_prs"), } } } From bb173815bb5b60964e9eac0e6b42a299dc1f1284 Mon Sep 17 00:00:00 2001 From: apiraino Date: Wed, 3 Apr 2024 12:34:42 +0200 Subject: [PATCH 2/3] Ensure user capacity is respected on PR assignment This check is specifically used when assigning from the Github web UI --- src/handlers.rs | 4 +-- src/handlers/pr_tracking.rs | 56 +++++++++++++++++++++++++++++++++---- 2 files changed, 53 insertions(+), 7 deletions(-) diff --git a/src/handlers.rs b/src/handlers.rs index f3045a55..3bc63498 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -159,7 +159,7 @@ macro_rules! issue_handlers { // Handle events that happened on issues // -// This is for events that happen only on issues (e.g. label changes). +// This is for events that happen only on issues or pull requests (e.g. label changes or assignments). // Each module in the list must contain the functions `parse_input` and `handle_input`. issue_handlers! { assign, @@ -280,7 +280,7 @@ macro_rules! command_handlers { // // This is for handlers for commands parsed by the `parser` crate. // Each variant of `parser::command::Command` must be in this list, -// preceded by the module containing the coresponding `handle_command` function +// preceded by the module containing the corresponding `handle_command` function command_handlers! { assign: Assign, glacier: Glacier, diff --git a/src/handlers/pr_tracking.rs b/src/handlers/pr_tracking.rs index 2e30ab36..b72172c7 100644 --- a/src/handlers/pr_tracking.rs +++ b/src/handlers/pr_tracking.rs @@ -1,19 +1,23 @@ //! This module updates the PR workqueue of the Rust project contributors +//! Runs after a PR has been assigned or unassigned //! //! Purpose: //! -//! - Adds the PR to the workqueue of one team member (when the PR has been assigned) -//! - Removes the PR from the workqueue of one team member (when the PR is unassigned or closed) +//! - Adds the PR to the workqueue of one team member (after the PR has been assigned) +//! - Removes the PR from the workqueue of one team member (after the PR has been unassigned or closed) use crate::{ config::ReviewPrefsConfig, db::notifications::record_username, github::{IssuesAction, IssuesEvent}, handlers::Context, + ReviewPrefs, }; use anyhow::Context as _; use tokio_postgres::Client as DbClient; +use super::assign::{FindReviewerError, REVIEWER_HAS_NO_CAPACITY, SELF_ASSIGN_HAS_NO_CAPACITY}; + pub(super) struct ReviewPrefsInput {} pub(super) async fn parse_input( @@ -49,7 +53,7 @@ pub(super) async fn handle_input<'a>( ) -> anyhow::Result<()> { let db_client = ctx.db.get().await; - // extract the assignee matching the assignment or unassignment enum variants or return and ignore this handler + // extract the assignee or ignore this handler and return let IssuesEvent { action: IssuesAction::Assigned { assignee } | IssuesAction::Unassigned { assignee }, .. @@ -66,18 +70,60 @@ pub(super) async fn handle_input<'a>( if matches!(event.action, IssuesAction::Unassigned { .. }) { delete_pr_from_workqueue(&db_client, assignee.id.unwrap(), event.issue.number) .await - .context("Failed to remove PR from workqueue")?; + .context("Failed to remove PR from work queue")?; } + // This handler is reached also when assigning a PR using the Github UI + // (i.e. from the "Assignees" dropdown menu). + // We need to also check assignee availability here. if matches!(event.action, IssuesAction::Assigned { .. }) { + let work_queue = has_user_capacity(&db_client, &assignee.login) + .await + .context("Failed to retrieve user work queue"); + + // if user has no capacity, revert the PR assignment (GitHub has already assigned it) + // and post a comment suggesting what to do + if let Err(_) = work_queue { + event + .issue + .remove_assignees(&ctx.github, crate::github::Selection::One(&assignee.login)) + .await?; + + let msg = if assignee.login.to_lowercase() == event.issue.user.login.to_lowercase() { + SELF_ASSIGN_HAS_NO_CAPACITY.replace("{username}", &assignee.login) + } else { + REVIEWER_HAS_NO_CAPACITY.replace("{username}", &assignee.login) + }; + event.issue.post_comment(&ctx.github, &msg).await?; + } + upsert_pr_into_workqueue(&db_client, assignee.id.unwrap(), event.issue.number) .await - .context("Failed to add PR to workqueue")?; + .context("Failed to add PR to work queue")?; } Ok(()) } +pub async fn has_user_capacity( + db: &crate::db::PooledClient, + assignee: &str, +) -> anyhow::Result { + let q = " +SELECT username, r.* +FROM review_prefs r +JOIN users ON users.user_id = r.user_id +WHERE username = $1 +AND CARDINALITY(r.assigned_prs) < max_assigned_prs;"; + let rec = db.query_one(q, &[&assignee]).await; + if let Err(_) = rec { + return Err(FindReviewerError::ReviewerHasNoCapacity { + username: assignee.to_string(), + }); + } + Ok(rec.unwrap().into()) +} + /// Add a PR to the workqueue of a team member. /// Ensures no accidental PR duplicates. async fn upsert_pr_into_workqueue( From 80303014f21a46dda2d4688854834c1a4396a879 Mon Sep 17 00:00:00 2001 From: apiraino Date: Wed, 3 Apr 2024 12:41:21 +0200 Subject: [PATCH 3/3] PR assignment based on work queue availability 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. --- src/handlers/assign.rs | 145 ++++++++++++++++++++++++++++++++++++----- src/lib.rs | 7 +- 2 files changed, 133 insertions(+), 19 deletions(-) diff --git a/src/handlers/assign.rs b/src/handlers/assign.rs index 3a91e501..cacf722d 100644 --- a/src/handlers/assign.rs +++ b/src/handlers/assign.rs @@ -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. @@ -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 _}; @@ -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)] @@ -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: @jackh726 @apiraino"; + fn on_vacation_msg(user: &str) -> String { ON_VACATION_WARNING.replace("{username}", user) } @@ -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 @@ -282,13 +309,14 @@ async fn determine_assignee( config: &AssignConfig, diff: &[FileDiff], ) -> anyhow::Result<(Option, 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 @@ -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 {}, \ @@ -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() @@ -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!( @@ -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?; @@ -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) => { @@ -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(()); } @@ -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) => { @@ -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), @@ -621,6 +677,11 @@ enum FindReviewerError { initial: Vec, filtered: Vec, }, + /// 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 {} @@ -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) + } } } } @@ -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, @@ -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> { + let usernames = candidates + .iter() + .map(|c| *c) + .collect::>() + .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 = 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, diff --git a/src/lib.rs b/src/lib.rs index 8aedf117..ad8e7ba4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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