/*
 * Licensed to the Apache Software Foundation (ASF) under one or more contributor license
 * agreements. See the NOTICE file distributed with this work for additional information regarding
 * copyright ownership. The ASF licenses this file to You 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 org.apache.geode.internal.cache.partitioned;

import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.MINUTES;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.apache.geode.admin.AdminDistributedSystemFactory.defineDistributedSystem;
import static org.apache.geode.admin.AdminDistributedSystemFactory.getDistributedSystem;
import static org.apache.geode.cache.EvictionAction.OVERFLOW_TO_DISK;
import static org.apache.geode.cache.EvictionAttributes.createLRUEntryAttributes;
import static org.apache.geode.cache.Region.SEPARATOR;
import static org.apache.geode.cache.RegionShortcut.PARTITION_OVERFLOW;
import static org.apache.geode.cache.RegionShortcut.PARTITION_PERSISTENT;
import static org.apache.geode.cache.RegionShortcut.PARTITION_PROXY;
import static org.apache.geode.cache.RegionShortcut.REPLICATE;
import static org.apache.geode.cache.RegionShortcut.REPLICATE_PERSISTENT;
import static org.apache.geode.distributed.ConfigurationProperties.SERIALIZABLE_OBJECT_FILTER;
import static org.apache.geode.test.awaitility.GeodeAwaitility.await;
import static org.apache.geode.test.dunit.IgnoredException.addIgnoredException;
import static org.apache.geode.test.dunit.Invoke.invokeInEveryVM;
import static org.apache.geode.test.dunit.VM.getVM;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatCode;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assertions.catchThrowable;
import static org.junit.Assert.assertEquals;

import java.io.Serializable;
import java.net.InetAddress;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Properties;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import junitparams.Parameters;
import junitparams.naming.TestCaseName;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;

import org.apache.geode.admin.AdminDistributedSystem;
import org.apache.geode.admin.AdminException;
import org.apache.geode.admin.DistributedSystemConfig;
import org.apache.geode.admin.internal.AdminDistributedSystemImpl;
import org.apache.geode.cache.Cache;
import org.apache.geode.cache.CacheClosedException;
import org.apache.geode.cache.DiskAccessException;
import org.apache.geode.cache.PartitionAttributesFactory;
import org.apache.geode.cache.Region;
import org.apache.geode.cache.RegionFactory;
import org.apache.geode.cache.execute.Function;
import org.apache.geode.cache.execute.FunctionContext;
import org.apache.geode.cache.execute.FunctionService;
import org.apache.geode.cache.persistence.ConflictingPersistentDataException;
import org.apache.geode.cache.persistence.PartitionOfflineException;
import org.apache.geode.cache.persistence.PersistentID;
import org.apache.geode.cache.persistence.RevokeFailedException;
import org.apache.geode.cache.persistence.RevokedPersistentDataException;
import org.apache.geode.distributed.DistributedSystemDisconnectedException;
import org.apache.geode.distributed.internal.ClusterDistributionManager;
import org.apache.geode.distributed.internal.DistributionMessage;
import org.apache.geode.distributed.internal.DistributionMessageObserver;
import org.apache.geode.distributed.internal.InternalDistributedSystem;
import org.apache.geode.distributed.internal.ReplyException;
import org.apache.geode.distributed.internal.membership.InternalDistributedMember;
import org.apache.geode.internal.cache.CacheClosingDistributionMessageObserver;
import org.apache.geode.internal.cache.DiskRegion;
import org.apache.geode.internal.cache.InitialImageOperation.RequestImageMessage;
import org.apache.geode.internal.cache.InternalCache;
import org.apache.geode.internal.cache.PartitionedRegion;
import org.apache.geode.internal.cache.PartitionedRegionDataStore;
import org.apache.geode.internal.cache.control.InternalResourceManager;
import org.apache.geode.internal.cache.control.InternalResourceManager.ResourceObserver;
import org.apache.geode.internal.cache.control.InternalResourceManager.ResourceObserverAdapter;
import org.apache.geode.internal.cache.persistence.PersistentMemberID;
import org.apache.geode.test.awaitility.GeodeAwaitility;
import org.apache.geode.test.dunit.AsyncInvocation;
import org.apache.geode.test.dunit.IgnoredException;
import org.apache.geode.test.dunit.SerializableRunnableIF;
import org.apache.geode.test.dunit.VM;
import org.apache.geode.test.dunit.rules.CacheRule;
import org.apache.geode.test.dunit.rules.DistributedDiskDirRule;
import org.apache.geode.test.dunit.rules.DistributedRule;
import org.apache.geode.test.junit.rules.serializable.SerializableTestName;
import org.apache.geode.test.junit.runners.GeodeParamsRunner;

/**
 * Tests the basic use cases for PR persistence.
 */
