Merge branch 'stable-2.16' into stable-3.0 * stable-2.16: Rollback ref update when global-refdb update failure Delegate canPerformGC functionality to the wrapped instance Change-Id: I5bb34b8dc8de51badaf8156a1e93d34f97333b16

diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/BatchRefUpdateValidator.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/BatchRefUpdateValidator.java
index cfe585d..7aa3328 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/BatchRefUpdateValidator.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/BatchRefUpdateValidator.java

@@ -57,7 +57,9 @@
   }
 
   public void executeBatchUpdateWithValidation(
-      BatchRefUpdate batchRefUpdate, NoParameterVoidFunction batchRefUpdateFunction)
+      BatchRefUpdate batchRefUpdate,
+      NoParameterVoidFunction batchRefUpdateFunction,
+      OneParameterVoidFunction<List<ReceiveCommand>> batchRefUpdateRollbackFunction)
       throws IOException {
     if (refEnforcement.getPolicy(projectName) == EnforcePolicy.IGNORED) {
       batchRefUpdateFunction.invoke();
@@ -65,7 +67,7 @@
     }
 
     try {
-      doExecuteBatchUpdate(batchRefUpdate, batchRefUpdateFunction);
+      doExecuteBatchUpdate(batchRefUpdate, batchRefUpdateFunction, batchRefUpdateRollbackFunction);
     } catch (IOException e) {
       logger.atWarning().withCause(e).log(
           "Failed to execute Batch Update on project %s", projectName);
@@ -76,7 +78,10 @@
   }
 
   private void doExecuteBatchUpdate(
-      BatchRefUpdate batchRefUpdate, NoParameterVoidFunction delegateUpdate) throws IOException {
+      BatchRefUpdate batchRefUpdate,
+      NoParameterVoidFunction delegateUpdate,
+      OneParameterVoidFunction<List<ReceiveCommand>> delegateUpdateRollback)
+      throws IOException {
 
     List<ReceiveCommand> commands = batchRefUpdate.getCommands();
     if (commands.isEmpty()) {
@@ -98,9 +103,18 @@
     }
 
     try (CloseableSet<AutoCloseable> locks = new CloseableSet<>()) {
-      refsToUpdate = compareAndGetLatestLocalRefs(refsToUpdate, locks);
+      final List<RefPair> finalRefsToUpdate = compareAndGetLatestLocalRefs(refsToUpdate, locks);
       delegateUpdate.invoke();
-      updateSharedRefDb(batchRefUpdate.getCommands().stream(), refsToUpdate);
+      try {
+        updateSharedRefDb(batchRefUpdate.getCommands().stream(), finalRefsToUpdate);
+      } catch (Exception e) {
+        List<ReceiveCommand> receiveCommands = batchRefUpdate.getCommands();
+        logger.atWarning().withCause(e).log(
+            String.format(
+                "Batch ref-update failing because of failure during the global refdb update. Set all commands Result to LOCK_FAILURE [%d]",
+                receiveCommands.size()));
+        rollback(delegateUpdateRollback, finalRefsToUpdate, receiveCommands);
+      }
     } catch (OutOfSyncException e) {
       List<ReceiveCommand> receiveCommands = batchRefUpdate.getCommands();
       logger.atWarning().withCause(e).log(
@@ -111,6 +125,24 @@
     }
   }
 
+  private void rollback(
+      OneParameterVoidFunction<List<ReceiveCommand>> delegateUpdateRollback,
+      List<RefPair> refsBeforeUpdate,
+      List<ReceiveCommand> receiveCommands)
+      throws IOException {
+    List<ReceiveCommand> rollbackCommands =
+        refsBeforeUpdate.stream()
+            .map(
+                refBeforeUpdate ->
+                    new ReceiveCommand(
+                        refBeforeUpdate.putValue,
+                        refBeforeUpdate.compareRef.getObjectId(),
+                        refBeforeUpdate.getName()))
+            .collect(Collectors.toList());
+    delegateUpdateRollback.invoke(rollbackCommands);
+    receiveCommands.forEach(command -> command.setResult(ReceiveCommand.Result.LOCK_FAILURE));
+  }
+
   private void updateSharedRefDb(Stream<ReceiveCommand> commandStream, List<RefPair> refsToUpdate)
       throws IOException {
     if (commandStream

diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdate.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdate.java
index c152b12..fd4f910 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdate.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdate.java

@@ -31,6 +31,7 @@
 public class MultiSiteBatchRefUpdate extends BatchRefUpdate {
 
   private final BatchRefUpdate batchRefUpdate;
+  private final BatchRefUpdate batchRefUpdateRollback;
   private final String project;
   private final BatchRefUpdateValidator.Factory batchRefValidatorFactory;
   private final RefDatabase refDb;
@@ -48,6 +49,7 @@
     this.refDb = refDb;
     this.project = project;
     this.batchRefUpdate = refDb.newBatchUpdate();
+    this.batchRefUpdateRollback = refDb.newBatchUpdate();
     this.batchRefValidatorFactory = batchRefValidatorFactory;
   }
 
@@ -68,6 +70,7 @@
 
   @Override
   public BatchRefUpdate setAllowNonFastForwards(boolean allow) {
+    batchRefUpdateRollback.setAllowNonFastForwards(allow);
     return batchRefUpdate.setAllowNonFastForwards(allow);
   }
 
@@ -78,6 +81,7 @@
 
   @Override
   public BatchRefUpdate setRefLogIdent(PersonIdent pi) {
+    batchRefUpdateRollback.setRefLogIdent(pi);
     return batchRefUpdate.setRefLogIdent(pi);
   }
 
@@ -93,6 +97,7 @@
 
   @Override
   public BatchRefUpdate setRefLogMessage(String msg, boolean appendStatus) {
+    batchRefUpdateRollback.setRefLogMessage(msg, appendStatus);
     return batchRefUpdate.setRefLogMessage(msg, appendStatus);
   }
 
@@ -103,6 +108,7 @@
 
   @Override
   public BatchRefUpdate setForceRefLog(boolean force) {
+    batchRefUpdateRollback.setForceRefLog(force);
     return batchRefUpdate.setForceRefLog(force);
   }
 
@@ -167,7 +173,10 @@
     batchRefValidatorFactory
         .create(project, refDb)
         .executeBatchUpdateWithValidation(
-            batchRefUpdate, () -> batchRefUpdate.execute(walk, monitor, options));
+            batchRefUpdate,
+            () -> batchRefUpdate.execute(walk, monitor, options),
+            (commands) ->
+                batchRefUpdateRollback.addCommand(commands).execute(walk, monitor, options));
   }
 
   @Override
@@ -175,7 +184,9 @@
     batchRefValidatorFactory
         .create(project, refDb)
         .executeBatchUpdateWithValidation(
-            batchRefUpdate, () -> batchRefUpdate.execute(walk, monitor));
+            batchRefUpdate,
+            () -> batchRefUpdate.execute(walk, monitor),
+            (commands) -> batchRefUpdateRollback.addCommand(commands).execute(walk, monitor));
   }
 
   @Override

diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteGitRepositoryManager.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteGitRepositoryManager.java
index 8ce1f59..594b4a0 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteGitRepositoryManager.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteGitRepositoryManager.java

@@ -51,6 +51,11 @@
   }
 
   @Override
+  public Boolean canPerformGC() {
+    return gitRepositoryManager.canPerformGC();
+  }
+
+  @Override
   public SortedSet<Project.NameKey> list() {
     return gitRepositoryManager.list();
   }

diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefUpdate.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefUpdate.java
index eadb633..20e12fe 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefUpdate.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefUpdate.java

@@ -16,6 +16,7 @@
 
 import com.google.inject.Inject;
 import com.google.inject.assistedinject.Assisted;
+import com.googlesource.gerrit.plugins.multisite.validation.RefUpdateValidator.NoParameterFunction;
 import java.io.IOException;
 import org.eclipse.jgit.lib.AnyObjectId;
 import org.eclipse.jgit.lib.ObjectId;
@@ -92,22 +93,34 @@
 
   @Override
   public Result update() throws IOException {
-    return refUpdateValidator.executeRefUpdate(refUpdateBase, refUpdateBase::update);
+    return refUpdateValidator.executeRefUpdate(
+        refUpdateBase,
+        refUpdateBase::update,
+        objectId -> rollback(objectId, refUpdateBase::update));
   }
 
   @Override
   public Result update(RevWalk rev) throws IOException {
-    return refUpdateValidator.executeRefUpdate(refUpdateBase, () -> refUpdateBase.update(rev));
+    return refUpdateValidator.executeRefUpdate(
+        refUpdateBase,
+        () -> refUpdateBase.update(rev),
+        objectId -> rollback(objectId, () -> refUpdateBase.update(rev)));
   }
 
   @Override
   public Result delete() throws IOException {
-    return refUpdateValidator.executeRefUpdate(refUpdateBase, refUpdateBase::delete);
+    return refUpdateValidator.executeRefUpdate(
+        refUpdateBase,
+        refUpdateBase::delete,
+        objectId -> rollback(objectId, refUpdateBase::update));
   }
 
   @Override
   public Result delete(RevWalk walk) throws IOException {
-    return refUpdateValidator.executeRefUpdate(refUpdateBase, () -> refUpdateBase.delete(walk));
+    return refUpdateValidator.executeRefUpdate(
+        refUpdateBase,
+        () -> refUpdateBase.delete(walk),
+        objectId -> rollback(objectId, () -> refUpdateBase.update(walk)));
   }
 
   @Override
