Take review preferences into account when determining reviewers

src/config.rs

@@ -91,6 +91,10 @@ pub(crate) struct PingTeamConfig {
     pub(crate) label: Option<String>,
 }
 
+#[derive(PartialEq, Eq, Debug, serde::Deserialize)]
+#[serde(deny_unknown_fields)]
+pub(crate) struct AssignReviewPrefsConfig {}
+
 #[derive(PartialEq, Eq, Debug, serde::Deserialize)]
 #[serde(deny_unknown_fields)]
 pub(crate) struct AssignConfig {
@@ -110,6 +114,9 @@ pub(crate) struct AssignConfig {
     pub(crate) owners: HashMap<String, Vec<String>>,
     #[serde(default)]
     pub(crate) users_on_vacation: HashSet<String>,
+    /// Should review preferences be taken into account when deciding who to assign to a PR?
+    #[serde(default)]
+    pub(crate) review_prefs: Option<AssignReviewPrefsConfig>,
 }
 
 impl AssignConfig {
@@ -590,6 +597,7 @@ mod tests {
                     adhoc_groups: HashMap::new(),
                     owners: HashMap::new(),
                     users_on_vacation: HashSet::from(["jyn514".into()]),
+                    review_prefs: None,
                 }),
                 note: Some(NoteConfig { _empty: () }),
                 ping: Some(PingConfig { teams: ping_teams }),
@@ -661,6 +669,7 @@ mod tests {
                     adhoc_groups: HashMap::new(),
                     owners: HashMap::new(),
                     users_on_vacation: HashSet::new(),
+                    review_prefs: None,
                 }),
                 note: None,
                 ping: None,

src/db/review_prefs.rs

@@ -2,6 +2,7 @@ use crate::db::users::record_username;
 use crate::github::{User, UserId};
 use anyhow::Context;
 use serde::Serialize;
+use std::collections::HashMap;
 
 #[derive(Debug, Serialize)]
 pub struct ReviewPrefs {
@@ -37,6 +38,50 @@ WHERE review_prefs.user_id = $1;";
     Ok(row.map(|r| r.into()))
 }
 
