Add migration mode to support missing entries

With migration mode, the ZkSharedRefDatabase will handle failures in
updating refs caused by a missing entry in Zookeeper as the expected
consequence of the Zookeeper storage not being aware of all the
objects already in the Git repo but not stored in it.

This is an approach alternative of having a batch process required
to initialise the Zookeeper storage with all existing objects in a
given Git repo.

Feature: Issue 10554
Change-Id: I26ae9ef76e288cfe4fa86b65a33b966ce43006a6
diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/Configuration.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/Configuration.java
index 8bb2a8b..88ab53e 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/Configuration.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/Configuration.java
@@ -25,6 +25,7 @@
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
 import com.googlesource.gerrit.plugins.multisite.forwarder.events.EventFamily;
+import com.googlesource.gerrit.plugins.multisite.validation.dfsrefdb.zookeeper.ZkSharedRefDatabase;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
@@ -422,6 +423,7 @@
     private final int DEFAULT_CAS_RETRY_POLICY_BASE_SLEEP_TIME_MS = 100;
     private final int DEFAULT_CAS_RETRY_POLICY_MAX_SLEEP_TIME_MS = 300;
     private final int DEFAULT_CAS_RETRY_POLICY_MAX_RETRIES = 3;
+    private final boolean DEFAULT_MIGRATE = false;
 
     static {
       CuratorFrameworkFactory.Builder b = CuratorFrameworkFactory.builder();
@@ -441,6 +443,7 @@
     public final String KEY_CAS_RETRY_POLICY_BASE_SLEEP_TIME_MS = "casRetryPolicyBaseSleepTimeMs";
     public final String KEY_CAS_RETRY_POLICY_MAX_SLEEP_TIME_MS = "casRetryPolicyMaxSleepTimeMs";
     public final String KEY_CAS_RETRY_POLICY_MAX_RETRIES = "casRetryPolicyMaxRetries";
+    public static final String KEY_MIGRATE = "migrate";
 
     private final String connectionString;
     private final String root;
@@ -452,6 +455,7 @@
     private final int casBaseSleepTimeMs;
     private final int casMaxSleepTimeMs;
     private final int casMaxRetries;
+    private final boolean migrate;
 
     private CuratorFramework build;
 
@@ -518,6 +522,8 @@
               KEY_CAS_RETRY_POLICY_MAX_RETRIES,
               DEFAULT_CAS_RETRY_POLICY_MAX_RETRIES);
 
+      migrate = getBoolean(cfg, SplitBrain.SECTION, SUBSECTION, KEY_MIGRATE, DEFAULT_MIGRATE);
+
       checkArgument(StringUtils.isNotEmpty(connectionString), "zookeeper.%s contains no servers");
     }
 
@@ -542,6 +548,12 @@
       return new BoundedExponentialBackoffRetry(
           casBaseSleepTimeMs, casMaxSleepTimeMs, casMaxRetries);
     }
+
+    public ZkSharedRefDatabase.OperationMode getOperationMode() {
+      return migrate
+          ? ZkSharedRefDatabase.OperationMode.MIGRATION
+          : ZkSharedRefDatabase.OperationMode.NORMAL;
+    }
   }
 
   public static class SplitBrain {
diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkSharedRefDatabase.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkSharedRefDatabase.java
index c1272a7..86be4f5 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkSharedRefDatabase.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkSharedRefDatabase.java
@@ -35,11 +35,21 @@
   private final CuratorFramework client;
   private final RetryPolicy retryPolicy;
 
+  private final OperationMode operationMode;
+
+  public enum OperationMode {
+    NORMAL,
+    MIGRATION;
+  }
+
   @Inject
   public ZkSharedRefDatabase(
-      CuratorFramework client, @Named("ZkLockRetryPolicy") RetryPolicy retryPolicy) {
+      CuratorFramework client,
+      @Named("ZkLockRetryPolicy") RetryPolicy retryPolicy,
+      OperationMode operationMode) {
     this.client = client;
     this.retryPolicy = retryPolicy;
+    this.operationMode = operationMode;
   }
 
   @Override
@@ -62,6 +72,15 @@
           distributedRefValue.compareAndSet(
               writeObjectId(oldRef.getObjectId()), writeObjectId(newValue));
 
+      if (!newDistributedValue.succeeded()
+          && operationMode == OperationMode.MIGRATION
+          && refNotInZk(projectName, oldRef, newRef)) {
+        logger.atInfo().log(
+            "Missing entry in ZK for ref at path '%'. Assuming this is because we are in migration mode and creating it",
+            pathFor(projectName, oldRef, newRef));
+
+        return distributedRefValue.initialize(writeObjectId(newRef.getObjectId()));
+      }
       return newDistributedValue.succeeded();
     } catch (Exception e) {
       logger.atWarning().withCause(e).log(
@@ -73,6 +92,10 @@
     }
   }
 
