/*
 * 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.control;

import static java.util.concurrent.TimeUnit.SECONDS;
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;
import static org.apache.geode.cache.RegionShortcut.PARTITION_PERSISTENT;
import static org.apache.geode.cache.RegionShortcut.REPLICATE;
import static org.apache.geode.distributed.ConfigurationProperties.ENFORCE_UNIQUE_HOST;
import static org.apache.geode.distributed.ConfigurationProperties.REDUNDANCY_ZONE;
import static org.apache.geode.test.awaitility.GeodeAwaitility.await;
import static org.apache.geode.test.dunit.Invoke.invokeInEveryVM;
import static org.apache.geode.test.dunit.VM.getVM;
import static org.apache.geode.test.dunit.VM.toArray;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.catchThrowable;
import static org.assertj.core.api.Assertions.within;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.anyBoolean;
import static org.mockito.Mockito.anyInt;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;

import java.io.File;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.BrokenBarrierException;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

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

import org.apache.geode.cache.AttributesFactory;
import org.apache.geode.cache.CacheClosedException;
import org.apache.geode.cache.CacheLoader;
import org.apache.geode.cache.CacheLoaderException;
import org.apache.geode.cache.CacheWriterException;
import org.apache.geode.cache.DiskStoreFactory;
import org.apache.geode.cache.EntryEvent;
import org.apache.geode.cache.EvictionAttributes;
import org.apache.geode.cache.LoaderHelper;
import org.apache.geode.cache.PartitionAttributesFactory;
import org.apache.geode.cache.Region;
import org.apache.geode.cache.RegionFactory;
import org.apache.geode.cache.asyncqueue.AsyncEvent;
import org.apache.geode.cache.asyncqueue.AsyncEventListener;
import org.apache.geode.cache.control.RebalanceOperation;
import org.apache.geode.cache.control.RebalanceResults;
import org.apache.geode.cache.control.ResourceManager;
import org.apache.geode.cache.partition.PartitionMemberInfo;
import org.apache.geode.cache.partition.PartitionRebalanceInfo;
import org.apache.geode.cache.partition.PartitionRegionHelper;
import org.apache.geode.cache.partition.PartitionRegionInfo;
import org.apache.geode.cache.util.CacheListenerAdapter;
import org.apache.geode.distributed.DistributedMember;
import org.apache.geode.distributed.internal.DistributionMessageObserver;
import org.apache.geode.distributed.internal.membership.InternalDistributedMember;
import org.apache.geode.internal.cache.BucketRegion;
import org.apache.geode.internal.cache.CacheClosingDistributionMessageObserver;
import org.apache.geode.internal.cache.ColocationHelper;
import org.apache.geode.internal.cache.DiskStoreImpl;
import org.apache.geode.internal.cache.InternalCache;
import org.apache.geode.internal.cache.PRHARedundancyProvider;
import org.apache.geode.internal.cache.PartitionedRegion;
import org.apache.geode.internal.cache.PartitionedRegionDataStore;
import org.apache.geode.internal.cache.SignalBounceOnRequestImageMessageObserver;
import org.apache.geode.internal.cache.control.InternalResourceManager.ResourceObserverAdapter;
import org.apache.geode.internal.cache.partitioned.BucketCountLoadProbe;
import org.apache.geode.internal.cache.partitioned.LoadProbe;
import org.apache.geode.test.awaitility.GeodeAwaitility;
import org.apache.geode.test.dunit.AsyncInvocation;
import org.apache.geode.test.dunit.SerializableRunnable;
import org.apache.geode.test.dunit.SerializableRunnableIF;
import org.apache.geode.test.dunit.VM;
import org.apache.geode.test.dunit.WaitCriterion;
import org.apache.geode.test.dunit.cache.CacheTestCase;
import org.apache.geode.test.dunit.rules.DistributedRestoreSystemProperties;
import org.apache.geode.test.junit.rules.ExecutorServiceRule;
import org.apache.geode.test.junit.runners.GeodeParamsRunner;
import org.apache.geode.util.internal.GeodeGlossary;

/**
 * TODO test colocated regions where buckets aren't created for all subregions
 *
 * <p>
 * TODO test colocated regions where members aren't consistent in which regions they have
 */
