Merge "Revert "Delete references to Zombie draft comments""
diff --git a/java/com/google/gerrit/server/CommentsUtil.java b/java/com/google/gerrit/server/CommentsUtil.java
index 3d3603f..ba9f6d6 100644
--- a/java/com/google/gerrit/server/CommentsUtil.java
+++ b/java/com/google/gerrit/server/CommentsUtil.java
@@ -450,6 +450,10 @@
   /**
    * Get NoteDb draft refs for a change.
    *
+   * <p>This is just a simple ref scan, so the results may potentially include refs for zombie draft
+   * comments. A zombie draft is one which has been published but the write to delete the draft ref
+   * from All-Users failed.
+   *
    * @param changeId change ID.
    * @return raw refs from All-Users repo.
    */
diff --git a/java/com/google/gerrit/server/notedb/ChangeNotes.java b/java/com/google/gerrit/server/notedb/ChangeNotes.java
index f2034af..2d9b014 100644
--- a/java/com/google/gerrit/server/notedb/ChangeNotes.java
+++ b/java/com/google/gerrit/server/notedb/ChangeNotes.java
@@ -32,6 +32,7 @@
 import com.google.common.collect.ListMultimap;
 import com.google.common.collect.Lists;
 import com.google.common.collect.MultimapBuilder;
+import com.google.common.collect.Multimaps;
 import com.google.common.collect.Ordering;
 import com.google.common.collect.Sets;
 import com.google.common.collect.Sets.SetView;
