From 124e2873e4cd4ae5fa712277fcf325bddf4b2e20 Mon Sep 17 00:00:00 2001 From: apiraino Date: Tue, 7 Jan 2025 12:20:56 +0100 Subject: [PATCH 1/3] Add work queue tracking in triagebot DB Signed-off-by: apiraino --- src/db.rs | 3 +++ src/lib.rs | 11 ++++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/db.rs b/src/db.rs index c74606cf..54f7a05f 100644 --- a/src/db.rs +++ b/src/db.rs @@ -344,4 +344,7 @@ 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 max_assigned_prs INTEGER DEFAULT NULL; +", ]; diff --git a/src/lib.rs b/src/lib.rs index 9a8e8569..ce44f61a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -131,17 +131,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 + ) } } @@ -152,6 +160,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 73653c581b4f715e20291b0d6eaf2440960e5a29 Mon Sep 17 00:00:00 2001 From: apiraino Date: Tue, 7 Jan 2025 12:22:06 +0100 Subject: [PATCH 2/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. Signed-off-by: apiraino --- src/handlers/assign.rs | 144 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 130 insertions(+), 14 deletions(-) diff --git a/src/handlers/assign.rs b/src/handlers/assign.rs index e6e67306..51159b10 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, WarnNonDefaultBranchException}, 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)] @@ -80,6 +84,27 @@ const NON_DEFAULT_BRANCH_EXCEPTION: &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) } @@ -287,6 +312,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 @@ -297,13 +324,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 @@ -317,7 +345,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 {}, \ @@ -327,7 +357,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() @@ -345,7 +377,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!( @@ -507,7 +539,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?; @@ -528,7 +573,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) => { @@ -539,7 +591,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(()); } @@ -597,6 +653,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) => { @@ -618,7 +675,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), @@ -636,6 +693,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 {} @@ -665,13 +727,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) + } } } } @@ -682,7 +754,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, @@ -708,14 +781,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, From 886116c43d44ab6fc620b5f983f2214641bc629d Mon Sep 17 00:00:00 2001 From: apiraino Date: Tue, 7 Jan 2025 12:22:46 +0100 Subject: [PATCH 3/3] Ensure user capacity is respected on PR assignment This check is specifically used when assigning from the Github web UI Signed-off-by: apiraino --- .env.sample | 6 +++- src/handlers.rs | 4 +-- src/handlers/assign.rs | 2 +- src/handlers/pr_tracking.rs | 56 +++++++++++++++++++++++++++++++++---- 4 files changed, 59 insertions(+), 9 deletions(-) diff --git a/.env.sample b/.env.sample index d0092bd7..e8f51e43 100644 --- a/.env.sample +++ b/.env.sample @@ -20,4 +20,8 @@ GITHUB_WEBHOOK_SECRET=MUST_BE_CONFIGURED # ZULIP_API_TOKEN=yyy # Authenticates inbound webhooks from Github -# ZULIP_TOKEN=xxx \ No newline at end of file +# ZULIP_TOKEN=xxx + +# Use another endpoint to retrieve teams of the Rust project (useful for local testing) +# default: https://team-api.infra.rust-lang.org/v1 +# TEAMS_API_URL=http://localhost:8080 diff --git a/src/handlers.rs b/src/handlers.rs index cb125dc8..f59b6295 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -207,7 +207,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, @@ -328,7 +328,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 51159b10..42afb8db 100644 --- a/src/handlers/assign.rs +++ b/src/handlers/assign.rs @@ -820,7 +820,7 @@ 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", +AND CARDINALITY(r.assigned_prs) < LEAST(COALESCE(max_assigned_prs,1000000))", usernames ); let result = db.query(&q, &[]).await.context("Select DB error")?; diff --git a/src/handlers/pr_tracking.rs b/src/handlers/pr_tracking.rs index 980067d5..11e1ed60 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, event.issue.number) .await - .context("Failed to remove PR from workqueue")?; + .context("Failed to remove PR from workq ueue")?; } + // 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, 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) < LEAST(COALESCE(max_assigned_prs,1000000));"; + 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(