@RunWith(GeodeParamsRunner.class)
@SuppressWarnings("serial,unused")
public class PersistentPartitionedRegionDistributedTest implements Serializable {

  private static final int NUM_BUCKETS = 15;

  private String partitionedRegionName;
  private String parentRegion1Name;
  private String parentRegion2Name;

  private VM vm0;
  private VM vm1;
  private VM vm2;
  private VM vm3;

  private static final long TIMEOUT_MILLIS = GeodeAwaitility.getTimeout().toMillis();

  @Rule
  public DistributedRule distributedRule = new DistributedRule();

  @Rule
  public CacheRule cacheRule =
      CacheRule.builder().addConfig(getDistributedSystemProperties()).build();

  @Rule
  public SerializableTestName testName = new SerializableTestName();

  @Rule
  public DistributedDiskDirRule diskDirRule = new DistributedDiskDirRule();

  @Before
  public void setUp() {
    vm0 = getVM(0);
    vm1 = getVM(1);
    vm2 = getVM(2);
    vm3 = getVM(3);

    String uniqueName = getClass().getSimpleName() + "-" + testName.getMethodName();
    partitionedRegionName = uniqueName + "-partitionedRegion";
    parentRegion1Name = "parent1";
    parentRegion2Name = "parent2";
  }

  @After
  public void tearDown() {
    invokeInEveryVM(() -> {
      InternalResourceManager.setResourceObserver(null);
      DistributionMessageObserver.setInstance(null);
    });
  }

  private Properties getDistributedSystemProperties() {
    Properties config = new Properties();
    config.setProperty(SERIALIZABLE_OBJECT_FILTER, TestFunction.class.getName());
    return config;
  }

  /**
   * A simple test case that we are actually persisting with a PR.
   */
  @Test
  public void recoversFromDisk() throws Exception {
    createPartitionedRegion(0, -1, 113, true);

    createData(0, 1, "a");

    Set<Integer> vm0Buckets = getBucketList();

    getCache().close();

    createPartitionedRegion(0, -1, 113, true);

    assertThat(getBucketList()).isEqualTo(vm0Buckets);

    checkData(0, 1, "a");

    getCache().getRegion(partitionedRegionName).localDestroyRegion();

    getCache().close();

    createPartitionedRegion(0, -1, 113, true);

    // Make sure the data is now missing
    checkData(0, 1, null);
  }

  @Test
  public void numberOfBucketsIsImmutable() throws Exception {
    createPartitionedRegion(0, 0, 5, true);
    createData(0, 5, "a");
    getCache().close();

    addIgnoredException(IllegalStateException.class);
    addIgnoredException(DiskAccessException.class);

    assertThatThrownBy(() -> createPartitionedRegion(0, 0, 2, true))
        .isInstanceOf(IllegalStateException.class).hasMessageContaining(
            String.format(
                "For partition region %s,total-num-buckets %s should not be changed. Previous configured number is %s.",
                SEPARATOR + partitionedRegionName, 2, 5));

    getCache().close();

    assertThatThrownBy(() -> createPartitionedRegion(0, 0, 10, true))
        .isInstanceOf(IllegalStateException.class).hasMessageContaining(
            String.format(
                "For partition region %s,total-num-buckets %s should not be changed. Previous configured number is %s.",
                SEPARATOR + partitionedRegionName, 10, 5));
  }

