/**
    *
    * 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.Assert.assertTrue;
  
  import java.io.IOException;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.atomic.AtomicInteger;
  
  import org.apache.hadoop.hbase.HBaseTestingUtility;
  import org.apache.hadoop.hbase.HConstants;
  import org.apache.hadoop.hbase.HRegionInfo;
  import org.apache.hadoop.hbase.TableName;
  import org.apache.hadoop.hbase.regionserver.HRegionServer;
  import org.apache.hadoop.hbase.testclassification.MasterTests;
  import org.apache.hadoop.hbase.testclassification.MediumTests;
  import org.apache.hadoop.hbase.util.Bytes;
  import org.junit.After;
  import org.junit.Before;
  import org.junit.Test;
  import org.junit.experimental.categories.Category;
  
  @Category({MasterTests.class, MediumTests.class})
  public class TestMasterBalanceThrottling {
    private static final HBaseTestingUtility TEST_UTIL = new HBaseTestingUtility();
    private static final byte[] FAMILYNAME = Bytes.toBytes("fam");
  
    @Before
    public void setupConfiguration() {
      TEST_UTIL.getConfiguration().set(HConstants.HBASE_MASTER_LOADBALANCER_CLASS,
          "org.apache.hadoop.hbase.master.balancer.SimpleLoadBalancer");
    }
  
    @After
    public void shutdown() throws Exception {
      TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_BALANCER_MAX_BALANCING,
        HConstants.DEFAULT_HBASE_BALANCER_PERIOD);
      TEST_UTIL.getConfiguration().setDouble(HConstants.HBASE_MASTER_BALANCER_MAX_RIT_PERCENT,
        HConstants.DEFAULT_HBASE_MASTER_BALANCER_MAX_RIT_PERCENT);
      TEST_UTIL.shutdownMiniCluster();
    }
  
    @Test(timeout = 60000)
    public void testThrottlingByBalanceInterval() throws Exception {
      // Use default config and start a cluster of two regionservers.
      TEST_UTIL.startMiniCluster(2);
  
      TableName tableName = createTable("testNoThrottling");
      final HMaster master = TEST_UTIL.getHBaseCluster().getMaster();
  
      // Default max balancing time is 300000 ms and there are 50 regions to balance
      // The balance interval is 6000 ms, much longger than the normal region in transition duration
      // So the master can balance the region one by one
      unbalance(master, tableName);
      AtomicInteger maxCount = new AtomicInteger(0);
      AtomicBoolean stop = new AtomicBoolean(false);
      Thread checker = startBalancerChecker(master, maxCount, stop);
      master.balance();
      stop.set(true);
      checker.interrupt();
      checker.join();
      assertTrue("max regions in transition: " + maxCount.get(), maxCount.get() == 1);
  
      TEST_UTIL.deleteTable(tableName);
    }
  
    @Test(timeout = 60000)
    public void testThrottlingByMaxRitPercent() throws Exception {
      // Set max balancing time to 500 ms and max percent of regions in transition to 0.05
      TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_BALANCER_MAX_BALANCING, 500);
      TEST_UTIL.getConfiguration().setDouble(HConstants.HBASE_MASTER_BALANCER_MAX_RIT_PERCENT, 0.05);
      TEST_UTIL.startMiniCluster(2);
  
      TableName tableName = createTable("testThrottlingByMaxRitPercent");
      final HMaster master = TEST_UTIL.getHBaseCluster().getMaster();
  
      unbalance(master, tableName);
      AtomicInteger maxCount = new AtomicInteger(0);
      AtomicBoolean stop = new AtomicBoolean(false);
      Thread checker = startBalancerChecker(master, maxCount, stop);
      master.balance();
      stop.set(true);
     checker.interrupt();
     checker.join();
     // The max number of regions in transition is 100 * 0.05 = 5
     assertTrue("max regions in transition: " + maxCount.get(), maxCount.get() == 5);
 
     TEST_UTIL.deleteTable(tableName);
   }
 
   private TableName createTable(String table) throws IOException {
     TableName tableName = TableName.valueOf(table);
     byte[] startKey = new byte[] { 0x00 };
     byte[] stopKey = new byte[] { 0x7f };
     TEST_UTIL.createTable(tableName, new byte[][] { FAMILYNAME }, 1, startKey, stopKey,
       100);
     return tableName;
   }
 
   private Thread startBalancerChecker(final HMaster master, final AtomicInteger maxCount,
       final AtomicBoolean stop) {
     Runnable checker = new Runnable() {
       @Override
       public void run() {
         while (!stop.get()) {
           maxCount.set(Math.max(maxCount.get(), master.getAssignmentManager().getRegionStates()
               .getRegionsInTransitionCount()));
           try {
             Thread.sleep(10);
           } catch (InterruptedException e) {
             e.printStackTrace();
           }
         }
       }
     };
     Thread thread = new Thread(checker);
     thread.start();
     return thread;
   }
 
   private void unbalance(HMaster master, TableName tableName) throws Exception {
     while (master.getAssignmentManager().getRegionStates().getRegionsInTransitionCount() > 0) {
       Thread.sleep(100);
     }
     HRegionServer biasedServer = TEST_UTIL.getMiniHBaseCluster().getRegionServer(0);
     for (HRegionInfo regionInfo : TEST_UTIL.getHBaseAdmin().getTableRegions(tableName)) {
       master.move(regionInfo.getEncodedNameAsBytes(),
         Bytes.toBytes(biasedServer.getServerName().getServerName()));
     }
     while (master.getAssignmentManager().getRegionStates().getRegionsInTransitionCount() > 0) {
       Thread.sleep(100);
     }
   }
 }