Merge branch '3.1'

Author: keith-turner

core/src/main/java/org/apache/accumulo/core/manager/balancer/BalanceParamsImpl.java

@@ -27,6 +27,7 @@
 import java.util.SortedMap;
 import java.util.stream.Collectors;
 
+import org.apache.accumulo.core.data.TableId;
 import org.apache.accumulo.core.data.TabletId;
 import org.apache.accumulo.core.dataImpl.KeyExtent;
 import org.apache.accumulo.core.dataImpl.TabletIdImpl;
@@ -46,11 +47,13 @@ public class BalanceParamsImpl implements TabletBalancer.BalanceParameters {
   private final Set<KeyExtent> thriftCurrentMigrations;
   private final Map<String,Set<TabletServerId>> tserverResourceGroups;
   private final DataLevel currentDataLevel;
+  private final Map<String,TableId> tablesToBalance;
 
   public static BalanceParamsImpl fromThrift(SortedMap<TabletServerId,TServerStatus> currentStatus,
       Map<String,Set<TServerInstance>> currentTServerGrouping,
       SortedMap<TServerInstance,TabletServerStatus> thriftCurrentStatus,
-      Set<KeyExtent> thriftCurrentMigrations, DataLevel currentLevel) {
+      Set<KeyExtent> thriftCurrentMigrations, DataLevel currentLevel,
+      Map<String,TableId> tablesToBalance) {
     Set<TabletId> currentMigrations = thriftCurrentMigrations.stream().map(TabletIdImpl::new)
         .collect(Collectors.toUnmodifiableSet());
 
@@ -62,33 +65,37 @@ public static BalanceParamsImpl fromThrift(SortedMap<TabletServerId,TServerStatu
     });
 
     return new BalanceParamsImpl(currentStatus, tserverGroups, currentMigrations, new ArrayList<>(),
-        thriftCurrentStatus, thriftCurrentMigrations, currentLevel);
+        thriftCurrentStatus, thriftCurrentMigrations, currentLevel, tablesToBalance);
   }
 
   public BalanceParamsImpl(SortedMap<TabletServerId,TServerStatus> currentStatus,
       Map<String,Set<TabletServerId>> currentGroups, Set<TabletId> currentMigrations,
-      List<TabletMigration> migrationsOut, DataLevel currentLevel) {
+      List<TabletMigration> migrationsOut, DataLevel currentLevel,
+      Map<String,TableId> tablesToBalance) {
     this.currentStatus = currentStatus;
     this.tserverResourceGroups = currentGroups;
     this.currentMigrations = currentMigrations;
     this.migrationsOut = migrationsOut;
     this.thriftCurrentStatus = null;
     this.thriftCurrentMigrations = null;
     this.currentDataLevel = currentLevel;
+    this.tablesToBalance = tablesToBalance;
   }
 
   private BalanceParamsImpl(SortedMap<TabletServerId,TServerStatus> currentStatus,
       Map<String,Set<TabletServerId>> currentGroups, Set<TabletId> currentMigrations,
       List<TabletMigration> migrationsOut,
       SortedMap<TServerInstance,TabletServerStatus> thriftCurrentStatus,
-      Set<KeyExtent> thriftCurrentMigrations, DataLevel currentLevel) {
+      Set<KeyExtent> thriftCurrentMigrations, DataLevel currentLevel,
+      Map<String,TableId> tablesToBalance) {
     this.currentStatus = currentStatus;
     this.tserverResourceGroups = currentGroups;
     this.currentMigrations = currentMigrations;
     this.migrationsOut = migrationsOut;
     this.thriftCurrentStatus = thriftCurrentStatus;
     this.thriftCurrentMigrations = thriftCurrentMigrations;
     this.currentDataLevel = currentLevel;
+    this.tablesToBalance = tablesToBalance;
   }
 
   @Override
@@ -130,4 +137,9 @@ public Map<String,Set<TabletServerId>> currentResourceGroups() {
   public String currentLevel() {
     return currentDataLevel.name();
   }
+
+  @Override
+  public Map<String,TableId> getTablesToBalance() {
+    return tablesToBalance;
+  }
 }

core/src/main/java/org/apache/accumulo/core/rpc/TraceProtocolFactory.java