  @Test
  public void missingDiskStoreCanBeRevoked() throws Exception {
    vm0.invoke(() -> createPartitionedRegion(1, -1, 113, true));
    vm1.invoke(() -> createPartitionedRegion(1, -1, 113, true));

    int numBuckets = 50;
    vm0.invoke(() -> createData(0, numBuckets, "a"));

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM1 = vm1.invoke(this::getBucketList);
    assertThat(bucketsOnVM1).isEqualTo(bucketsOnVM0);

    vm0.invoke(() -> getCache().close());
    vm1.invoke(() -> createData(0, numBuckets, "b"));

    vm1.invoke(() -> getCache().close());

    AsyncInvocation<Void> createPartitionedRegionOnVM0 =
        vm0.invokeAsync(() -> createPartitionedRegion(1, -1, 113, true));
    // Note: Make sure that vm0 is waiting for vm1 to recover
    // If VM(0) recovers early, that is a problem, because vm1 has newer data
    createPartitionedRegionOnVM0.join(SECONDS.toMillis(1));
    assertThat(createPartitionedRegionOnVM0.isAlive()).isTrue();

    vm2.invoke(() -> {
      DistributedSystemConfig config = defineDistributedSystem(getSystem(), "");
      AdminDistributedSystem adminDS = getDistributedSystem(config);
      adminDS.connect();
      try {
        adminDS.waitToBeConnected(MINUTES.toMillis(2));

        await().until(() -> {
          Set<PersistentID> missingIds = adminDS.getMissingPersistentMembers();
          if (missingIds.size() != 1) {
            return false;
          }
          for (PersistentID missingId : missingIds) {
            adminDS.revokePersistentMember(missingId.getUUID());
          }
          return true;
        });

      } finally {
        adminDS.disconnect();
      }
    });

    createPartitionedRegionOnVM0.await();

    assertThat(vm0.invoke(this::getBucketList)).isEqualTo(bucketsOnVM0);

    vm0.invoke(() -> {
      checkData(0, numBuckets, "a");
      createData(numBuckets, 113, "b");
      checkData(numBuckets, 113, "b");
    });

    addIgnoredException(RevokedPersistentDataException.class);

    vm1.invoke(() -> {
      assertThatThrownBy(() -> createPartitionedRegion(1, -1, 113, true))
          .isInstanceOf(RevokedPersistentDataException.class);
    });
  }

  /**
   * Test to make sure that we can recover from a complete system shutdown
   */
  @Test
  @Parameters({"0", "1"})
  @TestCaseName("{method}({params})")
  public void recoversFromDiskAfterCompleteShutdown(int redundancy) throws Exception {
    int numBuckets = 50;

    vm0.invoke(() -> createPartitionedRegion(redundancy, -1, numBuckets, true));
    vm1.invoke(() -> createPartitionedRegion(redundancy, -1, numBuckets, true));
    vm2.invoke(() -> createPartitionedRegion(redundancy, -1, numBuckets, true));

    vm0.invoke(() -> createData(0, numBuckets, "a"));

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM1 = vm1.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM2 = vm2.invoke(this::getBucketList);

    vm0.invoke(() -> getCache().close());
    vm1.invoke(() -> getCache().close());
    vm2.invoke(() -> getCache().close());

    AsyncInvocation<Void> createPartitionedRegionOnVM0 =
        vm0.invokeAsync(() -> createPartitionedRegion(redundancy, -1, numBuckets, true));
    AsyncInvocation<Void> createPartitionedRegionOnVM1 =
        vm1.invokeAsync(() -> createPartitionedRegion(redundancy, -1, numBuckets, true));
    AsyncInvocation<Void> createPartitionedRegionOnVM2 =
        vm2.invokeAsync(() -> createPartitionedRegion(redundancy, -1, numBuckets, true));

    createPartitionedRegionOnVM0.await();
    createPartitionedRegionOnVM1.await();
    createPartitionedRegionOnVM2.await();

    assertThat(vm0.invoke(this::getBucketList)).isEqualTo(bucketsOnVM0);
    assertThat(vm1.invoke(this::getBucketList)).isEqualTo(bucketsOnVM1);
    assertThat(vm2.invoke(this::getBucketList)).isEqualTo(bucketsOnVM2);

    vm0.invoke(() -> {
      checkData(0, numBuckets, "a");
      createData(numBuckets, 113, "b");
      checkData(numBuckets, 113, "b");
    });

    // I think the following is sort of a TODO...
    // TRAC #43476: CREATE TABLE freeze on a 2-member distributed system (with exact commands)
    // RegressionTest for bug 43476 - make sure a destroy cleans up proxy bucket regions.

    vm0.invoke(() -> getCache().getRegion(partitionedRegionName).localDestroyRegion());
    vm1.invoke(() -> getCache().getRegion(partitionedRegionName).localDestroyRegion());
    vm2.invoke(() -> getCache().getRegion(partitionedRegionName).localDestroyRegion());

    createPartitionedRegionOnVM0 =
        vm0.invokeAsync(() -> createPartitionedRegion(redundancy, -1, numBuckets, true));
    createPartitionedRegionOnVM1 =
        vm1.invokeAsync(() -> createPartitionedRegion(redundancy, -1, numBuckets, true));
    createPartitionedRegionOnVM2 =
        vm2.invokeAsync(() -> createPartitionedRegion(redundancy, -1, numBuckets, true));

    createPartitionedRegionOnVM0.await();
    createPartitionedRegionOnVM1.await();
    createPartitionedRegionOnVM2.await();

    vm0.invoke(() -> checkData(0, numBuckets, null));
  }