@RunWith(GeodeParamsRunner.class)
@SuppressWarnings("serial")
public class RebalanceOperationDistributedTest extends CacheTestCase {

  private static final long TIMEOUT_SECONDS = GeodeAwaitility.getTimeout().getSeconds();

  @Rule
  public DistributedRestoreSystemProperties restoreSystemProperties =
      new DistributedRestoreSystemProperties();

  @Rule
  public ExecutorServiceRule executorServiceRule = new ExecutorServiceRule();

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

  @Test
  @Parameters({"true", "false"})
  @TestCaseName("{method}(simulate={0})")
  public void testRecoverRedundancy(boolean simulate) {
    VM vm0 = getVM(0);
    VM vm1 = getVM(1);

    // Create the region in only 1 VM
    vm0.invoke(() -> createPartitionedRegion("region1", 1));

    // Create some buckets
    vm0.invoke(() -> {
      Region<Number, String> region = getCache().getRegion("region1");
      region.put(1, "A");
      region.put(2, "A");
      region.put(3, "A");
      region.put(4, "A");
      region.put(5, "A");
      region.put(6, "A");
    });

    // make sure we can tell that the buckets have low redundancy
    vm0.invoke(() -> validateRedundancy("region1", 6, 0, 6));

    // Create the region in the other VM (should have no effect)
    vm1.invoke(() -> createPartitionedRegion("region1", 1));

    // Make sure we still have low redundancy
    vm0.invoke(() -> validateRedundancy("region1", 6, 0, 6));

    // Now simulate a rebalance
    vm0.invoke(() -> {
      InternalResourceManager manager = getCache().getInternalResourceManager();
      RebalanceResults results = doRebalance(simulate, manager);

      assertThat(results.getTotalBucketCreatesCompleted()).isEqualTo(6);
      assertThat(results.getTotalPrimaryTransfersCompleted()).isEqualTo(3);
      assertThat(results.getTotalBucketTransferBytes()).isEqualTo(0);
      assertThat(results.getTotalBucketTransfersCompleted()).isEqualTo(0);

      Set<PartitionRebalanceInfo> detailSet = results.getPartitionRebalanceDetails();
      assertThat(detailSet).hasSize(1);

      PartitionRebalanceInfo details = detailSet.iterator().next();
      assertThat(details.getBucketCreatesCompleted()).isEqualTo(6);
      assertThat(details.getPrimaryTransfersCompleted()).isEqualTo(3);
      assertThat(details.getBucketTransferBytes()).isEqualTo(0);
      assertThat(details.getBucketTransfersCompleted()).isEqualTo(0);

      Set<PartitionMemberInfo> afterDetails = details.getPartitionMemberDetailsAfter();
      assertThat(afterDetails).hasSize(2);

      for (PartitionMemberInfo memberDetails : afterDetails) {
        assertThat(memberDetails.getBucketCount()).isEqualTo(6);
        assertThat(memberDetails.getPrimaryCount()).isEqualTo(3);
      }

      if (!simulate) {
        validateStatistics(manager, results);
      }
    });

    if (simulate) {
      // Make sure the simulation didn't do anything
      vm0.invoke(() -> validateRedundancy("region1", 6, 0, 6));
    } else {
      vm0.invoke(() -> validateRedundancy("region1", 6, 1, 0));
      vm1.invoke(() -> validateRedundancy("region1", 6, 1, 0));
    }
  }

