Don't reimplement comment threads for attention set

As we extracted the logic for comment threads with Ie2713ebdb, we can
use it for attention set. There was also a corner case which wasn't
covered by the logic which was previously used. We add a test case for
it to avoid future regressions.

In AttentionSetIT, we also switched to a different approach for looking
up the ChangeData instance. That was necessary due to the different
types used in the new test but also had a different reason: There's
no guarantee that the ChangeData object returned by PushOneCommit.Result
is updated after the push happened. We skipped additional improvements
like looking up attention set updates via a cleaner/better way than a
ChangeData instance as that's much beyond the scope of this change.

Change-Id: I9d8091c3ad6575304de9a5698cb1c94e32421e7b

java/com/google/gerrit/server/CommentsUtil.java

@@ -21,7 +21,6 @@ import static java.util.stream.Collectors.toCollection;
 import static java.util.stream.Collectors.toList;
 
 import com.google.common.collect.ComparisonChain;
-import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Ordering;
 import com.google.gerrit.common.Nullable;
@@ -51,12 +50,8 @@ import java.io.IOException;
 import java.sql.Timestamp;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
-import java.util.Set;
-import java.util.stream.Collectors;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.Ref;
 import org.eclipse.jgit.lib.Repository;
@@ -354,43 +349,6 @@ public class CommentsUtil {
     update.deleteCommentByRewritingHistory(commentKey.uuid, newMessage);
   }
 