  /**
   * Note: This test is only valid for
   * {@link AdminDistributedSystem#revokePersistentMember(InetAddress, String)} which is deprecated
   * and does not exist in GFSH. GFSH only supports revoking by UUID which cannot be revoked
   * <bold>before</bold> a member starts up and is missing the diskStore identified by the UUID.
   */
  @Test
  public void missingDiskStoreCanBeRevokedBeforeStartingServer() {
    int numBuckets = 50;

    vm0.invoke(() -> createPartitionedRegion(1, -1, 113, true));
    vm1.invoke(() -> createPartitionedRegion(1, -1, 113, true));

    vm0.invoke(() -> createData(0, numBuckets, "a"));

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM1 = vm1.invoke(this::getBucketList);
    assertThat(bucketsOnVM1).isEqualTo(bucketsOnVM0);

    InetAddress vm1Address = vm1.invoke(() -> getSystem().getDistributedMember().getInetAddress());

    // This should fail with a revocation failed message
    try (IgnoredException ie = addIgnoredException(RevokeFailedException.class)) {
      vm2.invoke("revoke disk store should fail", () -> {
        assertThatThrownBy(() -> {
          DistributedSystemConfig config = defineDistributedSystem(getSystem(), "");
          AdminDistributedSystem adminDS = getDistributedSystem(config);
          adminDS.connect();

          try {
            adminDS.waitToBeConnected(MINUTES.toMillis(2));
            adminDS.revokePersistentMember(vm1Address, null);
          } finally {
            adminDS.disconnect();
          }
        }).isInstanceOf(ReplyException.class).hasCauseInstanceOf(RevokeFailedException.class);
      });
    }

    vm0.invoke(() -> getCache().close());
    vm1.invoke(() -> createData(0, numBuckets, "b"));

    String diskDirPathOnVM1 = diskDirRule.getDiskDirFor(vm1).getAbsolutePath();
    vm1.invoke(() -> getCache().close());

    vm0.invoke((SerializableRunnableIF) this::getCache);

    vm2.invoke(() -> {
      DistributedSystemConfig config = defineDistributedSystem(getSystem(), "");
      AdminDistributedSystem adminDS = getDistributedSystem(config);
      adminDS.connect();
      try {
        adminDS.waitToBeConnected(MINUTES.toMillis(2));
        adminDS.revokePersistentMember(vm1Address, diskDirPathOnVM1);
      } finally {
        adminDS.disconnect();
      }
    });

    vm0.invoke(() -> {
      createPartitionedRegion(1, -1, 113, true);
      assertThat(getBucketList()).isEqualTo(bucketsOnVM0);

      checkData(0, numBuckets, "a");
      createData(numBuckets, 113, "b");
      checkData(numBuckets, 113, "b");
    });

    try (IgnoredException ie = addIgnoredException(RevokedPersistentDataException.class)) {
      vm1.invoke(() -> {
        assertThatThrownBy(() -> createPartitionedRegion(1, -1, 113, true))
            .isInstanceOf(RevokedPersistentDataException.class);
      });
    }
  }

