Store draft PatchLineComments in Git notes

The draft comments for a user will be stored in their own unique ref
in the All-Users repository. The ref will look like
'refs/users/34/1234/comments-5678' where the user's accountId is 1234
and the change's changeId is 5678. Since these commits will live in a
different repository than the change at a different ref, I created two
new classes (ChangeDraftUpdate and DraftCommentNotes) to mimic the
workflow of ChangeNotes and ChangeUpdate.

The basic workflow is that when a draft comment is added, it will be
written to a note attached to the user's ref for the given change in
the All-Users repository. However, whenever a draft is published or
deleted, then it should be removed from that location and instead
written to the same place as the other published comments (in the
notes attached to the commits in the meta ref related to the given
change).

Additionally, I added a helper method to TestChanges to help construct
ChangeDraftUpdates for test. Finally, I added test coverage for
storing draft PatchLineComments in the notedb to ChangeNotesTest.

Change-Id: I178b039529937c2a65a0efed388ee6ca1043aaa7

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

@@ -77,6 +77,15 @@ public class PatchLineCommentsUtil {
     return commentsOnPs;
   }
 
+  // TODO(yyonas): Delete drafts if they already existed.
+  public void addPublishedComments(ReviewDb db, ChangeUpdate update,
+      Iterable<PatchLineComment> comments) throws OrmException {
+    for (PatchLineComment c : comments) {
+      update.upsertComment(c);
+    }
+    db.patchComments().upsert(comments);
+  }
+
   private static Collection<PatchLineComment> addCommentsInFile(
       Collection<PatchLineComment> commentsOnFile,
       Collection<PatchLineComment> allComments,
@@ -89,12 +98,4 @@ public class PatchLineCommentsUtil {
     }
     return commentsOnFile;
   }
-
-  public void addPublishedComments(ReviewDb db, ChangeUpdate update,
-      Iterable<PatchLineComment> comments) throws OrmException {
-    for (PatchLineComment c : comments) {
-      update.putComment(c);
-    }
-    db.patchComments().upsert(comments);
-  }
 }

gerrit-server/src/main/java/com/google/gerrit/server/git/VersionedMetaData.java

@@ -175,6 +175,7 @@ public abstract class VersionedMetaData {
     void write(CommitBuilder commit) throws IOException;
     void write(VersionedMetaData config, CommitBuilder commit) throws IOException;
     RevCommit createRef(String refName) throws IOException;
+    void removeRef(String refName) throws IOException;
     RevCommit commit() throws IOException;
     RevCommit commitAt(ObjectId revision) throws IOException;
     void close();
@@ -267,6 +268,24 @@ public abstract class VersionedMetaData {
         }
       }
 