@@ -222,7 +235,10 @@
 
   @Override
   public Result forceUpdate() throws IOException {
-    return refUpdateValidator.executeRefUpdate(refUpdateBase, refUpdateBase::forceUpdate);
+    return refUpdateValidator.executeRefUpdate(
+        refUpdateBase,
+        refUpdateBase::forceUpdate,
+        objectId -> rollback(objectId, refUpdateBase::forceUpdate));
   }
 
   @Override
@@ -234,4 +250,11 @@
   public void setCheckConflicting(boolean check) {
     refUpdateBase.setCheckConflicting(check);
   }
+
+  private Result rollback(ObjectId objectId, NoParameterFunction<Result> updateFunction)
+      throws IOException {
+    refUpdateBase.setExpectedOldObjectId(refUpdateBase.getNewObjectId());
+    refUpdateBase.setNewObjectId(objectId);
+    return updateFunction.invoke();
+  }
 }

diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/RefUpdateValidator.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/RefUpdateValidator.java
index 76c1a67..aec5619 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/RefUpdateValidator.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/RefUpdateValidator.java

@@ -34,6 +34,7 @@
 import org.eclipse.jgit.lib.Ref;
 import org.eclipse.jgit.lib.RefDatabase;
 import org.eclipse.jgit.lib.RefUpdate;
+import org.eclipse.jgit.lib.RefUpdate.Result;
 
 public class RefUpdateValidator {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
@@ -67,6 +68,14 @@
     void invoke() throws IOException;
   }
 