  /**
   * Test that we wait for missing data to come back if the redundancy was 0.
   */
  @Test
  public void waitsForMissingDiskStoreWhenRedundancyIsZero() {
    vm0.invoke(() -> createPartitionedRegion(0, -1, 113, true));
    vm1.invoke(() -> createPartitionedRegion(0, -1, 113, true));

    vm0.invoke(() -> createData(0, NUM_BUCKETS, "a"));

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM1 = vm1.invoke(this::getBucketList);

    int bucketOnVM0 = bucketsOnVM0.iterator().next();
    int bucketOnVM1 = bucketsOnVM1.iterator().next();

    vm1.invoke(() -> getCache().close());

    try (IgnoredException ie = addIgnoredException(PartitionOfflineException.class)) {
      vm0.invoke(() -> checkReadWriteOperationsWithOfflineMember(bucketOnVM0, bucketOnVM1));

      // Make sure that a newly created member is informed about the offline member
      vm2.invoke(() -> createPartitionedRegion(0, -1, 113, true));
      vm2.invoke(() -> checkReadWriteOperationsWithOfflineMember(bucketOnVM0, bucketOnVM1));
    }

    // This should work, because these are new buckets
    vm0.invoke(() -> createData(NUM_BUCKETS, 113, "a"));

    vm1.invoke(() -> createPartitionedRegion(0, -1, 113, true));

    // The data should be back online now.
    vm0.invoke(() -> checkData(0, 113, "a"));
  }

  /**
   * Test to make sure that we recreate a bucket if a member is destroyed
   */
  @Test
  public void recreatesBucketIfHostIsDestroyed() {
    vm0.invoke(() -> createPartitionedRegion(0, -1, 113, true));
    vm1.invoke(() -> createPartitionedRegion(0, -1, 113, true));

    vm0.invoke(() -> createData(0, NUM_BUCKETS, "a"));

    Set<Integer> bucketOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketOnVM1 = vm1.invoke(this::getBucketList);

    int bucketIdOnVM0 = bucketOnVM0.iterator().next();
    int bucketIdOnVM1 = bucketOnVM1.iterator().next();

    vm1.invoke(() -> getCache().getRegion(partitionedRegionName).localDestroyRegion());

    vm0.invoke(() -> {
      // This should work, because this bucket is still available.
      checkData(bucketIdOnVM0, bucketIdOnVM0 + 1, "a");

      // This should find that the data is missing, because we destroyed that bucket
      checkData(bucketIdOnVM1, bucketIdOnVM1 + 1, null);

      // We should be able to recreate that bucket
      createData(bucketIdOnVM1, bucketIdOnVM1 + 1, "b");
    });

    vm1.invoke(() -> createPartitionedRegion(0, -1, 113, true));

    vm0.invoke(() -> {
      // The data should still be available
      checkData(bucketIdOnVM0, bucketIdOnVM0 + 1, "a");
      checkData(bucketIdOnVM1, bucketIdOnVM1 + 1, "b");

      // This bucket should now be in vm0, because we recreated it there
      assertThat(getBucketList()).contains(bucketIdOnVM1);
    });
  }

  /**
   * Test to make sure that we recreate a bucket if a member is destroyed
   */
  @Test
  public void recreatesBucketIfHostIsDestroyedAndRedundancyIsOne() {
    vm0.invoke(() -> createPartitionedRegion(1, -1, 113, true));
    vm1.invoke(() -> createPartitionedRegion(1, -1, 113, true));

    vm0.invoke(() -> createData(0, NUM_BUCKETS, "a"));

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM1 = vm1.invoke(this::getBucketList);
    assertThat(bucketsOnVM1).isEqualTo(bucketsOnVM0);

    int bucketIdOnVM0 = bucketsOnVM0.iterator().next();
    vm1.invoke(() -> getCache().getRegion(partitionedRegionName).localDestroyRegion());

    // This should work, because this bucket is still available.
    vm0.invoke(() -> checkData(bucketIdOnVM0, bucketIdOnVM0 + 1, "a"));

    vm2.invoke(() -> createPartitionedRegion(1, -1, 113, true));

    Set<Integer> vm2Buckets = vm2.invoke(this::getBucketList);

    // VM 2 should have created a copy of all of the buckets
    assertThat(vm2Buckets).isEqualTo(bucketsOnVM0);
  }