  @Test
  @Parameters({"true", "false"})
  @TestCaseName("{method}(simulate={0})")
  public void testEnforceIP(boolean simulate) {
    VM vm0 = getVM(0);
    VM vm1 = getVM(1);

    for (VM vm : toArray(vm0, vm1)) {
      vm.invoke(() -> {
        Properties props = new Properties();
        props.setProperty(ENFORCE_UNIQUE_HOST, "true");
        getCache(props);
      });
    }

    // Create the region in only 1 VM
    vm0.invoke(() -> createPartitionedRegion("region1", 1));

    // Create some buckets
    vm0.invoke(() -> {
      Region<Number, String> region = getCache().getRegion("region1");
      region.put(1, "A");
      region.put(2, "A");
      region.put(3, "A");
      region.put(4, "A");
      region.put(5, "A");
      region.put(6, "A");
    });

    // make sure we can tell that the buckets have low redundancy
    vm0.invoke(() -> validateRedundancy("region1", 6, 0, 6));

    // Create the region in the other VM (should have no effect)
    vm1.invoke(() -> createPartitionedRegion("region1", 1));

    // Make sure we still have low redundancy
    vm0.invoke(() -> validateRedundancy("region1", 6, 0, 6));

    // Now simulate a rebalance
    vm0.invoke(() -> {
      InternalResourceManager manager = getCache().getInternalResourceManager();
      RebalanceResults results = doRebalance(simulate, manager);

      assertThat(results.getTotalBucketCreatesCompleted()).isEqualTo(0);
      assertThat(results.getTotalPrimaryTransfersCompleted()).isEqualTo(0);

      // We actually *will* transfer buckets, because that improves the balance
      assertThat(results.getTotalBucketTransfersCompleted()).isEqualTo(3);

      Set<PartitionRebalanceInfo> detailSet = results.getPartitionRebalanceDetails();
      assertThat(detailSet).hasSize(1);

      PartitionRebalanceInfo details = detailSet.iterator().next();
      assertThat(details.getBucketCreatesCompleted()).isEqualTo(0);
      assertThat(details.getPrimaryTransfersCompleted()).isEqualTo(0);
      assertThat(details.getBucketTransfersCompleted()).isEqualTo(3);

      if (!simulate) {
        validateStatistics(manager, results);
      }
    });

    // Make sure we still have low redundancy
    vm0.invoke(() -> validateRedundancy("region1", 6, 0, 6));
    vm1.invoke(() -> validateRedundancy("region1", 6, 0, 6));
  }

  /**
   * Test that we correctly use the redundancy-zone property to determine where to place redundant
   * copies of a buckets.
   */
  @Test
  @Parameters({"true", "false"})
  @TestCaseName("{method}(simulate={0})")
  public void testEnforceZone(boolean simulate) {
    VM vm0 = getVM(0);
    VM vm1 = getVM(1);
    VM vm2 = getVM(2);

    vm0.invoke(() -> setRedundancyZone("A"));
    vm1.invoke(() -> setRedundancyZone("A"));

    DistributedMember zoneBMember = vm2.invoke(() -> {
      setRedundancyZone("B");
      return getCache().getDistributedSystem().getDistributedMember();
    });

    // Create the region in only 1 VM
    vm0.invoke(() -> createPartitionedRegion("region1", 1));

    // Create some buckets
    vm0.invoke(() -> {
      Region<Number, String> region = getCache().getRegion("region1");
      region.put(1, "A");
      region.put(2, "A");
      region.put(3, "A");
      region.put(4, "A");
      region.put(5, "A");
      region.put(6, "A");
    });

    // make sure we can tell that the buckets have low redundancy
    vm0.invoke(() -> validateRedundancy("region1", 6, 0, 6));

    // Create the region in the other VMs (should have no effect)
    vm1.invoke(() -> createPartitionedRegion("region1", 1));
    vm2.invoke(() -> createPartitionedRegion("region1", 1));

    // Make sure we still have low redundancy
    vm0.invoke(() -> validateRedundancy("region1", 6, 0, 6));

    // Now simulate a rebalance
    vm0.invoke(() -> {
      InternalResourceManager manager = getCache().getInternalResourceManager();
      RebalanceResults results = doRebalance(simulate, manager);

      // We expect to satisfy redundancy with the zone B member
      assertThat(results.getTotalBucketCreatesCompleted()).isEqualTo(6);

      // 2 primaries will go to vm2, leaving vm0 and vm1 with 2 primaries each
      assertThat(results.getTotalPrimaryTransfersCompleted()).isEqualTo(2);

      // We actually *will* transfer 3 buckets to the other member in zone A, because that
      // improves
      // the balance
      assertThat(results.getTotalBucketTransfersCompleted()).isEqualTo(3);

      Set<PartitionRebalanceInfo> detailSet = results.getPartitionRebalanceDetails();
      assertThat(detailSet).hasSize(1);

      PartitionRebalanceInfo details = detailSet.iterator().next();
      assertThat(details.getBucketCreatesCompleted()).isEqualTo(6);
      assertThat(details.getPrimaryTransfersCompleted()).isEqualTo(2);
      assertThat(details.getBucketTransfersCompleted()).isEqualTo(3);

      Set<PartitionMemberInfo> afterDetails = details.getPartitionMemberDetailsAfter();
      for (PartitionMemberInfo info : afterDetails) {
        if (info.getDistributedMember().equals(zoneBMember)) {
          assertThat(info.getBucketCount()).isEqualTo(6);
        } else {
          assertThat(info.getBucketCount()).isEqualTo(3);
        }
        assertThat(info.getPrimaryCount()).isEqualTo(2);
      }

      // assertIndexDetailsEquals(0, details.getBucketTransferBytes());
      if (!simulate) {
        validateStatistics(manager, results);
      }
    });

    if (!simulate) {
      vm0.invoke(() -> validateBucketCount("region1", 3));
      vm1.invoke(() -> validateBucketCount("region1", 3));
      vm2.invoke(() -> validateBucketCount("region1", 6));
    }
  }

