Wrap batch refupdate for shared refdb validation

Intercept the batch refupdate for making sure  that none of them
is failing and all of them are successfully updated on the shared refdb.

Change-Id: I27583731811c3b93907a95a2522c16cadb5cb7ad
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
new file mode 100644
index 0000000..d93c744
--- /dev/null
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdate.java
@@ -0,0 +1,272 @@
+// Copyright (C) 2019 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// Copyright (C) 2018 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.googlesource.gerrit.plugins.multisite.validation;
+
+import static java.util.Comparator.comparing;
+
+import com.google.inject.Inject;
+import com.google.inject.assistedinject.Assisted;
+import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.SharedRefDatabase;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.eclipse.jgit.lib.BatchRefUpdate;
+import org.eclipse.jgit.lib.PersonIdent;
+import org.eclipse.jgit.lib.ProgressMonitor;
+import org.eclipse.jgit.lib.Ref;
+import org.eclipse.jgit.lib.RefDatabase;
+import org.eclipse.jgit.revwalk.RevWalk;
+import org.eclipse.jgit.transport.PushCertificate;
+import org.eclipse.jgit.transport.ReceiveCommand;
+import org.eclipse.jgit.util.time.ProposedTimestamp;
+
+public class MultiSiteBatchRefUpdate extends BatchRefUpdate {
+  private final BatchRefUpdate batchRefUpdate;
+  private final RefDatabase refDb;
+  private final SharedRefDatabase sharedRefDb;
+  private final String projectName;
+
+  public static class RefPair {
+    final Ref oldRef;
+    final Ref newRef;
+    final Exception exception;
+
+    RefPair(Ref oldRef, Ref newRef) {
+      this.oldRef = oldRef;
+      this.newRef = newRef;
+      this.exception = null;
+    }
+
+    RefPair(Ref newRef, Exception e) {
+      this.newRef = newRef;
+      this.oldRef = SharedRefDatabase.NULL_REF;
+      this.exception = e;
+    }
+
+    public boolean hasFailed() {
+      return exception != null;
+    }
+  }
+
+  public static interface Factory {
+    MultiSiteBatchRefUpdate create(String projectName, RefDatabase refDb);
+  }
+
+  @Inject
+  public MultiSiteBatchRefUpdate(
+      SharedRefDatabase sharedRefDb, @Assisted String projectName, @Assisted RefDatabase refDb) {
+    super(refDb);
+
+    this.sharedRefDb = sharedRefDb;
+    this.projectName = projectName;
+    this.refDb = refDb;
+    this.batchRefUpdate = refDb.newBatchUpdate();
+  }
+
+  @Override
+  public int hashCode() {
+    return batchRefUpdate.hashCode();
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    return batchRefUpdate.equals(obj);
+  }
+
+  @Override
+  public boolean isAllowNonFastForwards() {
+    return batchRefUpdate.isAllowNonFastForwards();
+  }
+
+  @Override
+  public BatchRefUpdate setAllowNonFastForwards(boolean allow) {
+    return batchRefUpdate.setAllowNonFastForwards(allow);
+  }
+
+  @Override
+  public PersonIdent getRefLogIdent() {
+    return batchRefUpdate.getRefLogIdent();
+  }
+
+  @Override
+  public BatchRefUpdate setRefLogIdent(PersonIdent pi) {
+    return batchRefUpdate.setRefLogIdent(pi);
+  }
+
+  @Override
+  public String getRefLogMessage() {
+    return batchRefUpdate.getRefLogMessage();
+  }
+
+  @Override
+  public boolean isRefLogIncludingResult() {
+    return batchRefUpdate.isRefLogIncludingResult();
+  }
+
+  @Override
+  public BatchRefUpdate setRefLogMessage(String msg, boolean appendStatus) {
+    return batchRefUpdate.setRefLogMessage(msg, appendStatus);
+  }
+
+  @Override
+  public BatchRefUpdate disableRefLog() {
+    return batchRefUpdate.disableRefLog();
+  }
+
+  @Override
+  public BatchRefUpdate setForceRefLog(boolean force) {
+    return batchRefUpdate.setForceRefLog(force);
+  }
+
+  @Override
+  public boolean isRefLogDisabled() {
+    return batchRefUpdate.isRefLogDisabled();
+  }
+
+  @Override
+  public BatchRefUpdate setAtomic(boolean atomic) {
+    return batchRefUpdate.setAtomic(atomic);
+  }
+
+  @Override
+  public boolean isAtomic() {
+    return batchRefUpdate.isAtomic();
+  }
+
+  @Override
+  public void setPushCertificate(PushCertificate cert) {
+    batchRefUpdate.setPushCertificate(cert);
+  }
+
+  @Override
+  public List<ReceiveCommand> getCommands() {
+    return batchRefUpdate.getCommands();
+  }
+
+  @Override
+  public BatchRefUpdate addCommand(ReceiveCommand cmd) {
+    return batchRefUpdate.addCommand(cmd);
+  }
+
+  @Override
+  public BatchRefUpdate addCommand(ReceiveCommand... cmd) {
+    return batchRefUpdate.addCommand(cmd);
+  }
+
+  @Override
+  public BatchRefUpdate addCommand(Collection<ReceiveCommand> cmd) {
+    return batchRefUpdate.addCommand(cmd);
+  }
+
+  @Override
+  public List<String> getPushOptions() {
+    return batchRefUpdate.getPushOptions();
+  }
+
+  @Override
+  public List<ProposedTimestamp> getProposedTimestamps() {
+    return batchRefUpdate.getProposedTimestamps();
+  }
+
+  @Override
+  public BatchRefUpdate addProposedTimestamp(ProposedTimestamp ts) {
+    return batchRefUpdate.addProposedTimestamp(ts);
+  }
+
+  @Override
+  public void execute(RevWalk walk, ProgressMonitor monitor, List<String> options)
+      throws IOException {
+    updateSharedRefDb(getRefsPairs());
+    batchRefUpdate.execute(walk, monitor, options);
+  }
+
+  @Override
+  public void execute(RevWalk walk, ProgressMonitor monitor) throws IOException {
+    updateSharedRefDb(getRefsPairs());
+    batchRefUpdate.execute(walk, monitor);
+  }
+
+  @Override
+  public String toString() {
+    return batchRefUpdate.toString();
+  }
+
+  private void updateSharedRefDb(Stream<RefPair> oldRefs) throws IOException {
+    List<RefPair> refsToUpdate =
+        oldRefs.sorted(comparing(RefPair::hasFailed).reversed()).collect(Collectors.toList());
+    if (refsToUpdate.get(0).hasFailed()) {
+      RefPair failedRef = refsToUpdate.get(0);
+      throw new IOException(
+          "Failed to fetch ref entries" + failedRef.newRef.getName(), failedRef.exception);
+    }
+
+    for (RefPair refPair : refsToUpdate) {
+      boolean compareAndPutResult =
+          sharedRefDb.compareAndPut(projectName, refPair.oldRef, refPair.newRef);
+      if (!compareAndPutResult) {
+        throw new IOException(
+            String.format(
+                "This repos is out of sync for project %s. old_ref=%s, new_ref=%s",
+                projectName, refPair.oldRef, refPair.newRef));
+      }
+    }
+  }
+
+  private Stream<RefPair> getRefsPairs() {
+    return batchRefUpdate.getCommands().stream().map(this::getRefPairForCommand);
+  }
+
+  private RefPair getRefPairForCommand(ReceiveCommand command) {
+    try {
+      switch (command.getType()) {
+        case CREATE:
+          return new RefPair(SharedRefDatabase.NULL_REF, getNewRef(command));
+
+        case UPDATE:
+        case UPDATE_NONFASTFORWARD:
+          return new RefPair(refDb.getRef(command.getRefName()), getNewRef(command));
+
+        case DELETE:
+          return new RefPair(refDb.getRef(command.getRefName()), SharedRefDatabase.NULL_REF);
+
+        default:
+          return new RefPair(
+              getNewRef(command),
+              new IllegalArgumentException("Unsupported command type " + command.getType()));
+      }
+    } catch (IOException e) {
+      return new RefPair(command.getRef(), e);
+    }
+  }
+
+  private Ref getNewRef(ReceiveCommand command) {
+    return sharedRefDb.newRef(command.getRefName(), command.getNewId());
+  }
+}
diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefDatabase.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefDatabase.java
index 5ed5854..d18bb77 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefDatabase.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefDatabase.java
@@ -41,6 +41,7 @@
 
 public class MultiSiteRefDatabase extends RefDatabase {
   private final MultiSiteRefUpdate.Factory refUpdateFactory;
+  private final MultiSiteBatchRefUpdate.Factory batchRefUpdateFactory;
   private final String projectName;
   private final RefDatabase refDatabase;
 
@@ -51,9 +52,11 @@
   @Inject
   public MultiSiteRefDatabase(
       MultiSiteRefUpdate.Factory refUpdateFactory,
+      MultiSiteBatchRefUpdate.Factory batchRefUpdateFactory,
       @Assisted String projectName,
       @Assisted RefDatabase refDatabase) {
     this.refUpdateFactory = refUpdateFactory;
+    this.batchRefUpdateFactory = batchRefUpdateFactory;
     this.projectName = projectName;
     this.refDatabase = refDatabase;
   }
@@ -100,7 +103,7 @@
 
   @Override
   public BatchRefUpdate newBatchUpdate() {
-    return refDatabase.newBatchUpdate();
+    return batchRefUpdateFactory.create(projectName, refDatabase);
   }
 
   @Override
diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/ValidationModule.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/ValidationModule.java
index 284ea2c..ea12e19 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/ValidationModule.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/ValidationModule.java
@@ -33,6 +33,7 @@
     factory(MultiSiteRepository.Factory.class);
     factory(MultiSiteRefDatabase.Factory.class);
     factory(MultiSiteRefUpdate.Factory.class);
+    factory(MultiSiteBatchRefUpdate.Factory.class);
 
     if (!disableGitRepositoryValidation) {
       bind(GitRepositoryManager.class).to(MultiSiteGitRepositoryManager.class);
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
new file mode 100644
index 0000000..c23135d
--- /dev/null
+++ b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteBatchRefUpdateTest.java
@@ -0,0 +1,102 @@
+// Copyright (C) 2019 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// Copyright (C) 2018 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.googlesource.gerrit.plugins.multisite.validation;
+
+import static org.mockito.Mockito.doReturn;
+
+import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.SharedRefDatabase;
+import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.zookeeper.RefFixture;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import org.eclipse.jgit.lib.BatchRefUpdate;
+import org.eclipse.jgit.lib.ObjectIdRef;
+import org.eclipse.jgit.lib.ProgressMonitor;
+import org.eclipse.jgit.lib.Ref;
+import org.eclipse.jgit.lib.RefDatabase;
+import org.eclipse.jgit.revwalk.RevWalk;
+import org.eclipse.jgit.transport.ReceiveCommand;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestName;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+@RunWith(MockitoJUnitRunner.class)
+public class MultiSiteBatchRefUpdateTest implements RefFixture {
+
+  @Mock SharedRefDatabase sharedRefDb;
+  @Mock BatchRefUpdate batchRefUpdate;
+  @Mock RefDatabase refDatabase;
+  @Mock RevWalk revWalk;
+  @Mock ProgressMonitor progressMonitor;
+
+  private final Ref oldRef =
+      new ObjectIdRef.Unpeeled(Ref.Storage.NETWORK, A_TEST_REF_NAME, AN_OBJECT_ID_1);
+  private final Ref newRef =
+      new ObjectIdRef.Unpeeled(Ref.Storage.NETWORK, A_TEST_REF_NAME, AN_OBJECT_ID_2);
+  ReceiveCommand receiveCommand =
+      new ReceiveCommand(oldRef.getObjectId(), newRef.getObjectId(), oldRef.getName());
+
+  private MultiSiteBatchRefUpdate multiSiteRefUpdate;
+
+  @Rule public TestName nameRule = new TestName();
+
+  @Override
+  public String testBranch() {
+    return "branch_" + nameRule.getMethodName();
+  }
+
+  private void setMockRequiredReturnValues() throws IOException {
+    doReturn(batchRefUpdate).when(refDatabase).newBatchUpdate();
+    doReturn(Arrays.asList(receiveCommand)).when(batchRefUpdate).getCommands();
+    doReturn(oldRef).when(refDatabase).getRef(A_TEST_REF_NAME);
+    doReturn(newRef).when(sharedRefDb).newRef(A_TEST_REF_NAME, AN_OBJECT_ID_2);
+
+    multiSiteRefUpdate = new MultiSiteBatchRefUpdate(sharedRefDb, A_TEST_PROJECT_NAME, refDatabase);
+  }
+
+  @Test
+  public void executeAndDelegateSuccessfullyWithNoExceptions() throws IOException {
+    setMockRequiredReturnValues();
+
+    // When compareAndPut against sharedDb succeeds
+    doReturn(true).when(sharedRefDb).compareAndPut(A_TEST_PROJECT_NAME, oldRef, newRef);
+    multiSiteRefUpdate.execute(revWalk, progressMonitor, Collections.emptyList());
+  }
+
+  @Test(expected = IOException.class)
+  public void executeAndFailsWithExceptions() throws IOException {
+    setMockRequiredReturnValues();
+
+    // When compareAndPut against sharedDb fails
+    doReturn(false).when(sharedRefDb).compareAndPut(A_TEST_PROJECT_NAME, oldRef, newRef);
+    multiSiteRefUpdate.execute(revWalk, progressMonitor, Collections.emptyList());
+  }
+}
diff --git a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefDatabaseTest.java b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefDatabaseTest.java
index 36d2082..9916199 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefDatabaseTest.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/MultiSiteRefDatabaseTest.java
@@ -46,6 +46,7 @@
   @Rule public TestName nameRule = new TestName();
 
   @Mock MultiSiteRefUpdate.Factory refUpdateFactoryMock;
+  @Mock MultiSiteBatchRefUpdate.Factory refBatchUpdateFactoryMock;
 
   @Mock RefDatabase refDatabaseMock;
 
@@ -60,7 +61,8 @@
   public void newUpdateShouldCreateMultiSiteRefUpdate() throws Exception {
     String refName = aBranchRef();
     MultiSiteRefDatabase multiSiteRefDb =
-        new MultiSiteRefDatabase(refUpdateFactoryMock, A_TEST_PROJECT_NAME, refDatabaseMock);
+        new MultiSiteRefDatabase(
+            refUpdateFactoryMock, refBatchUpdateFactoryMock, A_TEST_PROJECT_NAME, refDatabaseMock);
     doReturn(refUpdateMock).when(refDatabaseMock).newUpdate(refName, false);
 
     multiSiteRefDb.newUpdate(refName, false);