+  public interface OneParameterFunction<F, T> {
+    T invoke(F f) throws IOException;
+  }
+
+  public interface OneParameterVoidFunction<T> {
+    void invoke(T f) throws IOException;
+  }
+
   @Inject
   public RefUpdateValidator(
       SharedRefDatabaseWrapper sharedRefDb,
@@ -84,26 +93,16 @@
   }
 
   public RefUpdate.Result executeRefUpdate(
-      RefUpdate refUpdate, NoParameterFunction<RefUpdate.Result> refUpdateFunction)
+      RefUpdate refUpdate,
+      NoParameterFunction<RefUpdate.Result> refUpdateFunction,
+      OneParameterFunction<ObjectId, Result> rollbackFunction)
       throws IOException {
     if (isProjectVersionUpdate(refUpdate.getName())
         || refEnforcement.getPolicy(projectName) == EnforcePolicy.IGNORED) {
       return refUpdateFunction.invoke();
     }
 
-    try {
-      return doExecuteRefUpdate(refUpdate, refUpdateFunction);
-    } catch (SharedDbSplitBrainException e) {
-      validationMetrics.incrementSplitBrain();
-
-      logger.atWarning().withCause(e).log(
-          "Unable to execute ref-update on project=%s ref=%s",
-          projectName, refUpdate.getRef().getName());
-      if (refEnforcement.getPolicy(projectName) == EnforcePolicy.REQUIRED) {
-        throw e;
-      }
-    }
-    return null;
+    return doExecuteRefUpdate(refUpdate, refUpdateFunction, rollbackFunction);
   }
 
   private Boolean isProjectVersionUpdate(String refName) {
@@ -125,14 +124,27 @@
   }
 
   protected RefUpdate.Result doExecuteRefUpdate(
-      RefUpdate refUpdate, NoParameterFunction<RefUpdate.Result> refUpdateFunction)
+      RefUpdate refUpdate,
+      NoParameterFunction<Result> refUpdateFunction,
+      OneParameterFunction<ObjectId, Result> rollbackFunction)
       throws IOException {
     try (CloseableSet<AutoCloseable> locks = new CloseableSet<>()) {
       RefPair refPairForUpdate = newRefPairFrom(refUpdate);
       compareAndGetLatestLocalRef(refPairForUpdate, locks);
       RefUpdate.Result result = refUpdateFunction.invoke();
-      if (isSuccessful(result)) {
-        updateSharedDbOrThrowExceptionFor(refPairForUpdate);
+      try {
+        if (isSuccessful(result)) {
+          updateSharedDbOrThrowExceptionFor(refPairForUpdate);
+        }
+      } catch (Exception e) {
+        result = rollbackFunction.invoke(refPairForUpdate.compareRef.getObjectId());
+        if (isSuccessful(result)) {
+          result = RefUpdate.Result.LOCK_FAILURE;
+        }
+        logger.atSevere().withCause(e).log(
+            String.format(
+                "Failed to update global refdb, the local refdb has been rolled back: %s",
+                e.getMessage()));
       }
       return result;
     } catch (OutOfSyncException e) {
@@ -166,7 +178,7 @@
           projectName,
           refPair.getName(),
           e.getMessage());
-      throw new SharedDbSplitBrainException(errorMessage, e);
+      throw e;
     }
 
     if (!succeeded) {

diff --git a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/BatchRefUpdateValidatorTest.java b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/BatchRefUpdateValidatorTest.java
index 0960cf1..17e71ce 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/BatchRefUpdateValidatorTest.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/BatchRefUpdateValidatorTest.java

@@ -25,6 +25,7 @@
 import com.google.gerrit.metrics.DisabledMetricMaker;
 import com.google.gerrit.reviewdb.client.Project;
 import com.googlesource.gerrit.plugins.multisite.SharedRefDatabaseWrapper;
+import com.googlesource.gerrit.plugins.multisite.validation.RefUpdateValidator.OneParameterVoidFunction;
 import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.DefaultSharedRefEnforcement;
 import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.RefFixture;
 import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.SharedRefEnforcement;
@@ -43,6 +44,7 @@
 import org.eclipse.jgit.revwalk.RevCommit;
 import org.eclipse.jgit.revwalk.RevWalk;
 import org.eclipse.jgit.transport.ReceiveCommand;
+import org.eclipse.jgit.transport.ReceiveCommand.Result;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -64,6 +66,7 @@
   @Mock SharedRefDatabaseWrapper sharedRefDatabase;
 
   @Mock SharedRefEnforcement tmpRefEnforcement;
+  @Mock OneParameterVoidFunction<List<ReceiveCommand>> rollbackFunction;
 
   @Before
   public void setup() throws Exception {
@@ -90,7 +93,7 @@
     BatchRefUpdateValidator BatchRefUpdateValidator = newDefaultValidator(A_TEST_PROJECT_NAME);
 
     BatchRefUpdateValidator.executeBatchUpdateWithValidation(
-        batchRefUpdate, () -> execute(batchRefUpdate));
+        batchRefUpdate, () -> execute(batchRefUpdate), this::defaultRollback);
 
     verify(sharedRefDatabase, never())
         .compareAndPut(any(Project.NameKey.class), any(Ref.class), any(ObjectId.class));
@@ -105,7 +108,7 @@
     BatchRefUpdateValidator BatchRefUpdateValidator = newDefaultValidator(A_TEST_PROJECT_NAME);
 
     BatchRefUpdateValidator.executeBatchUpdateWithValidation(
-        batchRefUpdate, () -> execute(batchRefUpdate));
+        batchRefUpdate, () -> execute(batchRefUpdate), this::defaultRollback);
 
     verify(sharedRefDatabase, never())
         .compareAndPut(A_TEST_PROJECT_NAME_KEY, newRef(DRAFT_COMMENT, A.getId()), B.getId());
@@ -129,7 +132,9 @@
         .isUpToDate(A_TEST_PROJECT_NAME_KEY, newRef(AN_OUT_OF_SYNC_REF, AN_OBJECT_ID_1));
 
     batchRefUpdateValidator.executeBatchUpdateWithValidation(
-        batchRefUpdate, () -> execute(batchRefUpdate));
+        batchRefUpdate, () -> execute(batchRefUpdate), rollbackFunction);
+
+    verify(rollbackFunction, never()).invoke(any());
 
     final List<ReceiveCommand> commands = batchRefUpdate.getCommands();
     assertThat(commands.size()).isEqualTo(1);
@@ -162,6 +167,7 @@
   private BatchRefUpdate newBatchUpdate(List<ReceiveCommand> cmds) {
     BatchRefUpdate u = refdir.newBatchUpdate();
     u.addCommand(cmds);
+    cmds.forEach(c -> c.setResult(Result.OK));
     return u;
   }
 
