Merge branch '2.1' into 3.1

Author: keith-turner

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

@@ -20,10 +20,12 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 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;
@@ -42,38 +44,43 @@ public class BalanceParamsImpl implements TabletBalancer.BalanceParameters {
   private final SortedMap<TServerInstance,TabletServerStatus> thriftCurrentStatus;
   private final Set<KeyExtent> thriftCurrentMigrations;
   private final DataLevel currentDataLevel;
+  private final Map<String,TableId> tablesToBalance;
 
   public static BalanceParamsImpl fromThrift(SortedMap<TabletServerId,TServerStatus> currentStatus,
       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());
 
     return new BalanceParamsImpl(currentStatus, currentMigrations, new ArrayList<>(),
-        thriftCurrentStatus, thriftCurrentMigrations, currentLevel);
+        thriftCurrentStatus, thriftCurrentMigrations, currentLevel, tablesToBalance);
   }
 
   public BalanceParamsImpl(SortedMap<TabletServerId,TServerStatus> currentStatus,
-      Set<TabletId> currentMigrations, List<TabletMigration> migrationsOut,
-      DataLevel currentLevel) {
+      Set<TabletId> currentMigrations, List<TabletMigration> migrationsOut, DataLevel currentLevel,
+      Map<String,TableId> tablesToBalance) {
     this.currentStatus = currentStatus;
     this.currentMigrations = currentMigrations;
     this.migrationsOut = migrationsOut;
     this.thriftCurrentStatus = null;
     this.thriftCurrentMigrations = null;
     this.currentDataLevel = currentLevel;
+    this.tablesToBalance = tablesToBalance;
   }
 
   private BalanceParamsImpl(SortedMap<TabletServerId,TServerStatus> currentStatus,
       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.currentMigrations = currentMigrations;
     this.migrationsOut = migrationsOut;
     this.thriftCurrentStatus = thriftCurrentStatus;
     this.thriftCurrentMigrations = thriftCurrentMigrations;
     this.currentDataLevel = currentLevel;
+    this.tablesToBalance = tablesToBalance;
   }
 
   @Override
@@ -110,4 +117,9 @@ public void addMigration(KeyExtent extent, TServerInstance oldServer, TServerIns
   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

@@ -373,7 +373,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);
@@ -504,8 +504,8 @@ public long balance(BalanceParameters params) {
         continue;
       }
       ArrayList<TabletMigration> newMigrations = new ArrayList<>();
