Automatically rebuild out-of-date NoteDb drafts

Uses the same mechanism as when rebuilding changes, except checking
the state of the relevant drafts ref in All-Users.

The one place we couldn't do this easily is in
PatchLineCommentsUtil#draftsByAuthor, since it doesn't have a Change
object available for each change, and we don't want to go overboard in
reading the notes. Just bail and skip rebuilding in that particular
case, marking it as deprecated so it doesn't pick up any new users.

Change-Id: Iaeccc9e665eb6ee2ffa4583f5add8d0ba9968d70

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

@@ -18,7 +18,9 @@ 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.AcceptanceTestRequestScope;
 import com.google.gerrit.acceptance.PushOneCommit;
+import com.google.gerrit.acceptance.TestAccount;
 import com.google.gerrit.extensions.api.changes.DraftInput;
 import com.google.gerrit.extensions.restapi.ResourceNotFoundException;
 import com.google.gerrit.reviewdb.client.Change;
@@ -103,13 +105,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();
+    assertThat(unwrapDb().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();
+    assertThat(unwrapDb().changes().get(id).getNoteDbState()).isNotNull();
     checker.checkChanges(id);
   }
 
@@ -166,26 +168,14 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
     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);
-
+    putDraft(user, id, 1, "comment by user");
     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);
-
+    putDraft(admin, id, 2, "comment by admin");
     ObjectId adminDraftsId = getMetaRef(
         allUsers, RefNames.refsDraftComments(admin.getId(), id));
     assertThat(admin.getId().get()).isLessThan(user.getId().get());
@@ -194,9 +184,7 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
         + "," + admin.getId() + "=" + adminDraftsId.name()
         + "," + user.getId() + "=" + userDraftsId.name());
 
