Added multiple table subset balancing test

SimpleLoadBalancer did not have a multi-table balance test for a subset
of tables.
Added a test for balancing two of three total tables

Author: ddanielr

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

@@ -21,13 +21,11 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -248,32 +246,88 @@ public void testUnevenAssignment2() {
     for (int i = 0; i < 10; i++) {
       shortServer.getValue().tablets.add(makeTablet(table2, null, null));
     }
-    Set<TableId> migratedTables;
+    Map<TableId,Integer> migratedTables;
     // average is 58, with 2 at 59: we need 48 more moved to the short server
-    migratedTables = balanceTables(Map.of(table1, t1Id, table2, t2Id), 48);
-    assertEquals(1, migratedTables.size());
-    assertTrue(migratedTables.contains(t1Id));
+    migratedTables = balanceTables(Map.of(table1, t1Id, table2, t2Id));
+    assertEquals(48, migratedTables.get(t1Id));
 
     // ShortTable tablets are only on the short server: 9 should be migrated off
-    migratedTables = balanceTables(Map.of(table2, t2Id), 9);
+    migratedTables = balanceTables(Map.of(table2, t2Id));
     assertEquals(1, migratedTables.size());
-    assertTrue(migratedTables.contains(t2Id));
+    assertEquals(9, migratedTables.get(t2Id));
 
     // ShortServer now has 9 less LongTable tablets: 9 LongTable tablets should migrate back
-    migratedTables = balanceTables(Map.of(table1, t1Id), 9);
+    migratedTables = balanceTables(Map.of(table1, t1Id));
     assertEquals(1, migratedTables.size());
-    assertTrue(migratedTables.contains(t1Id));
+    assertEquals(9, migratedTables.get(t1Id));
 
     // Servers are now fully balanced. Attempt to balance both tables again
-    migratedTables = balanceTables(Map.of(table1, t1Id, table2, t2Id), 0);
+    migratedTables = balanceTables(Map.of(table1, t1Id, table2, t2Id));
+    assertEquals(0, migratedTables.size(), "Migrated Tables: " + migratedTables);
+
+  }
+
+  @Test
+  public void testBalanceSubset() {
+    // make 26 servers
+    for (char c : "abcdefghijklmnopqrstuvwxyz".toCharArray()) {
+      TabletServerId tsid = new TabletServerIdImpl("127.0.0.1", c, Character.toString(c));
+      FakeTServer fakeTServer = new FakeTServer();
+      servers.put(tsid, fakeTServer);
+    }
+    String table1 = "Long";
+    String table2 = "Short";
+    String table3 = "neverMigrate";
+
+    // Generate Table Ids
+    TableId t1Id = TableId.of(table1);
+    TableId t2Id = TableId.of(table2);
+    TableId t3Id = TableId.of(table3);
+
+    // put 10 tablets on 8 of the servers that should never migrate
+    int randServer = 0;
+    for (Entry<TabletServerId,FakeTServer> entry : servers.entrySet()) {
+      randServer++;
+      if (randServer % 3 == 0) {
+        for (int i = 0; i < 10; i++) {
+          entry.getValue().tablets.add(makeTablet(table3, null, null));
+        }
+      }
+    }
+
+    // put 60 tablets on 25 of them
+    List<Entry<TabletServerId,FakeTServer>> shortList = new ArrayList<>(servers.entrySet());
+    Entry<TabletServerId,FakeTServer> shortServer = shortList.remove(0);
+    for (int i = 0; i < 60; i++) {
+      for (Entry<TabletServerId,FakeTServer> entry : shortList) {
+        entry.getValue().tablets.add(makeTablet(table1, null, null));
+      }
+    }
+    // put 10 on the that short server:
+    for (int i = 0; i < 10; i++) {
+      shortServer.getValue().tablets.add(makeTablet(table2, null, null));
+    }
+
+    Map<TableId,Integer> migratedTables;
+    // average is 58, with 2 at 59: we need 48 more moved to the short server
+    migratedTables = balanceTables(Map.of(table1, t1Id, table2, t2Id));
+    assertEquals(1, migratedTables.size());
+    assertEquals(48, migratedTables.get(t1Id));
+
+    // Servers are now fully balanced. Attempt to balance both tables again
+    migratedTables = balanceTables(Map.of(table1, t1Id, table2, t2Id));
     assertEquals(0, migratedTables.size());
+
+    // Only balance the neverMigrate table
+    migratedTables = balanceTables(Map.of(table3, t3Id));
+    assertEquals(1, migratedTables.size(), "Too many tables migrated: " + migratedTables);
+    assertEquals(54, migratedTables.get(t3Id));
   }
 
-  private Set<TableId> balanceTables(Map<String,TableId> tablesToBalance, int expectedMigrations) {
+  private Map<TableId,Integer> balanceTables(Map<String,TableId> tablesToBalance) {
     TestSimpleLoadBalancer balancer = new TestSimpleLoadBalancer();
-    Set<TableId> migratedTables = new HashSet<>();
+    Map<TableId,Integer> migratedTables = new HashMap<>();
     Set<TabletId> migrations = Collections.emptySet();
-    int moved = 0;
     // balance until we can't balance no more!
     while (true) {
       List<TabletMigration> migrationsOut = new ArrayList<>();
@@ -284,15 +338,18 @@ private Set<TableId> balanceTables(Map<String,TableId> tablesToBalance, int expe
       }
       for (TabletMigration migration : migrationsOut) {
         if (servers.get(migration.getOldTabletServer()).tablets.remove(migration.getTablet())) {
-          moved++;
-          migratedTables.add(migration.getTablet().getTable());
+          TableId tableId = migration.getTablet().getTable();
+          Integer value = 0;
+          if (migratedTables.containsKey(tableId)) {
+            value = migratedTables.get(tableId);
+          }
+          migratedTables.put(tableId, value + 1);
         }
         last.remove(migration.getTablet());
         servers.get(migration.getNewTabletServer()).tablets.add(migration.getTablet());
         last.put(migration.getTablet(), migration.getNewTabletServer());
       }
     }
-    assertEquals(expectedMigrations, moved);
     return migratedTables;
   }