Store NoteDb migration state inside Change

Since a ReviewDb write might succeed but its corresponding NoteDb
write might fail, we would like to be able to detect and fix up such
problems after the fact. The most reliable way to do so is to store,
somewhere, the complete set of ref state across the change meta repo
and All-Users for a given Change entity. This effectively needs to
be in Change, or at least in ReviewDb, so that we can write it in the
same transaction that stores the rest of the Change mutation.

Split up NoteDbUpdateManager into two phases, "stage" and "execute".
Stage the change before committing the ReviewDb transaction, and use
the results of staging to update the noteDbState field in Change.

The state is abstracted as an immutable data type NoteDbChangeState,
and NoteDbUpdateManager computes a "Delta" over the type. This
abstraction helps with compile-time safety and testability. However,
for simplicity in the Change record, we just serialize the results
into a simple, parsimonious string.

Change-Id: I2cccda1ccb54062aa1e80ecc82994614ef6d9598

gerrit-acceptance-tests/src/test/java/com/google/gerrit/acceptance/server/notedb/ChangeRebuilderIT.java

@@ -14,27 +14,45 @@
 
 package com.google.gerrit.acceptance.server.notedb;
 
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.TruthJUnit.assume;
 
 import com.google.gerrit.acceptance.AbstractDaemonTest;
 import com.google.gerrit.acceptance.PushOneCommit;
+import com.google.gerrit.extensions.api.changes.DraftInput;
 import com.google.gerrit.extensions.restapi.ResourceNotFoundException;
 import com.google.gerrit.reviewdb.client.Change;
+import com.google.gerrit.reviewdb.client.Project;
+import com.google.gerrit.reviewdb.client.RefNames;
+import com.google.gerrit.reviewdb.server.ReviewDb;
 import com.google.gerrit.server.change.Rebuild;
+import com.google.gerrit.server.config.AllUsersName;
+import com.google.gerrit.server.notedb.ChangeNoteUtil;
+import com.google.gerrit.server.schema.DisabledChangesReviewDbWrapper;
 import com.google.gerrit.testutil.NoteDbChecker;
 import com.google.gerrit.testutil.NoteDbMode;
 import com.google.inject.Inject;
+import com.google.inject.Provider;
 
+import org.eclipse.jgit.lib.ObjectId;
+import org.eclipse.jgit.lib.Ref;
+import org.eclipse.jgit.lib.Repository;
 import org.junit.Before;
 import org.junit.Test;
 
 public class ChangeRebuilderIT extends AbstractDaemonTest {
+  @Inject
+  private AllUsersName allUsers;
+
   @Inject
   private NoteDbChecker checker;
 
   @Inject
   private Rebuild rebuildHandler;
 
+  @Inject
+  private Provider<ReviewDb> dbProvider;
+
   @Before
   public void setUp() {
     assume().that(NoteDbMode.readWrite()).isFalse();
@@ -68,11 +86,13 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
 
     // First write doesn't create the ref, but rebuilding works.
     checker.assertNoChangeRef(project, id);
+    assertThat(db.changes().get(id).getNoteDbState()).isNull();
     checker.rebuildAndCheckChanges(id);
 
     // Now that there is a ref, writes are "turned on" for this change, and
     // NoteDb stays up to date without explicit rebuilding.
     gApi.changes().id(id.get()).topic(name("new-topic"));
+    assertThat(db.changes().get(id).getNoteDbState()).isNotNull();
     checker.checkChanges(id);
   }
 
@@ -117,4 +137,69 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
     checker.assertNoChangeRef(project, id1);
     checker.rebuildAndCheckChanges(id1);
   }
