Skip to content

Commit a55a953

Browse files
committed
Store pull request titles in the workqueue
This allows us to display more useful output in `work show`, and it also opens up opportunities to add more interesting PR data to be stored in the future.
1 parent 1a5e499 commit a55a953

File tree

3 files changed

+74
-21
lines changed

3 files changed

+74
-21
lines changed

src/handlers/assign/tests/tests_candidates.rs

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
use super::super::*;
44
use crate::db::review_prefs::{upsert_review_prefs, RotationMode};
55
use crate::github::{PullRequestNumber, User};
6-
use crate::handlers::pr_tracking::ReviewerWorkqueue;
6+
use crate::handlers::pr_tracking::{AssignedPullRequest, ReviewerWorkqueue};
77
use crate::tests::github::{issue, user};
88
use crate::tests::{run_db_test, TestContext};
99

@@ -13,7 +13,7 @@ struct AssignCtx {
1313
teams: Teams,
1414
config: AssignConfig,
1515
issue: Issue,
16-
reviewer_workqueue: HashMap<UserId, HashSet<PullRequestNumber>>,
16+
reviewer_workqueue: HashMap<UserId, HashMap<PullRequestNumber, AssignedPullRequest>>,
1717
}
1818

1919
impl AssignCtx {
@@ -50,7 +50,16 @@ impl AssignCtx {
5050
}
5151

5252
fn assign_prs(mut self, user_id: UserId, count: u64) -> Self {
53-
let prs: HashSet<PullRequestNumber> = (0..count).collect();
53+
let prs = (0..count)
54+
.map(|pr_number| {
55+
(
56+
pr_number,
57+
AssignedPullRequest {
58+
title: format!("PR {pr_number}"),
59+
},
60+
)
61+
})
62+
.collect();
5463
self.reviewer_workqueue.insert(user_id, prs);
5564
self
5665
}

src/handlers/pr_tracking.rs

Lines changed: 54 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,22 +18,29 @@ use octocrab::models::IssueState;
1818
use octocrab::params::pulls::Sort;
1919
use octocrab::params::{Direction, State};
2020
use octocrab::Octocrab;
21-
use std::collections::{HashMap, HashSet};
21+
use std::collections::HashMap;
2222
use tokio::sync::RwLockWriteGuard;
2323
use tracing as log;
2424

25+
#[derive(Clone, Debug)]
26+
pub struct AssignedPullRequest {
27+
pub title: String,
28+
}
29+
2530
/// Maps users to a set of currently assigned open non-draft pull requests.
2631
/// We store this map in memory, rather than in the DB, because it can get desynced when webhooks
2732
/// are missed.
2833
/// It is thus reloaded when triagebot starts and also periodically, so it is not needed to store it
2934
/// in the DB.
3035
#[derive(Debug, Default)]
3136
pub struct ReviewerWorkqueue {
32-
reviewers: HashMap<UserId, HashSet<PullRequestNumber>>,
37+
reviewers: HashMap<UserId, HashMap<PullRequestNumber, AssignedPullRequest>>,
3338
}
3439

3540
impl ReviewerWorkqueue {
36-
pub fn new(reviewers: HashMap<UserId, HashSet<PullRequestNumber>>) -> Self {
41+
pub fn new(
42+
reviewers: HashMap<UserId, HashMap<PullRequestNumber, AssignedPullRequest>>,
43+
) -> Self {
3744
Self { reviewers }
3845
}
3946

@@ -112,6 +119,10 @@ pub(super) async fn handle_input<'a>(
112119
return Ok(());
113120
}
114121

122+
let assigned_pr = AssignedPullRequest {
123+
title: pr.title.clone(),
124+
};
125+
115126
match input {
116127
// The PR was assigned to a specific user, and it is waiting for a review.
117128
ReviewPrefsInput::Assigned { assignee } => {
@@ -120,7 +131,7 @@ pub(super) async fn handle_input<'a>(
120131
assignee.login
121132
);
122133

123-
upsert_pr_into_user_queue(&mut workqueue, assignee.id, pr_number);
134+
upsert_pr_into_user_queue(&mut workqueue, assignee.id, pr_number, assigned_pr);
124135
}
125136
ReviewPrefsInput::Unassigned { assignee } => {
126137
log::info!(
@@ -139,7 +150,12 @@ pub(super) async fn handle_input<'a>(
139150
// receive.
140151
ReviewPrefsInput::OtherChange => {
141152
for assignee in &event.issue.assignees {
142-
if upsert_pr_into_user_queue(&mut workqueue, assignee.id, pr_number) {
153+
if upsert_pr_into_user_queue(
154+
&mut workqueue,
155+
assignee.id,
156+
pr_number,
157+
assigned_pr.clone(),
158+
) {
143159
log::info!("Adding PR {pr_number} to workqueue of {}.", assignee.login);
144160
}
145161
}
@@ -155,10 +171,11 @@ pub async fn load_workqueue(client: &Octocrab) -> anyhow::Result<ReviewerWorkque
155171
let prs = retrieve_pull_request_assignments("rust-lang", "rust", &client).await?;
156172

157173
// Aggregate PRs by user
158-
let aggregated: HashMap<UserId, HashSet<PullRequestNumber>> =
159-
prs.into_iter().fold(HashMap::new(), |mut acc, (user, pr)| {
174+
let aggregated: HashMap<UserId, HashMap<PullRequestNumber, AssignedPullRequest>> = prs
175+
.into_iter()
176+
.fold(HashMap::new(), |mut acc, (user, pr_number, pr)| {
160177
let prs = acc.entry(user.id).or_default();
161-
prs.insert(pr);
178+
prs.insert(pr_number, pr);
162179
acc
163180
});
164181
tracing::debug!("PR assignments\n{aggregated:?}");
@@ -174,7 +191,7 @@ pub async fn retrieve_pull_request_assignments(
174191
owner: &str,
175192
repository: &str,
176193
client: &Octocrab,
177-
) -> anyhow::Result<Vec<(User, PullRequestNumber)>> {
194+
) -> anyhow::Result<Vec<(User, PullRequestNumber, AssignedPullRequest)>> {
178195
let mut assignments = vec![];
179196

180197
// We use the REST API to fetch open pull requests, as it is much (~5-10x)
@@ -224,6 +241,9 @@ pub async fn retrieve_pull_request_assignments(
224241
id: (*user.id).into(),
225242
},
226243
pr.number,
244+
AssignedPullRequest {
245+
title: pr.title.clone().unwrap_or_default(),
246+
},
227247
));
228248
}
229249
}
@@ -234,7 +254,10 @@ pub async fn retrieve_pull_request_assignments(
234254
}
235255

236256
/// Get pull request assignments for a team member
237-
pub async fn get_assigned_prs(ctx: &Context, user_id: UserId) -> HashSet<PullRequestNumber> {
257+
pub async fn get_assigned_prs(
258+
ctx: &Context,
259+
user_id: UserId,
260+
) -> HashMap<PullRequestNumber, AssignedPullRequest> {
238261
ctx.workqueue
239262
.read()
240263
.await
@@ -245,15 +268,22 @@ pub async fn get_assigned_prs(ctx: &Context, user_id: UserId) -> HashSet<PullReq
245268
}
246269

247270
/// Add a PR to the workqueue of a team member.
271+
/// Updates data of the pull request if it already was in the workqueue.
248272
/// Ensures no accidental PR duplicates.
249273
///
250274
/// Returns true if the PR was actually inserted.
251275
fn upsert_pr_into_user_queue(
252276
workqueue: &mut RwLockWriteGuard<ReviewerWorkqueue>,
253277
user_id: UserId,
254278
pr: PullRequestNumber,
279+
assigned_pr: AssignedPullRequest,
255280
) -> bool {
256-
workqueue.reviewers.entry(user_id).or_default().insert(pr)
281+
workqueue
282+
.reviewers
283+
.entry(user_id)
284+
.or_default()
285+
.insert(pr, assigned_pr)
286+
.is_none()
257287
}
258288

259289
/// Delete a PR from the workqueue of a team member.
@@ -270,7 +300,7 @@ fn delete_pr_from_user_queue(
270300
/// Delete a PR from the workqueue completely.
271301
fn delete_pr_from_all_queues(workqueue: &mut ReviewerWorkqueue, pr: PullRequestNumber) {
272302
for queue in workqueue.reviewers.values_mut() {
273-
queue.retain(|pr_number| *pr_number != pr);
303+
queue.retain(|pr_number, _| *pr_number != pr);
274304
}
275305
}
276306

@@ -308,7 +338,9 @@ mod tests {
308338
use crate::config::Config;
309339
use crate::github::{Issue, IssuesAction, IssuesEvent, Repository, User};
310340
use crate::github::{Label, PullRequestNumber};
311-
use crate::handlers::pr_tracking::{handle_input, parse_input, upsert_pr_into_user_queue};
341+
use crate::handlers::pr_tracking::{
342+
handle_input, parse_input, upsert_pr_into_user_queue, AssignedPullRequest,
343+
};
312344
use crate::tests::github::{default_test_user, issue, pull_request, user};
313345
use crate::tests::{run_db_test, TestContext};
314346

@@ -504,7 +536,7 @@ mod tests {
504536
.get(&user.id)
505537
.cloned()
506538
.unwrap_or_default()
507-
.into_iter()
539+
.into_keys()
508540
.collect::<Vec<_>>();
509541
assigned.sort();
510542
assert_eq!(assigned, expected_prs);
@@ -514,7 +546,14 @@ mod tests {
514546
{
515547
let mut workqueue = ctx.handler_ctx().workqueue.write().await;
516548
for &pr in prs {
517-
upsert_pr_into_user_queue(&mut workqueue, user.id, pr);
549+
upsert_pr_into_user_queue(
550+
&mut workqueue,
551+
user.id,
552+
pr,
553+
AssignedPullRequest {
554+
title: format!("PR {pr}"),
555+
},
556+
);
518557
}
519558
}
520559
check_assigned_prs(&ctx, user, prs).await;

src/zulip.rs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,7 @@ async fn workqueue_commands(
289289
.await
290290
.into_iter()
291291
.collect::<Vec<_>>();
292-
assigned_prs.sort();
292+
assigned_prs.sort_by_key(|(pr_number, _)| *pr_number);
293293

294294
let review_prefs = get_review_prefs(&db_client, gh_id)
295295
.await
@@ -309,9 +309,14 @@ async fn workqueue_commands(
309309

310310
let prs = assigned_prs
311311
.iter()
312-
.map(|pr| format!("- [#{pr}](https://github.com/rust-lang/rust/pull/{pr})"))
312+
.map(|(pr_number, pr)| {
313+
format!(
314+
"- [#{pr_number}](https://github.com/rust-lang/rust/pull/{pr_number}) {}",
315+
pr.title
316+
)
317+
})
313318
.collect::<Vec<String>>()
314-
.join(", ");
319+
.join("\n");
315320
let mut response = format!(
316321
"`rust-lang/rust` PRs in your review queue ({} {}):\n{prs}\n",
317322
assigned_prs.len(),

0 commit comments

Comments
 (0)