Wrap ref updates in the correct RefUpdateContext.

Bug: Google b/261005032
Release-Notes: skip
Forward-Compatible: checked
Change-Id: If48c444ad172038896f43076ea8f1e03cf32afa1
diff --git a/java/com/google/gerrit/plugins/checks/acceptance/testsuite/CheckerOperationsImpl.java b/java/com/google/gerrit/plugins/checks/acceptance/testsuite/CheckerOperationsImpl.java
index c242f68..a65bf95 100644
--- a/java/com/google/gerrit/plugins/checks/acceptance/testsuite/CheckerOperationsImpl.java
+++ b/java/com/google/gerrit/plugins/checks/acceptance/testsuite/CheckerOperationsImpl.java
@@ -16,6 +16,7 @@
 
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.gerrit.testing.TestActionRefUpdateContext.testRefAction;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.stream.Collectors.toList;
 import static org.eclipse.jgit.lib.Constants.OBJ_BLOB;
@@ -301,7 +302,7 @@
             TestRepository<Repository> testRepo = new TestRepository<>(repo)) {
           RefUpdate ru = testRepo.getRepository().updateRef(checkerUuid.toRefName(), true);
           ru.setForceUpdate(true);
-          ru.delete();
+          testRefAction(() -> ru.delete());
         }
       }
     }
diff --git a/java/com/google/gerrit/plugins/checks/db/CheckerRefMigration.java b/java/com/google/gerrit/plugins/checks/db/CheckerRefMigration.java
index 40eb858..e539fe7 100644
--- a/java/com/google/gerrit/plugins/checks/db/CheckerRefMigration.java
+++ b/java/com/google/gerrit/plugins/checks/db/CheckerRefMigration.java
@@ -13,11 +13,15 @@
 // limitations under the License.
 package com.google.gerrit.plugins.checks.db;
 
+import static com.google.gerrit.server.update.context.RefUpdateContext.RefUpdateType.OFFLINE_OPERATION;
+import static com.google.gerrit.server.update.context.RefUpdateContext.RefUpdateType.PLUGIN;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.gerrit.entities.Project;
 import com.google.gerrit.pgm.init.api.AllProjectsNameOnInitProvider;
 import com.google.gerrit.plugins.checks.CheckerRef;
 import com.google.gerrit.server.git.GitRepositoryManager;
+import com.google.gerrit.server.update.context.RefUpdateContext;
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
 import java.io.IOException;
@@ -49,7 +53,9 @@
   }
 
   public void migrate() throws Exception {
-    try (Repository repo = repoManager.openRepository(allProjectsName)) {
+    try (Repository repo = repoManager.openRepository(allProjectsName);
+        RefUpdateContext pluginCtx = RefUpdateContext.open(PLUGIN);
+        RefUpdateContext ctx = RefUpdateContext.open(OFFLINE_OPERATION)) {
 
       // This part is specifically for cases where the rename failed half-way last time.
       Ref ref = repo.exactRef(TMP_REF);
@@ -68,7 +74,8 @@
           ex.getMessage()
               + " Ensure that refs/tmp/checker-migration doesn't exist. Also,"
               + "ensure that refs/meta/checkers/ (with trailing '/') has been migrated to refs/meta/checkers (without "
-              + "trailing '/'). Consider documentation for updating refs manually: https://git-scm.com/docs/git-update-ref");
+              + "trailing '/'). Consider documentation for updating refs manually: https://git-scm.com/docs/git-update-ref",
+          ex);
     }
   }
 
diff --git a/java/com/google/gerrit/plugins/checks/db/NoteDbCheckersUpdate.java b/java/com/google/gerrit/plugins/checks/db/NoteDbCheckersUpdate.java
index e84e129..45f0952 100644
--- a/java/com/google/gerrit/plugins/checks/db/NoteDbCheckersUpdate.java
+++ b/java/com/google/gerrit/plugins/checks/db/NoteDbCheckersUpdate.java
@@ -14,6 +14,9 @@
 
 package com.google.gerrit.plugins.checks.db;
 
+import static com.google.gerrit.server.update.context.RefUpdateContext.RefUpdateType.PLUGIN;
+import static com.google.gerrit.server.update.context.RefUpdateContext.RefUpdateType.VERSIONED_META_DATA_CHANGE;
+
 import com.google.common.base.Throwables;
 import com.google.gerrit.entities.Project;
 import com.google.gerrit.exceptions.DuplicateKeyException;
