Parse ChangeNotes from staged results if auto-rebuilding fails

Loading the change screen is really pessimal when it comes to
rebuilding out-of-date changes in NoteDb, because it causes 10 or so
request handlers to try to rebuild the change concurrently. Even on a
backend that's more performant than googlesource.com's, this can lead
to contention when trying to update the same NoteDb ref.

A previous fix in I0b80c975 handles some failure modes by rechecking
the change after a failure during rebuilding, but that wasn't enough.
Go one step further and just serve up the ChangeNotes from data the
caller tried but failed to store in NoteDb. For a given state in
ReviewDb (i.e. a given ChangeBundle), the exact set of git objects
produced by rebuilding is completely deterministic, so we can say with
confidence that this is what we should have returned. Of course, it's
still possible that multiple concurrent read-and-rebuild requests will
race with concurrent write requests, so different requests may have
rebuilt from different snapshots of ReviewDb. But this is a completely
normal and acceptable race: if a write happens in the middle of a
sequence of concurrent reads, some reads will observe the old state
and some will observe the new state.

Change-Id: I2797d4fbc59fe9e0507db21ed53c69625df1eb07

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

@@ -24,7 +24,6 @@ import static org.eclipse.jgit.lib.Constants.OBJ_BLOB;
 import com.google.common.collect.ImmutableList;
 import com.google.gerrit.acceptance.AbstractDaemonTest;
 import com.google.gerrit.acceptance.AcceptanceTestRequestScope;
-import com.google.gerrit.acceptance.GerritConfig;
 import com.google.gerrit.acceptance.PushOneCommit;
 import com.google.gerrit.acceptance.TestAccount;
 import com.google.gerrit.common.TimeUtil;
@@ -53,6 +52,7 @@ import com.google.gerrit.server.notedb.ChangeNotes;
 import com.google.gerrit.server.notedb.NoteDbChangeState;
 import com.google.gerrit.server.notedb.TestChangeRebuilderWrapper;
 import com.google.gerrit.server.schema.DisabledChangesReviewDbWrapper;
+import com.google.gerrit.testutil.ConfigSuite;
 import com.google.gerrit.testutil.NoteDbChecker;
 import com.google.gerrit.testutil.NoteDbMode;
 import com.google.gerrit.testutil.TestChanges;
@@ -60,6 +60,7 @@ import com.google.gerrit.testutil.TestTimeUtil;
 import com.google.inject.Inject;
 import com.google.inject.Provider;
 
+import org.eclipse.jgit.lib.Config;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.ObjectInserter;
 import org.eclipse.jgit.lib.Ref;
@@ -76,6 +77,13 @@ import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 public class ChangeRebuilderIT extends AbstractDaemonTest {
+  @ConfigSuite.Default
+  public static Config defaultConfig() {
+    Config cfg = new Config();
+    cfg.setBoolean("noteDb", null, "testRebuilderWrapper", true);
+    return cfg;
+  }
+
   @Inject
   private AllUsersName allUsers;
 
@@ -358,7 +366,6 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
   }
 
   @Test
-  @GerritConfig(name = "noteDb.testRebuilderWrapper", value = "true")
   public void rebuildIgnoresErrorIfChangeIsUpToDateAfter() throws Exception {
     setNotesMigration(true, true);
 
@@ -387,6 +394,85 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
     assertThat(actual.differencesFrom(expected)).isEmpty();
   }
 
