Re-apply: Resolve draft comment reference ambiguity for imported changes

Due to the complexity of `stable-3.9` merge to `master` the change
"Resolve draft comment reference ambiguity for imported changes`
(I108671395) was reverted in the merge commit.

This is re-applying that change separetly to `master`.

Bug: Issue 325309574
Release-Notes: skip
Forward-Compatible: checked
Change-Id: Ic481a75d036c657a34144a9c1baed522250218d7

4 files changed, 70 insertions, 12 deletions

diff --git a/java/com/google/gerrit/server/notedb/ChangeDraftNotesUpdate.java b/java/com/google/gerrit/server/notedb/ChangeDraftNotesUpdate.java
index 1edc7bd6a4..4bb347abdc 100644
--- a/java/com/google/gerrit/server/notedb/ChangeDraftNotesUpdate.java
+++ b/java/com/google/gerrit/server/notedb/ChangeDraftNotesUpdate.java
@@ -40,6 +40,7 @@ import com.google.gerrit.server.experiments.ExperimentFeaturesConstants;
 import com.google.gerrit.server.git.GitRepositoryManager;
 import com.google.gerrit.server.logging.Metadata;
 import com.google.gerrit.server.logging.TraceContext;
+import com.google.gerrit.server.query.change.ChangeNumberVirtualIdAlgorithm;
 import com.google.gerrit.server.update.BatchUpdateListener;
 import com.google.inject.Inject;
 import com.google.inject.assistedinject.Assisted;
@@ -73,6 +74,8 @@ import org.eclipse.jgit.transport.ReceiveCommand;
  * <p>This class is not thread safe.
  */
 public class ChangeDraftNotesUpdate extends AbstractChangeUpdate implements ChangeDraftUpdate {
+  private final ChangeNumberVirtualIdAlgorithm virtualIdFunc;
+
   public interface Factory extends ChangeDraftUpdateFactory {
     @Override
     ChangeDraftNotesUpdate create(
@@ -234,6 +237,7 @@ public class ChangeDraftNotesUpdate extends AbstractChangeUpdate implements Chan
       AllUsersName allUsers,
       ChangeNoteUtil noteUtil,
       ExperimentFeatures experimentFeatures,
+      @Nullable ChangeNumberVirtualIdAlgorithm virtualIdFunc,
       @Assisted ChangeNotes notes,
       @Assisted("effective") Account.Id accountId,
       @Assisted("real") Account.Id realAccountId,
@@ -242,6 +246,7 @@ public class ChangeDraftNotesUpdate extends AbstractChangeUpdate implements Chan
     super(noteUtil, serverIdent, notes, null, accountId, realAccountId, authorIdent, when);
     this.draftsProject = allUsers;
     this.experimentFeatures = experimentFeatures;
+    this.virtualIdFunc = virtualIdFunc;
   }
 
   @AssistedInject
@@ -250,6 +255,7 @@ public class ChangeDraftNotesUpdate extends AbstractChangeUpdate implements Chan
       AllUsersName allUsers,
       ChangeNoteUtil noteUtil,
       ExperimentFeatures experimentFeatures,
+      @Nullable ChangeNumberVirtualIdAlgorithm virtualIdFunc,
       @Assisted Change change,
       @Assisted("effective") Account.Id accountId,
       @Assisted("real") Account.Id realAccountId,
@@ -258,6 +264,7 @@ public class ChangeDraftNotesUpdate extends AbstractChangeUpdate implements Chan
     super(noteUtil, serverIdent, null, change, accountId, realAccountId, authorIdent, when);
     this.draftsProject = allUsers;
     this.experimentFeatures = experimentFeatures;
+    this.virtualIdFunc = virtualIdFunc;
   }
 
   @Override
@@ -320,6 +327,7 @@ public class ChangeDraftNotesUpdate extends AbstractChangeUpdate implements Chan
             draftsProject,
             noteUtil,
             experimentFeatures,
+            virtualIdFunc,
             new Change(getChange()),
             accountId,
             realAccountId,
@@ -430,7 +438,7 @@ public class ChangeDraftNotesUpdate extends AbstractChangeUpdate implements Chan
 
   @Override
   protected String getRefName() {
-    return RefNames.refsDraftComments(getId(), accountId);
+    return RefNames.refsDraftComments(getVirtualId(), accountId);
   }
 
   @Override
@@ -447,4 +455,11 @@ public class ChangeDraftNotesUpdate extends AbstractChangeUpdate implements Chan
   public boolean isEmpty() {
     return delete.isEmpty() && put.isEmpty();
   }
+
+  private Change.Id getVirtualId() {
+    Change change = getChange();
+    return virtualIdFunc == null
+        ? change.getId()
+        : virtualIdFunc.apply(change.getServerId(), change.getId());
+  }
 }