-  /**
-   * Gets all of the {@link HumanComment} in the comment threads that received a reply.
-   *
-   * @param changeNotes notes of this change.
-   * @param newComments set of all the new comments added on the change by the current user.
-   * @return set of all comments in the comments thread that received a reply.
-   */
-  public Set<HumanComment> getAllHumanCommentsInCommentThreads(
-      ChangeNotes changeNotes, ImmutableSet<HumanComment> newComments) {
-    Map<String, HumanComment> uuidToComment =
-        publishedHumanCommentsByChange(changeNotes).stream()
-            .collect(Collectors.toMap(c -> c.key.uuid, c -> c));
-
-    // Copy the set so that it won't be mutated.
-    List<HumanComment> toTraverse = new ArrayList<>(newComments);
-    Set<String> seen = new HashSet<>();
-    Set<HumanComment> allCommentsInCommentThreads = new HashSet<>();
-    while (!toTraverse.isEmpty()) {
-      HumanComment current = toTraverse.remove(0);
-      allCommentsInCommentThreads.add(current);
-
-      if (current.parentUuid != null) {
-        HumanComment parent = uuidToComment.get(current.parentUuid);
-        if (parent == null) {
-          // If we can't find the parent within the human comments, the parent must be a robot
-          // comment and can be ignored.
-          continue;
-        }
-        if (!seen.contains(current.parentUuid)) {
-          toTraverse.add(parent);
-          seen.add(current.parentUuid);
-        }
-      }
-    }
-    return allCommentsInCommentThreads;
-  }
-
   private static List<HumanComment> commentsOnFile(
       Collection<HumanComment> allComments, String file) {
     List<HumanComment> result = new ArrayList<>(allComments.size());

java/com/google/gerrit/server/change/CommentThreads.java

@@ -17,6 +17,7 @@ package com.google.gerrit.server.change;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static java.util.stream.Collectors.groupingBy;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import com.google.gerrit.entities.Comment;
@@ -24,6 +25,7 @@ import java.util.Comparator;
 import java.util.Map;
 import java.util.PriorityQueue;
 import java.util.Queue;
+import java.util.function.Function;
 
 /**
  * Identifier of comment threads.
@@ -40,31 +42,74 @@ import java.util.Queue;
  */
 public class CommentThreads<T extends Comment> {
 
-  private final Iterable<T> comments;
+  private final ImmutableMap<String, T> commentPerUuid;
+  private final Map<String, ImmutableSet<T>> childrenPerParent;
 
-  private CommentThreads(Iterable<T> comments) {
+  public CommentThreads(
-    this.comments = comments;
+      ImmutableMap<String, T> commentPerUuid, Map<String, ImmutableSet<T>> childrenPerParent) {
+    this.commentPerUuid = commentPerUuid;
+    this.childrenPerParent = childrenPerParent;
   }
 
   public static <T extends Comment> CommentThreads<T> forComments(Iterable<T> comments) {
-    return new CommentThreads<>(comments);
+    ImmutableMap<String, T> commentPerUuid =
-  }
-
-  public ImmutableSet<CommentThread<T>> getThreads() {
-    ImmutableSet<String> commentUuids =
-        Streams.stream(comments).map(comment -> comment.key.uuid).collect(toImmutableSet());
-
-    ImmutableSet<T> roots =
         Streams.stream(comments)
-            .filter(
+            .distinct()
-                comment -> comment.parentUuid == null || !commentUuids.contains(comment.parentUuid))
+            .collect(ImmutableMap.toImmutableMap(comment -> comment.key.uuid, Function.identity()));
-            .collect(toImmutableSet());
 
     Map<String, ImmutableSet<T>> childrenPerParent =
-        Streams.stream(comments)
+        commentPerUuid.values().stream()
             .filter(comment -> comment.parentUuid != null)
             .collect(groupingBy(comment -> comment.parentUuid, toImmutableSet()));
+    return new CommentThreads<>(commentPerUuid, childrenPerParent);
+  }
 
+  /**
+   * Returns all comments organized into threads.
+   *
+   * <p>Comments appear only once.
+   */
+  public ImmutableSet<CommentThread<T>> getThreads() {
+    ImmutableSet<T> roots =
+        commentPerUuid.values().stream().filter(this::isRoot).collect(toImmutableSet());
+
+    return buildThreadsOf(roots);
+  }
+
+  /**
+   * Returns only the comment threads to which the specified comments are a reply.
+   *
+   * <p>If the specified child comments are part of the comments originally provided to {@link
+   * CommentThreads#forComments(Iterable)}, they will also appear in the returned comment threads.
+   * They don't need to be part of the originally provided comments, though, but should refer to one
+   * of these comments via their {@link Comment#parentUuid}. Child comments not referring to any
+   * known comments will be ignored.
+   *
+   * @param childComments comments for which the matching threads should be determined
+   * @return threads to which the provided child comments are a reply
+   */
+  public ImmutableSet<CommentThread<T>> getThreadsForChildren(Iterable<? extends T> childComments) {
+    ImmutableSet<T> relevantRoots =
+        Streams.stream(childComments)
+            .map(this::findRoot)
+            .filter(root -> commentPerUuid.containsKey(root.key.uuid))
+            .collect(toImmutableSet());
+    return buildThreadsOf(relevantRoots);
+  }
+
+  private T findRoot(T comment) {
+    T current = comment;
+    while (!isRoot(current)) {
+      current = commentPerUuid.get(current.parentUuid);
+    }
+    return current;
+  }
+
+  private boolean isRoot(T current) {
+    return current.parentUuid == null || !commentPerUuid.containsKey(current.parentUuid);
+  }
+
+  private ImmutableSet<CommentThread<T>> buildThreadsOf(ImmutableSet<T> roots) {
     return roots.stream()
         .map(root -> buildCommentThread(root, childrenPerParent))
         .collect(toImmutableSet());

java/com/google/gerrit/server/restapi/change/ReplyAttentionSetUpdates.java

@@ -17,6 +17,8 @@ package com.google.gerrit.server.restapi.change;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Sets;
+import com.google.common.collect.Sets.SetView;
 import com.google.gerrit.entities.Account;
 import com.google.gerrit.entities.HumanComment;
 import com.google.gerrit.entities.PatchSet;
@@ -32,6 +34,8 @@ import com.google.gerrit.server.account.AccountResolver;
 import com.google.gerrit.server.account.ServiceUserClassifier;
 import com.google.gerrit.server.change.AddToAttentionSetOp;
 import com.google.gerrit.server.change.AttentionSetUnchangedOp;
+import com.google.gerrit.server.change.CommentThread;
+import com.google.gerrit.server.change.CommentThreads;
 import com.google.gerrit.server.change.RemoveFromAttentionSetOp;
 import com.google.gerrit.server.notedb.ChangeNotes;
 import com.google.gerrit.server.notedb.ChangeUpdate;
@@ -44,6 +48,7 @@ import com.google.gerrit.server.util.time.TimeUtil;
 import com.google.inject.Inject;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -208,19 +213,22 @@ public class ReplyAttentionSetUpdates {
   /** Adds all authors of all comment threads that received a reply during this update */
   private void addAllAuthorsOfCommentThreads(
       BatchUpdate bu, ChangeNotes changeNotes, ImmutableSet<HumanComment> allNewComments) {
-    Set<HumanComment> allCommentsInCommentThreads =
+    List<HumanComment> publishedComments = commentsUtil.publishedHumanCommentsByChange(changeNotes);
-        commentsUtil.getAllHumanCommentsInCommentThreads(changeNotes, allNewComments);
+    ImmutableSet<CommentThread<HumanComment>> repliedToCommentThreads =
-    // Copy the set to make it mutable, so that we can delete users that were already added.
+        CommentThreads.forComments(publishedComments).getThreadsForChildren(allNewComments);
-    Set<Account.Id> possibleUsersToAdd =
-        new HashSet<>(approvalsUtil.getReviewers(changeNotes).all());
 
-    for (HumanComment comment : allCommentsInCommentThreads) {
+    ImmutableSet<Account.Id> repliedToUsers =
-      Account.Id author = comment.author.getId();
+        repliedToCommentThreads.stream()
-      if (possibleUsersToAdd.contains(author)) {
+            .map(CommentThread::comments)
-        addToAttentionSet(
+            .flatMap(Collection::stream)
-            bu, changeNotes, author, "Someone else replied on a comment you posted", false);
+            .map(comment -> comment.author.getId())
-        possibleUsersToAdd.remove(author);
+            .collect(toImmutableSet());
-      }
+    ImmutableSet<Account.Id> possibleUsersToAdd = approvalsUtil.getReviewers(changeNotes).all();
+    SetView<Account.Id> usersToAdd = Sets.intersection(possibleUsersToAdd, repliedToUsers);
+
+    for (Account.Id user : usersToAdd) {
+      addToAttentionSet(
+          bu, changeNotes, user, "Someone else replied on a comment you posted", false);
     }
   }
 

javatests/com/google/gerrit/acceptance/rest/change/AttentionSetIT.java

@@ -22,14 +22,19 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
+import com.google.common.truth.Correspondence;
 import com.google.gerrit.acceptance.AbstractDaemonTest;
 import com.google.gerrit.acceptance.NoHttpd;
 import com.google.gerrit.acceptance.PushOneCommit;
 import com.google.gerrit.acceptance.TestAccount;
 import com.google.gerrit.acceptance.UseClockStep;
+import com.google.gerrit.acceptance.testsuite.account.AccountOperations;
+import com.google.gerrit.acceptance.testsuite.change.ChangeOperations;
 import com.google.gerrit.acceptance.testsuite.request.RequestScopeOperations;
 import com.google.gerrit.common.Nullable;
+import com.google.gerrit.entities.Account;
 import com.google.gerrit.entities.AttentionSetUpdate;
+import com.google.gerrit.entities.Change;
 import com.google.gerrit.entities.Patch;
 import com.google.gerrit.extensions.api.changes.AddReviewerInput;
 import com.google.gerrit.extensions.api.changes.AttentionSetInput;
@@ -39,12 +44,17 @@ import com.google.gerrit.extensions.api.changes.ReviewInput;
 import com.google.gerrit.extensions.client.ReviewerState;
 import com.google.gerrit.extensions.client.Side;
 import com.google.gerrit.extensions.restapi.BadRequestException;
+import com.google.gerrit.server.query.change.ChangeData;
+import com.google.gerrit.server.query.change.InternalChangeQuery;
 import com.google.gerrit.server.util.time.TimeUtil;
 import com.google.gerrit.testing.FakeEmailSender;
 import com.google.gerrit.testing.TestCommentHelper;
+import com.google.gerrit.truth.NullAwareCorrespondence;
 import com.google.inject.Inject;
+import com.google.inject.Provider;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.function.LongSupplier;
@@ -56,9 +66,13 @@ import org.junit.Test;
 @UseClockStep(clockStepUnit = TimeUnit.MINUTES)
 public class AttentionSetIT extends AbstractDaemonTest {
 
+  @Inject private ChangeOperations changeOperations;
+  @Inject private AccountOperations accountOperations;
   @Inject private RequestScopeOperations requestScopeOperations;
+
   @Inject private FakeEmailSender email;
   @Inject private TestCommentHelper testCommentHelper;
+  @Inject private Provider<InternalChangeQuery> changeQueryProvider;
 
   /** Simulates a fake clock. Uses second granularity. */
   private static class FakeClock implements LongSupplier {
@@ -165,7 +179,8 @@ public class AttentionSetIT extends AbstractDaemonTest {
     assertThat(emailBody)
         .contains(
             user.fullName()
-                + " removed themselves from the attention set of this change.\n The reason is: removed.");
+                + " removed themselves from the attention set of this change.\n"
+                + " The reason is: removed.");
   }
 
   @Test
@@ -611,7 +626,8 @@ public class AttentionSetIT extends AbstractDaemonTest {
 
     assertThat(exception.getMessage())
         .isEqualTo(
-            "user can not be added/removed twice, and can not be added and removed at the same time");
+            "user can not be added/removed twice, and can not be added and removed at the same"
+                + " time");
   }
 
   @Test
@@ -627,7 +643,8 @@ public class AttentionSetIT extends AbstractDaemonTest {
 
     assertThat(exception.getMessage())
         .isEqualTo(
-            "user can not be added/removed twice, and can not be added and removed at the same time");
+            "user can not be added/removed twice, and can not be added and removed at the same"
+                + " time");
   }
 
   @Test
@@ -663,7 +680,8 @@ public class AttentionSetIT extends AbstractDaemonTest {
 
     assertThat(exception.getMessage())
         .isEqualTo(
-            "user can not be added/removed twice, and can not be added and removed at the same time");
+            "user can not be added/removed twice, and can not be added and removed at the same"
+                + " time");
   }
 
   @Test
@@ -957,6 +975,64 @@ public class AttentionSetIT extends AbstractDaemonTest {
         .isEqualTo("Someone else replied on a comment you posted");
   }
 
+  @Test
+  public void reviewAddsAllUsersInCommentThreadEvenOfDifferentChildBranch() throws Exception {
+    Account.Id changeOwner = accountOperations.newAccount().create();
+    Change.Id changeId = changeOperations.newChange().owner(changeOwner).create();
+    Account.Id user1 = accountOperations.newAccount().create();
+    Account.Id user2 = accountOperations.newAccount().create();
+    Account.Id user3 = accountOperations.newAccount().create();
+    Account.Id user4 = accountOperations.newAccount().create();
+    // Add users as reviewers.
+    gApi.changes().id(changeId.get()).addReviewer(user1.toString());
+    gApi.changes().id(changeId.get()).addReviewer(user2.toString());
+    gApi.changes().id(changeId.get()).addReviewer(user3.toString());
+    gApi.changes().id(changeId.get()).addReviewer(user4.toString());
+    // Add a comment thread with branches. Such threads occur if people reply in parallel without
+    // having seen/loaded the reply of another person.
+    String root =
+        changeOperations.change(changeId).currentPatchset().newComment().author(user1).create();
+    String sibling1 =
+        changeOperations
+            .change(changeId)
+            .currentPatchset()
+            .newComment()
+            .author(user2)
+            .parentUuid(root)
+            .create();
+    String sibling2 =
+        changeOperations
+            .change(changeId)
+            .currentPatchset()
+            .newComment()
+            .author(user3)
+            .parentUuid(root)
+            .create();
+    changeOperations
+        .change(changeId)
+        .currentPatchset()
+        .newComment()
+        .author(user4)
+        .parentUuid(sibling2)
+        .create();
+    // Clear the attention set. Necessary as we used Gerrit APIs above which affect the attention
+    // set.
+    AttentionSetInput clearAttention = new AttentionSetInput("clear attention set");
+    gApi.changes().id(changeId.get()).attention(user1.toString()).remove(clearAttention);
+    gApi.changes().id(changeId.get()).attention(user2.toString()).remove(clearAttention);
+    gApi.changes().id(changeId.get()).attention(user3.toString()).remove(clearAttention);
+    gApi.changes().id(changeId.get()).attention(user4.toString()).remove(clearAttention);
+
+    requestScopeOperations.setApiUser(changeOwner);
+    // Simulate that this reply is a child of sibling1 and thus parallel to sibling2 and its child.
+    gApi.changes().id(changeId.get()).current().review(reviewInReplyToComment(sibling1));
+
+    List<AttentionSetUpdate> attentionSetUpdates = getAttentionSetUpdates(changeId);
+    assertThat(attentionSetUpdates)
+        .comparingElementsUsing(hasAccount())
+        .containsExactly(user1, user2, user3, user4);
+  }
+
   @Test
   public void reviewAddsAllUsersInCommentThreadWhenPostedAsDraft() throws Exception {
     PushOneCommit.Result r = createChange();
@@ -1281,11 +1357,20 @@ public class AttentionSetIT extends AbstractDaemonTest {
 
   private List<AttentionSetUpdate> getAttentionSetUpdatesForUser(
       PushOneCommit.Result r, TestAccount account) {
-    return r.getChange().attentionSet().stream()
+    return getAttentionSetUpdates(r.getChange().getId()).stream()
-        .filter(a -> a.account().get() == account.id().get())
+        .filter(a -> a.account().equals(account.id()))
         .collect(Collectors.toList());
   }
 
+  private List<AttentionSetUpdate> getAttentionSetUpdates(Change.Id changeId) {
+    List<ChangeData> changeData = changeQueryProvider.get().byLegacyChangeId(changeId);
+    if (changeData.size() != 1) {
+      throw new IllegalStateException(
+          String.format("Not exactly one change found for ID %s.", changeId));
+    }
+    return new ArrayList<>(Iterables.getOnlyElement(changeData).attentionSet());
+  }
+
   private ReviewInput reviewWithComment() {
     return reviewInReplyToComment(null);
   }
@@ -1301,4 +1386,8 @@ public class AttentionSetIT extends AbstractDaemonTest {
     reviewInput.comments = ImmutableMap.of(Patch.COMMIT_MSG, ImmutableList.of(comment));
     return reviewInput;
   }
+
+  private Correspondence<AttentionSetUpdate, Account.Id> hasAccount() {
+    return NullAwareCorrespondence.transforming(AttentionSetUpdate::account, "hasAccount");
+  }
 }

javatests/com/google/gerrit/server/change/CommentThreadsTest.java

@@ -174,6 +174,90 @@ public class CommentThreadsTest {
     assertThat(commentThreads).isEqualTo(expectedThreads);
   }
 
+  @Test
+  public void specificThreadsCanBeRequestedByTheirReply() {
+    HumanComment thread1Root = createComment("thread1Root");
+    HumanComment thread2Root = createComment("thread2Root");
+
+    HumanComment thread1Reply = asReply(createComment("thread1Reply"), "thread1Root");
+
+    ImmutableList<HumanComment> comments = ImmutableList.of(thread1Root, thread2Root, thread1Reply);
+    ImmutableSet<CommentThread<HumanComment>> commentThreads =
+        CommentThreads.forComments(comments).getThreadsForChildren(ImmutableList.of(thread1Reply));
+
+    ImmutableSet<CommentThread<HumanComment>> expectedThreads =
+        ImmutableSet.of(toThread(thread1Root, thread1Reply));
+    assertThat(commentThreads).isEqualTo(expectedThreads);
+  }
+
+  @Test
+  public void requestedThreadsDoNotNeedToContainReply() {
+    HumanComment thread1Root = createComment("thread1Root");
+    HumanComment thread2Root = createComment("thread2Root");
+
+    HumanComment thread1Reply = asReply(createComment("thread1Reply"), "thread1Root");
+
+    ImmutableList<HumanComment> comments = ImmutableList.of(thread1Root, thread2Root);
+    ImmutableSet<CommentThread<HumanComment>> commentThreads =
+        CommentThreads.forComments(comments).getThreadsForChildren(ImmutableList.of(thread1Reply));
+
+    ImmutableSet<CommentThread<HumanComment>> expectedThreads =
+        ImmutableSet.of(toThread(thread1Root));
+    assertThat(commentThreads).isEqualTo(expectedThreads);
+  }
+
+  @Test
+  public void completeThreadCanBeRequestedByReplyToRootComment() {
+    HumanComment root = createComment("root");
+    HumanComment child = asReply(createComment("child"), "root");
+
+    HumanComment reply = asReply(createComment("reply"), "root");
+
+    ImmutableList<HumanComment> comments = ImmutableList.of(root, child);
+    ImmutableSet<CommentThread<HumanComment>> commentThreads =
+        CommentThreads.forComments(comments).getThreadsForChildren(ImmutableList.of(reply));
+
+    ImmutableSet<CommentThread<HumanComment>> expectedThreads =
+        ImmutableSet.of(toThread(root, child));
+    assertThat(commentThreads).isEqualTo(expectedThreads);
+  }
+
+  @Test
+  public void completeThreadWithBranchesCanBeRequestedByReplyToIntermediateComment() {
+    HumanComment root = writtenOn(createComment("root"), new Timestamp(1));
+    HumanComment sibling1 = writtenOn(asReply(createComment("sibling1"), "root"), new Timestamp(2));
+    HumanComment sibling2 = writtenOn(asReply(createComment("sibling2"), "root"), new Timestamp(3));
+    HumanComment sibling1Child =
+        writtenOn(asReply(createComment("sibling1Child"), "sibling1"), new Timestamp(4));
+    HumanComment sibling2Child =
+        writtenOn(asReply(createComment("sibling2Child"), "sibling2"), new Timestamp(5));
+
+    HumanComment reply = asReply(createComment("sibling1"), "root");
+
+    ImmutableList<HumanComment> comments =
+        ImmutableList.of(root, sibling1, sibling2, sibling1Child, sibling2Child);
+    ImmutableSet<CommentThread<HumanComment>> commentThreads =
+        CommentThreads.forComments(comments).getThreadsForChildren(ImmutableList.of(reply));
+
+    ImmutableSet<CommentThread<HumanComment>> expectedThreads =
+        ImmutableSet.of(toThread(root, sibling1, sibling2, sibling1Child, sibling2Child));
+    assertThat(commentThreads).isEqualTo(expectedThreads);
+  }
+
+  @Test
+  public void requestedThreadsAreEmptyIfReplyDoesNotReferToAThread() {
+    HumanComment root = createComment("root");
+
+    HumanComment reply = asReply(createComment("reply"), "invalid");
+
+    ImmutableList<HumanComment> comments = ImmutableList.of(root);
+    ImmutableSet<CommentThread<HumanComment>> commentThreads =
+        CommentThreads.forComments(comments).getThreadsForChildren(ImmutableList.of(reply));
+
+    ImmutableSet<CommentThread<HumanComment>> expectedThreads = ImmutableSet.of();
+    assertThat(commentThreads).isEqualTo(expectedThreads);
+  }
+
   private static HumanComment createComment(String commentUuid) {
     return new HumanComment(
         new Key(commentUuid, "myFile", 1),