+/// Returns a set of review preferences for all passed usernames.
+/// Usernames are matched regardless of case.
+///
+/// Usernames that are not present in the resulting map have no review preferences configured
+/// in the database.
+pub async fn get_review_prefs_batch<'a>(
+    db: &tokio_postgres::Client,
+    users: &[&'a str],
+) -> anyhow::Result<HashMap<&'a str, ReviewPrefs>> {
+    // We need to make sure that we match users regardless of case, but at the
+    // same time we need to return the originally-cased usernames in the final hashmap.
+    // At the same time, we can't depend on the order of results returned by the DB.
+    // So we need to do some additional bookkeeping here.
+    let lowercase_map: HashMap<String, &str> = users
+        .iter()
+        .map(|name| (name.to_lowercase(), *name))
+        .collect();
+    let lowercase_users: Vec<&str> = lowercase_map.keys().map(|s| s.as_str()).collect();
+
+    // The id/user_id/max_assigned_prs columns have to match the names used in
+    // `From<tokio_postgres::row::Row> for ReviewPrefs`.
+    let query = "
+SELECT lower(u.username) AS username, r.id AS id, r.user_id AS user_id, r.max_assigned_prs AS max_assigned_prs
+FROM review_prefs AS r
+JOIN users AS u ON u.user_id = r.user_id
+WHERE lower(u.username) = ANY($1);";
+
+    Ok(db
+        .query(query, &[&lowercase_users])
+        .await
+        .context("Error retrieving review preferences from usernames")?
+        .into_iter()
+        .map(|row| {
+            // Map back from the lowercase username to the original username.
+            let username_lower: &str = row.get("username");
+            let username = lowercase_map
+                .get(username_lower)
+                .expect("Lowercase username not found");
+            let prefs: ReviewPrefs = row.into();
+            (*username, prefs)
+        })
+        .collect())
+}
+
 /// Updates review preferences of the specified user, or creates them
 /// if they do not exist yet.
 pub async fn upsert_review_prefs(
@@ -67,17 +112,14 @@ mod tests {
     use crate::db::review_prefs::{get_review_prefs, upsert_review_prefs};
     use crate::db::users::get_user;
     use crate::tests::github::user;
-    use crate::tests::run_test;
+    use crate::tests::run_db_test;
 
     #[tokio::test]
     async fn insert_prefs_create_user() {
-        run_test(|ctx| async {
-            let db = ctx.db_client().await;
-
+        run_db_test(|ctx| async {
             let user = user("Martin", 1);
-            upsert_review_prefs(&db, user.clone(), Some(1)).await?;
-
-            assert_eq!(get_user(&db, user.id).await?.unwrap(), user);
+            upsert_review_prefs(&ctx.db_client(), user.clone(), Some(1)).await?;
+            assert_eq!(get_user(&ctx.db_client(), user.id).await?.unwrap(), user);
 
             Ok(ctx)
         })
@@ -86,12 +128,13 @@ mod tests {
 
     #[tokio::test]
     async fn insert_max_assigned_prs() {
-        run_test(|ctx| async {
-            let db = ctx.db_client().await;
-
-            upsert_review_prefs(&db, user("Martin", 1), Some(5)).await?;
+        run_db_test(|ctx| async {
+            upsert_review_prefs(&ctx.db_client(), user("Martin", 1), Some(5)).await?;
             assert_eq!(
-                get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
+                get_review_prefs(&ctx.db_client(), 1)
+                    .await?
+                    .unwrap()
+                    .max_assigned_prs,
                 Some(5)
             );
 
@@ -102,8 +145,8 @@ mod tests {
 
     #[tokio::test]
     async fn update_max_assigned_prs() {
-        run_test(|ctx| async {
-            let db = ctx.db_client().await;
+        run_db_test(|ctx| async {
+            let db = ctx.db_client();
 
             upsert_review_prefs(&db, user("Martin", 1), Some(5)).await?;
             assert_eq!(

src/db/users.rs

@@ -41,12 +41,12 @@ WHERE user_id = $1;",
 #[cfg(test)]
 mod tests {
     use crate::db::users::{get_user, record_username};
-    use crate::tests::run_test;
+    use crate::tests::run_db_test;
 
     #[tokio::test]
     async fn update_username_on_conflict() {
-        run_test(|ctx| async {
-            let db = ctx.db_client().await;
+        run_db_test(|ctx| async {
+            let db = ctx.db_client();
 
             record_username(&db, 1, "Foo").await?;
             record_username(&db, 1, "Bar").await?;

src/handlers/assign.rs

@@ -20,6 +20,9 @@
 //! `assign.owners` config, it will auto-select an assignee based on the files
 //! the PR modifies.
 
+use crate::db::review_prefs::get_review_prefs_batch;
+use crate::github::UserId;
+use crate::handlers::pr_tracking::ReviewerWorkqueue;
 use crate::{
     config::AssignConfig,
     github::{self, Event, FileDiff, Issue, IssuesAction, Selection},
@@ -33,6 +36,8 @@ use rand::seq::IteratorRandom;
 use rust_team_data::v1::Teams;
 use std::collections::{HashMap, HashSet};
 use std::fmt;
+use std::sync::Arc;
+use tokio::sync::RwLock;
 use tokio_postgres::Client as DbClient;
 use tracing as log;
 
@@ -265,7 +270,16 @@ async fn determine_assignee(
     let teams = crate::team_data::teams(&ctx.github).await?;
     if let Some(name) = find_assign_command(ctx, event) {
         // User included `r?` in the opening PR body.
-        match find_reviewer_from_names(&db_client, &teams, config, &event.issue, &[name]).await {
+        match find_reviewer_from_names(
+            &db_client,
+            ctx.workqueue.clone(),
+            &teams,
+            config,
+            &event.issue,
+            &[name],
+        )
+        .await
+        {
             Ok(assignee) => return Ok((Some(assignee), true)),
             Err(e) => {
                 event
@@ -279,8 +293,15 @@ 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(&db_client, &teams, config, &event.issue, &candidates)
-                .await
+            match find_reviewer_from_names(
+                &db_client,
+                ctx.workqueue.clone(),
+                &teams,
+                config,
+                &event.issue,
+                &candidates,
+            )
+            .await
             {
                 Ok(assignee) => return Ok((Some(assignee), false)),
                 Err(FindReviewerError::TeamNotFound(team)) => log::warn!(
@@ -292,7 +313,9 @@ async fn determine_assignee(
                     e @ FindReviewerError::NoReviewer { .. }
                     | e @ FindReviewerError::ReviewerIsPrAuthor { .. }
                     | e @ FindReviewerError::ReviewerAlreadyAssigned { .. }
-                    | e @ FindReviewerError::ReviewerOnVacation { .. },
+                    | e @ FindReviewerError::ReviewerOnVacation { .. }
+                    | e @ FindReviewerError::DatabaseError(_)
+                    | e @ FindReviewerError::ReviewerAtMaxCapacity { .. },
                 ) => log::trace!(
                     "no reviewer could be determined for PR {}: {e}",
                     event.issue.global_id()
@@ -310,7 +333,16 @@ async fn determine_assignee(
     }
 
     if let Some(fallback) = config.adhoc_groups.get("fallback") {
-        match find_reviewer_from_names(&db_client, &teams, config, &event.issue, fallback).await {
+        match find_reviewer_from_names(
+            &db_client,
+            ctx.workqueue.clone(),
+            &teams,
+            config,
+            &event.issue,
+            fallback,
+        )
+        .await
+        {
             Ok(assignee) => return Ok((Some(assignee), false)),
             Err(e) => {
                 log::trace!(
@@ -488,6 +520,7 @@ pub(super) async fn handle_command(
         let db_client = ctx.db.get().await;
         let assignee = match find_reviewer_from_names(
             &db_client,
+            ctx.workqueue.clone(),
             &teams,
             config,
             issue,
@@ -612,6 +645,10 @@ pub enum FindReviewerError {
     ReviewerIsPrAuthor { username: String },
     /// Requested reviewer is already assigned to that PR
     ReviewerAlreadyAssigned { username: String },
+    /// Data required for assignment could not be loaded from the DB.
+    DatabaseError(String),
+    /// The reviewer has too many PRs alreayd assigned.
+    ReviewerAtMaxCapacity { username: String },
 }
 
 impl std::error::Error for FindReviewerError {}
@@ -654,6 +691,17 @@ impl fmt::Display for FindReviewerError {
                     REVIEWER_ALREADY_ASSIGNED.replace("{username}", username)
                 )
             }
+            FindReviewerError::DatabaseError(error) => {
+                write!(f, "Database error: {error}")
+            }
+            FindReviewerError::ReviewerAtMaxCapacity { username } => {
+                write!(
+                    f,
+                    r"`{username}` has too many PRs assigned to them.
+
+Please select a different reviewer.",
+                )
+            }
         }
     }
 }
@@ -665,7 +713,8 @@ impl fmt::Display for FindReviewerError {
 /// auto-assign groups, or rust-lang team names. It must have at least one
 /// entry.
 async fn find_reviewer_from_names(
-    _db: &DbClient,
+    db: &DbClient,
+    workqueue: Arc<RwLock<ReviewerWorkqueue>>,
     teams: &Teams,
     config: &AssignConfig,
     issue: &Issue,
@@ -678,7 +727,10 @@ async fn find_reviewer_from_names(
         }
     }
 
-    let candidates = candidate_reviewers_from_names(teams, config, issue, names)?;
+    let candidates =
+        candidate_reviewers_from_names(db, workqueue, teams, config, issue, names).await?;
+    assert!(!candidates.is_empty());
+
     // This uses a relatively primitive random choice algorithm.
     // GitHub's CODEOWNERS supports much more sophisticated options, such as:
     //
@@ -782,19 +834,23 @@ fn expand_teams_and_groups(
 
 /// Returns a list of candidate usernames (from relevant teams) to choose as a reviewer.
 /// If not reviewer is available, returns an error.
-fn candidate_reviewers_from_names<'a>(
+async fn candidate_reviewers_from_names<'a>(
+    db: &DbClient,
+    workqueue: Arc<RwLock<ReviewerWorkqueue>>,
     teams: &'a Teams,
     config: &'a AssignConfig,
     issue: &Issue,
     names: &'a [String],
 ) -> Result<HashSet<String>, FindReviewerError> {
+    // Step 1: expand teams and groups into candidate names
     let (expanded, expansion_happened) = expand_teams_and_groups(teams, issue, config, names)?;
     let expanded_count = expanded.len();
 
     // Set of candidate usernames to choose from.
     // We go through each expanded candidate and store either success or an error for them.
     let mut candidates: Vec<Result<String, FindReviewerError>> = Vec::new();
 
+    // Step 2: pre-filter candidates based on checks that we can perform quickly
     for candidate in expanded {
         let name_lower = candidate.to_lowercase();
         let is_pr_author = name_lower == issue.user.login.to_lowercase();
@@ -831,9 +887,50 @@ fn candidate_reviewers_from_names<'a>(
     }
     assert_eq!(candidates.len(), expanded_count);
 
-    let valid_candidates: HashSet<String> = candidates
+    if config.review_prefs.is_some() {
+        // Step 3: gather potential usernames to form a DB query for review preferences
+        let usernames: Vec<String> = candidates
+            .iter()
+            .filter_map(|res| res.as_deref().ok().map(|s| s.to_string()))
+            .collect();
+        let usernames: Vec<&str> = usernames.iter().map(|s| s.as_str()).collect();
+        let review_prefs = get_review_prefs_batch(db, &usernames)
+            .await
+            .context("cannot fetch review preferences")
+            .map_err(|e| FindReviewerError::DatabaseError(e.to_string()))?;
+
+        let workqueue = workqueue.read().await;
+
+        // Step 4: check review preferences
+        candidates = candidates
+            .into_iter()
+            .map(|username| {
+                // Only consider candidates that did not have an earlier error
+                let username = username?;
+
+                // If no review prefs were found, we assume the default unlimited
+                // review capacity.
+                let Some(review_prefs) = review_prefs.get(username.as_str()) else {
+                    return Ok(username);
+                };
+                let Some(capacity) = review_prefs.max_assigned_prs else {
+                    return Ok(username);
+                };
+                let assigned_prs = workqueue.assigned_pr_count(review_prefs.user_id as UserId);
+                // Can we assign one more PR?
+                if (assigned_prs as i32) < capacity {
+                    Ok(username)
+                } else {
+                    Err(FindReviewerError::ReviewerAtMaxCapacity { username })
+                }
+            })
+            .collect();
+    }
+    assert_eq!(candidates.len(), expanded_count);
+
+    let valid_candidates: HashSet<&str> = candidates
         .iter()
-        .filter_map(|res| res.as_ref().ok().cloned())
+        .filter_map(|res| res.as_deref().ok())
         .collect();
 
     if valid_candidates.is_empty() {
@@ -860,6 +957,9 @@ fn candidate_reviewers_from_names<'a>(
             })
         }
     } else {
-        Ok(valid_candidates)
+        Ok(valid_candidates
+            .into_iter()
+            .map(|s| s.to_string())
+            .collect())
     }
 }