+  private boolean refNotInZk(String projectName, Ref oldRef, Ref newRef) throws Exception {
+    return client.checkExists().forPath(pathFor(projectName, oldRef, newRef)) == null;
+  }
+
   static String pathFor(String projectName, Ref oldRef, Ref newRef) {
     return pathFor(projectName, MoreObjects.firstNonNull(oldRef.getName(), newRef.getName()));
   }
diff --git a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkValidationModule.java b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkValidationModule.java
index d4a3126..0b6b457 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkValidationModule.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkValidationModule.java
@@ -36,5 +36,7 @@
     bind(RetryPolicy.class)
         .annotatedWith(Names.named("ZkLockRetryPolicy"))
         .toInstance(cfg.getSplitBrain().getZookeeper().buildCasRetryPolicy());
+    bind(ZkSharedRefDatabase.OperationMode.class)
+        .toInstance(cfg.getSplitBrain().getZookeeper().getOperationMode());
   }
 }
diff --git a/src/main/resources/Documentation/config.md b/src/main/resources/Documentation/config.md
index ca3f87c..76f9ee9 100644
--- a/src/main/resources/Documentation/config.md
+++ b/src/main/resources/Documentation/config.md
@@ -211,6 +211,14 @@
 :   Configuration for the max number of retries to use to create the BoundedExponentialBackoffRetry policy
 used for the Compare and Swap operations on Zookeeper
     Defaults: 3
+
+```split-brain.zookeeper.migrate```
+:   Set to true when the plugin has been applied to an already existing module and
+there are no entries in Zookeeper for the existing refs. It will handle update failures
+caused by the old refs not existing forcing the creation of the new one
+    Defaults: false
+    
+
     
 #### Custom kafka properties:
 
diff --git a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/ValidationIT.java b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/ValidationIT.java
index fb63793..c65002a 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/ValidationIT.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/ValidationIT.java
@@ -70,7 +70,7 @@
 
     @Override
     protected void configure() {
-      ZookeeperTestContainerSupport zookeeperContainer = new ZookeeperTestContainerSupport();
+      ZookeeperTestContainerSupport zookeeperContainer = new ZookeeperTestContainerSupport(true);
       Configuration multiSiteConfig = zookeeperContainer.getConfig();
       bind(Configuration.class).toInstance(multiSiteConfig);
       install(new Module(multiSiteConfig, noteDb));
@@ -91,10 +91,8 @@
 
   @Test
   public void inSyncChangeValidatorShouldAcceptNewChange() throws Exception {
-    // FIXME: The code does not work for already existing refs (need a migration step for first
-    // run - T0). Using "refs/heads/master2" in this test for now
     final PushOneCommit.Result change =
-        createCommitAndPush(testRepo, "refs/heads/master2", "msg", "file", "content");
+        createCommitAndPush(testRepo, "refs/heads/master", "msg", "file", "content");
 
     change.assertOkStatus();
   }
diff --git a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkSharedRefDatabaseTest.java b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkSharedRefDatabaseTest.java
index b5ac1e4..88548f4 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkSharedRefDatabaseTest.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZkSharedRefDatabaseTest.java
@@ -47,9 +47,12 @@
 
   @Before
   public void setup() {
-    zookeeperContainer = new ZookeeperTestContainerSupport();
+    zookeeperContainer = new ZookeeperTestContainerSupport(false);
     zkSharedRefDatabase =
-        new ZkSharedRefDatabase(zookeeperContainer.getCurator(), new RetryNTimes(5, 30));
+        new ZkSharedRefDatabase(
+            zookeeperContainer.getCurator(),
+            new RetryNTimes(5, 30),
+            ZkSharedRefDatabase.OperationMode.NORMAL);
   }
 
   @After
@@ -59,12 +62,12 @@
 
   @Test
   public void shouldCompareAndCreateSuccessfully() throws Exception {
-  	Ref ref = refOf(AN_OBJECT_ID_1);
+    Ref ref = refOf(AN_OBJECT_ID_1);
 
-  	assertThat(zkSharedRefDatabase.compareAndCreate(A_TEST_PROJECT_NAME, ref)).isTrue();
+    assertThat(zkSharedRefDatabase.compareAndCreate(A_TEST_PROJECT_NAME, ref)).isTrue();
 
-  	assertThat(zookeeperContainer.readRefValueFromZk(A_TEST_PROJECT_NAME, ref))
-  			.isEqualTo(ref.getObjectId());
+    assertThat(zookeeperContainer.readRefValueFromZk(A_TEST_PROJECT_NAME, ref))
+        .isEqualTo(ref.getObjectId());
   }
 
   @Test
@@ -100,6 +103,24 @@
   }
 
   @Test