  /**
   * Test to make sure that we recreate a bucket if a member is revoked
   */
  @Test
  public void recreatesBucketIfDiskStoreIsRevokedAndRedundancyIsZero() {
    vm0.invoke(() -> createPartitionedRegion(0, -1, 113, true));
    vm1.invoke(() -> createPartitionedRegion(0, -1, 113, true));

    vm0.invoke(() -> createData(0, NUM_BUCKETS, "a"));

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM1 = vm1.invoke(this::getBucketList);

    int bucketOnVM0 = bucketsOnVM0.iterator().next();
    int bucketOnVM1 = bucketsOnVM1.iterator().next();

    vm1.invoke(() -> getCache().close());

    vm0.invoke(() -> {
      // This should work, because this bucket is still available.
      checkData(bucketOnVM0, bucketOnVM0 + 1, "a");

      try (IgnoredException ie = addIgnoredException("PartitionOfflineException", vm0)) {
        assertThatThrownBy(() -> checkData(bucketOnVM1, bucketOnVM1 + 1, "a"))
            .isInstanceOf(PartitionOfflineException.class);

        assertThatThrownBy(() -> checkData(bucketOnVM1, bucketOnVM1 + 1, "b"))
            .isInstanceOf(PartitionOfflineException.class);
      }

      // This should work, because these are new buckets
      createData(NUM_BUCKETS, 113, "a");
    });

    vm2.invoke(() -> revokeKnownMissingMembers(1));

    vm2.invoke(() -> {
      createPartitionedRegion(0, -1, 113, true);

      // We should be able to use that missing bucket now
      checkData(bucketOnVM1, bucketOnVM1 + 1, null);
      createData(bucketOnVM1, bucketOnVM1 + 1, "a");
      checkData(bucketOnVM1, bucketOnVM1 + 1, "a");
    });

    vm1.invoke(() -> {
      try (IgnoredException ie = addIgnoredException(RevokedPersistentDataException.class)) {
        assertThatThrownBy(() -> createPartitionedRegion(0, -1, 113, true))
            .isInstanceOf(RevokedPersistentDataException.class);
      }
    });
  }

  /**
   * Test to make sure that we recreate a bucket if a member is revoked
   */
  @Test
  public void recreatesBucketIfDiskStoreIsRevokedAndRedundancyIsOne() throws Exception {
    vm0.invoke(() -> createPartitionedRegion(1, -1, 113, true));
    vm1.invoke(() -> createPartitionedRegion(1, -1, 113, true));

    vm0.invoke(() -> createData(0, NUM_BUCKETS, "a"));

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM1 = vm1.invoke(this::getBucketList);
    assertThat(bucketsOnVM1).isEqualTo(bucketsOnVM0);

    vm1.invoke(() -> getCache().close());

    vm0.invoke(() -> {
      // This should work, because this bucket is still available.
      checkData(0, NUM_BUCKETS, "a");

      createData(0, NUM_BUCKETS, "b");
    });

    vm2.invoke(() -> {
      revokeKnownMissingMembers(1);

      // This should make a copy of all of the buckets, because we have revoked VM1.
      createPartitionedRegion(1, -1, 113, true);

      Set<Integer> bucketsOnVM2 = vm2.invoke(this::getBucketList);
      assertThat(bucketsOnVM2).isEqualTo(bucketsOnVM1);
    });

    vm1.invoke(() -> {
      try (IgnoredException ie = addIgnoredException(RevokedPersistentDataException.class)) {
        assertThatThrownBy(() -> createPartitionedRegion(1, -1, 113, true))
            .isInstanceOf(RevokedPersistentDataException.class);
      }
    });

    // Test that we can bounce vm0 and vm1, and still get a RevokedPersistentDataException
    // when vm1 tries to recover
    vm0.invoke(() -> getCache().close());
    vm2.invoke(() -> getCache().close());

    AsyncInvocation<Void> createPartitionedRegionOnVM0 =
        vm0.invokeAsync(() -> createPartitionedRegion(1, -1, 113, true));
    AsyncInvocation<Void> createPartitionedRegionOnVM2 =
        vm2.invokeAsync(() -> createPartitionedRegion(1, -1, 113, true));

    createPartitionedRegionOnVM0.await();
    createPartitionedRegionOnVM2.await();

    vm1.invoke(() -> {
      try (IgnoredException ie = addIgnoredException(RevokedPersistentDataException.class)) {
        assertThatThrownBy(() -> createPartitionedRegion(1, -1, 113, true))
            .isInstanceOf(RevokedPersistentDataException.class);
      }
    });

    // The data shouldn't be affected.
    vm2.invoke(() -> checkData(0, NUM_BUCKETS, "b"));
  }