+  @Test
+  public void rebuildReturnsCorrectResultEvenIfSavingToNoteDbFailed()
+      throws Exception {
+    setNotesMigration(true, true);
+
+    PushOneCommit.Result r = createChange();
+    Change.Id id = r.getPatchSetId().getParentKey();
+    assertChangeUpToDate(true, id);
+    ObjectId oldMetaId = getMetaRef(project, changeMetaRef(id));
+
+    // Make a ReviewDb change behind NoteDb's back.
+    setNotesMigration(false, false);
+    gApi.changes().id(id.get()).topic(name("a-topic"));
+    setInvalidNoteDbState(id);
+    assertChangeUpToDate(false, id);
+    assertThat(getMetaRef(project, changeMetaRef(id))).isEqualTo(oldMetaId);
+
+    // Force the next rebuild attempt to fail.
+    rebuilderWrapper.failNextUpdate();
+    setNotesMigration(true, true);
+    ChangeNotes notes = notesFactory.create(dbProvider.get(), project, id);
+
+    // Not up to date, but the actual returned state matches anyway.
+    assertChangeUpToDate(false, id);
+    assertThat(getMetaRef(project, changeMetaRef(id))).isEqualTo(oldMetaId);
+    ChangeBundle actual = ChangeBundle.fromNotes(plcUtil, notes);
+    ChangeBundle expected = ChangeBundle.fromReviewDb(unwrapDb(), id);
+    assertThat(actual.differencesFrom(expected)).isEmpty();
+    assertChangeUpToDate(false, id);
+
+    // Another rebuild attempt succeeds
+    notesFactory.create(dbProvider.get(), project, id);
+    assertThat(getMetaRef(project, changeMetaRef(id))).isNotEqualTo(oldMetaId);
+    assertChangeUpToDate(true, id);
+  }
+
+  @Test
+  public void rebuildReturnsCorrectDraftResultEvenIfSavingToNoteDbFailed()
+      throws Exception {
+    setNotesMigration(true, true);
+
+    PushOneCommit.Result r = createChange();
+    Change.Id id = r.getPatchSetId().getParentKey();
+    putDraft(user, id, 1, "comment by user");
+    assertChangeUpToDate(true, id);
+
+    ObjectId oldMetaId =
+        getMetaRef(allUsers, refsDraftComments(id, user.getId()));
+
+    // Add a draft behind NoteDb's back.
+    setNotesMigration(false, false);
+    putDraft(user, id, 1, "second comment by user");
+    setInvalidNoteDbState(id);
+    assertDraftsUpToDate(false, id, user);
+    assertThat(getMetaRef(allUsers, refsDraftComments(id, user.getId())))
+        .isEqualTo(oldMetaId);
+
+    // Force the next rebuild attempt to fail.
+    rebuilderWrapper.failNextUpdate();
+    setNotesMigration(true, true);
+    ChangeNotes notes = notesFactory.create(dbProvider.get(), project, id);
+    notes.getDraftComments(user.getId());
+    assertThat(getMetaRef(allUsers, refsDraftComments(id, user.getId())))
+        .isEqualTo(oldMetaId);
+
+    // Not up to date, but the actual returned state matches anyway.
+    assertDraftsUpToDate(false, id, user);
+    ChangeBundle actual = ChangeBundle.fromNotes(plcUtil, notes);
+    ChangeBundle expected = ChangeBundle.fromReviewDb(unwrapDb(), id);
+    assertThat(actual.differencesFrom(expected)).isEmpty();
+
+    // Another rebuild attempt succeeds
+    notesFactory.create(dbProvider.get(), project, id);
+    assertChangeUpToDate(true, id);
+    assertDraftsUpToDate(true, id, user);
+    assertThat(getMetaRef(allUsers, refsDraftComments(id, user.getId())))
+        .isNotEqualTo(oldMetaId);
+  }
+
   @Test
   public void rebuildAutomaticallyWhenDraftsOutOfDate() throws Exception {
     setNotesMigration(true, true);

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

@@ -53,7 +53,6 @@ import com.google.gerrit.reviewdb.server.ReviewDb;
 import com.google.gerrit.server.ReviewerSet;
 import com.google.gerrit.server.git.RefCache;
 import com.google.gerrit.server.git.RepoRefCache;
-import com.google.gerrit.server.notedb.NoteDbUpdateManager.Result;
 import com.google.gerrit.server.project.NoSuchChangeException;
 import com.google.gerrit.server.project.ProjectCache;
 import com.google.gerrit.server.query.change.ChangeData;
@@ -396,6 +395,7 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
   // notes easier.
   RevisionNoteMap revisionNoteMap;
 
+  private NoteDbUpdateManager.Result rebuildResult;
   private DraftCommentNotes draftCommentNotes;
 
   @VisibleForTesting
@@ -503,8 +503,8 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
       throws OrmException {
     if (draftCommentNotes == null ||
         !author.equals(draftCommentNotes.getAuthor())) {
-      draftCommentNotes =
-          new DraftCommentNotes(args, change, author, autoRebuild);
+      draftCommentNotes = new DraftCommentNotes(
+          args, change, author, autoRebuild, rebuildResult);
       draftCommentNotes.load();
     }
   }
@@ -593,63 +593,42 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
 
   private LoadHandle rebuildAndOpen(Repository repo, ObjectId oldId)
       throws IOException {
+    Change.Id cid = getChangeId();
+    ReviewDb db = args.db.get();
+    ChangeRebuilder rebuilder = args.rebuilder.get();
     try {
-      NoteDbChangeState newState;
-      try {
-        NoteDbUpdateManager.Result r =
-            args.rebuilder.get().rebuild(args.db.get(), getChangeId());
-        newState = r != null ? r.newState() : null;
-        repo.scanForRepoChanges();
-      } catch (IOException e) {
-        newState = recheckUpToDate(repo, e);
-      }
-      if (newState == null) {
+      NoteDbUpdateManager manager = rebuilder.stage(db, cid);
+      if (manager == null) {
         return super.openHandle(repo, oldId); // May be null in tests.
       }
+      NoteDbUpdateManager.Result r = manager.stageAndApplyDelta(change);
+      try {
+        rebuilder.execute(db, cid, manager);
+        repo.scanForRepoChanges();
+      } catch (OrmException | IOException e) {
+        // Rebuilding failed. Most likely cause is contention on one or more
+        // change refs; there are other types of errors that can happen during
+        // rebuilding, but generally speaking they should happen during stage(),
+        // not execute(). Assume that some other worker is going to successfully
+        // store the rebuilt state, which is deterministic given an input
+        // ChangeBundle.
+        //
+        // Parse notes from the staged result so we can return something useful
+        // to the caller instead of throwing.
+        rebuildResult = checkNotNull(r);
+        checkNotNull(r.newState());
+        checkNotNull(r.staged());
         return LoadHandle.create(
-          ChangeNotesCommit.newRevWalk(repo), newState.getChangeMetaId());
+            ChangeNotesCommit.newStagedRevWalk(
+                repo, r.staged().changeObjects()),
+            r.newState().getChangeMetaId());
+      }
+      return LoadHandle.create(
+          ChangeNotesCommit.newRevWalk(repo), r.newState().getChangeMetaId());
     } catch (NoSuchChangeException e) {
       return super.openHandle(repo, oldId);
-    } catch (OrmException | ConfigInvalidException e) {
+    } catch (OrmException e) {
       throw new IOException(e);
     }
   }
-
-  private NoteDbChangeState recheckUpToDate(Repository repo, IOException e)
-      throws IOException {
-    // Should only be non-null if auto-rebuilding disabled.
-    checkState(refs == null);
-    // An error during auto-rebuilding might be caused by LOCK_FAILURE or a
-    // similar contention issue, where another thread successfully rebuilt the
-    // change. Reread the change from ReviewDb and NoteDb and recheck the state.
-    Change newChange;
-    try {
-      newChange = unwrap(args.db.get()).changes().get(getChangeId());
-    } catch (OrmException e2) {
-      logRecheckError(e2);
-      throw e;
-    }
-    NoteDbChangeState newState = NoteDbChangeState.parse(newChange);
-    boolean upToDate;
-    try {
-      repo.scanForRepoChanges();
-      upToDate = NoteDbChangeState.isChangeUpToDate(
-          newState, new RepoRefCache(repo), getChangeId());
-    } catch (IOException e2) {
-      logRecheckError(e2);
-      throw e;
-    }
-    if (!upToDate) {
-      log.warn("Rechecked change {} after a rebuild error, but it was not up to"
-          + " date; rethrowing exception", getChangeId());
-      throw e;
-    }
-    change = new Change(newChange);
-    return newState;
-  }
-
-  private void logRecheckError(Throwable t) {
-    log.error("Error rechecking if change " + getChangeId()
-        + " is up to date; logging this exception but rethrowing original", t);
-  }
 }

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