diff --git a/java/com/google/gerrit/server/notedb/ChangeNotes.java b/java/com/google/gerrit/server/notedb/ChangeNotes.java
index 35add5fff3..a2c9c2f299 100644
--- a/java/com/google/gerrit/server/notedb/ChangeNotes.java
+++ b/java/com/google/gerrit/server/notedb/ChangeNotes.java
@@ -550,11 +550,21 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
   }
 
   public ImmutableList<HumanComment> getDraftComments(Account.Id author) {
-    return getDraftComments(author, null);
+    return getDraftComments(author, null, null);
   }
 
   public ImmutableList<HumanComment> getDraftComments(Account.Id author, @Nullable Ref ref) {
-    loadDraftComments(author, ref);
+    return getDraftComments(author, null, ref);
+  }
+
+  public ImmutableList<HumanComment> getDraftComments(
+      Account.Id author, @Nullable Change.Id virtualId) {
+    return getDraftComments(author, virtualId, null);
+  }
+
+  ImmutableList<HumanComment> getDraftComments(
+      Account.Id author, @Nullable Change.Id virtualId, @Nullable Ref ref) {
+    loadDraftComments(author, virtualId, ref);
     // Filter out any zombie draft comments. These are drafts that are also in
     // the published map, and arise when the update to All-Users to delete them
     // during the publish operation failed.
@@ -573,9 +583,10 @@ public class ChangeNotes extends AbstractChangeNotes<ChangeNotes> {
    * 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, @Nullable Ref ref) {
+  private void loadDraftComments(
+      Account.Id author, @Nullable Change.Id virtualId, @Nullable Ref ref) {
     if (draftCommentNotes == null || !author.equals(draftCommentNotes.getAuthor()) || ref != null) {
-      draftCommentNotes = new DraftCommentNotes(args, getChangeId(), author, ref);
+      draftCommentNotes = new DraftCommentNotes(args, getChangeId(), virtualId, author, ref);
       draftCommentNotes.load();
     }
   }
diff --git a/java/com/google/gerrit/server/notedb/DraftCommentNotes.java b/java/com/google/gerrit/server/notedb/DraftCommentNotes.java
index 186b49a8d8..639633e6e0 100644
--- a/java/com/google/gerrit/server/notedb/DraftCommentNotes.java
+++ b/java/com/google/gerrit/server/notedb/DraftCommentNotes.java
@@ -42,8 +42,15 @@ import org.eclipse.jgit.revwalk.RevCommit;
 public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
+  private final Change.Id virtualId;
+
   public interface Factory {
     DraftCommentNotes create(Change.Id changeId, Account.Id accountId);
+
+    DraftCommentNotes create(
+        @Assisted("changeId") Change.Id changeId,
+        @Assisted("virtualId") Change.Id virtualId,
+        Account.Id accountId);
   }
 
   private final Account.Id author;
@@ -54,11 +61,26 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
 
   @AssistedInject
   DraftCommentNotes(Args args, @Assisted Change.Id changeId, @Assisted Account.Id author) {
-    this(args, changeId, author, null);
+    this(args, changeId, null, author, null);
+  }
+
+  @AssistedInject
+  DraftCommentNotes(
+      Args args,
+      @Assisted("changeId") Change.Id changeId,
+      @Assisted("virtualId") Change.Id virtualId,
+      @Assisted Account.Id author) {
+    this(args, changeId, virtualId, author, null);
   }
 
-  DraftCommentNotes(Args args, Change.Id changeId, Account.Id author, @Nullable Ref ref) {
+  DraftCommentNotes(
+      Args args,
+      Change.Id changeId,
+      @Nullable Change.Id virtualId,
+      Account.Id author,
+      @Nullable Ref ref) {
     super(args, changeId, null);
+    this.virtualId = virtualId;
     this.author = requireNonNull(author);
     this.ref = ref;
     if (ref != null) {
@@ -94,7 +116,7 @@ public class DraftCommentNotes extends AbstractChangeNotes<DraftCommentNotes> {
 
   @Override
   protected String getRefName() {
-    return refsDraftComments(getChangeId(), author);
+    return refsDraftComments(virtualId != null ? virtualId : getChangeId(), author);
   }
 
   @Override
diff --git a/java/com/google/gerrit/server/notedb/DraftCommentsNotesReader.java b/java/com/google/gerrit/server/notedb/DraftCommentsNotesReader.java
index 27c59f968f..ea3dd0a33b 100644
--- a/java/com/google/gerrit/server/notedb/DraftCommentsNotesReader.java
+++ b/java/com/google/gerrit/server/notedb/DraftCommentsNotesReader.java
@@ -26,6 +26,7 @@ import com.google.gerrit.server.DraftCommentsReader;
 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.query.change.ChangeNumberVirtualIdAlgorithm;
 import com.google.inject.Singleton;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -43,15 +44,18 @@ public class DraftCommentsNotesReader implements DraftCommentsReader {
   private final DraftCommentNotes.Factory draftCommentNotesFactory;
   private final GitRepositoryManager repoManager;
   private final AllUsersName allUsers;
+  private final ChangeNumberVirtualIdAlgorithm virtualIdAlgorithm;
 
   @Inject
   DraftCommentsNotesReader(
       DraftCommentNotes.Factory draftCommentNotesFactory,
       GitRepositoryManager repoManager,
-      AllUsersName allUsers) {
+      AllUsersName allUsers,
+      ChangeNumberVirtualIdAlgorithm virtualIdAlgorithm) {
     this.draftCommentNotesFactory = draftCommentNotesFactory;
     this.repoManager = repoManager;
     this.allUsers = allUsers;
+    this.virtualIdAlgorithm = virtualIdAlgorithm;
   }
 
   @Override
@@ -64,7 +68,7 @@ public class DraftCommentsNotesReader implements DraftCommentsReader {
 
   @Override
   public List<HumanComment> getDraftsByChangeAndDraftAuthor(ChangeNotes notes, Account.Id author) {
-    return sort(new ArrayList<>(notes.getDraftComments(author)));
+    return sort(new ArrayList<>(notes.getDraftComments(author, getVirtualId(notes))));
   }
 
   @Override
@@ -77,7 +81,7 @@ public class DraftCommentsNotesReader implements DraftCommentsReader {
   public List<HumanComment> getDraftsByPatchSetAndDraftAuthor(
       ChangeNotes notes, PatchSet.Id psId, Account.Id author) {
     return sort(
-        notes.load().getDraftComments(author).stream()
+        notes.load().getDraftComments(author, getVirtualId(notes)).stream()
             .filter(c -> c.key.patchSetId == psId.get())
             .collect(Collectors.toList()));
   }
@@ -136,7 +140,7 @@ public class DraftCommentsNotesReader implements DraftCommentsReader {
   private List<Ref> getDraftRefs(ChangeNotes notes) {
     try (Repository repo = repoManager.openRepository(allUsers)) {
       return repo.getRefDatabase()
-          .getRefsByPrefix(RefNames.refsDraftCommentsPrefix(notes.getChangeId()));
+          .getRefsByPrefix(RefNames.refsDraftCommentsPrefix(getVirtualId(notes)));
     } catch (IOException e) {
       throw new StorageException(e);
     }
@@ -145,4 +149,10 @@ public class DraftCommentsNotesReader implements DraftCommentsReader {
   private List<HumanComment> sort(List<HumanComment> comments) {
     return CommentsUtil.sort(comments);
   }
+
+  private Change.Id getVirtualId(ChangeNotes notes) {
+    return virtualIdAlgorithm == null
+        ? notes.getChangeId()
+        : virtualIdAlgorithm.apply(notes.getServerId(), notes.getChangeId());
+  }
 }