  @Test
  public void testEnforceZoneWithMultipleRegions() {
    VM vm0 = getVM(0);
    VM vm1 = getVM(1);
    VM vm2 = getVM(2);

    vm0.invoke(() -> setRedundancyZone("A"));
    vm1.invoke(() -> setRedundancyZone("A"));

    DistributedMember zoneBMember = vm2.invoke(() -> {
      setRedundancyZone("B");
      return getCache().getDistributedSystem().getDistributedMember();
    });

    vm0.invoke(() -> InternalResourceManager.setResourceObserver(new ParallelRecoveryObserver(2)));

    // Create the region in only 1 VM
    vm0.invoke(this::createTwoRegionsWithParallelRecoveryObserver);

    // Create some buckets
    vm0.invoke(() -> {
      doPuts("region1");
      doPuts("region2");
    });

    // make sure we can tell that the buckets have low redundancy
    vm0.invoke(() -> validateRedundancyOfTwoRegions("region1", "region2", 6, 0, 6));

    // Create the region in the other VMs (should have no effect)
    for (VM vm : toArray(vm1, vm2)) {
      vm.invoke(() -> {
        InternalResourceManager.setResourceObserver(new ParallelRecoveryObserver(2));
        createTwoRegionsWithParallelRecoveryObserver();
      });
    }

    // Make sure we still have low redundancy
    vm0.invoke(() -> validateRedundancyOfTwoRegions("region1", "region2", 6, 0, 6));

    // Now do a rebalance
    vm0.invoke(() -> {
      InternalResourceManager manager = getCache().getInternalResourceManager();
      RebalanceResults results = doRebalance(false, manager);

      // We expect to satisfy redundancy with the zone B member
      assertThat(results.getTotalBucketCreatesCompleted()).isEqualTo(12);

      // 2 primaries will go to vm2, leaving vm0 and vm1 with 2 primaries each
      assertThat(results.getTotalPrimaryTransfersCompleted()).isEqualTo(4);

      // We actually *will* transfer 3 buckets to the other member in zone A, because that
      // improves the balance
      assertThat(results.getTotalBucketTransfersCompleted()).isEqualTo(6);

      Set<PartitionRebalanceInfo> detailSet = results.getPartitionRebalanceDetails();
      assertThat(detailSet).hasSize(2);

      for (PartitionRebalanceInfo details : detailSet) {
        assertThat(details.getBucketCreatesCompleted()).isEqualTo(6);
        assertThat(details.getPrimaryTransfersCompleted()).isEqualTo(2);
        assertThat(details.getBucketTransfersCompleted()).isEqualTo(3);

        Set<PartitionMemberInfo> afterDetails = details.getPartitionMemberDetailsAfter();
        for (PartitionMemberInfo info : afterDetails) {
          if (info.getDistributedMember().equals(zoneBMember)) {
            assertThat(info.getBucketCount()).isEqualTo(6);
          } else {
            assertThat(info.getBucketCount()).isEqualTo(3);
          }
          assertThat(info.getPrimaryCount()).isEqualTo(2);
        }
      }

      validateStatistics(manager, results);
    });

    vm0.invoke(() -> {
      assertThat(((ParallelRecoveryObserver) InternalResourceManager.getResourceObserver())
          .isObserverCalled()).isTrue();
    });

    vm0.invoke(() -> validateBucketCount("region1", 3));
    vm1.invoke(() -> validateBucketCount("region1", 3));
    vm2.invoke(() -> validateBucketCount("region1", 6));

    vm0.invoke(() -> validateBucketCount("region2", 3));
    vm1.invoke(() -> validateBucketCount("region2", 3));
    vm2.invoke(() -> validateBucketCount("region2", 6));
  }

