/*
 * 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.hadoop.hbase.master;

import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.hadoop.hbase.HBaseTestingUtility;
import org.apache.hadoop.hbase.HConstants;
import org.apache.hadoop.hbase.ServerName;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.RegionInfo;
import org.apache.hadoop.hbase.master.assignment.AssignmentManager;
import org.apache.hadoop.hbase.master.balancer.LoadBalancerFactory;
import org.apache.hadoop.hbase.master.balancer.StochasticLoadBalancer;
import org.apache.hadoop.hbase.testclassification.MasterTests;
import org.apache.hadoop.hbase.testclassification.MediumTests;
import org.apache.hadoop.hbase.util.Bytes;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInfo;
import org.mockito.Mockito;
import org.mockito.invocation.InvocationOnMock;

@Tag(MasterTests.TAG)
@Tag(MediumTests.TAG)
public class TestMasterBalancerNPE {

  private static final HBaseTestingUtility TEST_UTIL = new HBaseTestingUtility();
  private static final byte[] FAMILYNAME = Bytes.toBytes("fam");
  private String testMethodName;

  @BeforeEach
  public void setTestMethod(TestInfo testInfo) {
    testMethodName = testInfo.getTestMethod().get().getName();
  }

  @BeforeEach
  public void setupConfiguration() {
    /**
     * Make {@link BalancerChore} not run,so does not disrupt the test.
     */
    HMaster.setDisableBalancerChoreForTest(true);
    TEST_UTIL.getConfiguration().set(HConstants.HBASE_MASTER_LOADBALANCER_CLASS,
      MyLoadBalancer.class.getName());
  }

  @AfterEach
  public void shutdown() throws Exception {
    HMaster.setDisableBalancerChoreForTest(false);
    TEST_UTIL.getConfiguration().set(HConstants.HBASE_MASTER_LOADBALANCER_CLASS,
      LoadBalancerFactory.getDefaultLoadBalancerClass().getName());
    TEST_UTIL.shutdownMiniCluster();
  }

  /**
   * This test is for HBASE-26712, to make the region is unassigned just before
   * {@link AssignmentManager#balance} is invoked on the region.
   */
  @Test
  public void testBalancerNPE() throws Exception {
    TEST_UTIL.startMiniCluster(2);
    TEST_UTIL.getAdmin().balancerSwitch(false, true);
    TableName tableName = createTable(testMethodName);
    final HMaster master = TEST_UTIL.getHBaseCluster().getMaster();
    List<RegionInfo> regionInfos = TEST_UTIL.getAdmin().getRegions(tableName);
    assertTrue(regionInfos.size() == 1);
    final ServerName serverName1 =
      TEST_UTIL.getMiniHBaseCluster().getRegionServer(0).getServerName();
    final ServerName serverName2 =
      TEST_UTIL.getMiniHBaseCluster().getRegionServer(1).getServerName();

    final RegionInfo regionInfo = regionInfos.get(0);

    MyLoadBalancer loadBalancer = (MyLoadBalancer) master.getLoadBalancer();
    MyLoadBalancer spiedLoadBalancer = Mockito.spy(loadBalancer);
    final AtomicReference<RegionPlan> regionPlanRef = new AtomicReference<RegionPlan>();

    /**
     * Mock {@link StochasticLoadBalancer#balanceTable} to return the {@link RegionPlan} to move the
     * only region to the other RegionServer.
     */
    Mockito.doAnswer((InvocationOnMock invocation) -> {
      @SuppressWarnings("unchecked")
      Map<ServerName, List<RegionInfo>> regionServerNameToRegionInfos =
        (Map<ServerName, List<RegionInfo>>) invocation.getArgument(1);

      List<ServerName> assignedRegionServerNames = new ArrayList<ServerName>();
      for (Map.Entry<ServerName, List<RegionInfo>> entry : regionServerNameToRegionInfos
        .entrySet()) {
        if (entry.getValue() != null) {
          entry.getValue().stream().forEach((reginInfo) -> {
            if (reginInfo.getTable().equals(tableName)) {
              assignedRegionServerNames.add(entry.getKey());
            }
          });
        }
      }
      assertTrue(assignedRegionServerNames.size() == 1);
      ServerName assignedRegionServerName = assignedRegionServerNames.get(0);
      ServerName notAssignedRegionServerName =
        assignedRegionServerName.equals(serverName1) ? serverName2 : serverName1;
      RegionPlan regionPlan =
        new RegionPlan(regionInfo, assignedRegionServerName, notAssignedRegionServerName);
      regionPlanRef.set(regionPlan);
      return Arrays.asList(regionPlan);
    }).when(spiedLoadBalancer).balanceTable(Mockito.eq(HConstants.ENSEMBLE_TABLE_NAME),
      Mockito.anyMap());

    AssignmentManager assignmentManager = master.getAssignmentManager();
    final AssignmentManager spiedAssignmentManager = Mockito.spy(assignmentManager);
    final CyclicBarrier cyclicBarrier = new CyclicBarrier(2);

    /**
     * Override {@link AssignmentManager#balance} to invoke real {@link AssignmentManager#balance}
     * after the region is successfully unassigned.
     */
    Mockito.doAnswer((InvocationOnMock invocation) -> {
      RegionPlan regionPlan = invocation.getArgument(0);
      RegionPlan referedRegionPlan = regionPlanRef.get();
      assertTrue(referedRegionPlan != null);
      if (referedRegionPlan.equals(regionPlan)) {
        /**
         * To make {@link AssignmentManager#unassign} could be invoked just before
         * {@link AssignmentManager#balance} is invoked.
         */
        cyclicBarrier.await();
        /**
         * After {@link AssignmentManager#unassign} is completed,we could invoke
         * {@link AssignmentManager#balance}.
         */
        cyclicBarrier.await();
      }
      /**
       * Before HBASE-26712,here may throw NPE.
       */
      return invocation.callRealMethod();
    }).when(spiedAssignmentManager).balance(Mockito.any());

    try {
      final AtomicReference<Throwable> exceptionRef = new AtomicReference<Throwable>(null);
      Thread unassignThread = new Thread(() -> {
        try {
          /**
           * To invoke {@link AssignmentManager#unassign} just before
           * {@link AssignmentManager#balance} is invoked.
           */
          cyclicBarrier.await();
          spiedAssignmentManager.unassign(regionInfo);
          assertTrue(spiedAssignmentManager.getRegionStates().getRegionAssignments().get(regionInfo)
              == null);
          /**
           * After {@link AssignmentManager#unassign} is completed,we could invoke
           * {@link AssignmentManager#balance}.
           */
          cyclicBarrier.await();
        } catch (Exception e) {
          exceptionRef.set(e);
        }
      });
      unassignThread.setName("UnassignThread");
      unassignThread.start();

      master.setLoadBalancer(spiedLoadBalancer);
      master.setAssignmentManager(spiedAssignmentManager);
      /**
       * enable balance
       */
      TEST_UTIL.getAdmin().balancerSwitch(true, false);
      /**
       * Before HBASE-26712,here invokes {@link AssignmentManager#balance(RegionPlan)} which may
       * throw NPE.
       */
      master.balanceOrUpdateMetrics();

      unassignThread.join();
      assertTrue(exceptionRef.get() == null);
    } finally {
      master.setLoadBalancer(loadBalancer);
      master.setAssignmentManager(assignmentManager);
    }
  }

  private TableName createTable(String table) throws IOException {
    TableName tableName = TableName.valueOf(table);
    TEST_UTIL.createTable(tableName, FAMILYNAME);
    return tableName;
  }

  /**
   * Define this class because the test needs to override
   * {@link StochasticLoadBalancer#balanceTable}, which is protected.
   */
  static class MyLoadBalancer extends StochasticLoadBalancer {
    @Override
    protected List<RegionPlan> balanceTable(TableName tableName,
      Map<ServerName, List<RegionInfo>> loadOfOneTable) {
      return super.balanceTable(tableName, loadOfOneTable);
    }
  }
}