+
+  @Test
+  public void noteDbChangeState() throws Exception {
+    notesMigration.setAllEnabled(true);
+    PushOneCommit.Result r = createChange();
+    Change.Id id = r.getPatchSetId().getParentKey();
+
+    ObjectId changeMetaId = getMetaRef(
+        project, ChangeNoteUtil.changeRefName(id));
+    assertThat(unwrapDb().changes().get(id).getNoteDbState()).isEqualTo(
+        changeMetaId.name());
+
+    DraftInput in = new DraftInput();
+    in.line = 1;
+    in.message = "comment by user";
+    in.path = PushOneCommit.FILE_NAME;
+    setApiUser(user);
+    gApi.changes().id(id.get()).current().createDraft(in);
+
+    ObjectId userDraftsId = getMetaRef(
+        allUsers, RefNames.refsDraftComments(user.getId(), id));
+    assertThat(unwrapDb().changes().get(id).getNoteDbState()).isEqualTo(
+        changeMetaId.name()
+        + "," + user.getId() + "=" + userDraftsId.name());
+
+    in = new DraftInput();
+    in.line = 2;
+    in.message = "comment by admin";
+    in.path = PushOneCommit.FILE_NAME;
+    setApiUser(admin);
+    gApi.changes().id(id.get()).current().createDraft(in);
+
+    ObjectId adminDraftsId = getMetaRef(
+        allUsers, RefNames.refsDraftComments(admin.getId(), id));
+    assertThat(admin.getId().get()).isLessThan(user.getId().get());
+    assertThat(unwrapDb().changes().get(id).getNoteDbState()).isEqualTo(
+        changeMetaId.name()
+        + "," + admin.getId() + "=" + adminDraftsId.name()
+        + "," + user.getId() + "=" + userDraftsId.name());
+
+    in.message = "revised comment by admin";
+    gApi.changes().id(id.get()).current().createDraft(in);
+
+    adminDraftsId = getMetaRef(
+        allUsers, RefNames.refsDraftComments(admin.getId(), id));
+    assertThat(unwrapDb().changes().get(id).getNoteDbState()).isEqualTo(
+        changeMetaId.name()
+        + "," + admin.getId() + "=" + adminDraftsId.name()
+        + "," + user.getId() + "=" + userDraftsId.name());
+  }
+
+  private ObjectId getMetaRef(Project.NameKey p, String name) throws Exception {
+    try (Repository repo = repoManager.openMetadataRepository(p)) {
+      Ref ref = repo.exactRef(name);
+      return ref != null ? ref.getObjectId() : null;
+    }
+  }
+
+  private ReviewDb unwrapDb() {
+    ReviewDb db = dbProvider.get();
+    if (db instanceof DisabledChangesReviewDbWrapper) {
+      db = ((DisabledChangesReviewDbWrapper) db).unsafeGetDelegate();
+    }
+    return db;
+  }
 }

gerrit-reviewdb/src/main/java/com/google/gerrit/reviewdb/client/Change.java