+  public void compareAndPutShouldDoAnInsertIfTheObjectionDoesNotExistAndInMigrationMode()
+      throws Exception {
+    zkSharedRefDatabase =
+        new ZkSharedRefDatabase(
+            zookeeperContainer.getCurator(),
+            new RetryNTimes(5, 30),
+            ZkSharedRefDatabase.OperationMode.MIGRATION);
+
+    Ref oldRef = refOf(AN_OBJECT_ID_1);
+    assertThat(
+            zkSharedRefDatabase.compareAndPut(A_TEST_PROJECT_NAME, oldRef, refOf(AN_OBJECT_ID_2)))
+        .isTrue();
+
+    assertThat(zookeeperContainer.readRefValueFromZk(A_TEST_PROJECT_NAME, oldRef))
+        .isEqualTo(AN_OBJECT_ID_2);
+  }
+
+  @Test
   public void shouldCompareAndRemoveSuccessfully() throws Exception {
     Ref oldRef = refOf(AN_OBJECT_ID_1);
     String projectName = A_TEST_PROJECT_NAME;
@@ -112,13 +133,12 @@
   @Test
   public void shouldReplaceTheRefWithATombstoneAfterCompareAndPutRemove() throws Exception {
     Ref oldRef = refOf(AN_OBJECT_ID_1);
-    String projectName = A_TEST_PROJECT_NAME;
 
-    zookeeperContainer.createRefInZk(projectName, oldRef);
+    zookeeperContainer.createRefInZk(A_TEST_PROJECT_NAME, oldRef);
 
-    assertThat(zkSharedRefDatabase.compareAndRemove(projectName, oldRef)).isTrue();
+    assertThat(zkSharedRefDatabase.compareAndRemove(A_TEST_PROJECT_NAME, oldRef)).isTrue();
 
-    assertThat(zookeeperContainer.readRefValueFromZk(projectName, oldRef))
+    assertThat(zookeeperContainer.readRefValueFromZk(A_TEST_PROJECT_NAME, oldRef))
         .isEqualTo(ObjectId.zeroId());
   }
 
diff --git a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZookeeperTestContainerSupport.java b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZookeeperTestContainerSupport.java
index 3d041ae..d479f7f 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZookeeperTestContainerSupport.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/multisite/validation/dfsrefdb/zookeeper/ZookeeperTestContainerSupport.java
@@ -68,7 +68,7 @@
   }
 
   @SuppressWarnings("resource")
-  public ZookeeperTestContainerSupport() {
+  public ZookeeperTestContainerSupport(boolean migrationMode) {
     container = new ZookeeperContainer().withExposedPorts(2181).waitingFor(Wait.forListeningPort());
     container.start();
     Integer zkHostPort = container.getMappedPort(2181);
@@ -76,6 +76,16 @@
     String connectString = "localhost:" + zkHostPort;
     splitBrainconfig.setBoolean("split-brain", null, "enabled", true);
     splitBrainconfig.setString("split-brain", "zookeeper", "connectString", connectString);
+    splitBrainconfig.setString(
+        "split-brain",
+        Configuration.Zookeeper.SUBSECTION,
+        Configuration.Zookeeper.KEY_CONNECT_STRING,
+        connectString);
+    splitBrainconfig.setBoolean(
+        "split-brain",
+        Configuration.Zookeeper.SUBSECTION,
+        Configuration.Zookeeper.KEY_MIGRATE,
+        migrationMode);
 
     configuration = new Configuration(splitBrainconfig);
     this.curator = configuration.getSplitBrain().getZookeeper().buildCurator();