@@ -169,4 +175,8 @@
   public String testBranch() {
     return "branch_" + nameRule.getMethodName();
   }
+
+  private void defaultRollback(List<ReceiveCommand> cmds) throws IOException {
+    // do nothing
+  }
 }

diff --git a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdateTest.java b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdateTest.java
index 187f686..622cf9d 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdateTest.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdateTest.java

@@ -25,10 +25,12 @@
 import static org.mockito.Mockito.when;
 
 import com.googlesource.gerrit.plugins.multisite.SharedRefDatabaseWrapper;
+import com.googlesource.gerrit.plugins.multisite.validation.RefUpdateValidator.OneParameterVoidFunction;
 import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.DefaultSharedRefEnforcement;
 import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.RefFixture;
 import java.io.IOException;
 import java.util.Collections;
+import java.util.List;
 import org.eclipse.jgit.lib.BatchRefUpdate;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.ObjectIdRef;
@@ -56,6 +58,7 @@
   @Mock RevWalk revWalk;
   @Mock ProgressMonitor progressMonitor;
   @Mock ValidationMetrics validationMetrics;
+  @Mock OneParameterVoidFunction<List<ReceiveCommand>> rollbackFunction;
 
   private final Ref oldRef =
       new ObjectIdRef.Unpeeled(Ref.Storage.NETWORK, A_TEST_REF_NAME, AN_OBJECT_ID_1);
@@ -130,7 +133,7 @@
     multiSiteRefUpdate = getMultiSiteBatchRefUpdateWithMockedValidator();
     doThrow(new IOException("IO Test Exception"))
         .when(batchRefUpdateValidator)
-        .executeBatchUpdateWithValidation(any(), any());
+        .executeBatchUpdateWithValidation(any(), any(), any());
 
     multiSiteRefUpdate.execute(revWalk, progressMonitor, Collections.emptyList());
   }

diff --git a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/RefUpdateValidatorTest.java b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/RefUpdateValidatorTest.java
index a9968c4..09e60d3 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/RefUpdateValidatorTest.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/RefUpdateValidatorTest.java

@@ -19,13 +19,16 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import com.gerritforge.gerrit.globalrefdb.GlobalRefDbSystemError;
 import com.google.gerrit.reviewdb.client.Project;
 import com.googlesource.gerrit.plugins.multisite.SharedRefDatabaseWrapper;
+import com.googlesource.gerrit.plugins.multisite.validation.RefUpdateValidator.OneParameterFunction;
 import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.DefaultSharedRefEnforcement;
 import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.RefFixture;
-import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.SharedDbSplitBrainException;
 import org.eclipse.jgit.lib.ObjectId;
 import org.eclipse.jgit.lib.Ref;
 import org.eclipse.jgit.lib.RefDatabase;
@@ -50,6 +53,8 @@
 
   @Mock RefUpdate refUpdate;
 
+  @Mock OneParameterFunction<ObjectId, Result> rollbackFunction;
+
   String refName;
   Ref oldUpdateRef;
   Ref newUpdateRef;
@@ -66,10 +71,10 @@
 
     doReturn(localRef).when(localRefDb).getRef(refName);
     doReturn(localRef).when(localRefDb).exactRef(refName);
-    doReturn(oldUpdateRef).when(refUpdate).getRef();
     doReturn(newUpdateRef.getObjectId()).when(refUpdate).getNewObjectId();
     doReturn(refName).when(refUpdate).getName();
     lenient().doReturn(oldUpdateRef.getObjectId()).when(refUpdate).getOldObjectId();
+    doReturn(Result.FAST_FORWARD).when(rollbackFunction).invoke(any());
 
     refUpdateValidator =
         new RefUpdateValidator(
@@ -96,7 +101,9 @@
         .when(sharedRefDb)
         .compareAndPut(A_TEST_PROJECT_NAME_KEY, localRef, newUpdateRef.getObjectId());
 
-    Result result = refUpdateValidator.executeRefUpdate(refUpdate, () -> RefUpdate.Result.NEW);
+    Result result =
+        refUpdateValidator.executeRefUpdate(
+            refUpdate, () -> RefUpdate.Result.NEW, this::defaultRollback);
 
     assertThat(result).isEqualTo(RefUpdate.Result.NEW);
   }
@@ -114,7 +121,9 @@
         .compareAndPut(A_TEST_PROJECT_NAME_KEY, localRef, ObjectId.zeroId());
     doReturn(localRef).doReturn(null).when(localRefDb).getRef(refName);
 
