Rewrite updating comments in NoteDb

The high-level API is simplified: we no longer ask callers to
distinguish between insert/update/upsert, and instead use upsert
everywhere, calling it "put". These distinctions were not
particularly important from the perspective of e.g. the REST API.

Although API simplification is a good thing regardless, this change
was primarily driven by the fact that we do not want to use the
eagerly-read ChangeNotes from within ChangeUpdate internals prior to
actually applying the update. The various checkArgument calls in
insert/update/upsert were trying to enforce, for example, that we only
insert changes that don't exist at the time the setter was called.

But in the context of NoteDbUpdateManager that executes multiple
ChangeUpdates in sequence, it doesn't make sense to look up a
comment in the old comment map based on the current branch tip when
the ChangeUpdate instance was created. Instead we should be looking at
the comment map that is the immediate output of the previous update in
the sequence. This distinction is precisely what was causing the test
in I064f004e to fail.

To solve this problem at the storage layer, instead of depending on
the NoteMap from the original ChangeNotes, read the RevisionNotes
from the current revision passed in from the NoteDbUpdateManager
(which may not have been flushed yet, but is still available to the
reader). Don't bother reading the whole ChangeNotes again, as that is
unnecessary. Encapsulate the RevisionNotes and the mutable NoteMap
instance itself in a small POJO, RevisionNoteMap. In the common case
of a single ChangeUpdate, we can even reuse the previously parsed
RevisionNoteMap from the ChangeNotes that was passed in; we only have
to re-read when doing multiple updates.

Mutations are done to the RevisionNoteMap using the same
RevisionNoteMapBuilder idiom we were already using for merging push
certificates with inline comments. Moreover, we can use the same
logic within ChangeDraftUpdate, obviating the need for
DraftCommentNotesParser and simplifying the parse codepath somewhat.

Finally, this change improves reliability and robustness around the
fact that we are writing ref updates to two distinct repositories
non-atomically, the change meta repo and All-Users. Order the updates
so that the write to the change repo happens first, taking advantage
of the fact that comments can only go from DRAFT to PUBLISHED, not
vice versa. If the change update succeeds and the All-Users update
fails, then we will have a zombie DRAFT comment left in All-Users
where there is a PUBLISHED comment in the change repo. We fix this up
in two ways: at read time, by explicitly filtering out any comments
from the draft list that are also published; and at write time, by
deleting all zombie draft comments any time we write any other draft
comment.

Add test cases for a few more combinations of comment updates,
including one that simulates a failed All-Users write.

Change-Id: I69c46ab4b55d4e40ad80754ee8528e99ae6ad4c5

gerrit-server/src/main/java/com/google/gerrit/server/PatchLineCommentsUtil.java

@@ -288,30 +288,14 @@ public class PatchLineCommentsUtil {
     return sort(comments);
   }
 