@@ -22,6 +22,7 @@ import com.google.common.collect.ListMultimap;
 import org.eclipse.jgit.errors.IncorrectObjectTypeException;
 import org.eclipse.jgit.errors.MissingObjectException;
 import org.eclipse.jgit.lib.AnyObjectId;
+import org.eclipse.jgit.lib.ObjectReader;
 import org.eclipse.jgit.lib.Repository;
 import org.eclipse.jgit.revwalk.FooterKey;
 import org.eclipse.jgit.revwalk.FooterLine;
@@ -44,11 +45,30 @@ public class ChangeNotesCommit extends RevCommit {
     return new ChangeNotesRevWalk(repo);
   }
 
+  public static ChangeNotesRevWalk newStagedRevWalk(Repository repo,
+      Iterable<InsertedObject> stagedObjs) {
+    final InMemoryInserter ins = new InMemoryInserter(repo);
+    for (InsertedObject obj : stagedObjs) {
+      ins.insert(obj);
+    }
+    return new ChangeNotesRevWalk(ins.newReader()) {
+      @Override
+      public void close() {
+        ins.close();
+        super.close();
+      }
+    };
+  }
+
   public static class ChangeNotesRevWalk extends RevWalk {
     private ChangeNotesRevWalk(Repository repo) {
       super(repo);
     }
 
+    private ChangeNotesRevWalk(ObjectReader reader) {
+      super(reader);
+    }
+
     @Override
     protected ChangeNotesCommit createCommit(AnyObjectId id) {
       return new ChangeNotesCommit(id);

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

@@ -63,4 +63,11 @@ public abstract class ChangeRebuilder {
       Project.NameKey project, Repository allUsersRepo)
       throws NoSuchChangeException, IOException, OrmException,
       ConfigInvalidException;
+
+  public abstract NoteDbUpdateManager stage(ReviewDb db, Change.Id changeId)
+      throws NoSuchChangeException, IOException, OrmException;
+
+  public abstract Result execute(ReviewDb db, Change.Id changeId,
+      NoteDbUpdateManager manager) throws NoSuchChangeException, OrmException,
+      IOException;
 }

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

@@ -165,9 +165,43 @@ public class ChangeRebuilderImpl extends ChangeRebuilder {
     return execute(db, changeId, manager);
   }
 
-  private Result execute(ReviewDb db, Change.Id changeId,
-      NoteDbUpdateManager manager)
-      throws NoSuchChangeException, OrmException, IOException {
+  private static class AbortUpdateException extends OrmRuntimeException {
+    private static final long serialVersionUID = 1L;
+
+    AbortUpdateException() {
+      super("aborted");
+    }
+  }
+
+  @Override
+  public Result rebuild(NoteDbUpdateManager manager,
+      ChangeBundle bundle) throws NoSuchChangeException, IOException,
+      OrmException, ConfigInvalidException {
+    Change change = new Change(bundle.getChange());
+    buildUpdates(manager, bundle);
+    return manager.stageAndApplyDelta(change);
+  }
+
+  @Override
+  public NoteDbUpdateManager stage(ReviewDb db, Change.Id changeId)
+      throws NoSuchChangeException, IOException, OrmException {
+    db = unwrapDb(db);
+    Change change = db.changes().get(changeId);
+    if (change == null) {
+      throw new NoSuchChangeException(changeId);
+    }
+    NoteDbUpdateManager manager =
+        updateManagerFactory.create(change.getProject());
+    buildUpdates(manager, ChangeBundle.fromReviewDb(db, changeId));
+    manager.stage();
+    return manager;
+  }
+
+  @Override
+  public Result execute(ReviewDb db, Change.Id changeId,
+      NoteDbUpdateManager manager) throws NoSuchChangeException, OrmException,
+      IOException {
+    db = unwrapDb(db);
     Change change = db.changes().get(changeId);
     if (change == null) {
       throw new NoSuchChangeException(changeId);
@@ -194,23 +228,6 @@ public class ChangeRebuilderImpl extends ChangeRebuilder {
     return r;
   }
 
-  private static class AbortUpdateException extends OrmRuntimeException {
-    private static final long serialVersionUID = 1L;
-
-    AbortUpdateException() {
-      super("aborted");
-    }
-  }
-
-  @Override
-  public Result rebuild(NoteDbUpdateManager manager,
-      ChangeBundle bundle) throws NoSuchChangeException, IOException,
-      OrmException, ConfigInvalidException {
-    Change change = new Change(bundle.getChange());
-    buildUpdates(manager, bundle);
-    return manager.stageAndApplyDelta(change);
-  }
-
   @Override
   public boolean rebuildProject(ReviewDb db,
       ImmutableMultimap<Project.NameKey, Change.Id> allChanges,

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

@@ -14,6 +14,8 @@
 
 package com.google.gerrit.server.notedb;
 
+import static com.google.common.base.Preconditions.checkNotNull;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableListMultimap;
@@ -25,6 +27,7 @@ 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.git.RepoRefCache;
+import com.google.gerrit.server.notedb.NoteDbUpdateManager.StagedResult;
 import com.google.gerrit.server.project.NoSuchChangeException;
 import com.google.gwtorm.server.OrmException;
 import com.google.inject.assistedinject.Assisted;
@@ -36,6 +39,7 @@ import org.eclipse.jgit.lib.ObjectReader;
 import org.eclipse.jgit.lib.Repository;
 import org.eclipse.jgit.notes.NoteMap;
 import org.eclipse.jgit.revwalk.RevCommit;
+import org.eclipse.jgit.transport.ReceiveCommand;
 
 import java.io.IOException;
 
@@ -52,6 +56,7 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
 
   private final Change change;
   private final Account.Id author;
+  private final NoteDbUpdateManager.Result rebuildResult;
 
   private ImmutableListMultimap<RevId, PatchLineComment> comments;
   private RevisionNoteMap revisionNoteMap;
@@ -61,7 +66,7 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
       Args args,
       @Assisted Change change,
       @Assisted Account.Id author) {
-    this(args, change, author, true);
+    this(args, change, author, true, null);
   }
 
   @AssistedInject
@@ -72,16 +77,19 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
     super(args, changeId, true);
     this.change = null;
     this.author = author;
+    this.rebuildResult = null;
   }
 
   DraftCommentNotes(
       Args args,
       Change change,
       Account.Id author,
-      boolean autoRebuild) {
+      boolean autoRebuild,
+      NoteDbUpdateManager.Result rebuildResult) {
     super(args, change.getId(), autoRebuild);
     this.change = change;
     this.author = author;
+    this.rebuildResult = rebuildResult;
   }
 
   RevisionNoteMap getRevisionNoteMap() {
@@ -146,7 +154,12 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
 
   @Override
   protected LoadHandle openHandle(Repository repo) throws IOException {
-    if (change != null && autoRebuild) {
+    if (rebuildResult != null) {
+      StagedResult sr = checkNotNull(rebuildResult.staged());
+      return LoadHandle.create(
+          ChangeNotesCommit.newStagedRevWalk(repo, sr.allUsersObjects()),
+          findNewId(sr.allUsersCommands(), getRefName()));
+    } else if (change != null && autoRebuild) {
       NoteDbChangeState state = NoteDbChangeState.parse(change);
       // Only check if this particular user's drafts are up to date, to avoid
       // reading unnecessary refs.
@@ -158,6 +171,16 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
     return super.openHandle(repo);
   }
 
+  private static ObjectId findNewId(
+      Iterable<ReceiveCommand> cmds, String refName) {
+    for (ReceiveCommand cmd : cmds) {
+      if (cmd.getRefName().equals(refName)) {
+        return cmd.getNewId();
+      }
+    }
+    return null;
+  }
+
   private LoadHandle rebuildAndOpen(Repository repo) throws IOException {
     try {
       NoteDbUpdateManager.Result r =

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

@@ -25,6 +25,7 @@ import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.ObjectInserter;
 import org.eclipse.jgit.lib.ObjectLoader;
 import org.eclipse.jgit.lib.ObjectReader;
+import org.eclipse.jgit.lib.Repository;
 import org.eclipse.jgit.transport.PackParser;
 
 import java.io.IOException;
@@ -36,11 +37,17 @@ import java.util.Set;
 
 class InMemoryInserter extends ObjectInserter {
   private final ObjectReader reader;
-  private final Map<ObjectId, InsertedObject> inserted;
+  private final Map<ObjectId, InsertedObject> inserted = new LinkedHashMap<>();
+  private final boolean closeReader;
 
   InMemoryInserter(ObjectReader reader) {
     this.reader = checkNotNull(reader);
-    inserted = new LinkedHashMap<>();
+    closeReader = false;
+  }
+
+  InMemoryInserter(Repository repo) {
+    this.reader = repo.newObjectReader();
+    closeReader = true;
   }
 
   @Override
@@ -59,7 +66,7 @@ class InMemoryInserter extends ObjectInserter {
     return insert(InsertedObject.create(type, data, off, len));
   }
 
-  private ObjectId insert(InsertedObject obj) {
+  ObjectId insert(InsertedObject obj) {
     inserted.put(obj.id(), obj);
     return obj.id();
   }
@@ -81,7 +88,9 @@ class InMemoryInserter extends ObjectInserter {
 
   @Override
   public void close() {
-    // Do nothing; this class owns no open resources.
+    if (closeReader) {
+      reader.close();
+    }
   }
 
   public ImmutableList<InsertedObject> getInsertedObjects() {

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

@@ -80,6 +80,17 @@ public class NoteDbModule extends FactoryModule {
             Repository allUsersRepo) {
           return false;
         }
+
+        @Override
+        public NoteDbUpdateManager stage(ReviewDb db, Change.Id changeId) {
+          return null;
+        }
+
+        @Override
+        public Result execute(ReviewDb db, Change.Id changeId,
+            NoteDbUpdateManager manager) {
+          return null;
+        }
       });
       bind(new TypeLiteral<Cache<ChangeNotesCache.Key, ChangeNotesState>>() {})
           .annotatedWith(Names.named(ChangeNotesCache.CACHE_NAME))

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

@@ -36,6 +36,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
 @Singleton
 public class TestChangeRebuilderWrapper extends ChangeRebuilder {
   private final ChangeRebuilderImpl delegate;
+  private final AtomicBoolean failNextUpdate;
   private final AtomicBoolean stealNextUpdate;
 
   @Inject
@@ -43,9 +44,14 @@ public class TestChangeRebuilderWrapper extends ChangeRebuilder {
       ChangeRebuilderImpl rebuilder) {
     super(schemaFactory);
     this.delegate = rebuilder;
+    this.failNextUpdate = new AtomicBoolean();
     this.stealNextUpdate = new AtomicBoolean();
   }
 
+  public void failNextUpdate() {
+    failNextUpdate.set(true);
+  }
+
   public void stealNextUpdate() {
     stealNextUpdate.set(true);
   }
@@ -54,6 +60,9 @@ public class TestChangeRebuilderWrapper extends ChangeRebuilder {
   public Result rebuild(ReviewDb db, Change.Id changeId)
       throws NoSuchChangeException, IOException, OrmException,
       ConfigInvalidException {
+    if (failNextUpdate.getAndSet(false)) {
+      throw new IOException("Update failed");
+    }
     Result result = delegate.rebuild(db, changeId);
     if (stealNextUpdate.getAndSet(false)) {
       throw new IOException("Update stolen");
@@ -77,6 +86,9 @@ public class TestChangeRebuilderWrapper extends ChangeRebuilder {
       Project.NameKey project, Repository allUsersRepo)
       throws NoSuchChangeException, IOException, OrmException,
       ConfigInvalidException {
+    if (failNextUpdate.getAndSet(false)) {
+      throw new IOException("Update failed");
+    }
     boolean result =
         delegate.rebuildProject(db, allChanges, project, allUsersRepo);
     if (stealNextUpdate.getAndSet(false)) {
@@ -84,4 +96,25 @@ public class TestChangeRebuilderWrapper extends ChangeRebuilder {
     }
     return result;
   }
+
+  @Override
+  public NoteDbUpdateManager stage(ReviewDb db, Change.Id changeId)
+      throws NoSuchChangeException, IOException, OrmException {
+    // Don't inspect stealNextUpdate; that happens in execute() below.
+    return delegate.stage(db, changeId);
+  }
+
+  @Override
+  public Result execute(ReviewDb db, Change.Id changeId,
+      NoteDbUpdateManager manager) throws NoSuchChangeException, OrmException,
+      IOException {
+    if (failNextUpdate.getAndSet(false)) {
+      throw new IOException("Update failed");
+    }
+    Result result = delegate.execute(db, changeId, manager);
+    if (stealNextUpdate.getAndSet(false)) {
+      throw new IOException("Update stolen");
+    }
+    return result;
+  }
 }