+      @Override
+      public void removeRef(String refName) throws IOException {
+        RefUpdate ru = db.updateRef(refName);
+        ru.setForceUpdate(true);
+        if (revision != null) {
+          ru.setExpectedOldObjectId(revision);
+        }
+        RefUpdate.Result result = ru.delete();
+        switch (result) {
+          case FORCED:
+            update.fireGitRefUpdatedEvent(ru);
+            return;
+          default:
+            throw new IOException("Cannot delete " + ru.getName() + " in "
+                + db.getDirectory() + ": " + ru.getResult());
+        }
+      }
+
       @Override
       public RevCommit commit() throws IOException {
         return commitAt(revision);

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

@@ -59,6 +59,10 @@ public abstract class AbstractChangeUpdate extends VersionedMetaData {
     this.when = when;
   }
 
+  public ChangeNotes getChangeNotes() {
+    return ctl.getNotes();
+  }
+
   public Change getChange() {
     return ctl.getChange();
   }
@@ -119,6 +123,11 @@ public abstract class AbstractChangeUpdate extends VersionedMetaData {
         return null;
       }
 
+      @Override
+      public void removeRef(String refName) {
+        // Do nothing.
+      }
+
       @Override
       public RevCommit commit() {
         return null;

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

@@ -0,0 +1,286 @@
+// 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 static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
+import static com.google.gerrit.server.notedb.CommentsInNotesUtil.getCommentPsId;
+
+import com.google.common.collect.Lists;
+import com.google.common.collect.Table;
+import com.google.gerrit.reviewdb.client.Account;
+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.Project;
+import com.google.gerrit.reviewdb.client.RefNames;
+import com.google.gerrit.reviewdb.client.RevId;
+import com.google.gerrit.server.GerritPersonIdent;
+import com.google.gerrit.server.IdentifiedUser;
+import com.google.gerrit.server.config.AllUsersName;
+import com.google.gerrit.server.git.GitRepositoryManager;
+import com.google.gerrit.server.git.MetaDataUpdate;
+import com.google.gerrit.server.project.ChangeControl;
+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.PersonIdent;
+import org.eclipse.jgit.notes.NoteMap;
+import org.eclipse.jgit.revwalk.RevCommit;
+
+import java.io.IOException;
+import java.util.Date;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * A single delta to apply atomically to a change.
+ * <p>
+ * This delta contains only draft comments on a single patch set of a change by
+ * a single author. This delta will become a single commit in the All-Users
+ * repository.
+ * <p>
+ * This class is not thread safe.
+ */
+public class ChangeDraftUpdate extends AbstractChangeUpdate {
+  public interface Factory {
+    ChangeDraftUpdate create(ChangeControl ctl, Date when);
+  }
+
+  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;
+  private List<PatchLineComment> deleteComments;
+
+  @AssistedInject
+  private ChangeDraftUpdate(
+      @GerritPersonIdent PersonIdent serverIdent,
+      GitRepositoryManager repoManager,
+      NotesMigration migration,
+      MetaDataUpdate.User updateFactory,
+      DraftCommentNotes.Factory draftNotesFactory,
+      AllUsersName allUsers,
+      CommentsInNotesUtil commentsUtil,
+      @Assisted ChangeControl ctl,
+      @Assisted Date when) throws OrmException {
+    super(migration, repoManager, updateFactory, ctl, serverIdent, when);
+    this.draftsProject = allUsers;
+    this.commentsUtil = commentsUtil;
+    checkState(ctl.getCurrentUser().isIdentifiedUser(),
+        "Current user must be identified");
+    IdentifiedUser user = (IdentifiedUser) ctl.getCurrentUser();
+    this.accountId = user.getAccountId();
+    this.changeNotes = getChangeNotes().load();
+    this.draftNotes = draftNotesFactory.create(ctl.getChange(),
+        accountId).load();
+
+    this.upsertComments = Lists.newArrayList();
+    this.deleteComments = Lists.newArrayList();
+  }
+
+  public void insertComment(PatchLineComment c) throws OrmException {
+    verifyComment(c);
+    checkArgument(c.getStatus() == Status.DRAFT,
+        "Cannot insert a published comment into a ChangeDraftUpdate");
+    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) {
+    verifyComment(c);
+    checkArgument(c.getStatus() == Status.DRAFT,
+        "Cannot upsert a published comment into a ChangeDraftUpdate");
+    upsertComments.add(c);
+  }
+
+  public void updateComment(PatchLineComment c) throws OrmException {
+    verifyComment(c);
+    checkArgument(c.getStatus() == Status.DRAFT,
+        "Cannot update a published comment into a ChangeDraftUpdate");
+    checkArgument(draftNotes.containsComment(c),
+        "Cannot update this comment because it didn't exist previously");
+    upsertComments.add(c);
+  }
+
+  public void deleteComment(PatchLineComment c) {
+    verifyComment(c);
+    checkArgument(draftNotes.containsComment(c), "Cannot delete this comment"
+        + " because it didn't previously exist as a draft");
+    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) {
+    if (draftNotes.containsComment(c)) {
+      verifyComment(c);
+      deleteComments.add(c);
+    }
+  }
+
+  private void verifyComment(PatchLineComment comment) {
+    checkState(psId != null,
+        "setPatchSetId must be called first");
+    checkArgument(getCommentPsId(comment).equals(psId),
+        "Comment on %s does not match configured patch set %s",
+        getCommentPsId(comment), psId);
+    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(AtomicBoolean removedAllComments)
+      throws OrmException, IOException {
+    if (isEmpty()) {
+      return null;
+    }
+
+    NoteMap noteMap = draftNotes.getNoteMap();
+    if (noteMap == null) {
+      noteMap = NoteMap.newEmptyMap();
+    }
+
+    Table<PatchSet.Id, String, PatchLineComment> baseDrafts =
+        draftNotes.getDraftBaseComments();
+    Table<PatchSet.Id, String, PatchLineComment> psDrafts =
+        draftNotes.getDraftPsComments();
+
+    // There is no need to rewrite the note for one of the sides of the patch
+    // set if all of the modifications were made to the comments of one side,
+    // so we set these flags to potentially save that extra work.
+    boolean baseSideChanged = false;
+    boolean revisionSideChanged = false;
+
+    // We must define these RevIds so that if this update deletes all
+    // remaining comments on a given side, then we can remove that note.
+    // However, if this update doesn't delete any comments, it is okay for these
+    // to be null because they won't be used.
+    RevId baseRevId = null;
+    RevId psRevId = null;
+
+    for (PatchLineComment c : deleteComments) {
+      if (c.getSide() == (short) 0) {
+        baseSideChanged = true;
+        baseRevId = c.getRevId();
+        baseDrafts.remove(psId, c.getKey().get());
+      } else {
+        revisionSideChanged = true;
+        psRevId = c.getRevId();
+        psDrafts.remove(psId, c.getKey().get());
+      }
+    }
+
+    for (PatchLineComment c : upsertComments) {
+      if (c.getSide() == (short) 0) {
+        baseSideChanged = true;
+        baseDrafts.put(psId, c.getKey().get(), c);
+      } else {
+        revisionSideChanged = true;
+        psDrafts.put(psId, c.getKey().get(), c);
+      }
+    }
+
+    List<PatchLineComment> newBaseDrafts =
+        Lists.newArrayList(baseDrafts.row(psId).values());
+    List<PatchLineComment> newPsDrafts =
+        Lists.newArrayList(psDrafts.row(psId).values());
+
+    updateNoteMap(baseSideChanged, noteMap, newBaseDrafts,
+        baseRevId);
+    updateNoteMap(revisionSideChanged, noteMap, newPsDrafts,
+        psRevId);
+
+    removedAllComments.set(baseDrafts.isEmpty() && psDrafts.isEmpty());
+
+    return noteMap.writeTree(inserter);
+  }
+
+  private void updateNoteMap(boolean changed, NoteMap noteMap,
+      List<PatchLineComment> comments, RevId commitId)
+      throws IOException, OrmException {
+    if (changed) {
+      if (comments.isEmpty()) {
+        commentsUtil.removeNote(noteMap, commitId);
+      } else {
+        commentsUtil.writeCommentsToNoteMap(noteMap, comments, inserter);
+      }
+    }
+  }
+
+  public RevCommit commit() throws IOException {
+    BatchMetaDataUpdate batch = openUpdate();
+    try {
+      CommitBuilder builder = new CommitBuilder();
+      if (migration.write()) {
+        AtomicBoolean removedAllComments = new AtomicBoolean();
+        ObjectId treeId = storeCommentsInNotes(removedAllComments);
+        if (treeId != null) {
+          if (removedAllComments.get()) {
+            batch.removeRef(getRefName());
+          } else {
+            builder.setTreeId(treeId);
+            batch.write(builder);
+          }
+        }
+      }
+      return batch.commit();
+    } catch (OrmException e) {
+      throw new IOException(e);
+    } finally {
+      batch.close();
+    }
+  }
+
+  @Override
+  protected Project.NameKey getProjectName() {
+    return draftsProject;
+  }
+
+  @Override
+  protected String getRefName() {
+    return RefNames.refsDraftComments(accountId, getChange().getId());
+  }
+
+  @Override
+  protected boolean onSave(CommitBuilder commit) throws IOException,
+      ConfigInvalidException {
+    if (isEmpty()) {
+      return false;
+    }
+    commit.setAuthor(newIdent(getUser().getAccount(), when));
+    commit.setCommitter(new PersonIdent(serverIdent, when));
+    commit.setMessage(String.format("Comment on patch set %d", psId.get()));
+    return true;
+  }
+
+  private boolean isEmpty() {
+    return deleteComments.isEmpty()
+        && upsertComments.isEmpty();
+  }
+}

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

@@ -19,6 +19,7 @@ import static com.google.gerrit.server.notedb.ChangeNoteUtil.FOOTER_PATCH_SET;
 import static com.google.gerrit.server.notedb.ChangeNoteUtil.FOOTER_STATUS;
 import static com.google.gerrit.server.notedb.ChangeNoteUtil.FOOTER_SUBMITTED_WITH;
 import static com.google.gerrit.server.notedb.ChangeNoteUtil.GERRIT_PLACEHOLDER_HOST;
+import static com.google.gerrit.server.notedb.CommentsInNotesUtil.getCommentPsId;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Enums;
@@ -49,8 +50,11 @@ import com.google.gerrit.reviewdb.client.PatchSet.Id;
 import com.google.gerrit.reviewdb.client.PatchSetApproval;
 import com.google.gerrit.reviewdb.client.PatchSetApproval.LabelId;
 import com.google.gerrit.reviewdb.client.Project;
+import com.google.gerrit.server.config.AllUsersName;
+import com.google.gerrit.server.config.AllUsersNameProvider;
 import com.google.gerrit.server.git.GitRepositoryManager;
 import com.google.gerrit.server.util.LabelVote;
+import com.google.gwtorm.server.OrmException;
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
 
@@ -134,15 +138,17 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
   @Singleton
   public static class Factory {
     private final GitRepositoryManager repoManager;
+    private final AllUsersNameProvider allUsersProvider;
 
     @VisibleForTesting
     @Inject
-    public Factory(GitRepositoryManager repoManager) {
+    public Factory(GitRepositoryManager repoManager, AllUsersNameProvider allUsersProvider) {
       this.repoManager = repoManager;
+      this.allUsersProvider = allUsersProvider;
     }
 
     public ChangeNotes create(Change change) {
-      return new ChangeNotes(repoManager, change);
+      return new ChangeNotes(repoManager, allUsersProvider, change);
     }
   }
 
@@ -492,9 +498,15 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
   private ImmutableListMultimap<PatchSet.Id, PatchLineComment> commentsForPS;
   NoteMap noteMap;
 
+  private final AllUsersName allUsers;
+  private DraftCommentNotes draftCommentNotes;
+
+  @Inject
   @VisibleForTesting
-  public ChangeNotes(GitRepositoryManager repoManager, Change change) {
+  public ChangeNotes(GitRepositoryManager repoManager,
+      AllUsersNameProvider allUsersProvider, Change change) {
     super(repoManager, change);
+    this.allUsers = allUsersProvider.get();
   }
 
   public ImmutableListMultimap<PatchSet.Id, PatchSetApproval> getApprovals() {
@@ -530,6 +542,55 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
     return commentsForPS;
   }
 
+  public Table<PatchSet.Id, String, PatchLineComment> getDraftBaseComments(
+      Account.Id author) throws OrmException {
+    loadDraftComments(author);
+    return draftCommentNotes.getDraftBaseComments();
+  }
+
+  public Table<PatchSet.Id, String, PatchLineComment> getDraftPsComments(
+      Account.Id author) throws OrmException {
+    loadDraftComments(author);
+    return draftCommentNotes.getDraftPsComments();
+  }
+
+  /**
+   * If draft comments have already been loaded for this author, then they will
+   * not be reloaded. However, this method will load the comments if no draft
+   * comments have been loaded or if the caller would like the drafts for
+   * another author.
+   */
+  private void loadDraftComments(Account.Id author)
+      throws OrmException {
+    if (draftCommentNotes == null ||
+        !author.equals(draftCommentNotes.getAuthor())) {
+      draftCommentNotes = new DraftCommentNotes(repoManager, allUsers,
+          getChange(), author);
+      draftCommentNotes.load();
+    }
+  }
+
+  public boolean containsComment(PatchLineComment c) throws OrmException {
+    if (containsCommentPublished(c)) {
+      return true;
+    }
+    loadDraftComments(c.getAuthor());
+    return draftCommentNotes.containsComment(c);
+  }
+
+  public boolean containsCommentPublished(PatchLineComment c) {
+    PatchSet.Id psId = getCommentPsId(c);
+    List<PatchLineComment> list = (c.getSide() == (short) 0)
+        ? getBaseComments().get(psId)
+        : getPatchSetComments().get(psId);
+    for (PatchLineComment l : list) {
+      if (c.getKey().equals(l.getKey())) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   /** @return the NoteMap */
   NoteMap getNoteMap() {
     return noteMap;

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

@@ -33,9 +33,12 @@ import com.google.gerrit.reviewdb.client.Change;
 import com.google.gerrit.reviewdb.client.PatchLineComment;
 import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.client.Project;
+import com.google.gerrit.reviewdb.client.PatchLineComment.Status;
 import com.google.gerrit.server.GerritPersonIdent;
 import com.google.gerrit.server.IdentifiedUser;
 import com.google.gerrit.server.account.AccountCache;
+import com.google.gerrit.server.config.AllUsersName;
+import com.google.inject.Provider;
 import com.google.gerrit.server.git.GitRepositoryManager;
 import com.google.gerrit.server.git.MetaDataUpdate;
 import com.google.gerrit.server.project.ChangeControl;
@@ -60,12 +63,14 @@ import java.util.List;
 import java.util.Map;
 
 /**
- * A single delta to apply atomically to a change.
+ * A delta to apply to a change.
  * <p>
- * This delta becomes a single commit on the notes branch, so there are
- * limitations on the set of modifications that can be handled in a single
- * update. In particular, there is a single author and timestamp for each
- * update.
+ * This delta will become two unique commits: one in the AllUsers repo that will
+ * contain the draft comments on this change and one in the notes branch that
+ * will contain approvals, reviewers, change status, subject, submit records,
+ * the change message, and published comments. There are limitations on the set
+ * of modifications that can be handled in a single update. In particular, there
+ * is a single author and timestamp for each update.
  * <p>
  * This class is not thread-safe.
  */
@@ -88,6 +93,10 @@ public class ChangeUpdate extends AbstractChangeUpdate {
   private List<PatchLineComment> commentsForBase;
   private List<PatchLineComment> commentsForPs;
   private String changeMessage;
+  private ChangeNotes notes;
+
+  private final ChangeDraftUpdate.Factory draftUpdateFactory;
+  private ChangeDraftUpdate draftUpdate;
 
   @AssistedInject
   private ChangeUpdate(
@@ -96,12 +105,16 @@ public class ChangeUpdate extends AbstractChangeUpdate {
       NotesMigration migration,
       AccountCache accountCache,
       MetaDataUpdate.User updateFactory,
+      DraftCommentNotes.Factory draftNotesFactory,
+      Provider<AllUsersName> allUsers,
+      ChangeDraftUpdate.Factory draftUpdateFactory,
       ProjectCache projectCache,
       IdentifiedUser user,
       @Assisted ChangeControl ctl,
       CommentsInNotesUtil commentsUtil) {
     this(serverIdent, repoManager, migration, accountCache, updateFactory,
-        projectCache, ctl, serverIdent.getWhen(), commentsUtil);
+        draftNotesFactory, allUsers, draftUpdateFactory, projectCache, ctl,
+        serverIdent.getWhen(), commentsUtil);
   }
 
   @AssistedInject
@@ -111,12 +124,15 @@ public class ChangeUpdate extends AbstractChangeUpdate {
       NotesMigration migration,
       AccountCache accountCache,
       MetaDataUpdate.User updateFactory,
+      DraftCommentNotes.Factory draftNotesFactory,
+      Provider<AllUsersName> allUsers,
+      ChangeDraftUpdate.Factory draftUpdateFactory,
       ProjectCache projectCache,
       @Assisted ChangeControl ctl,
       @Assisted Date when,
       CommentsInNotesUtil commentsUtil) {
-    this(serverIdent, repoManager, migration, accountCache, updateFactory, ctl,
-        when,
+    this(serverIdent, repoManager, migration, accountCache, updateFactory,
+        draftNotesFactory, allUsers, draftUpdateFactory, ctl, when,
         projectCache.get(getProjectName(ctl)).getLabelTypes().nameComparator(),
         commentsUtil);
   }
@@ -132,11 +148,15 @@ public class ChangeUpdate extends AbstractChangeUpdate {
       NotesMigration migration,
       AccountCache accountCache,
       MetaDataUpdate.User updateFactory,
+      DraftCommentNotes.Factory draftNotesFactory,
+      Provider<AllUsersName> allUsers,
+      ChangeDraftUpdate.Factory draftUpdateFactory,
       @Assisted ChangeControl ctl,
       @Assisted Date when,
       @Assisted Comparator<String> labelNameComparator,
       CommentsInNotesUtil commentsUtil) {
     super(migration, repoManager, updateFactory, ctl, serverIdent, when);
+    this.draftUpdateFactory = draftUpdateFactory;
     this.accountCache = accountCache;
     this.commentsUtil = commentsUtil;
     this.approvals = Maps.newTreeMap(labelNameComparator);
@@ -174,20 +194,132 @@ public class ChangeUpdate extends AbstractChangeUpdate {
     this.changeMessage = changeMessage;
   }
 
-  public void putComment(PatchLineComment comment) {
-    checkArgument(psId != null,
-        "setPatchSetId must be called before putComment");
-    checkArgument(getCommentPsId(comment).equals(psId),
-        "Comment on %s doesn't match previous patch set %s",
-        getCommentPsId(comment), psId);
-    checkArgument(comment.getRevId() != null);
-    if (comment.getSide() == 0) {
-      commentsForBase.add(comment);
+  public void insertComment(PatchLineComment comment) throws OrmException {
+    if (comment.getStatus() == Status.DRAFT) {
+      insertDraftComment(comment);
     } else {
-      commentsForPs.add(comment);
+      insertPublishedComment(comment);
     }
   }
 
+  public void upsertComment(PatchLineComment comment) throws OrmException {
+    if (comment.getStatus() == Status.DRAFT) {
+      upsertDraftComment(comment);
+    } else {
+      deleteDraftCommentIfPresent(comment);
+      upsertPublishedComment(comment);
+    }
+  }
+
+  public void updateComment(PatchLineComment comment) throws OrmException {
+    if (comment.getStatus() == Status.DRAFT) {
+      updateDraftComment(comment);
+    } else {
+      deleteDraftCommentIfPresent(comment);
+      updatePublishedComment(comment);
+    }
+  }
+
+  public void deleteComment(PatchLineComment comment) throws OrmException {
+    if (comment.getStatus() == 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();
+    }
+    checkArgument(!notes.containsComment(c),
+        "A comment already exists with the same key as the following comment,"
+        + " so we cannot insert this comment: %s", c);
+    if (c.getSide() == 0) {
+      commentsForBase.add(c);
+    } else {
+      commentsForPs.add(c);
+    }
+  }
+
+  private void insertDraftComment(PatchLineComment c) throws OrmException {
+    createDraftUpdateIfNull(c);
+    draftUpdate.insertComment(c);
+  }
+
+  private void upsertPublishedComment(PatchLineComment c) {
+    verifyComment(c);
+    if (c.getSide() == 0) {
+      commentsForBase.add(c);
+    } else {
+      commentsForPs.add(c);
+    }
+  }
+
+  private void upsertDraftComment(PatchLineComment c) throws OrmException {
+    createDraftUpdateIfNull(c);
+    draftUpdate.upsertComment(c);
+  }
+
+  private void updatePublishedComment(PatchLineComment c) throws OrmException {
+    verifyComment(c);
+    if (notes == null) {
+      notes = getChangeNotes().load();
+    }
+    checkArgument(!notes.containsCommentPublished(c),
+        "Cannot update a comment that has already been published and saved");
+    if (c.getSide() == 0) {
+      commentsForBase.add(c);
+    } else {
+      commentsForPs.add(c);
+    }
+  }
+
+  private void updateDraftComment(PatchLineComment c) throws OrmException {
+    createDraftUpdateIfNull(c);
+    draftUpdate.updateComment(c);
+  }
+
+  private void deleteDraftComment(PatchLineComment c) throws OrmException {
+    createDraftUpdateIfNull(c);
+    draftUpdate.deleteComment(c);
+  }
+
+  private void deleteDraftCommentIfPresent(PatchLineComment c)
+      throws OrmException {
+    createDraftUpdateIfNull(c);
+    draftUpdate.deleteCommentIfPresent(c);
+  }
+
+
+  private void createDraftUpdateIfNull(PatchLineComment c) throws OrmException {
+    if (draftUpdate == null) {
+      draftUpdate = draftUpdateFactory.create(ctl, when);
+      if (psId != null) {
+        draftUpdate.setPatchSetId(psId);
+      } else {
+        draftUpdate.setPatchSetId(getCommentPsId(c));
+      }
+    }
+  }
+
+  private void verifyComment(PatchLineComment c) {
+    checkArgument(psId != null,
+        "setPatchSetId must be called first");
+    checkArgument(getCommentPsId(c).equals(psId),
+        "Comment on %s doesn't match previous patch set %s",
+        getCommentPsId(c), psId);
+    checkArgument(c.getRevId() != null);
+    checkArgument(c.getStatus() == 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);
+
+  }
+
   public void putReviewer(Account.Id reviewer, ReviewerState type) {
     checkArgument(type != ReviewerState.REMOVED, "invalid ReviewerType");
     reviewers.put(reviewer, type);
@@ -243,6 +375,9 @@ public class ChangeUpdate extends AbstractChangeUpdate {
         }
       }
       batch.write(builder);
+      if (draftUpdate != null) {
+        draftUpdate.commit();
+      }
       RevCommit c = batch.commit();
       return c;
     } catch (OrmException e) {

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

@@ -0,0 +1,199 @@
+// 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 static com.google.gerrit.server.notedb.CommentsInNotesUtil.getCommentPsId;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.HashBasedTable;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Table;
+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.PatchLineComment.Status;
+import com.google.gerrit.reviewdb.client.PatchSet;
+import com.google.gerrit.reviewdb.client.Project;
+import com.google.gerrit.server.config.AllUsersName;
+import com.google.gerrit.server.config.AllUsersNameProvider;
+import com.google.gerrit.server.git.GitRepositoryManager;
+import com.google.inject.Inject;
+import com.google.inject.Singleton;
+
+import org.eclipse.jgit.errors.ConfigInvalidException;
+import org.eclipse.jgit.errors.RepositoryNotFoundException;
+import org.eclipse.jgit.lib.CommitBuilder;
+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;
+
+/**
+ * View of the draft comments for a single {@link Change} based on the log of
+ * its drafts branch.
+ */
+public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
+  @Singleton
+  public static class Factory {
+    private final GitRepositoryManager repoManager;
+    private final AllUsersName draftsProject;
+
+    @VisibleForTesting
+    @Inject
+    public Factory(GitRepositoryManager repoManager,
+        AllUsersNameProvider allUsers) {
+      this.repoManager = repoManager;
+      this.draftsProject = allUsers.get();
+    }
+
+    public DraftCommentNotes create(Change change, Account.Id accountId) {
+      return new DraftCommentNotes(repoManager, draftsProject, change,
+          accountId);
+    }
+  }
+
+  private static class Parser {
+    private final Change.Id changeId;
+    private final ObjectId tip;
+    private final RevWalk walk;
+    private final Repository repo;
+    private final Account.Id author;
+
+    private final Multimap<PatchSet.Id, PatchLineComment> draftBaseComments;
+    private final Multimap<PatchSet.Id, PatchLineComment> draftPsComments;
+    private NoteMap noteMap;
+
+    private Parser(Change change, RevWalk walk, ObjectId tip,
+        GitRepositoryManager repoManager, AllUsersName draftsProject,
+        Account.Id author) throws RepositoryNotFoundException, IOException {
+      this.changeId = change.getId();
+      this.walk = walk;
+      this.tip = tip;
+      this.repo = repoManager.openRepository(draftsProject);
+      this.author = author;
+
+      draftBaseComments = ArrayListMultimap.create();
+      draftPsComments = ArrayListMultimap.create();
+    }
+
+    private void parseDraftComments() throws IOException, ConfigInvalidException {
+      walk.markStart(walk.parseCommit(tip));
+      noteMap = CommentsInNotesUtil.parseCommentsFromNotes(repo,
+          RefNames.refsDraftComments(author, changeId),
+          walk, changeId, draftBaseComments,
+          draftPsComments, Status.DRAFT);
+    }
+  }
+
+  private final AllUsersName draftsProject;
+  private final Account.Id author;
+
+  private final Table<PatchSet.Id, String, PatchLineComment> draftBaseComments;
+  private final Table<PatchSet.Id, String, PatchLineComment> draftPsComments;
+  private NoteMap noteMap;
+
+  DraftCommentNotes(GitRepositoryManager repoManager,
+      AllUsersName draftsProject, Change change, Account.Id author) {
+    super(repoManager, change);
+    this.draftsProject = draftsProject;
+    this.author = author;
+
+    this.draftBaseComments = HashBasedTable.create();
+    this.draftPsComments = HashBasedTable.create();
+  }
+
+  public NoteMap getNoteMap() {
+    return noteMap;
+  }
+
+  public Account.Id getAuthor() {
+    return author;
+  }
+
+  /**
+   * @return a defensive copy of the table containing all draft comments
+   *    on this change with side == 0. The row key is the comment's PatchSet.Id,
+   *    the column key is the comment's UUID, and the value is the comment.
+   */
+  public Table<PatchSet.Id, String, PatchLineComment>
+      getDraftBaseComments() {
+    return HashBasedTable.create(draftBaseComments);
+  }
+
+  /**
+   * @return a defensive copy of the table containing all draft comments
+   *    on this change with side == 1. The row key is the comment's PatchSet.Id,
+   *    the column key is the comment's UUID, and the value is the comment.
+   */
+  public Table<PatchSet.Id, String, PatchLineComment>
+      getDraftPsComments() {
+    return HashBasedTable.create(draftPsComments);
+  }
+
+  public boolean containsComment(PatchLineComment c) {
+    Table<PatchSet.Id, String, PatchLineComment> t =
+        c.getSide() == (short) 0
+        ? getDraftBaseComments()
+        : getDraftPsComments();
+    return t.contains(getCommentPsId(c), c.getKey().get());
+  }
+
+  @Override
+  protected String getRefName() {
+    return RefNames.refsDraftComments(author, getChangeId());
+  }
+
+  @Override
+  protected void onLoad() throws IOException, ConfigInvalidException {
+    ObjectId rev = getRevision();
+    if (rev == null) {
+      return;
+    }
+
+    RevWalk walk = new RevWalk(reader);
+    Parser parser = new Parser(getChange(), walk, rev, repoManager,
+        draftsProject, author);
+    parser.parseDraftComments();
+
+    buildCommentTable(draftBaseComments, parser.draftBaseComments);
+    buildCommentTable(draftPsComments, parser.draftPsComments);
+    noteMap = parser.noteMap;
+  }
+
+  @Override
+  protected boolean onSave(CommitBuilder commit) throws IOException,
+      ConfigInvalidException {
+    throw new UnsupportedOperationException(
+        getClass().getSimpleName() + " is read-only");
+  }
+
+  @Override
+  protected Project.NameKey getProjectName() {
+    return draftsProject;
+  }
+
+  private void buildCommentTable(
+      Table<PatchSet.Id, String, PatchLineComment> commentTable,
+      Multimap<PatchSet.Id, PatchLineComment> allComments) {
+    for (PatchLineComment c : allComments.values()) {
+      commentTable.put(getCommentPsId(c), c.getKey().get(), c);
+    }
+  }
+
+}

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

@@ -20,5 +20,6 @@ public class NoteDbModule extends FactoryModule {
   @Override
   public void configure() {
     factory(ChangeUpdate.Factory.class);
+    factory(ChangeDraftUpdate.Factory.class);
   }
 }