-    Result result = refUpdateValidator.executeRefUpdate(refUpdate, () -> RefUpdate.Result.FORCED);
+    Result result =
+        refUpdateValidator.executeRefUpdate(
+            refUpdate, () -> RefUpdate.Result.FORCED, this::defaultRollback);
 
     assertThat(result).isEqualTo(RefUpdate.Result.FORCED);
   }
@@ -133,7 +142,9 @@
         .compareAndPut(A_TEST_PROJECT_NAME_KEY, localNullRef, newUpdateRef.getObjectId());
     doReturn(localNullRef).doReturn(newUpdateRef).when(localRefDb).getRef(refName);
 
-    Result result = refUpdateValidator.executeRefUpdate(refUpdate, () -> RefUpdate.Result.NEW);
+    Result result =
+        refUpdateValidator.executeRefUpdate(
+            refUpdate, () -> RefUpdate.Result.NEW, this::defaultRollback);
 
     assertThat(result).isEqualTo(RefUpdate.Result.NEW);
   }
@@ -147,13 +158,15 @@
     doReturn(true).when(sharedRefDb).exists(A_TEST_PROJECT_NAME_KEY, refName);
     doReturn(false).when(sharedRefDb).isUpToDate(A_TEST_PROJECT_NAME_KEY, localRef);
 
-    Result result = refUpdateValidator.executeRefUpdate(refUpdate, () -> RefUpdate.Result.NEW);
+    Result result =
+        refUpdateValidator.executeRefUpdate(
+            refUpdate, () -> RefUpdate.Result.NEW, this::defaultRollback);
 
     assertThat(result).isEqualTo(Result.LOCK_FAILURE);
   }
 
-  @Test(expected = SharedDbSplitBrainException.class)
-  public void shouldTrowSplitBrainWhenLocalRefDbIsUpToDateButFinalCompareAndPutIsFailing()
+  @Test
+  public void shouldRollbackWhenLocalRefDbIsUpToDateButFinalCompareAndPutIsFailing()
       throws Exception {
     lenient()
         .doReturn(false)
@@ -168,7 +181,11 @@
         .when(sharedRefDb)
         .compareAndPut(A_TEST_PROJECT_NAME_KEY, localRef, newUpdateRef.getObjectId());
 
-    refUpdateValidator.executeRefUpdate(refUpdate, () -> RefUpdate.Result.NEW);
+    Result result =
+        refUpdateValidator.executeRefUpdate(refUpdate, () -> Result.NEW, rollbackFunction);
+
+    verify(rollbackFunction, times(1)).invoke(any());
+    assertThat(result).isEqualTo(Result.LOCK_FAILURE);
   }
 
   @Test
@@ -180,10 +197,32 @@
     doReturn(true).when(sharedRefDb).isUpToDate(A_TEST_PROJECT_NAME_KEY, localRef);
 
     Result result =
-        refUpdateValidator.executeRefUpdate(refUpdate, () -> RefUpdate.Result.LOCK_FAILURE);
+        refUpdateValidator.executeRefUpdate(
+            refUpdate, () -> Result.LOCK_FAILURE, this::defaultRollback);
 
     verify(sharedRefDb, never())
         .compareAndPut(any(Project.NameKey.class), any(Ref.class), any(ObjectId.class));
     assertThat(result).isEqualTo(RefUpdate.Result.LOCK_FAILURE);
   }
+
+  @Test
+  public void shouldRollbackRefUpdateWhenCompareAndPutIsFailing() throws Exception {
+    lenient()
+        .doReturn(false)
+        .when(sharedRefDb)
+        .isUpToDate(any(Project.NameKey.class), any(Ref.class));
+    doReturn(true).when(sharedRefDb).isUpToDate(A_TEST_PROJECT_NAME_KEY, localRef);
+
+    when(sharedRefDb.compareAndPut(any(Project.NameKey.class), any(Ref.class), any(ObjectId.class)))
+        .thenThrow(GlobalRefDbSystemError.class);
+    when(rollbackFunction.invoke(any())).thenReturn(Result.LOCK_FAILURE);
+
+    refUpdateValidator.executeRefUpdate(refUpdate, () -> Result.NEW, rollbackFunction);
+
+    verify(rollbackFunction, times(1)).invoke(any());
+  }
+
+  private Result defaultRollback(ObjectId objectId) {
+    return Result.NO_CHANGE;
+  }
 }