@@ -35,25 +35,32 @@
 public class TraceProtocolFactory extends TCompactProtocol.Factory {
   private static final long serialVersionUID = 1L;
 
+  private static class TraceProtocol extends TCompactProtocol {
+
+    private Span span = null;
+    private Scope scope = null;
+
+    public TraceProtocol(TTransport transport) {
+      super(transport);
+    }
+
+    @Override
+    public void writeMessageBegin(TMessage message) throws TException {
+      span = TraceUtil.startClientRpcSpan(this.getClass(), message.name);
+      scope = span.makeCurrent();
+      super.writeMessageBegin(message);
+    }
+
+    @Override
+    public void writeMessageEnd() throws TException {
+      super.writeMessageEnd();
+      scope.close();
+      span.end();
+    }
+  }
+
   @Override
   public TProtocol getProtocol(TTransport trans) {
-    return new TCompactProtocol(trans) {
-      private Span span = null;
-      private Scope scope = null;
-
-      @Override
-      public void writeMessageBegin(TMessage message) throws TException {
-        span = TraceUtil.startClientRpcSpan(this.getClass(), message.name);
-        scope = span.makeCurrent();
-        super.writeMessageBegin(message);
-      }
-
-      @Override
-      public void writeMessageEnd() throws TException {
-        super.writeMessageEnd();
-        scope.close();
-        span.end();
-      }
-    };
+    return new TraceProtocol(trans);
   }
 }

core/src/main/java/org/apache/accumulo/core/spi/balancer/BalancerEnvironment.java