@@ -470,6 +470,10 @@ public final class Change {
   @Column(id = 18, notNull = false)
   protected String submissionId;
 
+  /** @see com.google.gerrit.server.notedb.NoteDbChangeState */
+  @Column(id = 101, notNull = false, length = Integer.MAX_VALUE)
+  protected String noteDbState;
+
   protected Change() {
   }
 
@@ -498,6 +502,7 @@ public final class Change {
     originalSubject = other.originalSubject;
     submissionId = other.submissionId;
     topic = other.topic;
+    noteDbState = other.noteDbState;
   }
 
   /** Legacy 32 bit integer identity for a change. */
@@ -633,6 +638,14 @@ public final class Change {
     this.topic = topic;
   }
 
+  public String getNoteDbState() {
+    return noteDbState;
+  }
+
+  public void setNoteDbState(String state) {
+    noteDbState = state;
+  }
+
   @Override
   public String toString() {
     return new StringBuilder(getClass().getSimpleName())

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

@@ -40,6 +40,7 @@ import com.google.gerrit.server.index.change.ChangeIndexer;
 import com.google.gerrit.server.notedb.ChangeDelete;
 import com.google.gerrit.server.notedb.ChangeNotes;
 import com.google.gerrit.server.notedb.ChangeUpdate;
+import com.google.gerrit.server.notedb.NoteDbChangeState;
 import com.google.gerrit.server.notedb.NoteDbUpdateManager;
 import com.google.gerrit.server.notedb.NotesMigration;
 import com.google.gerrit.server.project.ChangeControl;
@@ -556,15 +557,30 @@ public class BatchUpdate implements AutoCloseable {
         Change.Id id = e.getKey();
         db.changes().beginTransaction(id);
         ChangeContext ctx;
+        NoteDbUpdateManager updateManager = null;
         boolean dirty = false;
         try {
           ctx = newChangeContext(id);
+          // Call updateChange on each op.
           for (Op op : e.getValue()) {
             dirty |= op.updateChange(ctx);
           }
           if (!dirty) {
             return;
           }
+
+          // Stage the NoteDb update and store its state in the Change.
+          if (!ctx.deleted && notesMigration.writeChanges()) {
+            updateManager = updateManagerFactory.create(ctx.getProject());
+            for (ChangeUpdate u : ctx.updates.values()) {
+              updateManager.add(u);
+            }
+            NoteDbChangeState.applyDelta(
+                ctx.getChange(),
+                updateManager.stage().get(id));
+          }
+
+          // Bump lastUpdatedOn or rowVersion and commit.
           if (newChanges.containsKey(id)) {
             db.changes().insert(bumpLastUpdatedOn(ctx));
           } else if (ctx.saved) {
@@ -579,20 +595,20 @@ public class BatchUpdate implements AutoCloseable {
           db.rollback();
         }
 
-        if (ctx.deleted) {
+        // Execute NoteDb updates after committing ReviewDb updates.
-          if (notesMigration.writeChanges()) {
+        if (notesMigration.writeChanges()) {
+          if (updateManager != null) {
+            updateManager.execute();
+          }
+          if (ctx.deleted) {
             new ChangeDelete(plcUtil, getRepository(), ctx.getNotes()).delete();
           }
+        }
+
+        // Reindex changes.
+        if (ctx.deleted) {
           indexFutures.add(indexer.deleteAsync(id));
         } else {
-          if (notesMigration.writeChanges()) {
-            NoteDbUpdateManager manager =
-                updateManagerFactory.create(ctx.getProject());
-            for (ChangeUpdate u : ctx.updates.values()) {
-              manager.add(u);
-            }
-            manager.execute();
-          }
           indexFutures.add(indexer.indexAsync(ctx.getProject(), id));
         }
       }

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

@@ -24,6 +24,7 @@ import org.eclipse.jgit.lib.Repository;
 import org.eclipse.jgit.transport.ReceiveCommand;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.Map;
@@ -120,4 +121,9 @@ public class ChainedReceiveCommands {
       bru.addCommand(cmd);
     }
   }
+
+  /** @return an unmodifiable view of commands. */
+  public Map<String, ReceiveCommand> getCommands() {
+    return Collections.unmodifiableMap(commands);
+  }
 }

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

@@ -215,7 +215,12 @@ public class ChangeBundle {
     // Initialization-time checks that the column set hasn't changed since the
     // last time this file was updated.
     checkColumns(Change.Id.class, 1);
-    checkColumns(Change.class, 1, 2, 3, 4, 5, 7, 8, 10, 12, 13, 14, 17, 18);
+
+    checkColumns(Change.class,
+        1, 2, 3, 4, 5, 7, 8, 10, 12, 13, 14, 17, 18,
+        // TODO(dborowitz): It's potentially possible to compare noteDbState in
+        // the Change with the state implied by a ChangeNotes.
+        101);
     checkColumns(ChangeMessage.Key.class, 1, 2);
     checkColumns(ChangeMessage.class, 1, 2, 3, 4, 5);
     checkColumns(PatchSet.Id.class, 1, 2);
@@ -286,7 +291,8 @@ public class ChangeBundle {
     Change a = bundleA.change;
     Change b = bundleB.change;
     String desc = a.getId().equals(b.getId()) ? describe(a.getId()) : "Changes";
-    diffColumns(diffs, Change.class, desc, bundleA, a, bundleB, b);
+    diffColumnsExcluding(diffs, Change.class, desc, bundleA, a, bundleB, b,
+        "rowVersion", "noteDbState");
   }
 
   private static void diffChangeMessages(List<String> diffs,

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

@@ -201,6 +201,9 @@ public class ChangeUpdate extends AbstractChangeUpdate {
     NoteDbUpdateManager updateManager =
         updateManagerFactory.create(getProjectName());
     updateManager.add(this);
+    NoteDbChangeState.applyDelta(
+        getChange(),
+        updateManager.stage().get(ctl.getId()));
     updateManager.execute();
     return getResult();
   }

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

@@ -0,0 +1,202 @@
+// 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 static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Optional;
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import com.google.gerrit.reviewdb.client.Account;
+import com.google.gerrit.reviewdb.client.Change;
+import com.google.gerrit.reviewdb.client.RefNames;
+import com.google.gerrit.reviewdb.server.ReviewDbUtil;
+
+import org.eclipse.jgit.lib.ObjectId;
+import org.eclipse.jgit.lib.Ref;
+import org.eclipse.jgit.lib.Repository;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * The state of all relevant NoteDb refs across all repos corresponding to a
+ * given Change entity.
+ * <p>
+ * Stored serialized in the {@code Change#noteDbState} field, and used to
+ * determine whether the state in NoteDb is out of date.
+ * <p>
+ * Serialized in the form:
+ * <pre>
+ *   [meta-sha],[account1]=[drafts-sha],[account2]=[drafts-sha]...
+ * </pre>
+ * in numeric account ID order, with hex SHA-1s for human readability.
+ */
+public class NoteDbChangeState {
+  @AutoValue
+  public abstract static class Delta {
+    static Delta create(Change.Id changeId, Optional<ObjectId> newChangeMetaId,
+        Map<Account.Id, ObjectId> newDraftIds) {
+      if (newDraftIds == null) {
+        newDraftIds = ImmutableMap.of();
+      }
+      return new AutoValue_NoteDbChangeState_Delta(
+          changeId,
+          newChangeMetaId,
+          ImmutableMap.copyOf(newDraftIds));
+    }
+
+    abstract Change.Id changeId();
+    abstract Optional<ObjectId> newChangeMetaId();
+    abstract ImmutableMap<Account.Id, ObjectId> newDraftIds();
+  }
+
+  static NoteDbChangeState parse(Change c) {
+    return parse(c.getId(), c.getNoteDbState());
+  }
+
+  @VisibleForTesting
+  static NoteDbChangeState parse(Change.Id id, String str) {
+    if (str == null) {
+      return null;
+    }
+    List<String> parts = Splitter.on(',').splitToList(str);
+    checkArgument(!parts.isEmpty(),
+        "invalid state string for change %s: %s", id, str);
+    ObjectId changeMetaId = ObjectId.fromString(parts.get(0));
+    Map<Account.Id, ObjectId> draftIds =
+        Maps.newHashMapWithExpectedSize(parts.size() - 1);
+    Splitter s = Splitter.on('=');
+    for (int i = 1; i < parts.size(); i++) {
+      String p = parts.get(i);
+      List<String> draftParts = s.splitToList(p);
+      checkArgument(draftParts.size() == 2,
+          "invalid draft state part for change %s: %s", id, p);
+      draftIds.put(Account.Id.parse(draftParts.get(0)),
+          ObjectId.fromString(draftParts.get(1)));
+    }
+    return new NoteDbChangeState(id, changeMetaId, draftIds);
+  }
+
+  public static void applyDelta(Change change, Delta delta) {
+    if (delta == null) {
+      return;
+    }
+    String oldStr = change.getNoteDbState();
+    if (oldStr == null && !delta.newChangeMetaId().isPresent()) {
+      // Neither an old nor a new meta ID was present, most likely because we
+      // aren't writing a NoteDb graph at all for this change at this point. No
+      // point in proceeding.
+      return;
+    }
+    NoteDbChangeState oldState = parse(change.getId(), oldStr);
+
+    ObjectId changeMetaId;
+    if (delta.newChangeMetaId().isPresent()) {
+      changeMetaId = delta.newChangeMetaId().get();
+      if (changeMetaId.equals(ObjectId.zeroId())) {
+        change.setNoteDbState(null);
+        return;
+      }
+    } else {
+      changeMetaId = oldState.changeMetaId;
+    }
+
+    Map<Account.Id, ObjectId> draftIds = new HashMap<>();
+    if (oldState != null) {
+      draftIds.putAll(oldState.draftIds);
+    }
+    for (Map.Entry<Account.Id, ObjectId> e : delta.newDraftIds().entrySet()) {
+      if (e.getValue().equals(ObjectId.zeroId())) {
+        draftIds.remove(e.getKey());
+      } else {
+        draftIds.put(e.getKey(), e.getValue());
+      }
+    }
+
+    change.setNoteDbState(toString(changeMetaId, draftIds));
+  }
+
+  private static String toString(ObjectId changeMetaId,
+      Map<Account.Id, ObjectId> draftIds) {
+    List<Account.Id> accountIds = Lists.newArrayList(draftIds.keySet());
+    Collections.sort(accountIds, ReviewDbUtil.intKeyOrdering());
+    StringBuilder sb = new StringBuilder(changeMetaId.name());
+    for (Account.Id id : accountIds) {
+      sb.append(',')
+          .append(id.get())
+          .append('=')
+          .append(draftIds.get(id).name());
+    }
+    return sb.toString();
+  }
+
+  private final Change.Id changeId;
+  private final ObjectId changeMetaId;
+  private final ImmutableMap<Account.Id, ObjectId> draftIds;
+
+  NoteDbChangeState(Change.Id changeId, ObjectId changeMetaId,
+      Map<Account.Id, ObjectId> draftIds) {
+    this.changeId = checkNotNull(changeId);
+    this.changeMetaId = checkNotNull(changeMetaId);
+    this.draftIds = ImmutableMap.copyOf(draftIds);
+  }
+
+  public boolean isChangeUpToDate(Repository changeRepo) throws IOException {
+    Ref ref = changeRepo.exactRef(ChangeNoteUtil.changeRefName(changeId));
+    if (ref == null) {
+      return changeMetaId.equals(ObjectId.zeroId());
+    }
+    return ref.getObjectId().equals(changeMetaId);
+  }
+
+  public boolean areDraftsUpToDate(Repository draftsRepo, Account.Id accountId)
+      throws IOException {
+    Ref ref = draftsRepo.exactRef(
+        RefNames.refsDraftComments(accountId, changeId));
+    if (ref == null) {
+      return !draftIds.containsKey(accountId);
+    }
+    return ref.getObjectId().equals(draftIds.get(accountId));
+  }
+
+  @VisibleForTesting
+  Change.Id getChangeId() {
+    return changeId;
+  }
+
+  @VisibleForTesting
+  ObjectId getChangeMetaId() {
+    return changeMetaId;
+  }
+
+  @VisibleForTesting
+  ImmutableMap<Account.Id, ObjectId> getDraftIds() {
+    return draftIds;
+  }
+
+  @Override
+  public String toString() {
+    return toString(changeMetaId, draftIds);
+  }
+}

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

@@ -27,6 +27,14 @@ class NoteDbMetrics {
   /** End-to-end latency for writing a collection of updates. */
   final Timer1<NoteDbTable> updateLatency;
 
+  /**
+   * The portion of {@link #updateLatency} due to preparing the sequence of
+   * updates.
+   * <p>
+   * May include some I/O (e.g. reading old refs), but excludes writes.
+   */
+  final Timer1<NoteDbTable> stageUpdateLatency;
+
   /**
    * End-to-end latency for reading changes from NoteDb, including reading
    * ref(s) and parsing.
@@ -50,6 +58,13 @@ class NoteDbMetrics {
             .setUnit(Units.MILLISECONDS),
         view);
 
+    stageUpdateLatency = metrics.newTimer(
+        "notedb/stage_update_latency",
+        new Description("Latency for staging updates to NoteDb by table")
+            .setCumulative()
+            .setUnit(Units.MICROSECONDS),
+        view);
+
     readLatency = metrics.newTimer(
         "notedb/read_latency",
         new Description("NoteDb read latency by table")

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

@@ -18,11 +18,17 @@ 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.reviewdb.client.RefNames.REFS_DRAFT_COMMENTS;
 import static com.google.gerrit.server.notedb.NoteDbTable.CHANGES;
 
+import com.google.common.base.Optional;
 import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.HashBasedTable;
 import com.google.common.collect.ListMultimap;
+import com.google.common.collect.Table;
 import com.google.gerrit.metrics.Timer1;
+import com.google.gerrit.reviewdb.client.Account;
+import com.google.gerrit.reviewdb.client.Change;
 import com.google.gerrit.reviewdb.client.Project;
 import com.google.gerrit.server.config.AllUsersName;
 import com.google.gerrit.server.git.ChainedReceiveCommands;
@@ -41,8 +47,20 @@ import org.eclipse.jgit.transport.ReceiveCommand;
 
 import java.io.IOException;
 import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 
+/**
+ * Object to manage a single sequence of updates to NoteDb.
+ * <p>
+ * Instances are one-time-use. Handles updating both the change meta repo and
+ * the All-Users meta repo for any affected changes, with proper ordering.
+ * <p>
+ * To see the state that would be applied prior to executing the full sequence
+ * of updates, use {@link #stage()}.
+ */
 public class NoteDbUpdateManager {
   public interface Factory {
     NoteDbUpdateManager create(Project.NameKey projectName);
@@ -84,6 +102,7 @@ public class NoteDbUpdateManager {
 
   private OpenRepo changeRepo;
   private OpenRepo allUsersRepo;
+  private Map<Change.Id, NoteDbChangeState.Delta> staged;
 
   @AssistedInject
   NoteDbUpdateManager(GitRepositoryManager repoManager,
@@ -178,6 +197,7 @@ public class NoteDbUpdateManager {
     checkArgument(update.getProjectName().equals(projectName),
       "update for project %s cannot be added to manager for project %s",
       update.getProjectName(), projectName);
+    checkState(staged == null, "cannot add new update after staging");
     changeUpdates.put(update.getRefName(), update);
     ChangeDraftUpdate du = update.getDraftUpdate();
     if (du != null) {
@@ -186,19 +206,87 @@ public class NoteDbUpdateManager {
   }
 
   public void add(ChangeDraftUpdate draftUpdate) {
+    checkState(staged == null, "cannot add new update after staging");
     draftUpdates.put(draftUpdate.getRefName(), draftUpdate);
   }
 
+  /**
+   * Stage updates in the manager's internal list of commands.
+   *
+   * @return map of the state that would get written to the applicable repo(s)
+   *     for each affected change.
+   * @throws OrmException if a database layer error occurs.
+   * @throws IOException if a storage layer error occurs.
+   */
+  public Map<Change.Id, NoteDbChangeState.Delta> stage()
+      throws OrmException, IOException {
+    if (staged != null) {
+      return staged;
+    }
+    try (Timer1.Context timer = metrics.stageUpdateLatency.start(CHANGES)) {
+      staged = new HashMap<>();
+      if (isEmpty()) {
+        return staged;
+      }
+
+      initChangeRepo();
+      if (!draftUpdates.isEmpty()) {
+        initAllUsersRepo();
+      }
+      addCommands();
+
+      Table<Change.Id, Account.Id, ObjectId> allDraftIds = getDraftIds();
+      Set<Change.Id> changeIds = new HashSet<>();
+      for (ReceiveCommand cmd : changeRepo.cmds.getCommands().values()) {
+        Change.Id changeId = Change.Id.fromRef(cmd.getRefName());
+        changeIds.add(changeId);
+        Optional<ObjectId> metaId = Optional.of(cmd.getNewId());
+        staged.put(
+            changeId,
+            NoteDbChangeState.Delta.create(
+                changeId, metaId, allDraftIds.rowMap().remove(changeId)));
+      }
+
+      for (Map.Entry<Change.Id, Map<Account.Id, ObjectId>> e
+          : allDraftIds.rowMap().entrySet()) {
+        // If a change remains in the table at this point, it means we are
+        // updating its drafts but not the change itself.
+        staged.put(
+            e.getKey(),
+            NoteDbChangeState.Delta.create(
+                e.getKey(), Optional.<ObjectId>absent(), e.getValue()));
+      }
+
+      return staged;
+    }
+  }
+
+  private Table<Change.Id, Account.Id, ObjectId> getDraftIds() {
+    Table<Change.Id, Account.Id, ObjectId> draftIds = HashBasedTable.create();
+    if (allUsersRepo == null) {
+      return draftIds;
+    }
+    for (ReceiveCommand cmd : allUsersRepo.cmds.getCommands().values()) {
+      String r = cmd.getRefName();
+      if (r.startsWith(REFS_DRAFT_COMMENTS)) {
+        Account.Id accountId =
+            Account.Id.fromRefPart(r.substring(REFS_DRAFT_COMMENTS.length()));
+        checkDraftRef(accountId != null, r);
+        int s = r.lastIndexOf('/');
+        checkDraftRef(s >= 0 && s < r.length() - 1, r);
+        Change.Id changeId = Change.Id.parse(r.substring(s + 1));
+        draftIds.put(changeId, accountId, cmd.getNewId());
+      }
+    }
+    return draftIds;
+  }
+
   public void execute() throws OrmException, IOException {
     if (isEmpty()) {
       return;
     }
     try (Timer1.Context timer = metrics.updateLatency.start(CHANGES)) {
-      initChangeRepo();
+      stage();
-      if (!draftUpdates.isEmpty()) {
-        initAllUsersRepo();
-      }
-      addCommands();
 
       // ChangeUpdates must execute before ChangeDraftUpdates.
       //
@@ -286,4 +374,8 @@ public class NoteDbUpdateManager {
     }
     return updates.iterator().next().allowWriteToNewRef();
   }
+
+  private static void checkDraftRef(boolean condition, String refName) {
+    checkState(condition, "invalid draft ref: %s", refName);
+  }
 }

gerrit-server/src/main/java/com/google/gerrit/server/schema/SchemaVersion.java

@@ -32,7 +32,7 @@ import java.util.List;
 /** A version of the database schema. */
 public abstract class SchemaVersion {
   /** The current schema version. */
-  public static final Class<Schema_120> C = Schema_120.class;
+  public static final Class<Schema_121> C = Schema_121.class;
 
   public static int getBinaryVersion() {
     return guessVersion(C);

gerrit-server/src/main/java/com/google/gerrit/server/schema/Schema_121.java

@@ -0,0 +1,25 @@
+// 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.schema;
+
+import com.google.inject.Inject;
+import com.google.inject.Provider;
+
+public class Schema_121 extends SchemaVersion {
+  @Inject
+  Schema_121(Provider<Schema_120> prior) {
+    super(prior);
+  }
+}

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

@@ -0,0 +1,156 @@
+// 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 static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.gerrit.server.notedb.NoteDbChangeState.applyDelta;
+import static com.google.gerrit.server.notedb.NoteDbChangeState.parse;
+import static org.eclipse.jgit.lib.ObjectId.zeroId;
+
+import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableMap;
+import com.google.gerrit.reviewdb.client.Account;
+import com.google.gerrit.reviewdb.client.Change;
+import com.google.gerrit.reviewdb.client.Project;
+import com.google.gerrit.server.notedb.NoteDbChangeState.Delta;
+import com.google.gerrit.testutil.TestChanges;
+import com.google.gwtorm.client.KeyUtil;
+import com.google.gwtorm.server.StandardKeyEncoder;
+
+import org.eclipse.jgit.lib.ObjectId;
+import org.junit.Test;
+
+/** Unit tests for {@link NoteDbChangeState}. */
+public class NoteDbChangeStateTest {
+  static {
+    KeyUtil.setEncoderImpl(new StandardKeyEncoder());
+  }
+
+  ObjectId SHA1 =
+      ObjectId.fromString("deadbeefdeadbeefdeadbeefdeadbeefdeadbeef");
+  ObjectId SHA2 =
+      ObjectId.fromString("abcd1234abcd1234abcd1234abcd1234abcd1234");
+  ObjectId SHA3 =
+      ObjectId.fromString("badc0feebadc0feebadc0feebadc0feebadc0fee");
+
+  @Test
+  public void parseWithoutDrafts() {
+    NoteDbChangeState state = parse(new Change.Id(1), SHA1.name());
+
+    assertThat(state.getChangeId()).isEqualTo(new Change.Id(1));
+    assertThat(state.getChangeMetaId()).isEqualTo(SHA1);
+    assertThat(state.getDraftIds()).isEmpty();
+
+    assertThat(state.toString()).isEqualTo(SHA1.name());
+  }
+
+  @Test
+  public void parseWithDrafts() {
+    NoteDbChangeState state = parse(
+        new Change.Id(1),
+        SHA1.name() + ",2003=" + SHA2.name() + ",1001=" + SHA3.name());
+
+    assertThat(state.getChangeId()).isEqualTo(new Change.Id(1));
+    assertThat(state.getChangeMetaId()).isEqualTo(SHA1);
+    assertThat(state.getDraftIds()).containsExactly(
+        new Account.Id(1001), SHA3,
+        new Account.Id(2003), SHA2);
+
+    assertThat(state.toString()).isEqualTo(
+        SHA1.name() + ",1001=" + SHA3.name() + ",2003=" + SHA2.name());
+  }
+
+  @Test
+  public void applyDeltaToNullWithNoNewMetaId() {
+    Change c = newChange();
+    assertThat(c.getNoteDbState()).isNull();
+    applyDelta(c, Delta.create(c.getId(), noMetaId(), noDrafts()));
+    assertThat(c.getNoteDbState()).isNull();
+
+    applyDelta(c, Delta.create(c.getId(), noMetaId(),
+          drafts(new Account.Id(1001), zeroId())));
+    assertThat(c.getNoteDbState()).isNull();
+  }
+
+  @Test
+  public void applyDeltaToMetaId() {
+    Change c = newChange();
+    applyDelta(c, Delta.create(c.getId(), metaId(SHA1), noDrafts()));
+    assertThat(c.getNoteDbState()).isEqualTo(SHA1.name());
+
+    applyDelta(c, Delta.create(c.getId(), metaId(SHA2), noDrafts()));
+    assertThat(c.getNoteDbState()).isEqualTo(SHA2.name());
+
+    // No-op delta.
+    applyDelta(c, Delta.create(c.getId(), noMetaId(), noDrafts()));
+    assertThat(c.getNoteDbState()).isEqualTo(SHA2.name());
+
+    // Set to zero clears the field.
+    applyDelta(c, Delta.create(c.getId(), metaId(zeroId()), noDrafts()));
+    assertThat(c.getNoteDbState()).isNull();
+  }
+
+  @Test
+  public void applyDeltaToDrafts() {
+    Change c = newChange();
+    applyDelta(c, Delta.create(c.getId(), metaId(SHA1),
+          drafts(new Account.Id(1001), SHA2)));
+    assertThat(c.getNoteDbState()).isEqualTo(
+        SHA1.name() + ",1001=" + SHA2.name());
+
+    applyDelta(c, Delta.create(c.getId(), noMetaId(),
+          drafts(new Account.Id(2003), SHA3)));
+    assertThat(c.getNoteDbState()).isEqualTo(
+        SHA1.name() + ",1001=" + SHA2.name() + ",2003=" + SHA3.name());
+
+    applyDelta(c, Delta.create(c.getId(), noMetaId(),
+          drafts(new Account.Id(2003), zeroId())));
+    assertThat(c.getNoteDbState()).isEqualTo(
+        SHA1.name() + ",1001=" + SHA2.name());
+
+    applyDelta(c, Delta.create(c.getId(), metaId(SHA3), noDrafts()));
+    assertThat(c.getNoteDbState()).isEqualTo(
+        SHA3.name() + ",1001=" + SHA2.name());
+  }
+
+  private static Change newChange() {
+    return TestChanges.newChange(
+        new Project.NameKey("project"), new Account.Id(12345));
+  }
+
+  // Static factory methods to avoid type arguments when using as method args.
+
+  private static Optional<ObjectId> noMetaId() {
+    return Optional.absent();
+  }
+
+  private static Optional<ObjectId> metaId(ObjectId id) {
+    return Optional.of(id);
+  }
+
+  private static ImmutableMap<Account.Id, ObjectId> noDrafts() {
+    return ImmutableMap.of();
+  }
+
+  private static ImmutableMap<Account.Id, ObjectId> drafts(Object... args) {
+    checkArgument(args.length % 2 == 0);
+    ImmutableMap.Builder<Account.Id, ObjectId> b = ImmutableMap.builder();
+    for (int i = 0; i < args.length / 2; i++) {
+      b.put((Account.Id) args[2*i], (ObjectId) args[2*i+1]);
+    }
+    return b.build();
+  }
+}