  /**
   * Test to make sure that we recreate a bucket if a member is revoked, and that we do it
   * immediately if recovery delay is set to 0.
   */
  @Test
  public void recreatesBucketImmediatelyIfDiskStoreIsRevokedAndRecoveryDelayIsZero()
      throws Exception {
    vm0.invoke(() -> createPartitionedRegion(1, 0, 113, true));
    vm1.invoke(() -> createPartitionedRegion(1, 0, 113, true));

    vm0.invoke(() -> createData(0, NUM_BUCKETS, "a"));

    // This should do nothing because we have satisfied redundancy.
    vm2.invoke(() -> createPartitionedRegion(1, 0, 113, true));
    assertThat(vm2.invoke(this::getBucketList)).isEmpty();

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsLost = vm1.invoke(this::getBucketList);

    vm1.invoke(() -> getCache().close());

    // VM2 should pick up the slack

    await().until(() -> {
      Set<Integer> vm2Buckets = vm2.invoke(this::getBucketList);
      return bucketsLost.equals(vm2Buckets);
    });

    vm0.invoke(() -> createData(0, NUM_BUCKETS, "b"));

    // VM1 should recover, but it shouldn't host the bucket anymore
    vm1.invoke(() -> createPartitionedRegion(1, 0, 113, true));

    // The data shouldn't be affected.
    vm1.invoke(() -> checkData(0, NUM_BUCKETS, "b"));

    // restart everything, and make sure it comes back correctly.

    vm1.invoke(() -> getCache().close());
    vm0.invoke(() -> getCache().close());
    vm2.invoke(() -> getCache().close());

    AsyncInvocation<Void> createPartitionedRegionOnVM1 =
        vm1.invokeAsync(() -> createPartitionedRegion(1, 0, 113, true));
    AsyncInvocation<Void> createPartitionedRegionOnVM0 =
        vm0.invokeAsync(() -> createPartitionedRegion(1, 0, 113, true));

    // Make sure we wait for vm2, because it's got the latest copy of the bucket
    createPartitionedRegionOnVM1.join(SECONDS.toMillis(1));
    assertThat(createPartitionedRegionOnVM1.isAlive()).isTrue();

    AsyncInvocation<Void> createPartitionedRegionOnVM2 =
        vm2.invokeAsync(() -> createPartitionedRegion(1, 0, 113, true));

    createPartitionedRegionOnVM2.await();
    createPartitionedRegionOnVM0.await();
    createPartitionedRegionOnVM1.await();

    // The data shouldn't be affected.
    vm1.invoke(() -> checkData(0, NUM_BUCKETS, "b"));
    assertThat(vm1.invoke(this::getBucketList)).isEmpty();
    assertThat(vm0.invoke(this::getBucketList)).isEqualTo(bucketsOnVM0);
    assertThat(vm2.invoke(this::getBucketList)).isEqualTo(bucketsOnVM0);
  }

  /**
   * Test the with redundancy 1, we restore the same buckets when the missing member comes back
   * online.
   */
  @Test
  public void restoresSameBucketsWhenMissingDiskStoreReturnsOnline() {
    vm0.invoke(() -> createPartitionedRegion(1, 0, 113, true));
    vm1.invoke(() -> createPartitionedRegion(1, 0, 113, true));

    vm0.invoke(() -> createData(0, NUM_BUCKETS, "a"));

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM1 = vm1.invoke(this::getBucketList);

    vm1.invoke(() -> getCache().close());

    // This should work, because this bucket is still available.
    vm0.invoke(() -> {
      checkData(0, NUM_BUCKETS, "a");

      Region<Integer, String> region = getCache().getRegion(partitionedRegionName);
      for (int i = 0; i < NUM_BUCKETS / 2; i++) {
        region.destroy(i);
      }

      createData(NUM_BUCKETS / 2, NUM_BUCKETS, "b");
    });

    // This shouldn't create any buckets, because we know there are offline copies
    vm2.invoke(() -> createPartitionedRegion(1, 0, 113, true));

    Set<Integer> bucketsOnVM2 = vm2.invoke(this::getBucketList);
    assertThat(bucketsOnVM2).isEmpty();

    vm1.invoke(() -> {
      createPartitionedRegion(1, 0, 113, true);

      // The data should be back online now. vm1 should have received the latest copy of the data.
      checkData(0, NUM_BUCKETS / 2, null);
      checkData(NUM_BUCKETS / 2, NUM_BUCKETS, "b");
    });

    // Make sure we restored the buckets in the right place
    assertThat(vm0.invoke(this::getBucketList)).isEqualTo(bucketsOnVM0);
    assertThat(vm1.invoke(this::getBucketList)).isEqualTo(bucketsOnVM1);
    assertThat(vm2.invoke(this::getBucketList)).isEmpty();
  }