@@ -40,6 +40,11 @@ public interface BalancerEnvironment extends ServiceEnvironment {
    * Many Accumulo plugins are given table IDs as this is what Accumulo uses internally to identify
    * tables. This provides a mapping of table names to table IDs for the purposes of translating
    * and/or enumerating the existing tables.
+   *
+   * <p>
+   * This returns all tables that exists in the system. Each request to balance should limit itself
+   * to {@link TabletBalancer.BalanceParameters#getTablesToBalance()} and not balance everything
+   * returned by this.
    */
   Map<String,TableId> getTableIdMap();
 

core/src/main/java/org/apache/accumulo/core/spi/balancer/HostRegexTableLoadBalancer.java

@@ -379,7 +379,7 @@ public void getAssignments(AssignmentParameters params) {
   public long balance(BalanceParameters params) {
     long minBalanceTime = 20_000;
     // Iterate over the tables and balance each of them
-    Map<String,TableId> tableIdMap = environment.getTableIdMap();
+    Map<String,TableId> tableIdMap = params.getTablesToBalance();
     Map<TableId,String> tableIdToTableName = tableIdMap.entrySet().stream()
         .collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
     tableIdToTableName.keySet().forEach(this::checkTableConfig);
@@ -510,8 +510,9 @@ public long balance(BalanceParameters params) {
         continue;
       }
       ArrayList<TabletMigration> newMigrations = new ArrayList<>();
-      getBalancerForTable(tableId).balance(new BalanceParamsImpl(currentView,
-          params.currentResourceGroups(), migrations, newMigrations, DataLevel.of(tableId)));
+      getBalancerForTable(tableId)
+          .balance(new BalanceParamsImpl(currentView, params.currentResourceGroups(), migrations,
+              newMigrations, DataLevel.of(tableId), Map.of(tableName, tableId)));
 
       if (newMigrations.isEmpty()) {
         tableToTimeSinceNoMigrations.remove(tableId);

core/src/main/java/org/apache/accumulo/core/spi/balancer/TableLoadBalancer.java

@@ -124,7 +124,8 @@ protected TabletBalancer getBalancerForTable(TableId tableId) {
       }
 
       if (balancer == null) {
-        log.info("Using balancer {} for table {}", SimpleLoadBalancer.class.getName(), tableId);
+        log.info("Creating balancer {} limited to balancing table {}",
+            SimpleLoadBalancer.class.getName(), tableId);
         balancer = new SimpleLoadBalancer(tableId);
       }
       perTableBalancers.put(tableId, balancer);
@@ -210,9 +211,10 @@ public boolean needsReassignment(CurrentAssignment currentAssignment) {
   @Override
   public long balance(BalanceParameters params) {
     long minBalanceTime = 5_000;
-    // Iterate over the tables and balance each of them
     final DataLevel currentDataLevel = DataLevel.valueOf(params.currentLevel());
-    for (TableId tableId : environment.getTableIdMap().values()) {
+    for (Entry<String,TableId> entry : params.getTablesToBalance().entrySet()) {
+      String tableName = entry.getKey();
+      TableId tableId = entry.getValue();
       final String tableResourceGroup = getResourceGroupNameForTable(tableId);
       // get the group of tservers for this table
       SortedMap<TabletServerId,TServerStatus> groupedTServers = getCurrentSetForTable(
@@ -222,9 +224,10 @@ public long balance(BalanceParameters params) {
         continue;
       }
       ArrayList<TabletMigration> newMigrations = new ArrayList<>();
-      long tableBalanceTime = getBalancerForTable(tableId)
-          .balance(new BalanceParamsImpl(groupedTServers, params.currentResourceGroups(),
-              params.currentMigrations(), newMigrations, currentDataLevel));
+      long tableBalanceTime =
+          getBalancerForTable(tableId).balance(new BalanceParamsImpl(groupedTServers,
+              params.currentResourceGroups(), params.currentMigrations(), newMigrations,
+              currentDataLevel, Map.of(tableName, tableId)));
       if (tableBalanceTime < minBalanceTime) {
         minBalanceTime = tableBalanceTime;
       }

core/src/main/java/org/apache/accumulo/core/spi/balancer/TabletBalancer.java

@@ -24,6 +24,7 @@
 import java.util.SortedMap;
 
 import org.apache.accumulo.core.conf.Property;
+import org.apache.accumulo.core.data.TableId;
 import org.apache.accumulo.core.data.TabletId;
 import org.apache.accumulo.core.spi.balancer.data.TServerStatus;
 import org.apache.accumulo.core.spi.balancer.data.TabletMigration;
@@ -114,6 +115,15 @@ interface BalanceParameters {
      * @since 2.1.4
      */
     String currentLevel();
+
+    /**
+     * This is the set of tables the balancer should consider. Balancing any tables outside of this
+     * set will be ignored and result in an error in the logs.
+     *
+     * @return map of table names to table ids that should be balanced.
+     * @since 2.1.4
+     */
+    Map<String,TableId> getTablesToBalance();
   }
 
   /**

core/src/test/java/org/apache/accumulo/core/spi/balancer/GroupBalancerTest.java

@@ -107,10 +107,11 @@ protected int getMaxMigrations() {
         }
       };
 
-      balance(balancer, maxMigrations, tid);
+      balance(balancer, maxMigrations, tid, Map.of("1", tid));
     }
 
-    public void balance(TabletBalancer balancer, int maxMigrations, TableId tid) {
+    public void balance(TabletBalancer balancer, int maxMigrations, TableId tid,
+        Map<String,TableId> tablesToBalance) {
 
       while (true) {
         Set<TabletId> migrations = new HashSet<>();
@@ -124,7 +125,7 @@ public void balance(TabletBalancer balancer, int maxMigrations, TableId tid) {
 
         balancer.balance(new BalanceParamsImpl(current,
             Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, current.keySet()), migrations,
-            migrationsOut, DataLevel.of(tid)));
+            migrationsOut, DataLevel.of(tid), tablesToBalance));
 
         assertTrue(migrationsOut.size() <= (maxMigrations + 5),
             "Max Migration exceeded " + maxMigrations + " " + migrationsOut.size());

core/src/test/java/org/apache/accumulo/core/spi/balancer/HostRegexTableLoadBalancerReconfigurationTest.java

@@ -112,7 +112,7 @@ public void testConfigurationChanges() {
     UtilWaitThread.sleep(3000);
     this.balance(new BalanceParamsImpl(Collections.unmodifiableSortedMap(allTabletServers),
         Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, allTabletServers.keySet()), migrations,
-        migrationsOut, DataLevel.USER));
+        migrationsOut, DataLevel.USER, tables));
     assertEquals(0, migrationsOut.size());
     // Change property, simulate call by TableConfWatcher
 
@@ -125,7 +125,7 @@ public void testConfigurationChanges() {
     this.lastOOBCheckTimes.put(DataLevel.USER, System.currentTimeMillis() / 2);
     this.balance(new BalanceParamsImpl(Collections.unmodifiableSortedMap(allTabletServers),
         Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, allTabletServers.keySet()), migrations,
-        migrationsOut, DataLevel.USER));
+        migrationsOut, DataLevel.USER, tables));
     assertEquals(5, migrationsOut.size());
     for (TabletMigration migration : migrationsOut) {
       assertTrue(migration.getNewTabletServer().getHost().startsWith("192.168.0.1")

core/src/test/java/org/apache/accumulo/core/spi/balancer/HostRegexTableLoadBalancerTest.java

@@ -97,7 +97,7 @@ public void testBalance() {
     SortedMap<TabletServerId,TServerStatus> current = createCurrent(15);
     long wait = this.balance(new BalanceParamsImpl(current,
         Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, current.keySet()), migrations, migrationsOut,
-        DataLevel.USER));
+        DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // should balance four tablets in one of the tables before reaching max
     assertEquals(4, migrationsOut.size());
@@ -110,7 +110,7 @@ public void testBalance() {
     SortedMap<TabletServerId,TServerStatus> current2 = createCurrent(15);
     wait = this.balance(new BalanceParamsImpl(current2,
         Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, current2.keySet()), migrations, migrationsOut,
-        DataLevel.USER));
+        DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // should balance four tablets in one of the other tables before reaching max
     assertEquals(4, migrationsOut.size());
@@ -123,7 +123,7 @@ public void testBalance() {
     SortedMap<TabletServerId,TServerStatus> current3 = createCurrent(15);
     wait = this.balance(new BalanceParamsImpl(current3,
         Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, current3.keySet()), migrations, migrationsOut,
-        DataLevel.USER));
+        DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // should balance four tablets in one of the other tables before reaching max
     assertEquals(4, migrationsOut.size());
@@ -136,7 +136,7 @@ public void testBalance() {
     SortedMap<TabletServerId,TServerStatus> current4 = createCurrent(15);
     wait = this.balance(new BalanceParamsImpl(current4,
         Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, current4.keySet()), migrations, migrationsOut,
-        DataLevel.USER));
+        DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // no more balancing to do
     assertEquals(0, migrationsOut.size());
@@ -154,7 +154,7 @@ public void testBalanceWithTooManyOutstandingMigrations() {
     SortedMap<TabletServerId,TServerStatus> current = createCurrent(15);
     long wait = this.balance(new BalanceParamsImpl(current,
         Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, current.keySet()), migrations, migrationsOut,
-        DataLevel.USER));
+        DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // no migrations should have occurred as 10 is the maxOutstandingMigrations
     assertEquals(0, migrationsOut.size());
@@ -510,7 +510,7 @@ public void testOutOfBoundsTablets() {
     SortedMap<TabletServerId,TServerStatus> current = createCurrent(15);
     this.balance(new BalanceParamsImpl(current,
         Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, current.keySet()), migrations, migrationsOut,
-        DataLevel.USER));
+        DataLevel.USER, environment.getTableIdMap()));
     assertEquals(2, migrationsOut.size());
   }
 

core/src/test/java/org/apache/accumulo/core/spi/balancer/SimpleLoadBalancerTest.java

@@ -206,7 +206,7 @@ public void testUnevenAssignment() {
       SortedMap<TabletServerId,TServerStatus> tservers = getAssignments(servers);
       balancer.balance(new BalanceParamsImpl(tservers,
           Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, tservers.keySet()), migrations,
-          migrationsOut, DataLevel.USER));
+          migrationsOut, DataLevel.USER, Map.of()));
       if (migrationsOut.isEmpty()) {
         break;
       }
@@ -251,7 +251,7 @@ public void testUnevenAssignment2() {
       SortedMap<TabletServerId,TServerStatus> tservers = getAssignments(servers);
       balancer.balance(new BalanceParamsImpl(tservers,
           Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, tservers.keySet()), migrations,
-          migrationsOut, DataLevel.USER));
+          migrationsOut, DataLevel.USER, Map.of()));
       if (migrationsOut.isEmpty()) {
         break;
       }

core/src/test/java/org/apache/accumulo/core/spi/balancer/TableLoadBalancerTest.java

@@ -147,15 +147,15 @@ public void test() {
     tls.init(environment);
     tls.balance(
         new BalanceParamsImpl(state, Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, state.keySet()),
-            migrations, migrationsOut, DataLevel.USER));
+            migrations, migrationsOut, DataLevel.USER, tableIdMap));
     assertEquals(0, migrationsOut.size());
 
     state.put(mkts("10.0.0.2", 2345, "0x02030405"), status());
     tls = new TableLoadBalancer();
     tls.init(environment);
     tls.balance(
         new BalanceParamsImpl(state, Map.of(Constants.DEFAULT_RESOURCE_GROUP_NAME, state.keySet()),
-            migrations, migrationsOut, DataLevel.USER));
+            migrations, migrationsOut, DataLevel.USER, tableIdMap));
     int count = 0;
     Map<TableId,Integer> movedByTable = new HashMap<>();
     movedByTable.put(TableId.of(t1Id), 0);