  @Test
  @Parameters({"true", "false"})
  @TestCaseName("{method}(simulate={0})")
  public void testRecoverRedundancyBalancing(boolean simulate) {
    VM vm0 = getVM(0);
    VM vm1 = getVM(1);
    VM vm2 = getVM(2);

    DistributedMember member1 = vm0.invoke(() -> {
      createPartitionedRegion("region1", 1, 200);
      return getCache().getDistributedSystem().getDistributedMember();
    });

    vm0.invoke(() -> {
      Region<Number, String> region = getCache().getRegion("region1");
      for (int i = 0; i < 12; i++) {
        region.put(i, "A");
      }
    });

    vm0.invoke(() -> validateRedundancy("region1", 12, 0, 12));

    // Now create the region in 2 more VMs with half the localMaxMemory
    vm1.invoke(() -> createPartitionedRegion("region1", 1, 100));
    vm2.invoke(() -> createPartitionedRegion("region1", 1, 100));

    vm0.invoke(() -> validateRedundancy("region1", 12, 0, 12));

    // Now simulate a rebalance
    vm0.invoke(() -> {
      InternalResourceManager manager = getCache().getInternalResourceManager();
      RebalanceResults results = doRebalance(simulate, manager);

      assertThat(results.getTotalBucketCreatesCompleted()).isEqualTo(12);
      assertThat(results.getTotalPrimaryTransfersCompleted()).isEqualTo(6);
      assertThat(results.getTotalBucketTransferBytes()).isEqualTo(0);
      assertThat(results.getTotalBucketTransfersCompleted()).isEqualTo(0);

      Set<PartitionRebalanceInfo> detailSet = results.getPartitionRebalanceDetails();
      assertThat(detailSet).hasSize(1);

      PartitionRebalanceInfo details = detailSet.iterator().next();
      assertThat(details.getBucketCreatesCompleted()).isEqualTo(12);
      assertThat(details.getPrimaryTransfersCompleted()).isEqualTo(6);
      assertThat(details.getBucketTransferBytes()).isEqualTo(0);
      assertThat(details.getBucketTransfersCompleted()).isEqualTo(0);

      Set<PartitionMemberInfo> afterDetails = details.getPartitionMemberDetailsAfter();
      assertThat(afterDetails).hasSize(3);

      for (PartitionMemberInfo memberDetails : afterDetails) {
        // We have 1 member with a size of 200 and two members with size 100
        if (memberDetails.getDistributedMember().equals(member1)) {
          assertThat(memberDetails.getBucketCount()).isEqualTo(12);
          assertThat(memberDetails.getPrimaryCount()).isEqualTo(6);
        } else {
          assertThat(memberDetails.getBucketCount()).isEqualTo(6);
          assertThat(memberDetails.getPrimaryCount()).isEqualTo(3);
        }
      }

      if (!simulate) {
        validateStatistics(manager, results);
      }
    });

    if (!simulate) {
      for (VM vm : toArray(vm0, vm1, vm2)) {
        vm.invoke(() -> validateRedundancy("region1", 12, 1, 0));
      }
    }
  }

