Improve RoutingNode invariant assertion efficiency

[BalanceUnbalancedClusterTests take 10 or 11 minutes][1] in CI. The
reason is because of O(N^2) assertion behavior. This change moves the
invariant assertions out of the individual muting calls up to the
RoutingNodes level where it can do the assertion after bulk mutations.
This change is safe because assertions are only moved out of
package-private methods that are only called by RoutingNodes; the public
API of RoutingNode still enforces all the same invariants.

[1]: https://build.ci.opensearch.org/job/gradle-check/74280/testReport/org.opensearch.cluster.routing.allocation/BalanceUnbalancedClusterTests/

Signed-off-by: Andrew Ross <andrross@amazon.com>

Author: andrross

server/src/main/java/org/opensearch/cluster/routing/RoutingNode.java

@@ -269,7 +269,6 @@ public Collection<ShardRouting> getInitializingShards() {
      * @param shard Shard to create on this Node
      */
     void add(ShardRouting shard) {
-        assert invariant();
         if (shards.put(shard) != null) {
             throw new IllegalStateException(
                 "Trying to add a shard "
@@ -290,11 +289,9 @@ void add(ShardRouting shard) {
             relocatingShardsBucket.add(shard);
         }
         shardsByIndex.computeIfAbsent(shard.index(), k -> new LinkedHashSet<>()).add(shard);
-        assert invariant();
     }
 
     void update(ShardRouting oldShard, ShardRouting newShard) {
-        assert invariant();
         if (shards.containsKey(oldShard.shardId()) == false) {
             // Shard was already removed by routing nodes iterator
             // TODO: change caller logic in RoutingNodes so that this check can go away
@@ -320,11 +317,9 @@ void update(ShardRouting oldShard, ShardRouting newShard) {
             relocatingShardsBucket.add(newShard);
         }
         shardsByIndex.computeIfAbsent(newShard.index(), k -> new LinkedHashSet<>()).add(newShard);
-        assert invariant();
     }
 
     void remove(ShardRouting shard) {
-        assert invariant();
         ShardRouting previousValue = shards.remove(shard.shardId());
         assert previousValue == shard : "expected shard " + previousValue + " but was " + shard;
         if (shard.initializing()) {
@@ -338,7 +333,6 @@ void remove(ShardRouting shard) {
         if (shardsByIndex.get(shard.index()).isEmpty()) {
             shardsByIndex.remove(shard.index());
         }
-        assert invariant();
     }
 
     /**
@@ -502,7 +496,7 @@ public boolean isEmpty() {
         return shards.isEmpty();
     }
 
-    private boolean invariant() {
+    boolean invariant() {
 
         // initializingShards must consistent with that in shards
         Collection<ShardRouting> shardRoutingsInitializing = StreamSupport.stream(shards.spliterator(), false)

server/src/main/java/org/opensearch/cluster/routing/RoutingNodes.java

@@ -174,6 +174,7 @@ public RoutingNodes(ClusterState clusterState, boolean readOnly) {
                 }
             }
         }
+        assert nodesToShards.values().stream().allMatch(RoutingNode::invariant);
     }
 
     private void addRecovery(ShardRouting routing) {
@@ -563,6 +564,7 @@ public ShardRouting initializeShard(
         assert unassignedShard.unassigned() : "expected an unassigned shard " + unassignedShard;
         ShardRouting initializedShard = unassignedShard.initialize(nodeId, existingAllocationId, expectedSize);
         node(nodeId).add(initializedShard);
+        assert node(nodeId).invariant();
         inactiveShardCount++;
         if (initializedShard.primary()) {
             inactivePrimaryCount++;
@@ -591,6 +593,7 @@ public Tuple<ShardRouting, ShardRouting> relocateShard(
         ShardRouting target = source.getTargetRelocatingShard();
         updateAssigned(startedShard, source);
         node(target.currentNodeId()).add(target);
+        assert node(target.currentNodeId()).invariant();
         assignedShardsAdd(target);
         addRecovery(target);
         changes.relocationStarted(startedShard, target);
@@ -872,6 +875,7 @@ private ShardRouting promoteActiveReplicaShardToPrimary(ShardRouting replicaShar
     private void remove(ShardRouting shard) {
         assert shard.unassigned() == false : "only assigned shards can be removed here (" + shard + ")";
         node(shard.currentNodeId()).remove(shard);
+        assert node(shard.currentNodeId()).invariant();
         if (shard.initializing() && shard.relocatingNodeId() == null) {
             inactiveShardCount--;
             assert inactiveShardCount >= 0;
@@ -951,6 +955,7 @@ private void updateAssigned(ShardRouting oldShard, ShardRouting newShard) {
             + " by shard assigned to same node but was "
             + newShard;
         node(oldShard.currentNodeId()).update(oldShard, newShard);
+        assert node(oldShard.currentNodeId()).invariant();
         List<ShardRouting> shardsWithMatchingShardId = assignedShards.computeIfAbsent(oldShard.shardId(), k -> new ArrayList<>());
         int previousShardIndex = shardsWithMatchingShardId.indexOf(oldShard);
         assert previousShardIndex >= 0 : "shard to update " + oldShard + " does not exist in list of assigned shards";