From 15539b76de88a95a70dd52d1435a35aafb39411b Mon Sep 17 00:00:00 2001 From: apiraino Date: Tue, 16 Apr 2024 19:48:40 +0200 Subject: [PATCH] PR assignment based on work queue availability (take #2) This is a version of #1786 that fixes the two following bugs: There were 2 bugs: 1. the initial migration wasn't formatted correctly and the new schema changes weren't applied 2. The SELECT to find a reviewer would return wrong results when a team member had NULL in the table `review_prefs.max_assigned_prs` --- src/db.rs | 8 +- src/handlers.rs | 4 +- src/handlers/assign.rs | 145 ++++++++++++++++++++++++++++++++---- src/handlers/pr_tracking.rs | 56 ++++++++++++-- src/lib.rs | 11 ++- 5 files changed, 199 insertions(+), 25 deletions(-) diff --git a/src/db.rs b/src/db.rs index 1da96372..91282157 100644 --- a/src/db.rs +++ b/src/db.rs @@ -332,7 +332,11 @@ CREATE table review_prefs ( assigned_prs INT[] NOT NULL DEFAULT array[]::INT[] );", " -CREATE EXTENSION intarray; -CREATE UNIQUE INDEX review_prefs_user_id ON review_prefs(user_id); +CREATE EXTENSION IF NOT EXISTS intarray;", + " +CREATE UNIQUE INDEX IF NOT EXISTS review_prefs_user_id ON review_prefs(user_id); ", + " +ALTER TABLE review_prefs ADD COLUMN IF NOT EXISTS max_assigned_prs INTEGER DEFAULT NULL; +", ]; 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/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/handlers/pr_tracking.rs b/src/handlers/pr_tracking.rs index 2e30ab36..bb024173 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 OR max_assigned_prs IS NULL );"; + 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( diff --git a/src/lib.rs b/src/lib.rs index bf3df35a..ad8e7ba4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -132,17 +132,25 @@ 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 = match self.max_assigned_prs { + Some(max) => format!("{}", max), + None => 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 +161,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"), } } }