@@ -522,7 +523,12 @@
   public ImmutableListMultimap<ObjectId, HumanComment> getDraftComments(
       Account.Id author, @Nullable Ref ref) {
     loadDraftComments(author, ref);
-    return draftCommentNotes.getComments();
+    // Filter out any zombie draft comments. These are drafts that are also in
+    // the published map, and arise when the update to All-Users to delete them
+    // during the publish operation failed.
+    return ImmutableListMultimap.copyOf(
+        Multimaps.filterEntries(
+            draftCommentNotes.getComments(), e -> !getCommentKeys().contains(e.getValue().key)));
   }
 
   public ImmutableListMultimap<ObjectId, RobotComment> getRobotComments() {
diff --git a/java/com/google/gerrit/server/query/change/ChangeData.java b/java/com/google/gerrit/server/query/change/ChangeData.java
index 6e14a13..4112579 100644
--- a/java/com/google/gerrit/server/query/change/ChangeData.java
+++ b/java/com/google/gerrit/server/query/change/ChangeData.java
@@ -1373,7 +1373,14 @@
       draftsByUser = new HashMap<>();
       for (Ref ref : commentsUtil.getDraftRefs(notes().getChangeId())) {
         Account.Id account = Account.Id.fromRefSuffix(ref.getName());
-        if (account != null) {
+        if (account != null
+            // Double-check that any drafts exist for this user after
+            // filtering out zombies. If some but not all drafts in the ref
+            // were zombies, the returned Ref still includes those zombies;
+            // this is suboptimal, but is ok for the purposes of
+            // draftsByUser(), and easier than trying to rebuild the change at
+            // this point.
+            && !notes().getDraftComments(account, ref).isEmpty()) {
           draftsByUser.put(account, ref.getObjectId());
         }
       }
diff --git a/javatests/com/google/gerrit/server/notedb/ChangeNotesTest.java b/javatests/com/google/gerrit/server/notedb/ChangeNotesTest.java
index 4e7b3f3..c524c94 100644
--- a/javatests/com/google/gerrit/server/notedb/ChangeNotesTest.java
+++ b/javatests/com/google/gerrit/server/notedb/ChangeNotesTest.java
@@ -76,6 +76,8 @@
 import org.junit.Test;
 
 public class ChangeNotesTest extends AbstractChangeNotesTest {
+  @Inject private DraftCommentNotes.Factory draftNotesFactory;
+
   @Inject private ChangeNoteJson changeNoteJson;
 
   @Test
@@ -2978,6 +2980,86 @@
   }
 
   @Test
+  public void filterOutAndFixUpZombieDraftComments() throws Exception {
+    Change c = newChange();
+    ObjectId commitId1 = ObjectId.fromString("abcd1234abcd1234abcd1234abcd1234abcd1234");
+    CommentRange range = new CommentRange(1, 1, 2, 1);
+    PatchSet.Id ps1 = c.currentPatchSetId();
+    short side = (short) 1;
+
+    ChangeUpdate update = newUpdate(c, otherUser);
+    Timestamp now = TimeUtil.nowTs();
+    HumanComment comment1 =
+        newComment(
+            ps1,
+            "file1",
+            "uuid1",
+            range,
+            range.getEndLine(),
+            otherUser,
+            null,
+            now,
+            "comment on ps1",
+            side,
+            commitId1,
+            false);
+    HumanComment comment2 =
+        newComment(
+            ps1,
+            "file2",
+            "uuid2",
+            range,
+            range.getEndLine(),
+            otherUser,
+            null,
+            now,
+            "another comment",
+            side,
+            commitId1,
+            false);
+    update.putComment(HumanComment.Status.DRAFT, comment1);
+    update.putComment(HumanComment.Status.DRAFT, comment2);
+    update.commit();
+
+    String refName = refsDraftComments(c.getId(), otherUserId);
+    ObjectId oldDraftId = exactRefAllUsers(refName);
+
+    update = newUpdate(c, otherUser);
+    update.setPatchSetId(ps1);
+    update.putComment(HumanComment.Status.PUBLISHED, comment2);
+    update.commit();
+    assertThat(exactRefAllUsers(refName)).isNotNull();
+    assertThat(exactRefAllUsers(refName)).isNotEqualTo(oldDraftId);
+
+    // Re-add draft version of comment2 back to draft ref without updating
+    // change ref. Simulates the case where deleting the draft failed
+    // non-atomically after adding the published comment succeeded.
+    ChangeDraftUpdate draftUpdate = newUpdate(c, otherUser).createDraftUpdateIfNull();
+    draftUpdate.putComment(comment2);
+    try (NoteDbUpdateManager manager = updateManagerFactory.create(c.getProject())) {
+      manager.add(draftUpdate);
+      manager.execute();
+    }
+
+    // Looking at drafts directly shows the zombie comment.
+    DraftCommentNotes draftNotes = draftNotesFactory.create(c.getId(), otherUserId);
+    assertThat(draftNotes.load().getComments().get(commitId1)).containsExactly(comment1, comment2);
+
+    // Zombie comment is filtered out of drafts via ChangeNotes.
+    ChangeNotes notes = newNotes(c);
+    assertThat(notes.getDraftComments(otherUserId).get(commitId1)).containsExactly(comment1);
+    assertThat(notes.getHumanComments().get(commitId1)).containsExactly(comment2);
+
+    update = newUpdate(c, otherUser);
+    update.setPatchSetId(ps1);
+    update.putComment(HumanComment.Status.PUBLISHED, comment1);
+    update.commit();
+
+    // Updating an unrelated comment causes the zombie comment to get fixed up.
+    assertThat(exactRefAllUsers(refName)).isNull();
+  }
+
+  @Test
   public void updateCommentsInSequentialUpdates() throws Exception {
     Change c = newChange();
     CommentRange range = new CommentRange(1, 1, 2, 1);
diff --git a/javatests/com/google/gerrit/server/query/change/AbstractQueryChangesTest.java b/javatests/com/google/gerrit/server/query/change/AbstractQueryChangesTest.java
index 2663853..9ebee9c 100644
--- a/javatests/com/google/gerrit/server/query/change/AbstractQueryChangesTest.java
+++ b/javatests/com/google/gerrit/server/query/change/AbstractQueryChangesTest.java
@@ -69,6 +69,7 @@
 import com.google.gerrit.extensions.api.changes.DraftInput;
 import com.google.gerrit.extensions.api.changes.HashtagsInput;
 import com.google.gerrit.extensions.api.changes.ReviewInput;
+import com.google.gerrit.extensions.api.changes.ReviewInput.DraftHandling;
 import com.google.gerrit.extensions.api.changes.ReviewerInput;
 import com.google.gerrit.extensions.api.groups.GroupInput;
 import com.google.gerrit.extensions.api.projects.ConfigInput;
@@ -2333,6 +2334,44 @@
   }
 
   @Test
+  public void byHasDraftExcludesZombieDrafts() throws Exception {
+    Project.NameKey project = Project.nameKey("repo");
+    TestRepository<Repo> repo = createProject(project.get());
+    Change change = insert(repo, newChange(repo));
+    Change.Id id = change.getId();
+
+    DraftInput in = new DraftInput();
+    in.line = 1;
+    in.message = "nit: trailing whitespace";
+    in.path = Patch.COMMIT_MSG;
+    gApi.changes().id(id.get()).current().createDraft(in);
+
+    assertQuery("has:draft", change);
+    assertQuery("commentby:" + userId);
+
+    try (TestRepository<Repo> allUsers =
+        new TestRepository<>(repoManager.openRepository(allUsersName))) {
+      Ref draftsRef = allUsers.getRepository().exactRef(RefNames.refsDraftComments(id, userId));
+      assertThat(draftsRef).isNotNull();
+
+      ReviewInput rin = ReviewInput.dislike();
+      rin.drafts = DraftHandling.PUBLISH_ALL_REVISIONS;
+      gApi.changes().id(id.get()).current().review(rin);
+
+      assertQuery("has:draft");
+      assertQuery("commentby:" + userId, change);
+      assertThat(allUsers.getRepository().exactRef(draftsRef.getName())).isNull();
+
+      // Re-add drafts ref and ensure it gets filtered out during indexing.
+      allUsers.update(draftsRef.getName(), draftsRef.getObjectId());
+      assertThat(allUsers.getRepository().exactRef(draftsRef.getName())).isNotNull();
+    }
+
+    indexer.index(project, id);
+    assertQuery("has:draft");
+  }
+
+  @Test
   public void byStarredBy() throws Exception {
     TestRepository<Repo> repo = createProject("repo");
     Change change1 = insert(repo, newChange(repo));