  @Test
  public void testRecoverRedundancyBalancingIfCreateBucketFails() {
    VM vm0 = getVM(0);
    VM vm1 = getVM(1);
    VM vm2 = getVM(2);

    DistributedMember member1 = vm0.invoke(() -> {
      createPartitionedRegion("region1", 1, 100);
      return getCache().getDistributedSystem().getDistributedMember();
    });

    vm0.invoke(() -> {
      Region<Number, String> region = getCache().getRegion("region1");
      for (int i = 0; i < 1; i++) {
        region.put(i, "A");
      }
    });

    vm0.invoke(() -> validateRedundancy("region1", 1, 0, 1));

    // Now create the region in 2 more VMs
    // Let localMaxMemory(VM1) > localMaxMemory(VM2)
    // so that redundant bucket will always be attempted on VM1
    DistributedMember member2 = vm1.invoke(() -> {
      createPartitionedRegion("region1", 1, 100);
      return getCache().getDistributedSystem().getDistributedMember();
    });
    DistributedMember member3 = vm2.invoke(() -> {
      createPartitionedRegion("region1", 1, 90);
      return getCache().getDistributedSystem().getDistributedMember();
    });

    vm0.invoke(() -> validateRedundancy("region1", 1, 0, 1));

    // Inject mock PRHARedundancyProvider to simulate createBucketFailures
    vm0.invoke(() -> {
      InternalCache cache = spy(getCache());

      PartitionedRegion origRegion = (PartitionedRegion) cache.getRegion("region1");
      PartitionedRegion spyRegion = spy(origRegion);
      PRHARedundancyProvider redundancyProvider =
          spy(new PRHARedundancyProvider(spyRegion, mock(InternalResourceManager.class)));

      // return the spied region when ever getPartitionedRegions() is invoked
      Set<PartitionedRegion> parRegions = cache.getPartitionedRegions();
      parRegions.remove(origRegion);
      parRegions.add(spyRegion);

      doReturn(parRegions).when(cache).getPartitionedRegions();
      doReturn(redundancyProvider).when(spyRegion).getRedundancyProvider();

      // simulate create bucket fails on member2 and test if it creates on member3
      doReturn(false).when(redundancyProvider).createBackupBucketOnMember(anyInt(),
          eq((InternalDistributedMember) member2), anyBoolean(), anyBoolean(), any(),
          anyBoolean());

      // Now simulate a rebalance
      // Create operationImpl and not factory as we need spied cache to be passed to operationImpl
      RegionFilter filter = new FilterByPath(null, null);
      RebalanceOperationImpl operation = new RebalanceOperationImpl(cache, false, filter);
      operation.start();

      RebalanceResults results = operation.getResults(TIMEOUT_SECONDS, SECONDS);

      assertThat(results.getTotalBucketCreatesCompleted()).isEqualTo(1);
      assertThat(results.getTotalPrimaryTransfersCompleted()).isEqualTo(0);
      assertThat(results.getTotalBucketTransferBytes()).isEqualTo(0);
      assertThat(results.getTotalBucketTransfersCompleted()).isEqualTo(0);

      Set<PartitionRebalanceInfo> detailSet = results.getPartitionRebalanceDetails();
      assertThat(detailSet).hasSize(1);

      PartitionRebalanceInfo details = detailSet.iterator().next();
      assertThat(details.getBucketCreatesCompleted()).isEqualTo(1);
      assertThat(details.getPrimaryTransfersCompleted()).isEqualTo(0);
      assertThat(details.getBucketTransferBytes()).isEqualTo(0);
      assertThat(details.getBucketTransfersCompleted()).isEqualTo(0);

      Set<PartitionMemberInfo> afterDetails = details.getPartitionMemberDetailsAfter();
      assertThat(afterDetails).hasSize(3);
      for (PartitionMemberInfo memberDetails : afterDetails) {
        if (memberDetails.getDistributedMember().equals(member1)) {
          assertThat(memberDetails.getBucketCount()).isEqualTo(1);
          assertThat(memberDetails.getPrimaryCount()).isEqualTo(1);
        } else if (memberDetails.getDistributedMember().equals(member2)) {
          assertThat(memberDetails.getBucketCount()).isEqualTo(0);
          assertThat(memberDetails.getPrimaryCount()).isEqualTo(0);
        } else if (memberDetails.getDistributedMember().equals(member3)) {
          assertThat(memberDetails.getBucketCount()).isEqualTo(1);
          assertThat(memberDetails.getPrimaryCount()).isEqualTo(0);
        }
      }

      ResourceManagerStats stats = cache.getInternalResourceManager().getStats();

      assertThat(stats.getRebalancesInProgress()).isEqualTo(0);
      assertThat(stats.getRebalancesCompleted()).isEqualTo(1);
      assertThat(stats.getRebalanceBucketCreatesInProgress()).isEqualTo(0);
      assertThat(stats.getRebalanceBucketCreatesCompleted())
          .isEqualTo(results.getTotalBucketCreatesCompleted());
      assertThat(stats.getRebalanceBucketCreatesFailed()).isEqualTo(1);
    });

    for (VM vm : toArray(vm0, vm1, vm2)) {
      vm.invoke(() -> validateRedundancy("region1", 1, 1, 0));
    }
  }