  @Test
  public void bucketCanBeMovedMultipleTimes() throws Exception {
    int redundancy = 0;

    vm0.invoke(() -> {
      createPartitionedRegion(redundancy, -1, 113, true);
      createData(0, 2, "a");
    });

    vm1.invoke(() -> createPartitionedRegion(redundancy, -1, 113, true));
    vm2.invoke(() -> createPartitionedRegion(redundancy, -1, 113, true));

    Set<Integer> bucketsOnVM0 = vm0.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM1 = vm1.invoke(this::getBucketList);
    Set<Integer> bucketsOnVM2 = vm2.invoke(this::getBucketList);

    InternalDistributedMember memberVM0 = vm0.invoke(this::getInternalDistributedMember);
    InternalDistributedMember memberVM1 = vm1.invoke(this::getInternalDistributedMember);
    InternalDistributedMember memberVM2 = vm2.invoke(this::getInternalDistributedMember);

    vm1.invoke(() -> moveBucketToMe(0, memberVM0));
    vm2.invoke(() -> moveBucketToMe(0, memberVM1));
    vm0.invoke(() -> createData(113, 114, "a"));
    vm0.invoke(() -> moveBucketToMe(0, memberVM2));

    vm0.invoke(() -> createData(226, 227, "a"));

    assertThat(vm0.invoke(this::getBucketList)).isEqualTo(bucketsOnVM0);
    assertThat(vm1.invoke(this::getBucketList)).isEqualTo(bucketsOnVM1);
    assertThat(vm2.invoke(this::getBucketList)).isEqualTo(bucketsOnVM2);

    vm0.invoke(() -> getCache().close());
    vm1.invoke(() -> getCache().close());
    vm2.invoke(() -> getCache().close());

    AsyncInvocation<Void> createPartitionedRegionOnVM0 =
        vm0.invokeAsync(() -> createPartitionedRegion(redundancy, -1, 113, true));
    AsyncInvocation<Void> createPartitionedRegionOnVM1 =
        vm1.invokeAsync(() -> createPartitionedRegion(redundancy, -1, 113, true));
    AsyncInvocation<Void> createPartitionedRegionOnVM2 =
        vm2.invokeAsync(() -> createPartitionedRegion(redundancy, -1, 113, true));

    createPartitionedRegionOnVM0.await();
    createPartitionedRegionOnVM1.await();
    createPartitionedRegionOnVM2.await();

    assertThat(vm2.invoke(this::getBucketList)).isEqualTo(bucketsOnVM2);
    assertThat(vm1.invoke(this::getBucketList)).isEqualTo(bucketsOnVM1);
    assertThat(vm0.invoke(this::getBucketList)).isEqualTo(bucketsOnVM0);

    vm0.invoke(() -> {
      checkData(0, 2, "a");
      checkData(113, 114, "a");
      checkData(226, 227, "a");
    });
  }

  @Test
  public void recoversAfterCleanStop() throws Exception {
    vm0.invoke(() -> createPartitionedRegion(1, -1, 113, true));
    vm1.invoke(() -> createPartitionedRegion(1, -1, 113, true));

    vm0.invoke(() -> createData(0, 1, "a"));

    vm1.invoke(() -> fakeCleanShutdown(0));
    vm0.invoke(() -> fakeCleanShutdown(0));

    AsyncInvocation<Void> createPartitionedRegionOnVM0 =
        vm0.invokeAsync(() -> createPartitionedRegion(1, -1, 113, true));
    // Note: Make sure that vm0 is waiting for vm1 and vm2 to recover
    // If VM(0) recovers early, that is a problem, because we can now longer do a clean restart
    createPartitionedRegionOnVM0.join(SECONDS.toMillis(1));
    assertThat(createPartitionedRegionOnVM0.isAlive()).isTrue();

    AsyncInvocation<Void> createPartitionedRegionOnVM1 =
        vm1.invokeAsync(() -> createPartitionedRegion(1, -1, 113, true));

    createPartitionedRegionOnVM0.await();
    createPartitionedRegionOnVM1.await();

    vm0.invoke(() -> checkData(0, 1, "a"));
    vm1.invoke(() -> checkData(0, 1, "a"));

    vm0.invoke(() -> checkRecoveredFromDisk(0, true));
    vm1.invoke(() -> checkRecoveredFromDisk(0, true));