-      getBalancerForTable(tableId).balance(
-          new BalanceParamsImpl(currentView, migrations, newMigrations, DataLevel.of(tableId)));
+      getBalancerForTable(tableId).balance(new BalanceParamsImpl(currentView, 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

@@ -98,7 +98,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);
@@ -124,13 +125,14 @@ public void getAssignments(AssignmentParameters params) {
   @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();
       ArrayList<TabletMigration> newMigrations = new ArrayList<>();
-      long tableBalanceTime =
-          getBalancerForTable(tableId).balance(new BalanceParamsImpl(params.currentStatus(),
-              params.currentMigrations(), newMigrations, currentDataLevel));
+      long tableBalanceTime = getBalancerForTable(tableId)
+          .balance(new BalanceParamsImpl(params.currentStatus(), 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;
@@ -102,6 +103,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

@@ -106,10 +106,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<>();
@@ -121,8 +122,8 @@ public void balance(TabletBalancer balancer, int maxMigrations, TableId tid) {
               new org.apache.accumulo.core.manager.thrift.TabletServerStatus()));
         }
 
-        balancer
-            .balance(new BalanceParamsImpl(current, migrations, migrationsOut, DataLevel.of(tid)));
+        balancer.balance(new BalanceParamsImpl(current, migrations, 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

@@ -108,7 +108,7 @@ public void testConfigurationChanges() {
     // getOnlineTabletsForTable
     UtilWaitThread.sleep(3000);
     this.balance(new BalanceParamsImpl(Collections.unmodifiableSortedMap(allTabletServers),
-        migrations, migrationsOut, DataLevel.USER));
+        migrations, migrationsOut, DataLevel.USER, tables));
     assertEquals(0, migrationsOut.size());
     // Change property, simulate call by TableConfWatcher
 
@@ -120,7 +120,7 @@ public void testConfigurationChanges() {
     // populate the map with an older time value
     this.lastOOBCheckTimes.put(DataLevel.USER, System.currentTimeMillis() / 2);
     this.balance(new BalanceParamsImpl(Collections.unmodifiableSortedMap(allTabletServers),
-        migrations, migrationsOut, DataLevel.USER));
+        migrations, 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

@@ -94,7 +94,7 @@ public void testBalance() {
     List<TabletMigration> migrationsOut = new ArrayList<>();
     long wait =
         this.balance(new BalanceParamsImpl(Collections.unmodifiableSortedMap(createCurrent(15)),
-            migrations, migrationsOut, DataLevel.USER));
+            migrations, migrationsOut, DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // should balance four tablets in one of the tables before reaching max
     assertEquals(4, migrationsOut.size());
@@ -105,7 +105,7 @@ public void testBalance() {
     }
     migrationsOut.clear();
     wait = this.balance(new BalanceParamsImpl(Collections.unmodifiableSortedMap(createCurrent(15)),
-        migrations, migrationsOut, DataLevel.USER));
+        migrations, migrationsOut, DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // should balance four tablets in one of the other tables before reaching max
     assertEquals(4, migrationsOut.size());
@@ -116,7 +116,7 @@ public void testBalance() {
     }
     migrationsOut.clear();
     wait = this.balance(new BalanceParamsImpl(Collections.unmodifiableSortedMap(createCurrent(15)),
-        migrations, migrationsOut, DataLevel.USER));
+        migrations, migrationsOut, DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // should balance four tablets in one of the other tables before reaching max
     assertEquals(4, migrationsOut.size());
@@ -127,7 +127,7 @@ public void testBalance() {
     }
     migrationsOut.clear();
     wait = this.balance(new BalanceParamsImpl(Collections.unmodifiableSortedMap(createCurrent(15)),
-        migrations, migrationsOut, DataLevel.USER));
+        migrations, migrationsOut, DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // no more balancing to do
     assertEquals(0, migrationsOut.size());
@@ -144,7 +144,7 @@ public void testBalanceWithTooManyOutstandingMigrations() {
     migrations.addAll(tableTablets.get(BAR.getTableName()));
     long wait =
         this.balance(new BalanceParamsImpl(Collections.unmodifiableSortedMap(createCurrent(15)),
-            migrations, migrationsOut, DataLevel.USER));
+            migrations, migrationsOut, DataLevel.USER, environment.getTableIdMap()));
     assertEquals(20000, wait);
     // no migrations should have occurred as 10 is the maxOutstandingMigrations
     assertEquals(0, migrationsOut.size());
@@ -491,8 +491,8 @@ public void testOutOfBoundsTablets() {
     init(DEFAULT_TABLE_PROPERTIES);
     Set<TabletId> migrations = new HashSet<>();
     List<TabletMigration> migrationsOut = new ArrayList<>();
-    this.balance(
-        new BalanceParamsImpl(createCurrent(15), migrations, migrationsOut, DataLevel.USER));
+    this.balance(new BalanceParamsImpl(createCurrent(15), migrations, migrationsOut, DataLevel.USER,
+        environment.getTableIdMap()));
     assertEquals(2, migrationsOut.size());
   }
 

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

@@ -204,7 +204,7 @@ public void testUnevenAssignment() {
     while (true) {
       List<TabletMigration> migrationsOut = new ArrayList<>();
       balancer.balance(new BalanceParamsImpl(getAssignments(servers), migrations, migrationsOut,
-          DataLevel.USER));
+          DataLevel.USER, Map.of()));
       if (migrationsOut.isEmpty()) {
         break;
       }
@@ -247,7 +247,7 @@ public void testUnevenAssignment2() {
     while (true) {
       List<TabletMigration> migrationsOut = new ArrayList<>();
       balancer.balance(new BalanceParamsImpl(getAssignments(servers), migrations, migrationsOut,
-          DataLevel.USER));
+          DataLevel.USER, Map.of()));
       if (migrationsOut.isEmpty()) {
         break;
       }

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

@@ -141,13 +141,15 @@ public void test() {
     List<TabletMigration> migrationsOut = new ArrayList<>();
     TableLoadBalancer tls = new TableLoadBalancer();
     tls.init(environment);
-    tls.balance(new BalanceParamsImpl(state, migrations, migrationsOut, DataLevel.USER));
+    tls.balance(
+        new BalanceParamsImpl(state, 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, migrations, migrationsOut, DataLevel.USER));
+    tls.balance(
+        new BalanceParamsImpl(state, migrations, migrationsOut, DataLevel.USER, tableIdMap));
     int count = 0;
     Map<TableId,Integer> movedByTable = new HashMap<>();
     movedByTable.put(TableId.of(t1Id), 0);