  @Test
  @Parameters({"true", "false"})
  @TestCaseName("{method}(simulate={0})")
  public void testRecoverRedundancyColocatedRegions(boolean simulate) {
    VM vm0 = getVM(0);
    VM vm1 = getVM(1);
    VM vm2 = getVM(2);

    DistributedMember member1 = vm0.invoke(() -> {
      createPartitionedRegion("region1", 1, 200);
      return getCache().getDistributedSystem().getDistributedMember();
    });

    vm0.invoke(() -> createPartitionedRegion("region2", 200, "region1"));

    // Create some buckets.
    vm0.invoke(() -> {
      Region<Number, String> region = getCache().getRegion("region1");
      Region<Number, String> region2 = getCache().getRegion("region2");
      for (int i = 0; i < 12; i++) {
        region.put(i, "A");
        region2.put(i, "A");
      }
    });

    // check to make sure our redundancy is impaired
    vm0.invoke(() -> validateRedundancyOfTwoRegions("region1", "region2", 12, 0, 12));

    // Now create the region in 2 more vms, each which
    // has local max memory of 1/2 that of the original VM.
    vm1.invoke(() -> createPartitionedRegion("region1", 1, 100));
    vm2.invoke(() -> createPartitionedRegion("region1", 1, 100));
    vm1.invoke(() -> createPartitionedRegion("region2", 100, "region1"));
    vm2.invoke(() -> createPartitionedRegion("region2", 100, "region1"));

    vm0.invoke(() -> validateRedundancyOfTwoRegions("region1", "region2", 12, 0, 12));

    // Now simulate a rebalance
    vm0.invoke(() -> {
      InternalResourceManager manager = getCache().getInternalResourceManager();
      RebalanceResults results = doRebalance(simulate, manager);

      assertThat(results.getTotalBucketCreatesCompleted()).isEqualTo(24);
      assertThat(results.getTotalPrimaryTransfersCompleted()).isEqualTo(12);
      assertThat(results.getTotalBucketTransferBytes()).isEqualTo(0);
      assertThat(results.getTotalBucketTransfersCompleted()).isEqualTo(0);

      Set<PartitionRebalanceInfo> detailSet = results.getPartitionRebalanceDetails();
      assertThat(detailSet).hasSize(2);

      for (PartitionRebalanceInfo details : detailSet) {
        assertThat(details.getBucketCreatesCompleted()).isEqualTo(12);
        assertThat(details.getPrimaryTransfersCompleted()).isEqualTo(6);
        assertThat(details.getBucketTransferBytes()).isEqualTo(0);
        assertThat(details.getBucketTransfersCompleted()).isEqualTo(0);

        Set<PartitionMemberInfo> afterDetails = details.getPartitionMemberDetailsAfter();
        assertThat(afterDetails).hasSize(3);

        for (PartitionMemberInfo memberDetails : afterDetails) {
          if (memberDetails.getDistributedMember().equals(member1)) {
            assertThat(memberDetails.getBucketCount()).isEqualTo(12);
            assertThat(memberDetails.getPrimaryCount()).isEqualTo(6);
          } else {
            assertThat(memberDetails.getBucketCount()).isEqualTo(6);
            assertThat(memberDetails.getPrimaryCount()).isEqualTo(3);
          }
        }
        if (!simulate) {
          validateStatistics(manager, results);
        }
      }
    });

    if (!simulate) {
      for (VM vm : toArray(vm0, vm1, vm2)) {
        vm.invoke(() -> {
          PartitionedRegion region1 = (PartitionedRegion) getCache().getRegion("region1");
          PartitionedRegion region2 = (PartitionedRegion) getCache().getRegion("region2");

          PartitionRegionInfo details =
              PartitionRegionHelper.getPartitionRegionInfo(getCache().getRegion("region1"));
          assertThat(details.getCreatedBucketCount()).isEqualTo(12);
          assertThat(details.getActualRedundantCopies()).isEqualTo(1);
          assertThat(details.getLowRedundancyBucketCount()).isEqualTo(0);

          details = PartitionRegionHelper.getPartitionRegionInfo(getCache().getRegion("region2"));
          assertThat(details.getCreatedBucketCount()).isEqualTo(12);
          assertThat(details.getActualRedundantCopies()).isEqualTo(1);
          assertThat(details.getLowRedundancyBucketCount()).isEqualTo(0);

          assertThat(region2.getLocalPrimaryBucketsListTestOnly())
              .isEqualTo(region1.getLocalPrimaryBucketsListTestOnly());

          assertThat(region2.getLocalBucketsListTestOnly())
              .isEqualTo(region1.getLocalBucketsListTestOnly());
        });
      }
    }
  }