@@ -32,6 +35,7 @@
 import com.google.gerrit.server.git.GitRepositoryManager;
 import com.google.gerrit.server.git.meta.MetaDataUpdate;
 import com.google.gerrit.server.update.RetryHelper;
+import com.google.gerrit.server.update.context.RefUpdateContext;
 import com.google.inject.assistedinject.Assisted;
 import com.google.inject.assistedinject.AssistedInject;
 import java.io.IOException;
@@ -181,12 +185,16 @@
       CheckersByRepositoryNotes checkersByRepositoryNotes)
       throws IOException {
     BatchRefUpdate batchRefUpdate = allProjectsRepo.getRefDatabase().newBatchUpdate();
-    try (MetaDataUpdate metaDataUpdate =
-        metaDataUpdateFactory.create(allProjectsName, allProjectsRepo, batchRefUpdate)) {
-      checkerConfig.commit(metaDataUpdate);
-      checkersByRepositoryNotes.commit(metaDataUpdate);
+    try (RefUpdateContext pluginCtx = RefUpdateContext.open(PLUGIN)) {
+      try (RefUpdateContext ctx = RefUpdateContext.open(VERSIONED_META_DATA_CHANGE)) {
+        try (MetaDataUpdate metaDataUpdate =
+            metaDataUpdateFactory.create(allProjectsName, allProjectsRepo, batchRefUpdate)) {
+          checkerConfig.commit(metaDataUpdate);
+          checkersByRepositoryNotes.commit(metaDataUpdate);
+        }
+        RefUpdateUtil.executeChecked(batchRefUpdate, allProjectsRepo);
+      }
     }
-    RefUpdateUtil.executeChecked(batchRefUpdate, allProjectsRepo);
 
     gitRefUpdated.fire(
         allProjectsName, batchRefUpdate, currentUser.map(user -> user.state()).orElse(null));
diff --git a/java/com/google/gerrit/plugins/checks/db/NoteDbChecksUpdate.java b/java/com/google/gerrit/plugins/checks/db/NoteDbChecksUpdate.java
index 47de092..11e822d 100644
--- a/java/com/google/gerrit/plugins/checks/db/NoteDbChecksUpdate.java
+++ b/java/com/google/gerrit/plugins/checks/db/NoteDbChecksUpdate.java
@@ -15,6 +15,8 @@
 package com.google.gerrit.plugins.checks.db;
 
 import static com.google.gerrit.plugins.checks.CheckerRef.checksRef;
+import static com.google.gerrit.server.update.context.RefUpdateContext.RefUpdateType.CHANGE_MODIFICATION;
+import static com.google.gerrit.server.update.context.RefUpdateContext.RefUpdateType.PLUGIN;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.eclipse.jgit.lib.Constants.OBJ_BLOB;
 
@@ -36,6 +38,7 @@
 import com.google.gerrit.server.git.GitRepositoryManager;
 import com.google.gerrit.server.notedb.ChangeNoteUtil;
 import com.google.gerrit.server.update.RetryHelper;
+import com.google.gerrit.server.update.context.RefUpdateContext;
 import com.google.inject.assistedinject.Assisted;
 import com.google.inject.assistedinject.AssistedInject;
 import java.io.ByteArrayOutputStream;
@@ -171,45 +174,49 @@
     if (operation == Operation.CREATE) {
       assertCheckerIsPresent(checkKey.checkerUuid());
     }
+    try (RefUpdateContext pluginCtx = RefUpdateContext.open(PLUGIN)) {
+      try (RefUpdateContext ctx = RefUpdateContext.open(CHANGE_MODIFICATION)) {
+        try (Repository repo = repoManager.openRepository(checkKey.repository());
+            ObjectInserter objectInserter = repo.newObjectInserter();
+            RevWalk rw = new RevWalk(repo)) {
+          Ref checkRef = repo.getRefDatabase().exactRef(checksRef(checkKey.patchSet().changeId()));
+          ObjectId parent = checkRef == null ? ObjectId.zeroId() : checkRef.getObjectId();
+          CommitBuilder cb;
+          String message;
+          if (operation == Operation.CREATE) {
+            message = "Insert check " + checkKey.checkerUuid();
+            cb = commitBuilder(message, parent);
+          } else {
+            message = "Update check " + checkKey.checkerUuid();
+            cb = commitBuilder(message, parent);
+          }
 
-    try (Repository repo = repoManager.openRepository(checkKey.repository());
-        ObjectInserter objectInserter = repo.newObjectInserter();
-        RevWalk rw = new RevWalk(repo)) {
-      Ref checkRef = repo.getRefDatabase().exactRef(checksRef(checkKey.patchSet().changeId()));
-      ObjectId parent = checkRef == null ? ObjectId.zeroId() : checkRef.getObjectId();
-      CommitBuilder cb;
-      String message;
-      if (operation == Operation.CREATE) {
-        message = "Insert check " + checkKey.checkerUuid();
-        cb = commitBuilder(message, parent);
-      } else {
-        message = "Update check " + checkKey.checkerUuid();
-        cb = commitBuilder(message, parent);
+          boolean dirty =
+              updateNotesMap(
+                  checkKey, checkUpdate, repo, rw, objectInserter, parent, cb, operation);
+          if (!dirty) {
+            // This update is a NoOp, so omit writing a commit with the same tree.
+            return readSingleCheck(checkKey, repo, rw, checkRef.getObjectId());
+          }
+
+          ObjectId newCommitId = objectInserter.insert(cb);
+          objectInserter.flush();
+
+          String refName = CheckerRef.checksRef(checkKey.patchSet().changeId());
+          RefUpdate refUpdate = repo.updateRef(refName);
+          refUpdate.setExpectedOldObjectId(parent);
+          refUpdate.setNewObjectId(newCommitId);
+          refUpdate.setRefLogIdent(personIdent);
+          refUpdate.setRefLogMessage(message, false);
+          refUpdate.update();
+          RefUpdateUtil.checkResult(refUpdate);
+
+          combinedCheckStateCache.updateIfNecessary(checkKey.repository(), checkKey.patchSet());
+          gitRefUpdated.fire(
+              checkKey.repository(), refUpdate, currentUser.map(user -> user.state()).orElse(null));
+          return readSingleCheck(checkKey, repo, rw, newCommitId);
+        }
       }
-
-      boolean dirty =
-          updateNotesMap(checkKey, checkUpdate, repo, rw, objectInserter, parent, cb, operation);
-      if (!dirty) {
-        // This update is a NoOp, so omit writing a commit with the same tree.
-        return readSingleCheck(checkKey, repo, rw, checkRef.getObjectId());
-      }
-
-      ObjectId newCommitId = objectInserter.insert(cb);
-      objectInserter.flush();
-
-      String refName = CheckerRef.checksRef(checkKey.patchSet().changeId());
-      RefUpdate refUpdate = repo.updateRef(refName);
-      refUpdate.setExpectedOldObjectId(parent);
-      refUpdate.setNewObjectId(newCommitId);
-      refUpdate.setRefLogIdent(personIdent);
-      refUpdate.setRefLogMessage(message, false);
-      refUpdate.update();
-      RefUpdateUtil.checkResult(refUpdate);
-
-      combinedCheckStateCache.updateIfNecessary(checkKey.repository(), checkKey.patchSet());
-      gitRefUpdated.fire(
-          checkKey.repository(), refUpdate, currentUser.map(user -> user.state()).orElse(null));
-      return readSingleCheck(checkKey, repo, rw, newCommitId);
     }
   }
 
diff --git a/javatests/com/google/gerrit/plugins/checks/acceptance/CheckerRefsIT.java b/javatests/com/google/gerrit/plugins/checks/acceptance/CheckerRefsIT.java
index e9cadf3..dfeeb11 100644
--- a/javatests/com/google/gerrit/plugins/checks/acceptance/CheckerRefsIT.java
+++ b/javatests/com/google/gerrit/plugins/checks/acceptance/CheckerRefsIT.java
@@ -21,6 +21,7 @@
 import static com.google.gerrit.acceptance.testsuite.project.TestProjectUpdate.allowLabel;
 import static com.google.gerrit.server.group.SystemGroupBackend.REGISTERED_USERS;
 import static com.google.gerrit.testing.GerritJUnit.assertThrows;
+import static com.google.gerrit.testing.TestActionRefUpdateContext.testRefAction;
 
 import com.google.gerrit.acceptance.PushOneCommit;
 import com.google.gerrit.acceptance.SkipProjectClone;
@@ -329,7 +330,7 @@
           updateFactory.create(project, identifiedUserFactory.create(admin.id()), TimeUtil.now())) {
         bu.setRepository(git, rw, oi);
         bu.insertChange(ins);
-        bu.execute();
+        testRefAction(() -> bu.execute());
       }
       return changeId.toString();
     }