-    in.message = "revised comment by admin";
-    gApi.changes().id(id.get()).current().createDraft(in);
-
+    putDraft(admin, id, 2, "revised comment by admin");
     adminDraftsId = getMetaRef(
         allUsers, RefNames.refsDraftComments(admin.getId(), id));
     assertThat(unwrapDb().changes().get(id).getNoteDbState()).isEqualTo(
@@ -211,19 +199,19 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
 
     PushOneCommit.Result r = createChange();
     Change.Id id = r.getPatchSetId().getParentKey();
-    assertUpToDate(true, id);
+    assertChangeUpToDate(true, id);
 
     // Make a ReviewDb change behind NoteDb's back and ensure it's detected.
     notesMigration.setAllEnabled(false);
     gApi.changes().id(id.get()).topic(name("a-topic"));
     setInvalidNoteDbState(id);
-    assertUpToDate(false, id);
+    assertChangeUpToDate(false, id);
 
     // On next NoteDb read, the change is transparently rebuilt.
     notesMigration.setAllEnabled(true);
     assertThat(gApi.changes().id(id.get()).info().topic)
         .isEqualTo(name("a-topic"));
-    assertUpToDate(true, id);
+    assertChangeUpToDate(true, id);
 
     // Check that the bundles are equal.
     ChangeBundle actual = ChangeBundle.fromNotes(
@@ -232,6 +220,29 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
     assertThat(actual.differencesFrom(expected)).isEmpty();
   }
 
+  @Test
+  public void rebuildAutomaticallyWhenDraftsOutOfDate() throws Exception {
+    notesMigration.setAllEnabled(true);
+    setApiUser(user);
+
+    PushOneCommit.Result r = createChange();
+    Change.Id id = r.getPatchSetId().getParentKey();
+    putDraft(user, id, 1, "comment");
+    assertDraftsUpToDate(true, id, user);
+
+    // Make a ReviewDb change behind NoteDb's back and ensure it's detected.
+    notesMigration.setAllEnabled(false);
+    putDraft(user, id, 1, "comment");
+    setInvalidNoteDbState(id);
+    assertDraftsUpToDate(false, id, user);
+
+    // On next NoteDb read, the drafts are transparently rebuilt.
+    notesMigration.setAllEnabled(true);
+    assertThat(gApi.changes().id(id.get()).current().drafts())
+        .containsKey(PushOneCommit.FILE_NAME);
+    assertDraftsUpToDate(true, id, user);
+  }
+
   private void setInvalidNoteDbState(Change.Id id) throws Exception {
     ReviewDb db = unwrapDb();
     Change c = db.changes().get(id);
@@ -243,7 +254,8 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
     db.changes().update(Collections.singleton(c));
   }
 
-  private void assertUpToDate(boolean expected, Change.Id id) throws Exception {
+  private void assertChangeUpToDate(boolean expected, Change.Id id)
+      throws Exception {
     try (Repository repo = repoManager.openMetadataRepository(project)) {
       Change c = unwrapDb().changes().get(id);
       assertThat(c).isNotNull();
@@ -253,6 +265,18 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
     }
   }
 
+  private void assertDraftsUpToDate(boolean expected, Change.Id changeId,
+      TestAccount account) throws Exception {
+    try (Repository repo = repoManager.openMetadataRepository(allUsers)) {
+      Change c = unwrapDb().changes().get(changeId);
+      assertThat(c).isNotNull();
+      assertThat(c.getNoteDbState()).isNotNull();
+      NoteDbChangeState state = NoteDbChangeState.parse(c);
+      assertThat(state.areDraftsUpToDate(repo, account.getId()))
+          .isEqualTo(expected);
+    }
+  }
+
   private ObjectId getMetaRef(Project.NameKey p, String name) throws Exception {
     try (Repository repo = repoManager.openMetadataRepository(p)) {
       Ref ref = repo.exactRef(name);
@@ -260,6 +284,20 @@ public class ChangeRebuilderIT extends AbstractDaemonTest {
     }
   }
 
+  private void putDraft(TestAccount account, Change.Id id, int line, String msg)
+      throws Exception {
+    DraftInput in = new DraftInput();
+    in.line = line;
+    in.message = msg;
+    in.path = PushOneCommit.FILE_NAME;
+    AcceptanceTestRequestScope.Context old = setApiUser(account);
+    try {
+      gApi.changes().id(id.get()).current().createDraft(in);
+    } finally {
+      atrScope.set(old);
+    }
+  }
+
   private ReviewDb unwrapDb() {
     ReviewDb db = dbProvider.get();
     if (db instanceof DisabledChangesReviewDbWrapper) {

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

@@ -272,6 +272,7 @@ public class PatchLineCommentsUtil {
     return sort(comments);
   }
 
+  @Deprecated // To be used only by HasDraftByLegacyPredicate.
   public List<PatchLineComment> draftByAuthor(ReviewDb db,
       Account.Id author) throws OrmException {
     if (!migration.readChanges()) {
@@ -283,8 +284,12 @@ public class PatchLineCommentsUtil {
     List<PatchLineComment> comments = Lists.newArrayList();
     for (String refName : refNames) {
       Change.Id changeId = Change.Id.parse(refName);
-      comments.addAll(
-          draftFactory.create(changeId, author).load().getComments().values());
+      // Avoid loading notes for all affected changes just to be able to auto-
+      // rebuild. This is only used in a corner case in the search codepath, so
+      // returning slightly stale values is ok.
+      DraftCommentNotes notes =
+          draftFactory.createWithAutoRebuildingDisabled(changeId, author);
+      comments.addAll(notes.load().getComments().values());
     }
     return sort(comments);
   }

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

@@ -14,6 +14,7 @@
 
 package com.google.gerrit.server.notedb;
 
+import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.gerrit.server.notedb.NoteDbTable.CHANGES;
 
 import com.google.auto.value.AutoValue;
@@ -77,15 +78,21 @@ public abstract class AbstractChangeNotes<T> {
   public static abstract class LoadHandle implements AutoCloseable {
     public static LoadHandle create(RevWalk walk, ObjectId id) {
       return new AutoValue_AbstractChangeNotes_LoadHandle(
-          walk, id != null ? id.copy() : null);
+          checkNotNull(walk), id != null ? id.copy() : null);
     }
 
-    public abstract RevWalk walk();
+    public static LoadHandle missing() {
+      return new AutoValue_AbstractChangeNotes_LoadHandle(null, null);
+    }
+
+    @Nullable public abstract RevWalk walk();
     @Nullable public abstract ObjectId id();
 
     @Override
     public void close() {
-      walk().close();
+      if (walk() != null) {
+        walk().close();
+      }
     }
   }
 

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

@@ -491,7 +491,7 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
       throws OrmException {
     if (draftCommentNotes == null ||
         !author.equals(draftCommentNotes.getAuthor())) {
-      draftCommentNotes = new DraftCommentNotes(args, getChangeId(), author);
+      draftCommentNotes = new DraftCommentNotes(args, change, author);
       draftCommentNotes.load();
     }
   }

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

@@ -24,14 +24,18 @@ import com.google.gerrit.reviewdb.client.PatchLineComment;
 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.project.NoSuchChangeException;
+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.ObjectId;
 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.revwalk.RevWalk;
 
 import java.io.IOException;
 
@@ -41,20 +45,34 @@ import java.io.IOException;
  */
 public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
   public interface Factory {
-    DraftCommentNotes create(Change.Id changeId, Account.Id accountId);
+    DraftCommentNotes create(Change change, Account.Id accountId);
+    DraftCommentNotes createWithAutoRebuildingDisabled(
+        Change.Id changeId, Account.Id accountId);
   }
 
+  private final Change change;
   private final Account.Id author;
 
   private ImmutableListMultimap<RevId, PatchLineComment> comments;
   private RevisionNoteMap revisionNoteMap;
 
+  @AssistedInject
+  DraftCommentNotes(
+      Args args,
+      @Assisted Change change,
+      @Assisted Account.Id author) {
+    super(args, change.getId());
+    this.change = change;
+    this.author = author;
+  }
+
   @AssistedInject
   DraftCommentNotes(
       Args args,
       @Assisted Change.Id changeId,
       @Assisted Account.Id author) {
     super(args, changeId);
+    this.change = null;
     this.author = author;
   }
 
@@ -118,6 +136,36 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
     return args.allUsers;
   }
 
+  @Override
+  protected LoadHandle openHandle(Repository repo) throws IOException {
+    if (change != null) {
+      NoteDbChangeState state = NoteDbChangeState.parse(change);
+      // Only check if this particular user's drafts are up to date, to avoid
+      // reading unnecessary refs.
+      if (state == null || !state.areDraftsUpToDate(repo, author)) {
+        return rebuildAndOpen(repo);
+      }
+    }
+    return super.openHandle(repo);
+  }
+
+  private LoadHandle rebuildAndOpen(Repository repo) throws IOException {
+    try {
+      NoteDbChangeState newState =
+          args.rebuilder.get().rebuild(args.db.get(), getChangeId());
+      if (newState == null) {
+        return super.openHandle(repo); // May be null in tests.
+      }
+      ObjectId draftsId = newState.getDraftIds().get(author);
+      repo.scanForRepoChanges();
+      return LoadHandle.create(new RevWalk(repo), draftsId);
+    } catch (NoSuchChangeException e) {
+      return super.openHandle(repo);
+    } catch (OrmException | ConfigInvalidException e) {
+      throw new IOException(e);
+    }
+  }
+
   @VisibleForTesting
   NoteMap getNoteMap() {
     return revisionNoteMap != null ? revisionNoteMap.noteMap : null;

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

@@ -1808,8 +1808,7 @@ public class ChangeNotesTest extends AbstractChangeNotesTest {
     manager.execute();
 
     // Looking at drafts directly shows the zombie comment.
-    DraftCommentNotes draftNotes =
-        draftNotesFactory.create(c.getId(), otherUserId);
+    DraftCommentNotes draftNotes = draftNotesFactory.create(c, otherUserId);
     assertThat(draftNotes.load().getComments().get(rev1))
         .containsExactly(comment1, comment2);