  @Test
  @Parameters({"true", "false"})
  @TestCaseName("{method}(simulate={0})")
  public void testRecoverRedundancyParallelAsyncEventQueue(boolean simulate)
      throws SecurityException {
    if (simulate) {
      invokeInEveryVM(
          () -> System.setProperty(GeodeGlossary.GEMFIRE_PREFIX + "LOG_REBALANCE", "true"));
    }

    VM vm0 = getVM(0);
    VM vm1 = getVM(1);
    VM vm2 = getVM(2);

    DistributedMember member1 = vm0.invoke(() -> {
      createPRRegionWithAsyncQueue(200);
      return getCache().getDistributedSystem().getDistributedMember();
    });

    // Create some buckets. Put enough data to cause the queue to overflow (more than 1 MB)
    vm0.invoke(() -> {
      Region<Number, String> region = getCache().getRegion("region1");
      for (int i = 0; i < 12; i++) {
        region.put(i, "A", new byte[1024 * 512]);
      }

      // The async event queue uses asynchronous writes.
      // Flush the default disk store to make sure all values have overflowed
      getCache().findDiskStore(null).flush();
    });

    // check to make sure our redundancy is impaired

    vm0.invoke(() -> {
      Region region1 = getCache().getRegion("region1");
      PartitionedRegion region2 =
          ColocationHelper.getColocatedChildRegions((PartitionedRegion) region1).get(0);

      validateRedundancyOfTwoRegions(region1, region2, 12, 0, 12);

      assertThat(getCache().getAsyncEventQueue("parallelQueue").size()).isEqualTo(12);
    });

    // Create the region on two more members, each with 1/2 of the memory
    vm1.invoke(() -> createPRRegionWithAsyncQueue(100));
    vm2.invoke(() -> createPRRegionWithAsyncQueue(100));

    vm0.invoke(() -> {
      Region region1 = getCache().getRegion("region1");
      PartitionedRegion region2 =
          ColocationHelper.getColocatedChildRegions((PartitionedRegion) region1).get(0);

      validateRedundancyOfTwoRegions(region1, region2, 12, 0, 12);

      assertThat(getCache().getAsyncEventQueue("parallelQueue").size()).isEqualTo(12);
    });

    // Now simulate a rebalance
    vm0.invoke(() -> {
      InternalResourceManager manager = getCache().getInternalResourceManager();
      RebalanceResults results = doRebalance(simulate, manager);

      assertThat(results.getTotalBucketCreatesCompleted()).isEqualTo(24);
      assertThat(results.getTotalPrimaryTransfersCompleted()).isEqualTo(12);
      assertThat(results.getTotalBucketTransferBytes()).isEqualTo(0);
      assertThat(results.getTotalBucketTransfersCompleted()).isEqualTo(0);

      Set<PartitionRebalanceInfo> detailSet = results.getPartitionRebalanceDetails();
      assertThat(detailSet).hasSize(2);

      for (PartitionRebalanceInfo details : detailSet) {
        assertThat(details.getBucketCreatesCompleted()).isEqualTo(12);