-  public void insertComments(ReviewDb db, ChangeUpdate update,
+  public void putComments(ReviewDb db, ChangeUpdate update,
       Iterable<PatchLineComment> comments) throws OrmException {
     for (PatchLineComment c : comments) {
-      update.insertComment(c);
+      update.putComment(c);
-    }
-    db.patchComments().insert(comments);
-  }
-
-  public void upsertComments(ReviewDb db, ChangeUpdate update,
-      Iterable<PatchLineComment> comments) throws OrmException {
-    for (PatchLineComment c : comments) {
-      update.upsertComment(c);
     }
     db.patchComments().upsert(comments);
   }
 
-  public void updateComments(ReviewDb db, ChangeUpdate update,
-      Iterable<PatchLineComment> comments) throws OrmException {
-    for (PatchLineComment c : comments) {
-      update.updateComment(c);
-    }
-    db.patchComments().update(comments);
-  }
-
   public void deleteComments(ReviewDb db, ChangeUpdate update,
       Iterable<PatchLineComment> comments) throws OrmException {
     for (PatchLineComment c : comments) {

gerrit-server/src/main/java/com/google/gerrit/server/change/CreateDraftComment.java

@@ -124,7 +124,7 @@ public class CreateDraftComment implements RestModifyView<RevisionResource, Draf
       comment.setRange(in.range);
       setCommentRevId(
           comment, patchListCache, ctx.getChange(), ps);
-      plcUtil.insertComments(
+      plcUtil.putComments(
           ctx.getDb(), ctx.getUpdate(psId), Collections.singleton(comment));
       return true;
     }

gerrit-server/src/main/java/com/google/gerrit/server/change/PostReview.java

@@ -476,7 +476,7 @@ public class PostReview implements RestModifyView<RevisionResource, ReviewInput>
       // TODO(dborowitz): Currently doesn't work for PUBLISH_ALL_REVISIONS with
       // notedb.
       plcUtil.deleteComments(ctx.getDb(), u, del);
-      plcUtil.upsertComments(ctx.getDb(), u, ups);
+      plcUtil.putComments(ctx.getDb(), u, ups);
       comments.addAll(ups);
       return !del.isEmpty() || !ups.isEmpty();
     }

gerrit-server/src/main/java/com/google/gerrit/server/change/PutDraftComment.java

@@ -142,14 +142,14 @@ public class PutDraftComment implements RestModifyView<DraftCommentResource, Dra
             ctx.getUser().getAccountId(),
             comment.getParentUuid(), ctx.getWhen());
         setCommentRevId(comment, patchListCache, ctx.getChange(), ps);
-        plcUtil.insertComments(ctx.getDb(), update,
+        plcUtil.putComments(ctx.getDb(), update,
             Collections.singleton(update(comment, in)));
       } else {
         if (comment.getRevId() == null) {
           setCommentRevId(
               comment, patchListCache, ctx.getChange(), ps);
         }
-        plcUtil.updateComments(ctx.getDb(), update,
+        plcUtil.putComments(ctx.getDb(), update,
             Collections.singleton(update(comment, in)));
       }
       return true;

gerrit-server/src/main/java/com/google/gerrit/server/notedb/AbstractChangeUpdate.java

@@ -138,7 +138,7 @@ public abstract class AbstractChangeUpdate {
       return null;
     }
     ObjectId z = ObjectId.zeroId();
-    CommitBuilder cb = applyImpl(ins);
+    CommitBuilder cb = applyImpl(rw, ins, curr);
     if (cb == null) {
       result = z;
       return z; // Impl intends to delete the ref.
@@ -174,10 +174,8 @@ public abstract class AbstractChangeUpdate {
    * @throws OrmException if a Gerrit-level error occurred.
    * @throws IOException if a lower-level error occurred.
    */
-  // TODO(dborowitz): ChangeUpdate needs to be able to reread its ChangeNotes at
+  protected abstract CommitBuilder applyImpl(RevWalk rw, ObjectInserter ins,
-  // the old SHA-1, which would imply passing curr here.
+      ObjectId curr) throws OrmException, IOException;
-  protected abstract CommitBuilder applyImpl(ObjectInserter ins)
-      throws OrmException, IOException;
 
   ObjectId getResult() {
     return result;

gerrit-server/src/main/java/com/google/gerrit/server/notedb/ChangeDraftUpdate.java

@@ -14,13 +14,13 @@
 
 package com.google.gerrit.server.notedb;
 
+import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
-import static com.google.gerrit.server.notedb.CommentsInNotesUtil.addCommentToMap;
+import static org.eclipse.jgit.lib.Constants.OBJ_BLOB;
 
-import com.google.common.collect.ListMultimap;
+import com.google.auto.value.AutoValue;
-import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import com.google.gerrit.reviewdb.client.Account;
 import com.google.gerrit.reviewdb.client.PatchLineComment;
@@ -37,20 +37,20 @@ import com.google.gwtorm.server.OrmException;
 import com.google.inject.assistedinject.Assisted;
 import com.google.inject.assistedinject.AssistedInject;
 
+import org.eclipse.jgit.errors.ConfigInvalidException;
 import org.eclipse.jgit.lib.CommitBuilder;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.ObjectInserter;
 import org.eclipse.jgit.lib.PersonIdent;
 import org.eclipse.jgit.notes.NoteMap;
+import org.eclipse.jgit.revwalk.RevWalk;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Date;
 import java.util.HashMap;
-import java.util.List;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * A single delta to apply atomically to a change.
@@ -66,14 +66,23 @@ public class ChangeDraftUpdate extends AbstractChangeUpdate {
     ChangeDraftUpdate create(ChangeControl ctl, Date when);
   }
 
+  @AutoValue
+  static abstract class Key {
+    abstract RevId revId();
+    abstract PatchLineComment.Key key();
+  }
+
+  private static Key key(PatchLineComment c) {
+    return new AutoValue_ChangeDraftUpdate_Key(c.getRevId(), c.getKey());
+  }
+
   private final AllUsersName draftsProject;
   private final Account.Id accountId;
   private final CommentsInNotesUtil commentsUtil;
-  private final ChangeNotes changeNotes;
-  private final DraftCommentNotes draftNotes;
 
-  private List<PatchLineComment> upsertComments;
+  // TODO: can go back to a list?
-  private List<PatchLineComment> deleteComments;
+  private Map<Key, PatchLineComment> put;
+  private Set<Key> delete;
 
   @AssistedInject
   private ChangeDraftUpdate(
@@ -81,11 +90,10 @@ public class ChangeDraftUpdate extends AbstractChangeUpdate {
       @AnonymousCowardName String anonymousCowardName,
       GitRepositoryManager repoManager,
       NotesMigration migration,
-      DraftCommentNotes.Factory draftNotesFactory,
       AllUsersName allUsers,
       CommentsInNotesUtil commentsUtil,
       @Assisted ChangeControl ctl,
-      @Assisted Date when) throws OrmException {
+      @Assisted Date when) {
     super(migration, repoManager, ctl, serverIdent, anonymousCowardName, when);
     this.draftsProject = allUsers;
     this.commentsUtil = commentsUtil;
@@ -93,154 +101,117 @@ public class ChangeDraftUpdate extends AbstractChangeUpdate {
         "Current user must be identified");
     IdentifiedUser user = ctl.getUser().asIdentifiedUser();
     this.accountId = user.getAccountId();
-    this.changeNotes = getChangeNotes().load();
-    this.draftNotes = draftNotesFactory.create(ctl.getId(),
-        user.getAccountId());
 
-    this.upsertComments = Lists.newArrayList();
+    this.put = new HashMap<>();
-    this.deleteComments = Lists.newArrayList();
+    this.delete = new HashSet<>();
   }
 
-  public void insertComment(PatchLineComment c) throws OrmException {
+  public void putComment(PatchLineComment c) {
     verifyComment(c);
     checkArgument(c.getStatus() == PatchLineComment.Status.DRAFT,
         "Cannot insert a published comment into a ChangeDraftUpdate");
-    if (migration.readChanges()) {
+    put.put(key(c), c);
-      checkArgument(!changeNotes.containsComment(c),
-          "A comment already exists with the same key,"
-          + " so the following comment cannot be inserted: %s", c);
-    }
-    upsertComments.add(c);
   }
 
-  public void upsertComment(PatchLineComment c) {
+  public void deleteComment(PatchLineComment c) {
     verifyComment(c);
-    checkArgument(c.getStatus() == PatchLineComment.Status.DRAFT,
+    delete.add(key(c));
-        "Cannot upsert a published comment into a ChangeDraftUpdate");
-    upsertComments.add(c);
   }
 
-  public void updateComment(PatchLineComment c) throws OrmException {
+  public void deleteComment(RevId revId, PatchLineComment.Key key) {
-    verifyComment(c);
+    delete.add(new AutoValue_ChangeDraftUpdate_Key(revId, key));
-    checkArgument(c.getStatus() == PatchLineComment.Status.DRAFT,
-        "Cannot update a published comment into a ChangeDraftUpdate");
-    // Here, we check to see if this comment existed previously as a draft.
-    // However, this could cause a race condition if there is a delete and an
-    // update operation happening concurrently (or two deletes) and they both
-    // believe that the comment exists. If a delete happens first, then
-    // the update will fail. However, this is an acceptable risk since the
-    // caller wanted the comment deleted anyways, so the end result will be the
-    // same either way.
-    if (migration.readChanges()) {
-      checkArgument(draftNotes.load().containsComment(c),
-          "Cannot update this comment because it didn't exist previously");
-    }
-    upsertComments.add(c);
-  }
-
-  public void deleteComment(PatchLineComment c) throws OrmException {
-    verifyComment(c);
-    // See the comment above about potential race condition.
-    if (migration.readChanges()) {
-      checkArgument(draftNotes.load().containsComment(c),
-          "Cannot delete this comment because it didn't previously exist as a"
-          + " draft");
-    }
-    if (migration.writeChanges()) {
-      if (draftNotes.load().containsComment(c)) {
-        deleteComments.add(c);
-      }
-    }
-  }
-
-  /**
-   * Deletes a PatchLineComment from the list of drafts only if it existed
-   * previously as a draft. If it wasn't a draft previously, this is a no-op.
-   */
-  public void deleteCommentIfPresent(PatchLineComment c) throws OrmException {
-    if (draftNotes.load().containsComment(c)) {
-      verifyComment(c);
-      deleteComments.add(c);
-    }
   }
 
   private void verifyComment(PatchLineComment comment) {
-    if (migration.writeChanges()) {
-      checkArgument(comment.getRevId() != null);
-    }
     checkArgument(comment.getAuthor().equals(accountId),
         "The author for the following comment does not match the author of"
         + " this ChangeDraftUpdate (%s): %s", accountId, comment);
   }
 
   /** @return the tree id for the updated tree */
-  private ObjectId storeCommentsInNotes(ObjectInserter inserter,
+  private ObjectId storeCommentsInNotes(RevWalk rw, ObjectInserter ins,
-      AtomicBoolean removedAllComments) throws OrmException, IOException {
+      ObjectId curr) throws ConfigInvalidException, OrmException, IOException {
-    NoteMap noteMap = draftNotes.load().getNoteMap();
+    RevisionNoteMap rnm = getRevisionNoteMap(rw, curr);
-    if (noteMap == null) {
+    Set<RevId> updatedRevs =
-      noteMap = NoteMap.newEmptyMap();
+        Sets.newHashSetWithExpectedSize(rnm.revisionNotes.size());
+    RevisionNoteBuilder.Cache cache = new RevisionNoteBuilder.Cache(rnm);
+
+    for (PatchLineComment c : put.values()) {
+      if (!delete.contains(key(c))) {
+        cache.get(c.getRevId()).putComment(c);
+      }
+    }
+    for (Key k : delete) {
+      cache.get(k.revId()).deleteComment(k.key());
     }
 
-    Map<RevId, List<PatchLineComment>> allComments = new HashMap<>();
+    Map<RevId, RevisionNoteBuilder> builders = cache.getBuilders();
-
     boolean hasComments = false;
-    int n = deleteComments.size() + upsertComments.size();
+    for (Map.Entry<RevId, RevisionNoteBuilder> e : builders.entrySet()) {
-    Set<RevId> updatedRevs = Sets.newHashSetWithExpectedSize(n);
+      updatedRevs.add(e.getKey());
-    Set<PatchLineComment.Key> updatedKeys = Sets.newHashSetWithExpectedSize(n);
+      ObjectId id = ObjectId.fromString(e.getKey().get());
-    for (PatchLineComment c : deleteComments) {
+      byte[] data = e.getValue().build(commentsUtil);
-      allComments.put(c.getRevId(), new ArrayList<PatchLineComment>());
+      if (data.length == 0) {
-      updatedRevs.add(c.getRevId());
+        rnm.noteMap.remove(id);
-      updatedKeys.add(c.getKey());
+      } else {
-    }
-
-    for (PatchLineComment c : upsertComments) {
-      hasComments = true;
-      addCommentToMap(allComments, c);
-      updatedRevs.add(c.getRevId());
-      updatedKeys.add(c.getKey());
-    }
-
-    // Re-add old comments for updated revisions so the new note contents
-    // includes both old and new comments merged in the right order.
-    //
-    // writeCommentsToNoteMap doesn't touch notes for SHA-1s that are not
-    // mentioned in the input map, so by omitting comments for those revisions,
-    // we avoid the work of having to re-serialize identical comment data for
-    // those revisions.
-    ListMultimap<RevId, PatchLineComment> existing =
-        draftNotes.getComments();
-    for (Map.Entry<RevId, PatchLineComment> e : existing.entries()) {
-      PatchLineComment c = e.getValue();
-      if (updatedRevs.contains(c.getRevId())
-          && !updatedKeys.contains(c.getKey())) {
         hasComments = true;
-        addCommentToMap(allComments, e.getValue());
+        ObjectId dataBlob = ins.insert(OBJ_BLOB, data);
+        rnm.noteMap.set(id, dataBlob);
       }
     }
 
-    // If we touched every revision and there are no comments left, set the flag
+    // If we touched every revision and there are no comments left, tell the
-    // for the caller to delete the entire ref.
+    // caller to delete the entire ref.
-    boolean touchedAllRevs = updatedRevs.equals(existing.keySet());
+    boolean touchedAllRevs = updatedRevs.equals(rnm.revisionNotes.keySet());
     if (touchedAllRevs && !hasComments) {
-      removedAllComments.set(true);
       return null;
     }
 
-    commentsUtil.writeCommentsToNoteMap(noteMap, allComments, inserter);
+    return rnm.noteMap.writeTree(ins);
-    return noteMap.writeTree(inserter);
+  }
+
+  private RevisionNoteMap getRevisionNoteMap(RevWalk rw, ObjectId curr)
+      throws ConfigInvalidException, OrmException, IOException {
+    if (migration.readChanges()) {
+      // If reading from changes is enabled, then the old DraftCommentNotes
+      // already parsed the revision notes. We can reuse them as long as the ref
+      // hasn't advanced.
+      DraftCommentNotes draftNotes =
+          ctl.getNotes().load().getDraftCommentNotes();
+      if (draftNotes != null) {
+        ObjectId idFromNotes =
+            firstNonNull(draftNotes.getRevision(), ObjectId.zeroId());
+        if (idFromNotes.equals(curr)) {
+          return checkNotNull(ctl.getNotes().revisionNoteMap);
+        }
+      }
+    }
+    NoteMap noteMap;
+    if (!curr.equals(ObjectId.zeroId())) {
+      noteMap = NoteMap.read(rw.getObjectReader(), rw.parseCommit(curr));
+    } else {
+      noteMap = NoteMap.newEmptyMap();
+    }
+    // Even though reading from changes might not be enabled, we need to
+    // parse any existing revision notes so we can merge them.
+    return RevisionNoteMap.parse(
+        ctl.getId(), rw.getObjectReader(), noteMap, true);
   }
 
   @Override
-  protected CommitBuilder applyImpl(ObjectInserter ins)
+  protected CommitBuilder applyImpl(RevWalk rw, ObjectInserter ins,
-      throws OrmException, IOException {
+      ObjectId curr) throws OrmException, IOException {
     CommitBuilder cb = new CommitBuilder();
     cb.setMessage("Update draft comments");
-    AtomicBoolean removedAllComments = new AtomicBoolean();
+    try {
-    ObjectId treeId = storeCommentsInNotes(ins, removedAllComments);
+      ObjectId treeId = storeCommentsInNotes(rw, ins, curr);
-    if (removedAllComments.get()) {
+      if (treeId == null) {
-      return null; // Delete ref.
+        return null; // Delete ref.
+      }
+      cb.setTreeId(checkNotNull(treeId));
+    } catch (ConfigInvalidException e) {
+      throw new OrmException(e);
     }
-    cb.setTreeId(checkNotNull(treeId));
     return cb;
   }
 
@@ -256,7 +227,7 @@ public class ChangeDraftUpdate extends AbstractChangeUpdate {
 
   @Override
   public boolean isEmpty() {
-    return deleteComments.isEmpty()
+    return delete.isEmpty()
-        && upsertComments.isEmpty();
+        && put.isEmpty();
   }
 }

gerrit-server/src/main/java/com/google/gerrit/server/notedb/ChangeNotes.java

@@ -33,6 +33,8 @@ import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.ListMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Multimaps;
 import com.google.common.collect.Ordering;
 import com.google.common.primitives.Ints;
 import com.google.common.util.concurrent.AsyncFunction;
@@ -70,7 +72,6 @@ import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.PersonIdent;
 import org.eclipse.jgit.lib.Ref;
 import org.eclipse.jgit.lib.Repository;
-import org.eclipse.jgit.notes.NoteMap;
 import org.eclipse.jgit.revwalk.RevWalk;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -398,10 +399,9 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
   private ImmutableListMultimap<RevId, PatchLineComment> comments;
   private ImmutableSet<String> hashtags;
 
-  // Mutable note map state, only used by ChangeUpdate to make in-place editing
+  // Parsed note map state, used by ChangeUpdate to make in-place editing of
-  // of notes easier.
+  // notes easier.
-  NoteMap noteMap;
+  RevisionNoteMap revisionNoteMap;
-  Map<RevId, RevisionNote> revisionNotes;
 
   private final AllUsersName allUsers;
   private DraftCommentNotes draftCommentNotes;
@@ -477,7 +477,25 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
   public ImmutableListMultimap<RevId, PatchLineComment> getDraftComments(
       Account.Id author) throws OrmException {
     loadDraftComments(author);
-    return draftCommentNotes.getComments();
+    final Multimap<RevId, PatchLineComment> published = comments;
+    // Filter out any draft comments that also exist in the published map, in
+    // case the update to All-Users to delete them during the publish operation
+    // failed.
+    Multimap<RevId, PatchLineComment> filtered = Multimaps.filterEntries(
+        draftCommentNotes.getComments(),
+        new Predicate<Map.Entry<RevId, PatchLineComment>>() {
+          @Override
+          public boolean apply(Map.Entry<RevId, PatchLineComment> in) {
+            for (PatchLineComment c : published.get(in.getKey())) {
+              if (c.getKey().equals(in.getValue().getKey())) {
+                return false;
+              }
+            }
+            return true;
+          }
+        });
+    return ImmutableListMultimap.copyOf(
+        filtered);
   }
 
   /**
@@ -518,15 +536,6 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
     return false;
   }
 
-  /** @return the NoteMap */
-  NoteMap getNoteMap() {
-    return noteMap;
-  }
-
-  Map<RevId, RevisionNote> getRevisionNotes() {
-    return revisionNotes;
-  }
-
   @Override
   protected String getRefName() {
     return ChangeNoteUtil.changeRefName(getChangeId());
@@ -557,8 +566,7 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
       changeMessagesByPatchSet = parser.buildMessagesByPatchSet();
       allChangeMessages = parser.buildAllMessages();
       comments = ImmutableListMultimap.copyOf(parser.comments);
-      noteMap = parser.noteMap;
+      revisionNoteMap = parser.revisionNoteMap;
-      revisionNotes = parser.revisionNotes;
       change.setKey(new Change.Key(parser.changeId));
       change.setDest(new Branch.NameKey(project, parser.branch));
       change.setTopic(Strings.emptyToNull(parser.topic));

gerrit-server/src/main/java/com/google/gerrit/server/notedb/ChangeNotesParser.java

@@ -69,7 +69,6 @@ import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.ObjectReader;
 import org.eclipse.jgit.lib.PersonIdent;
 import org.eclipse.jgit.lib.Repository;
-import org.eclipse.jgit.notes.Note;
 import org.eclipse.jgit.notes.NoteMap;
 import org.eclipse.jgit.revwalk.FooterKey;
 import org.eclipse.jgit.revwalk.RevCommit;
@@ -101,7 +100,6 @@ class ChangeNotesParser implements AutoCloseable {
   final Multimap<RevId, PatchLineComment> comments;
   final TreeMap<PatchSet.Id, PatchSet> patchSets;
   final Map<PatchSet.Id, PatchSetState> patchSetStates;
-  final Map<RevId, RevisionNote> revisionNotes;
 
   String branch;
   Change.Status status;
@@ -115,7 +113,7 @@ class ChangeNotesParser implements AutoCloseable {
   String originalSubject;
   String submissionId;
   PatchSet.Id currentPatchSetId;
-  NoteMap noteMap;
+  RevisionNoteMap revisionNoteMap;
 
   private final Change.Id id;
   private final ObjectId tip;
@@ -142,7 +140,6 @@ class ChangeNotesParser implements AutoCloseable {
     comments = ArrayListMultimap.create();
     patchSets = Maps.newTreeMap(ReviewDbUtil.intKeyOrdering());
     patchSetStates = Maps.newHashMap();
-    revisionNotes = Maps.newHashMap();
   }
 
   @Override
@@ -503,19 +500,18 @@ class ChangeNotesParser implements AutoCloseable {
       throws IOException, ConfigInvalidException {
     ObjectReader reader = walk.getObjectReader();
     RevCommit tipCommit = walk.parseCommit(tip);
-    noteMap = NoteMap.read(reader, tipCommit);
+    revisionNoteMap = RevisionNoteMap.parse(
+        id, reader, NoteMap.read(reader, tipCommit), false);
+    Map<RevId, RevisionNote> rns = revisionNoteMap.revisionNotes;
 
-    for (Note note : noteMap) {
+    for (Map.Entry<RevId, RevisionNote> e : rns.entrySet()) {
-      RevisionNote rn = new RevisionNote(id, reader, note.getData());
+      for (PatchLineComment plc : e.getValue().comments) {
-      RevId revId = new RevId(note.name());
+        comments.put(e.getKey(), plc);
-      revisionNotes.put(revId, rn);
-      for (PatchLineComment plc : rn.comments) {
-        comments.put(revId, plc);
       }
     }
 
     for (PatchSet ps : patchSets.values()) {
-      RevisionNote rn = revisionNotes.get(ps.getRevision());
+      RevisionNote rn = rns.get(ps.getRevision());
       if (rn != null && rn.pushCert != null) {
         ps.setPushCertificate(rn.pushCert);
       }

gerrit-server/src/main/java/com/google/gerrit/server/notedb/ChangeRebuilder.java

@@ -418,14 +418,14 @@ public class ChangeRebuilder {
       if (c.getRevId() == null) {
         setCommentRevId(c, cache, change, ps);
       }
-      update.insertComment(c);
+      update.putComment(c);
     }
 
     void applyDraft(ChangeDraftUpdate draftUpdate) throws OrmException {
       if (c.getRevId() == null) {
         setCommentRevId(c, cache, change, ps);
       }
-      draftUpdate.insertComment(c);
+      draftUpdate.putComment(c);
     }
   }
 

gerrit-server/src/main/java/com/google/gerrit/server/notedb/ChangeUpdate.java

@@ -14,6 +14,7 @@
 
 package com.google.gerrit.server.notedb;
 
+import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
@@ -59,6 +60,7 @@ import com.google.gwtorm.server.OrmException;
 import com.google.inject.assistedinject.Assisted;
 import com.google.inject.assistedinject.AssistedInject;
 
+import org.eclipse.jgit.errors.ConfigInvalidException;
 import org.eclipse.jgit.lib.CommitBuilder;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.ObjectInserter;
@@ -71,7 +73,7 @@ import org.eclipse.jgit.revwalk.RevWalk;
 import java.io.IOException;
 import java.util.Comparator;
 import java.util.Date;
-import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -117,7 +119,6 @@ public class ChangeUpdate extends AbstractChangeUpdate {
   private ObjectId commit;
   private Set<String> hashtags;
   private String changeMessage;
-  private ChangeNotes notes;
   private PatchSetState psState;
   private Iterable<String> groups;
   private String pushCert;
@@ -205,13 +206,11 @@ public class ChangeUpdate extends AbstractChangeUpdate {
     return getResult();
   }
 
-  public void setChangeId(String changeId) throws OrmException {
+  public void setChangeId(String changeId) {
-    if (notes == null) {
+    String old = ctl.getChange().getKey().get();
-      notes = getChangeNotes().load();
+    checkArgument(old.equals(changeId),
-    }
-    checkArgument(notes.getChange().getKey().get().equals(changeId),
         "The Change-Id was already set to %s, so we cannot set this Change-Id: %s",
-        notes.getChange().getKey(), changeId);
+        old, changeId);
     this.changeId = changeId;
   }
 
@@ -270,120 +269,41 @@ public class ChangeUpdate extends AbstractChangeUpdate {
     this.changeMessage = changeMessage;
   }
 
-  public void insertComment(PatchLineComment comment) throws OrmException {
+  public void putComment(PatchLineComment c) {
-    if (comment.getStatus() == PatchLineComment.Status.DRAFT) {
-      insertDraftComment(comment);
-    } else {
-      insertPublishedComment(comment);
-    }
-  }
-
-  public void upsertComment(PatchLineComment comment) throws OrmException {
-    if (comment.getStatus() == PatchLineComment.Status.DRAFT) {
-      upsertDraftComment(comment);
-    } else {
-      deleteDraftCommentIfPresent(comment);
-      upsertPublishedComment(comment);
-    }
-  }
-
-  public void updateComment(PatchLineComment comment) throws OrmException {
-    if (comment.getStatus() == PatchLineComment.Status.DRAFT) {
-      updateDraftComment(comment);
-    } else {
-      deleteDraftCommentIfPresent(comment);
-      updatePublishedComment(comment);
-    }
-  }
-
-  public void deleteComment(PatchLineComment comment) throws OrmException {
-    if (comment.getStatus() == PatchLineComment.Status.DRAFT) {
-      deleteDraftComment(comment);
-    } else {
-      throw new IllegalArgumentException("Cannot delete a published comment.");
-    }
-  }
-
-  private void insertPublishedComment(PatchLineComment c) throws OrmException {
     verifyComment(c);
-    if (notes == null) {
-      notes = getChangeNotes().load();
-    }
-    if (migration.readChanges()) {
-      checkArgument(!notes.containsComment(c),
-          "A comment already exists with the same key as the following comment,"
-          + " so we cannot insert this comment: %s", c);
-    }
-    comments.add(c);
-  }
-
-  private void insertDraftComment(PatchLineComment c) throws OrmException {
     createDraftUpdateIfNull();
-    draftUpdate.insertComment(c);
+    if (c.getStatus() == PatchLineComment.Status.DRAFT) {
+      draftUpdate.putComment(c);
+    } else {
+      comments.add(c);
+      // Always delete the corresponding comment from drafts. Published comments
+      // are immutable, meaning in normal operation we only hit this path when
+      // publishing a comment. It's exactly in that case that we have to delete
+      // the draft.
+      draftUpdate.deleteComment(c);
+    }
   }
 
-  private void upsertPublishedComment(PatchLineComment c) throws OrmException {
+  public void deleteComment(PatchLineComment c) {
     verifyComment(c);
-    if (notes == null) {
+    if (c.getStatus() == PatchLineComment.Status.DRAFT) {
-      notes = getChangeNotes().load();
+      createDraftUpdateIfNull().deleteComment(c);
+    } else {
+      throw new IllegalArgumentException(
+          "Cannot delete published comment " + c);
     }
-    // This could allow callers to update a published comment if migration.write
-    // is on and migration.readComments is off because we will not be able to
-    // verify that the comment didn't already exist as a published comment
-    // since we don't have a ReviewDb.
-    if (migration.readChanges()) {
-      checkArgument(!notes.containsCommentPublished(c),
-          "Cannot update a comment that has already been published and saved");
-    }
-    comments.add(c);
   }
 
-  private void upsertDraftComment(PatchLineComment c) {
+  @VisibleForTesting
-    createDraftUpdateIfNull();
+  ChangeDraftUpdate createDraftUpdateIfNull() {
-    draftUpdate.upsertComment(c);
-  }
-
-  private void updatePublishedComment(PatchLineComment c) throws OrmException {
-    verifyComment(c);
-    if (notes == null) {
-      notes = getChangeNotes().load();
-    }
-    // See comment above in upsertPublishedComment() about potential risk with
-    // this check.
-    if (migration.readChanges()) {
-      checkArgument(!notes.containsCommentPublished(c),
-          "Cannot update a comment that has already been published and saved");
-    }
-    comments.add(c);
-  }
-
-  private void updateDraftComment(PatchLineComment c) throws OrmException {
-    createDraftUpdateIfNull();
-    draftUpdate.updateComment(c);
-  }
-
-  private void deleteDraftComment(PatchLineComment c) throws OrmException {
-    createDraftUpdateIfNull();
-    draftUpdate.deleteComment(c);
-  }
-
-  private void deleteDraftCommentIfPresent(PatchLineComment c)
-      throws OrmException {
-    createDraftUpdateIfNull();
-    draftUpdate.deleteCommentIfPresent(c);
-  }
-
-  private void createDraftUpdateIfNull() {
     if (draftUpdate == null) {
       draftUpdate = draftUpdateFactory.create(ctl, when);
     }
+    return draftUpdate;
   }
 
   private void verifyComment(PatchLineComment c) {
     checkArgument(c.getRevId() != null);
-    checkArgument(c.getStatus() == PatchLineComment.Status.PUBLISHED,
-        "Cannot add a draft comment to a ChangeUpdate. Use a ChangeDraftUpdate"
-        + " for draft comments");
     checkArgument(c.getAuthor().equals(getUser().getAccountId()),
         "The author for the following comment does not match the author of"
         + " this ChangeDraftUpdate (%s): %s", getUser().getAccountId(), c);
@@ -430,44 +350,93 @@ public class ChangeUpdate extends AbstractChangeUpdate {
   }
 
   /** @return the tree id for the updated tree */
-  private ObjectId storeRevisionNotes(ObjectInserter inserter)
+  private ObjectId storeRevisionNotes(RevWalk rw, ObjectInserter inserter,
-      throws OrmException, IOException {
+      ObjectId curr) throws ConfigInvalidException, OrmException, IOException {
-    ChangeNotes notes = ctl.getNotes().load();
-    NoteMap noteMap = notes.getNoteMap();
-    if (noteMap == null) {
-      noteMap = NoteMap.newEmptyMap();
-    }
     if (comments.isEmpty() && pushCert == null) {
       return null;
     }
+    RevisionNoteMap rnm = getRevisionNoteMap(rw, curr);
 
-    Map<RevId, RevisionNoteBuilder> builders = new HashMap<>();
+    RevisionNoteBuilder.Cache cache = new RevisionNoteBuilder.Cache(rnm);
     for (PatchLineComment c : comments) {
-      builder(builders, c.getRevId()).addComment(c);
+      cache.get(c.getRevId()).putComment(c);
     }
     if (pushCert != null) {
       checkState(commit != null);
-      builder(builders, new RevId(commit.name())).setPushCertificate(pushCert);
+      cache.get(new RevId(commit.name())).setPushCertificate(pushCert);
     }
+    Map<RevId, RevisionNoteBuilder> builders = cache.getBuilders();
+    checkComments(rnm.revisionNotes, builders);
 
     for (Map.Entry<RevId, RevisionNoteBuilder> e : builders.entrySet()) {
       ObjectId data = inserter.insert(
           OBJ_BLOB, e.getValue().build(commentsUtil));
-      noteMap.set(ObjectId.fromString(e.getKey().get()), data);
+      rnm.noteMap.set(ObjectId.fromString(e.getKey().get()), data);
     }
 
-    return noteMap.writeTree(inserter);
+    return rnm.noteMap.writeTree(inserter);
   }
 
-  private RevisionNoteBuilder builder(Map<RevId, RevisionNoteBuilder> builders,
+  private RevisionNoteMap getRevisionNoteMap(RevWalk rw, ObjectId curr)
-      RevId revId) {
+      throws ConfigInvalidException, OrmException, IOException {
-    RevisionNoteBuilder b = builders.get(revId);
+    if (migration.readChanges()) {
-    if (b == null) {
+      // If reading from changes is enabled, then the old ChangeNotes already
-      b = new RevisionNoteBuilder(
+      // parsed the revision notes. We can reuse them as long as the ref hasn't
-          getChangeNotes().getRevisionNotes().get(revId));
+      // advanced.
-      builders.put(revId, b);
+      ObjectId idFromNotes =
+          firstNonNull(ctl.getNotes().load().getRevision(), ObjectId.zeroId());
+      if (idFromNotes.equals(curr)) {
+        return checkNotNull(ctl.getNotes().revisionNoteMap);
+      }
+    }
+    NoteMap noteMap;
+    if (!curr.equals(ObjectId.zeroId())) {
+      noteMap = NoteMap.read(rw.getObjectReader(), rw.parseCommit(curr));
+    } else {
+      noteMap = NoteMap.newEmptyMap();
+    }
+    // Even though reading from changes might not be enabled, we need to
+    // parse any existing revision notes so we can merge them.
+    return RevisionNoteMap.parse(
+        ctl.getId(), rw.getObjectReader(), noteMap, false);
+  }
+
+  private void checkComments(Map<RevId, RevisionNote> existingNotes,
+      Map<RevId, RevisionNoteBuilder> toUpdate) throws OrmException {
+    // Prohibit various kinds of illegal operations on comments.
+    Set<PatchLineComment.Key> existing = new HashSet<>();
+    for (RevisionNote rn : existingNotes.values()) {
+      for (PatchLineComment c : rn.comments) {
+        existing.add(c.getKey());
+        if (draftUpdate != null) {
+          // Take advantage of an existing update on All-Users to prune any
+          // published comments from drafts. NoteDbUpdateManager takes care of
+          // ensuring that this update is applied before its dependent draft
+          // update.
+          //
+          // Deleting aggressively in this way, combined with filtering out
+          // duplicate published/draft comments in ChangeNotes#getDraftComments,
+          // makes up for the fact that updates between the change repo and
+          // All-Users are not atomic.
+          //
+          // TODO(dborowitz): We might want to distinguish between deleted
+          // drafts that we're fixing up after the fact by putting them in a
+          // separate commit. But note that we don't care much about the commit
+          // graph of the draft ref, particularly because the ref is completely
+          // deleted when all drafts are gone.
+          draftUpdate.deleteComment(c.getRevId(), c.getKey());
+        }
+      }
+    }
+
+    for (RevisionNoteBuilder b : toUpdate.values()) {
+      for (PatchLineComment c : b.put.values()) {
+        if (existing.contains(c.getKey())) {
+          throw new OrmException(
+              "Cannot update existing published comment: " + c);
+        }
+      }
     }
-    return b;
   }
 
   @Override
@@ -476,8 +445,8 @@ public class ChangeUpdate extends AbstractChangeUpdate {
   }
 
   @Override
-  protected CommitBuilder applyImpl(ObjectInserter ins)
+  protected CommitBuilder applyImpl(RevWalk rw, ObjectInserter ins,
-      throws OrmException, IOException {
+      ObjectId curr) throws OrmException, IOException {
     CommitBuilder cb = new CommitBuilder();
 
     int ps = psId != null ? psId.get() : getChange().currentPatchSetId().get();
@@ -580,9 +549,13 @@ public class ChangeUpdate extends AbstractChangeUpdate {
     }
 
     cb.setMessage(msg.toString());
-    ObjectId treeId = storeRevisionNotes(ins);
+    try {
-    if (treeId != null) {
+      ObjectId treeId = storeRevisionNotes(rw, ins, curr);
-      cb.setTreeId(treeId);
+      if (treeId != null) {
+        cb.setTreeId(treeId);
+      }
+    } catch (ConfigInvalidException e) {
+      throw new OrmException(e);
     }
     return cb;
   }

gerrit-server/src/main/java/com/google/gerrit/server/notedb/CommentsInNotesUtil.java

@@ -15,16 +15,12 @@
 package com.google.gerrit.server.notedb;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.gerrit.server.PatchLineCommentsUtil.PLC_ORDER;
 import static com.google.gerrit.server.notedb.ChangeNoteUtil.GERRIT_PLACEHOLDER_HOST;
 import static com.google.gerrit.server.notedb.ChangeNotes.parseException;
-import static com.google.gerrit.server.notedb.RevisionNote.MAX_NOTE_SZ;
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.eclipse.jgit.lib.Constants.OBJ_BLOB;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
-import com.google.common.collect.Multimap;
 import com.google.common.primitives.Ints;
 import com.google.gerrit.reviewdb.client.Account;
 import com.google.gerrit.reviewdb.client.Change;
@@ -42,16 +38,7 @@ import com.google.inject.Inject;
 import com.google.inject.Singleton;
 
 import org.eclipse.jgit.errors.ConfigInvalidException;
-import org.eclipse.jgit.lib.ObjectId;
-import org.eclipse.jgit.lib.ObjectInserter;
-import org.eclipse.jgit.lib.ObjectReader;
 import org.eclipse.jgit.lib.PersonIdent;
-import org.eclipse.jgit.lib.Ref;
-import org.eclipse.jgit.lib.Repository;
-import org.eclipse.jgit.notes.Note;
-import org.eclipse.jgit.notes.NoteMap;
-import org.eclipse.jgit.revwalk.RevCommit;
-import org.eclipse.jgit.revwalk.RevWalk;
 import org.eclipse.jgit.util.GitDateFormatter;
 import org.eclipse.jgit.util.GitDateFormatter.Format;
 import org.eclipse.jgit.util.GitDateParser;
@@ -60,18 +47,14 @@ import org.eclipse.jgit.util.QuotedString;
 import org.eclipse.jgit.util.RawParseUtils;
 
 import java.io.ByteArrayOutputStream;
-import java.io.IOException;
 import java.io.OutputStream;
 import java.io.OutputStreamWriter;
 import java.io.PrintWriter;
 import java.sql.Timestamp;
 import java.text.ParseException;
-import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Date;
 import java.util.List;
 import java.util.Locale;
-import java.util.Map;
 
 /**
  * Utility functions to parse PatchLineComments out of a note byte array and
@@ -89,33 +72,6 @@ public class CommentsInNotesUtil {
   private static final String REVISION = "Revision";
   private static final String UUID = "UUID";
 
-  public static NoteMap parseCommentsFromNotes(Repository repo, String refName,
-      RevWalk walk, Change.Id changeId,
-      Multimap<RevId, PatchLineComment> comments,
-      Status status)
-      throws IOException, ConfigInvalidException {
-    Ref ref = repo.getRefDatabase().exactRef(refName);
-    if (ref == null) {
-      return null;
-    }
-
-    ObjectReader reader = walk.getObjectReader();
-    RevCommit commit = walk.parseCommit(ref.getObjectId());
-    NoteMap noteMap = NoteMap.read(reader, commit);
-
-    for (Note note : noteMap) {
-      byte[] bytes =
-          reader.open(note.getData(), OBJ_BLOB).getCachedBytes(MAX_NOTE_SZ);
-      List<PatchLineComment> result =
-          parseNote(bytes, new MutableInteger(), changeId, status);
-      if (result == null || result.isEmpty()) {
-        continue;
-      }
-      comments.putAll(new RevId(note.name()), result);
-    }
-    return noteMap;
-  }
-
   public static List<PatchLineComment> parseNote(byte[] note, MutableInteger p,
       Change.Id changeId, Status status) throws ConfigInvalidException {
     if (p.value >= note.length) {
@@ -432,7 +388,7 @@ public class CommentsInNotesUtil {
    *            same side and must share the same patch set ID.
    * @param out output stream to write to.
    */
-  public void buildNote(List<PatchLineComment> comments, OutputStream out) {
+  void buildNote(List<PatchLineComment> comments, OutputStream out) {
     if (comments.isEmpty()) {
       return;
     }
@@ -514,51 +470,4 @@ public class CommentsInNotesUtil {
       }
     }
   }
-
-  /**
-   * Write comments for multiple revisions to a note map.
-   * <p>
-   * Mutates the map in-place. only notes for SHA-1s found as keys in the map
-   * are modified; all other notes are left untouched.
-   *
-   * @param noteMap note map to modify.
-   * @param allComments map of revision to all comments for that revision;
-   *     callers are responsible for reading the original comments and applying
-   *     any changes. Differs from a multimap in that present-but-empty values
-   *     are significant, and indicate the note for that SHA-1 should be
-   *     deleted.
-   * @param inserter object inserter for writing notes.
-   * @throws IOException if an error occurred.
-   */
-  public void writeCommentsToNoteMap(NoteMap noteMap,
-      Map<RevId, List<PatchLineComment>> allComments, ObjectInserter inserter)
-      throws IOException {
-    for (Map.Entry<RevId, List<PatchLineComment>> e : allComments.entrySet()) {
-      List<PatchLineComment> comments = e.getValue();
-      ObjectId commit = ObjectId.fromString(e.getKey().get());
-      if (comments.isEmpty()) {
-        noteMap.remove(commit);
-        continue;
-      }
-      Collections.sort(comments, PLC_ORDER);
-      // We allow comments for multiple commits to be written in the same
-      // update, even though the rest of the metadata update is associated with
-      // a single patch set.
-      byte[] note = buildNote(comments);
-      if (note != null && note.length > 0) {
-        noteMap.set(commit, inserter.insert(OBJ_BLOB, note));
-      }
-    }
-  }
-
-  static void addCommentToMap(Map<RevId, List<PatchLineComment>> map,
-      PatchLineComment c) {
-    List<PatchLineComment> list = map.get(c.getRevId());
-    if (list == null) {
-      list = new ArrayList<>();
-      map.put(c.getRevId(), list);
-    }
-    list.add(c);
-  }
-
 }

gerrit-server/src/main/java/com/google/gerrit/server/notedb/DraftCommentNotes.java

@@ -15,7 +15,9 @@
 package com.google.gerrit.server.notedb;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.Multimap;
 import com.google.gerrit.reviewdb.client.Account;
 import com.google.gerrit.reviewdb.client.Change;
 import com.google.gerrit.reviewdb.client.PatchLineComment;
@@ -31,6 +33,7 @@ import org.eclipse.jgit.errors.ConfigInvalidException;
 import org.eclipse.jgit.lib.CommitBuilder;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.notes.NoteMap;
+import org.eclipse.jgit.revwalk.RevCommit;
 import org.eclipse.jgit.revwalk.RevWalk;
 
 import java.io.IOException;
@@ -66,7 +69,7 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
   private final Account.Id author;
 
   private ImmutableListMultimap<RevId, PatchLineComment> comments;
-  private NoteMap noteMap;
+  private RevisionNoteMap revisionNoteMap;
 
   DraftCommentNotes(GitRepositoryManager repoManager, NotesMigration migration,
       AllUsersName draftsProject, Change.Id changeId, Account.Id author) {
@@ -75,8 +78,8 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
     this.author = author;
   }
 
-  public NoteMap getNoteMap() {
+  RevisionNoteMap getRevisionNoteMap() {
-    return noteMap;
+    return revisionNoteMap;
   }
 
   public Account.Id getAuthor() {
@@ -110,13 +113,17 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
       return;
     }
 
-    try (RevWalk walk = new RevWalk(reader);
+    try (RevWalk walk = new RevWalk(reader)) {
-        DraftCommentNotesParser parser = new DraftCommentNotesParser(
+      RevCommit tipCommit = walk.parseCommit(rev);
-          getChangeId(), walk, rev, repoManager, draftsProject, author)) {
+      revisionNoteMap = RevisionNoteMap.parse(
-      parser.parseDraftComments();
+          getChangeId(), reader, NoteMap.read(reader, tipCommit), true);
-
+      Multimap<RevId, PatchLineComment> cs = ArrayListMultimap.create();
-      comments = ImmutableListMultimap.copyOf(parser.comments);
+      for (RevisionNote rn : revisionNoteMap.revisionNotes.values()) {
-      noteMap = parser.noteMap;
+        for (PatchLineComment c : rn.comments) {
+          cs.put(c.getRevId(), c);
+        }
+      }
+      comments = ImmutableListMultimap.copyOf(cs);
     }
   }
 
@@ -136,4 +143,9 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
   public Project.NameKey getProjectName() {
     return draftsProject;
   }
+
+  @VisibleForTesting
+  NoteMap getNoteMap() {
+    return revisionNoteMap != null ? revisionNoteMap.noteMap : null;
+  }
 }

gerrit-server/src/main/java/com/google/gerrit/server/notedb/DraftCommentNotesParser.java

@@ -1,69 +0,0 @@
-// Copyright (C) 2014 The Android Open Source Project
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package com.google.gerrit.server.notedb;
-
-import com.google.common.collect.ArrayListMultimap;
-import com.google.common.collect.Multimap;
-import com.google.gerrit.reviewdb.client.Account;
-import com.google.gerrit.reviewdb.client.Change;
-import com.google.gerrit.reviewdb.client.PatchLineComment;
-import com.google.gerrit.reviewdb.client.RefNames;
-import com.google.gerrit.reviewdb.client.RevId;
-import com.google.gerrit.server.config.AllUsersName;
-import com.google.gerrit.server.git.GitRepositoryManager;
-
-import org.eclipse.jgit.errors.ConfigInvalidException;
-import org.eclipse.jgit.errors.RepositoryNotFoundException;
-import org.eclipse.jgit.lib.ObjectId;
-import org.eclipse.jgit.lib.Repository;
-import org.eclipse.jgit.notes.NoteMap;
-import org.eclipse.jgit.revwalk.RevWalk;
-
-import java.io.IOException;
-
-class DraftCommentNotesParser implements AutoCloseable {
-  final Multimap<RevId, PatchLineComment> comments;
-  NoteMap noteMap;
-
-  private final Change.Id changeId;
-  private final ObjectId tip;
-  private final RevWalk walk;
-  private final Repository repo;
-  private final Account.Id author;
-
-  DraftCommentNotesParser(Change.Id changeId, RevWalk walk, ObjectId tip,
-      GitRepositoryManager repoManager, AllUsersName draftsProject,
-      Account.Id author) throws RepositoryNotFoundException, IOException {
-    this.changeId = changeId;
-    this.walk = walk;
-    this.tip = tip;
-    this.repo = repoManager.openMetadataRepository(draftsProject);
-    this.author = author;
-
-    comments = ArrayListMultimap.create();
-  }
-
-  @Override
-  public void close() {
-    repo.close();
-  }
-
-  void parseDraftComments() throws IOException, ConfigInvalidException {
-    walk.markStart(walk.parseCommit(tip));
-    noteMap = CommentsInNotesUtil.parseCommentsFromNotes(repo,
-        RefNames.refsDraftComments(author, changeId),
-        walk, changeId, comments, PatchLineComment.Status.DRAFT);
-  }
-}

gerrit-server/src/main/java/com/google/gerrit/server/notedb/NoteDbUpdateManager.java

@@ -152,7 +152,11 @@ public class NoteDbUpdateManager {
   }
 
   private boolean isEmpty() {
-    return !migration.writeChanges() || changeUpdates.isEmpty();
+    if (!migration.writeChanges()) {
+      return true;
+    }
+    return changeUpdates.isEmpty()
+        && draftUpdates.isEmpty();
   }
 
   /**
@@ -189,8 +193,16 @@ public class NoteDbUpdateManager {
       }
       addCommands();
 
-      execute(allUsersRepo);
+      // ChangeUpdates must execute before ChangeDraftUpdates.
+      //
+      // ChangeUpdate will automatically delete draft comments for any published
+      // comments, but the updates to the two repos don't happen atomically.
+      // Thus if the change meta update succeeds and the All-Users update fails,
+      // we may have stale draft comments. Doing it in this order allows stale
+      // comments to be filtered out by ChangeNotes, reflecting the fact that
+      // comments can only go from DRAFT to PUBLISHED, not vice versa.
       execute(changeRepo);
+      execute(allUsersRepo);
     } finally {
       if (allUsersRepo != null) {
         allUsersRepo.close();

gerrit-server/src/main/java/com/google/gerrit/server/notedb/RevisionNote.java

@@ -63,14 +63,21 @@ class RevisionNote {
   final ImmutableList<PatchLineComment> comments;
   final String pushCert;
 
-  RevisionNote(Change.Id changeId, ObjectReader reader, ObjectId noteId)
+  RevisionNote(Change.Id changeId, ObjectReader reader, ObjectId noteId,
-      throws ConfigInvalidException, IOException {
+      boolean draftsOnly) throws ConfigInvalidException, IOException {
     byte[] bytes = reader.open(noteId, OBJ_BLOB).getCachedBytes(MAX_NOTE_SZ);
     MutableInteger p = new MutableInteger();
     trimLeadingEmptyLines(bytes, p);
-    pushCert = parsePushCert(changeId, bytes, p);
+    if (!draftsOnly) {
-    trimLeadingEmptyLines(bytes, p);
+      pushCert = parsePushCert(changeId, bytes, p);
-    comments = ImmutableList.copyOf(CommentsInNotesUtil.parseNote(
+      trimLeadingEmptyLines(bytes, p);
-        bytes, p, changeId, PatchLineComment.Status.PUBLISHED));
+    } else {
+      pushCert = null;
+    }
+    PatchLineComment.Status status = draftsOnly
+        ? PatchLineComment.Status.DRAFT
+        : PatchLineComment.Status.PUBLISHED;
+    comments = ImmutableList.copyOf(
+        CommentsInNotesUtil.parseNote(bytes, p, changeId, status));
   }
 }

gerrit-server/src/main/java/com/google/gerrit/server/notedb/RevisionNoteBuilder.java

@@ -14,35 +14,77 @@
 
 package com.google.gerrit.server.notedb;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.gerrit.server.PatchLineCommentsUtil.PLC_ORDER;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.gerrit.reviewdb.client.PatchLineComment;
+import com.google.gerrit.reviewdb.client.RevId;
 
 import java.io.ByteArrayOutputStream;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 class RevisionNoteBuilder {
-  private final Map<PatchLineComment.Key, PatchLineComment> comments;
+  static class Cache {
+    private final RevisionNoteMap revisionNoteMap;
+    private final Map<RevId, RevisionNoteBuilder> builders;
+
+    Cache(RevisionNoteMap revisionNoteMap) {
+      this.revisionNoteMap = revisionNoteMap;
+      this.builders = new HashMap<>();
+    }
+
+    RevisionNoteBuilder get(RevId revId) {
+      RevisionNoteBuilder b = builders.get(revId);
+      if (b == null) {
+        b = new RevisionNoteBuilder(
+            revisionNoteMap.revisionNotes.get(revId));
+        builders.put(revId, b);
+      }
+      return b;
+    }
+
+    Map<RevId, RevisionNoteBuilder> getBuilders() {
+      return Collections.unmodifiableMap(builders);
+    }
+  }
+
+  final List<PatchLineComment> baseComments;
+  final Map<PatchLineComment.Key, PatchLineComment> put;
+  final Set<PatchLineComment.Key> delete;
+
   private String pushCert;
 
   RevisionNoteBuilder(RevisionNote base) {
     if (base != null) {
-      comments = Maps.newHashMapWithExpectedSize(base.comments.size());
+      baseComments = base.comments;
-      for (PatchLineComment c : base.comments) {
+      put = Maps.newHashMapWithExpectedSize(base.comments.size());
-        addComment(c);
-      }
       pushCert = base.pushCert;
     } else {
-      comments = new HashMap<>();
+      baseComments = Collections.emptyList();
+      put = new HashMap<>();
       pushCert = null;
     }
+    delete = new HashSet<>();
   }
 
-  void addComment(PatchLineComment comment) {
+  void putComment(PatchLineComment comment) {
-    comments.put(comment.getKey(), comment);
+    checkArgument(!delete.contains(comment.getKey()),
+        "cannot both delete and put %s", comment.getKey());
+    put.put(comment.getKey(), comment);
+  }
+
+  void deleteComment(PatchLineComment.Key key) {
+    checkArgument(!put.containsKey(key), "cannot both delete and put %s", key);
+    delete.add(key);
   }
 
   void setPushCertificate(String pushCert) {
@@ -56,7 +98,16 @@ class RevisionNoteBuilder {
       out.write(certBytes, 0, trimTrailingNewlines(certBytes));
       out.write('\n');
     }
-    commentsUtil.buildNote(PLC_ORDER.sortedCopy(comments.values()), out);
+
+    List<PatchLineComment> all =
+        new ArrayList<>(baseComments.size() + put.size());
+    for (PatchLineComment c : Iterables.concat(baseComments, put.values())) {
+      if (!delete.contains(c.getKey())) {
+        all.add(c);
+      }
+    }
+    Collections.sort(all, PLC_ORDER);
+    commentsUtil.buildNote(all, out);
     return out.toByteArray();
   }
 

gerrit-server/src/main/java/com/google/gerrit/server/notedb/RevisionNoteMap.java

@@ -0,0 +1,51 @@
+// Copyright (C) 2016 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.gerrit.server.notedb;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.gerrit.reviewdb.client.Change;
+import com.google.gerrit.reviewdb.client.RevId;
+
+import org.eclipse.jgit.errors.ConfigInvalidException;
+import org.eclipse.jgit.lib.ObjectReader;
+import org.eclipse.jgit.notes.Note;
+import org.eclipse.jgit.notes.NoteMap;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+class RevisionNoteMap {
+  final NoteMap noteMap;
+  final ImmutableMap<RevId, RevisionNote> revisionNotes;
+
+  static RevisionNoteMap parse(Change.Id changeId, ObjectReader reader,
+      NoteMap noteMap, boolean draftsOnly)
+      throws ConfigInvalidException, IOException {
+    Map<RevId, RevisionNote> result = new HashMap<>();
+    for (Note note : noteMap) {
+      RevisionNote rn =
+          new RevisionNote(changeId, reader, note.getData(), draftsOnly);
+      result.put(new RevId(note.name()), rn);
+    }
+    return new RevisionNoteMap(noteMap, ImmutableMap.copyOf(result));
+  }
+
+  private RevisionNoteMap(NoteMap noteMap,
+      ImmutableMap<RevId, RevisionNote> revisionNotes) {
+    this.noteMap = noteMap;
+    this.revisionNotes = revisionNotes;
+  }
+}

gerrit-server/src/test/java/com/google/gerrit/server/notedb/AbstractChangeNotesTest.java

@@ -95,14 +95,18 @@ public class AbstractChangeNotesTest extends GerritBaseTests {
   protected RevWalk rw;
   protected TestRepository<InMemoryRepository> tr;
 
-  @Inject protected IdentifiedUser.GenericFactory userFactory;
+  @Inject
-  @Inject protected NoteDbUpdateManager.Factory updateManagerFactory;
+  protected IdentifiedUser.GenericFactory userFactory;
+
+  @Inject
+  protected NoteDbUpdateManager.Factory updateManagerFactory;
+
+  @Inject
+  protected AllUsersName allUsers;
 
   private Injector injector;
   private String systemTimeZone;
 
-  @Inject private AllUsersName allUsers;
-
   @Before
   public void setUp() throws Exception {
     setTimeForTesting();

gerrit-server/src/test/java/com/google/gerrit/server/notedb/ChangeNotesTest.java

@@ -43,12 +43,15 @@ import com.google.gerrit.reviewdb.client.PatchLineComment;
 import com.google.gerrit.reviewdb.client.PatchLineComment.Status;
 import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.client.PatchSetApproval;
+import com.google.gerrit.reviewdb.client.RefNames;
 import com.google.gerrit.reviewdb.client.RevId;
 import com.google.gwtorm.server.OrmException;
+import com.google.inject.Inject;
 
 import org.eclipse.jgit.lib.Constants;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.Ref;
+import org.eclipse.jgit.lib.Repository;
 import org.eclipse.jgit.notes.Note;
 import org.eclipse.jgit.notes.NoteMap;
 import org.eclipse.jgit.revwalk.RevCommit;
@@ -62,6 +65,9 @@ import java.util.List;
 import java.util.Map;
 
 public class ChangeNotesTest extends AbstractChangeNotesTest {
+  @Inject
+  private DraftCommentNotes.Factory draftNotesFactory;
+
   @Test
   public void approvalsOnePatchSet() throws Exception {
     Change c = newChange();
@@ -708,7 +714,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     update.setPatchSetState(PatchSetState.DRAFT);
     update.putApproval("Code-Review", (short) 1);
     update.setChangeMessage("This is a message");
-    update.insertComment(newPublishedComment(c.currentPatchSetId(), "a.txt",
+    update.putComment(newPublishedComment(c.currentPatchSetId(), "a.txt",
         "uuid1", new CommentRange(1, 2, 3, 4), 1, changeOwner, null,
         TimeUtil.nowTs(), "Comment", (short) 1, commit.name()));
     update.commit();
@@ -806,7 +812,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     update = newUpdate(c, changeOwner);
     update.setPatchSetId(psId2);
     Timestamp ts = TimeUtil.nowTs();
-    update.insertComment(newPublishedComment(psId2, "a.txt",
+    update.putComment(newPublishedComment(psId2, "a.txt",
         "uuid1", new CommentRange(1, 2, 3, 4), 1, changeOwner, null, ts,
         "Comment", (short) 1, commit.name()));
     update.commit();
@@ -883,7 +889,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid1, range1, range1.getEndLine(), otherUser, null, time1, message1,
         (short) 0, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update1.setPatchSetId(psId);
-    update1.upsertComment(comment1);
+    update1.putComment(comment1);
     updateManager.add(update1);
 
     ChangeUpdate update2 = newUpdate(c, otherUser);
@@ -1123,7 +1129,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid1, range1, range1.getEndLine(), otherUser, null, time1, message1,
         (short) 1, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
-    update.upsertComment(comment1);
+    update.putComment(comment1);
     update.commit();
 
     update = newUpdate(c, otherUser);
@@ -1132,7 +1138,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid2, range2, range2.getEndLine(), otherUser, null, time2, message2,
         (short) 1, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
-    update.upsertComment(comment2);
+    update.putComment(comment2);
     update.commit();
 
     update = newUpdate(c, otherUser);
@@ -1141,14 +1147,14 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid3, range3, range3.getEndLine(), otherUser, null, time3, message3,
         (short) 1, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
-    update.upsertComment(comment3);
+    update.putComment(comment3);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
 
     try (RevWalk walk = new RevWalk(repo)) {
       ArrayList<Note> notesInTree =
-          Lists.newArrayList(notes.getNoteMap().iterator());
+          Lists.newArrayList(notes.revisionNoteMap.noteMap.iterator());
       Note note = Iterables.getOnlyElement(notesInTree);
 
       byte[] bytes =
@@ -1202,7 +1208,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid1, range1, range1.getEndLine(), otherUser, null, time1, message1,
         (short) 0, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
-    update.upsertComment(comment1);
+    update.putComment(comment1);
     update.commit();
 
     update = newUpdate(c, otherUser);
@@ -1211,14 +1217,14 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid2, range2, range2.getEndLine(), otherUser, null, time2, message2,
         (short) 0, "abcd1234abcd1234abcd1234abcd1234abcd1234");
     update.setPatchSetId(psId);
-    update.upsertComment(comment2);
+    update.putComment(comment2);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
 
     try (RevWalk walk = new RevWalk(repo)) {
       ArrayList<Note> notesInTree =
-          Lists.newArrayList(notes.getNoteMap().iterator());
+          Lists.newArrayList(notes.revisionNoteMap.noteMap.iterator());
       Note note = Iterables.getOnlyElement(notesInTree);
 
       byte[] bytes =
@@ -1266,7 +1272,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         range, range.getEndLine(), otherUser, null, now, messageForBase,
         (short) 0, rev1);
     update.setPatchSetId(psId);
-    update.upsertComment(commentForBase);
+    update.putComment(commentForBase);
     update.commit();
 
     update = newUpdate(c, otherUser);
@@ -1275,7 +1281,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         range, range.getEndLine(), otherUser, null, now, messageForPS,
         (short) 1, rev2);
     update.setPatchSetId(psId);
-    update.upsertComment(commentForPS);
+    update.putComment(commentForPS);
     update.commit();
 
     assertThat(newNotes(c).getComments()).containsExactlyEntriesIn(
@@ -1302,7 +1308,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid1, range, range.getEndLine(), otherUser, null, timeForComment1,
         "comment 1", side, rev);
     update.setPatchSetId(psId);
-    update.upsertComment(comment1);
+    update.putComment(comment1);
     update.commit();
 
     update = newUpdate(c, otherUser);
@@ -1310,7 +1316,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid2, range, range.getEndLine(), otherUser, null, timeForComment2,
         "comment 2", side, rev);
     update.setPatchSetId(psId);
-    update.upsertComment(comment2);
+    update.putComment(comment2);
     update.commit();
 
     assertThat(newNotes(c).getComments()).containsExactlyEntriesIn(
@@ -1337,7 +1343,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid, range, range.getEndLine(), otherUser, null, now, "comment 1",
         side, rev);
     update.setPatchSetId(psId);
-    update.upsertComment(comment1);
+    update.putComment(comment1);
     update.commit();
 
     update = newUpdate(c, otherUser);
@@ -1345,7 +1351,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid, range, range.getEndLine(), otherUser, null, now, "comment 2",
         side, rev);
     update.setPatchSetId(psId);
-    update.upsertComment(comment2);
+    update.putComment(comment2);
     update.commit();
 
     assertThat(newNotes(c).getComments()).containsExactlyEntriesIn(
@@ -1371,7 +1377,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid, range, range.getEndLine(), otherUser, null, now, "comment on ps1",
         side, rev1);
     update.setPatchSetId(ps1);
-    update.upsertComment(comment1);
+    update.putComment(comment1);
     update.commit();
 
     incrementPatchSet(c);
@@ -1383,7 +1389,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid, range, range.getEndLine(), otherUser, null, now, "comment on ps2",
         side, rev2);
     update.setPatchSetId(ps2);
-    update.upsertComment(comment2);
+    update.putComment(comment2);
     update.commit();
 
     assertThat(newNotes(c).getComments()).containsExactlyEntriesIn(
@@ -1408,7 +1414,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         range.getEndLine(), otherUser, null, now, "comment on ps1", side,
         rev, Status.DRAFT);
     update.setPatchSetId(ps1);
-    update.insertComment(comment1);
+    update.putComment(comment1);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
@@ -1419,7 +1425,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     comment1.setStatus(Status.PUBLISHED);
     update = newUpdate(c, otherUser);
     update.setPatchSetId(ps1);
-    update.updateComment(comment1);
+    update.putComment(comment1);
     update.commit();
 
     notes = newNotes(c);
@@ -1451,8 +1457,8 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     PatchLineComment comment2 = newComment(psId, filename, uuid2,
         range2, range2.getEndLine(), otherUser, null, now, "other on ps1",
         side, rev, Status.DRAFT);
-    update.insertComment(comment1);
+    update.putComment(comment1);
-    update.insertComment(comment2);
+    update.putComment(comment2);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
@@ -1466,7 +1472,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     update = newUpdate(c, otherUser);
     update.setPatchSetId(psId);
     comment1.setStatus(Status.PUBLISHED);
-    update.updateComment(comment1);
+    update.putComment(comment1);
     update.commit();
 
     notes = newNotes(c);
@@ -1500,8 +1506,8 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         range2, range2.getEndLine(), otherUser, null, now, "comment on ps",
         (short) 1, rev2, Status.DRAFT);
 
-    update.insertComment(baseComment);
+    update.putComment(baseComment);
-    update.insertComment(psComment);
+    update.putComment(psComment);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
@@ -1517,8 +1523,8 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
 
     baseComment.setStatus(Status.PUBLISHED);
     psComment.setStatus(Status.PUBLISHED);
-    update.updateComment(baseComment);
+    update.putComment(baseComment);
-    update.updateComment(psComment);
+    update.putComment(psComment);
     update.commit();
 
     notes = newNotes(c);
@@ -1546,7 +1552,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         range.getEndLine(), otherUser, null, now, "comment on ps1", side,
         rev, Status.DRAFT);
     update.setPatchSetId(psId);
-    update.upsertComment(comment);
+    update.putComment(comment);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
@@ -1585,7 +1591,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid, range, range.getEndLine(), otherUser, null, now, "comment on ps1",
         side, rev1, Status.DRAFT);
     update.setPatchSetId(ps1);
-    update.upsertComment(comment1);
+    update.putComment(comment1);
     update.commit();
 
     incrementPatchSet(c);
@@ -1597,7 +1603,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         uuid, range, range.getEndLine(), otherUser, null, now, "comment on ps2",
         side, rev2, Status.DRAFT);
     update.setPatchSetId(ps2);
-    update.upsertComment(comment2);
+    update.putComment(comment2);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
@@ -1630,7 +1636,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         psId, "filename", uuid, null, 0, otherUser, null, now, messageForBase,
         (short) 0, rev);
     update.setPatchSetId(psId);
-    update.upsertComment(comment);
+    update.putComment(comment);
     update.commit();
 
     assertThat(newNotes(c).getComments()).containsExactlyEntriesIn(
@@ -1651,7 +1657,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
         psId, "filename", uuid, null, 1, otherUser, null, now, messageForBase,
         (short) 0, rev);
     update.setPatchSetId(psId);
-    update.upsertComment(comment);
+    update.putComment(comment);
     update.commit();
 
     assertThat(newNotes(c).getComments()).containsExactlyEntriesIn(
@@ -1659,7 +1665,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
   }
 
   @Test
-  public void updateCommentsForMultipleRevisions() throws Exception {
+  public void putCommentsForMultipleRevisions() throws Exception {
     Change c = newChange();
     String uuid = "uuid";
     String rev1 = "abcd1234abcd1234abcd1234abcd1234abcd1234";
@@ -1681,8 +1687,8 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     PatchLineComment comment2 = newComment(ps2, filename,
         uuid, range, range.getEndLine(), otherUser, null, now, "comment on ps2",
         side, rev2, Status.DRAFT);
-    update.upsertComment(comment1);
+    update.putComment(comment1);
-    update.upsertComment(comment2);
+    update.putComment(comment2);
     update.commit();
 
     ChangeNotes notes = newNotes(c);
@@ -1693,8 +1699,8 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     update.setPatchSetId(ps2);
     comment1.setStatus(Status.PUBLISHED);
     comment2.setStatus(Status.PUBLISHED);
-    update.upsertComment(comment1);
+    update.putComment(comment1);
-    update.upsertComment(comment2);
+    update.putComment(comment2);
     update.commit();
 
     notes = newNotes(c);
@@ -1702,6 +1708,44 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     assertThat(notes.getComments()).hasSize(2);
   }
 
+  @Test
+  public void publishSubsetOfCommentsOnRevision() throws Exception {
+    Change c = newChange();
+    RevId rev1 = new RevId("abcd1234abcd1234abcd1234abcd1234abcd1234");
+    CommentRange range = new CommentRange(1, 1, 2, 1);
+    PatchSet.Id ps1 = c.currentPatchSetId();
+    short side = (short) 1;
+
+    ChangeUpdate update = newUpdate(c, otherUser);
+    update.setPatchSetId(ps1);
+    Timestamp now = TimeUtil.nowTs();
+    PatchLineComment comment1 = newComment(ps1, "file1",
+        "uuid1", range, range.getEndLine(), otherUser, null, now, "comment1",
+        side, rev1.get(), Status.DRAFT);
+    PatchLineComment comment2 = newComment(ps1, "file2",
+        "uuid2", range, range.getEndLine(), otherUser, null, now, "comment2",
+        side, rev1.get(), Status.DRAFT);
+    update.putComment(comment1);
+    update.putComment(comment2);
+    update.commit();
+
+    ChangeNotes notes = newNotes(c);
+    assertThat(notes.getDraftComments(otherUserId).get(rev1))
+        .containsExactly(comment1, comment2);
+    assertThat(notes.getComments()).isEmpty();
+
+    update = newUpdate(c, otherUser);
+    update.setPatchSetId(ps1);
+    comment2.setStatus(Status.PUBLISHED);
+    update.putComment(comment2);
+    update.commit();
+
+    notes = newNotes(c);
+    assertThat(notes.getDraftComments(otherUserId).get(rev1))
+        .containsExactly(comment1);
+    assertThat(notes.getComments().get(rev1)).containsExactly(comment2);
+  }
+
   @Test
   public void updateWithServerIdent() throws Exception {
     Change c = newChange();
@@ -1718,9 +1762,84 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     update.putApproval("Code-Review", (short) 1);
   }
 
+  @Test
+  public void filterOutAndFixUpZombieDraftComments() throws Exception {
+    Change c = newChange();
+    RevId rev1 = new RevId("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();
+    PatchLineComment comment1 = newComment(ps1, "file1",
+        "uuid1", range, range.getEndLine(), otherUser, null, now, "comment on ps1",
+        side, rev1.get(), Status.DRAFT);
+    PatchLineComment comment2 = newComment(ps1, "file2",
+        "uuid2", range, range.getEndLine(), otherUser, null, now, "another comment",
+        side, rev1.get(), Status.DRAFT);
+    update.putComment(comment1);
+    update.putComment(comment2);
+    update.commit();
+
+
+    String refName = RefNames.refsDraftComments(otherUserId, c.getId());
+    ObjectId oldDraftId = exactRefAllUsers(refName);
+
+    update = newUpdate(c, otherUser);
+    update.setPatchSetId(ps1);
+    comment2.setStatus(Status.PUBLISHED);
+    update.putComment(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();
+    comment2.setStatus(Status.DRAFT);
+    draftUpdate.putComment(comment2);
+    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(rev1))
+        .containsExactly(comment1, comment2);
+
+    comment2.setStatus(Status.PUBLISHED); // Reset for later assertions.
+
+    // Zombie comment is filtered out of drafts via ChangeNotes.
+    ChangeNotes notes = newNotes(c);
+    assertThat(notes.getDraftComments(otherUserId).get(rev1))
+        .containsExactly(comment1);
+    assertThat(notes.getComments().get(rev1))
+        .containsExactly(comment2);
+
+    update = newUpdate(c, otherUser);
+    update.setPatchSetId(ps1);
+    comment1.setStatus(Status.PUBLISHED);
+    update.putComment(comment1);
+    update.commit();
+
+    // Updating an unrelated comment causes the zombie comment to get fixed up.
+    assertThat(exactRefAllUsers(refName)).isNull();
+  }
+
   private String readNote(ChangeNotes notes, ObjectId noteId) throws Exception {
-    ObjectId dataId = notes.getNoteMap().getNote(noteId).getData();
+    ObjectId dataId = notes.revisionNoteMap.noteMap.getNote(noteId).getData();
     return new String(
         rw.getObjectReader().open(dataId, OBJ_BLOB).getCachedBytes(), UTF_8);
   }
+
+  private ObjectId exactRefAllUsers(String refName) throws Exception {
+    try (Repository allUsersRepo = repoManager.openRepository(allUsers)) {
+      Ref ref = allUsersRepo.exactRef(refName);
+      return ref != null ? ref.getObjectId() : null;
+    }
+  }
 }