Cache WAGED partition-movement state maps

[sjainit]

Issues

• My PR addresses the following Helix issues and references them in the PR description:

No dedicated GitHub issue was filed; this is a self-contained performance optimization in the WAGED rebalancer's scoring hot path. Happy to open a tracking issue if maintainers prefer.

Description

• Here are some details about my PR, including screenshots of any UI changes:

What: Memoize the per-(resource, partition) instance-to-state map lookups used by the WAGED partition-movement soft constraints, caching them on ClusterContext.

Why: While scoring a replica, ConstraintBasedAlgorithm evaluates every candidate node in a parallelStream. For each node, PartitionMovementConstraint and BaselineInfluenceConstraint called getStateMap(replica, ...), which looked up the baseline and best-possible ResourceAssignment for the replica's (resource, partition). That result is independent of the candidate node, so the identical map was recomputed once per node — and each call allocated a throwaway new Partition(partitionName) key. For a cluster with N assignable nodes this repeated the same work N times per replica per constraint.

How: ClusterContext already owns the immutable baseline and best-possible assignments for the whole rebalance run. It now exposes getBaselineAssignmentStateMap(resource, partition) and getBestPossibleAssignmentStateMap(resource, partition), backed by ConcurrentHashMap caches (safe under the parallel scoring stream), so each map is resolved exactly once and reused across all nodes. The two movement constraints read these accessors, and the now-unused getStateMap helper is removed from AbstractPartitionMovementConstraint. Behavior is unchanged: absent resources/partitions still return Collections.emptyMap().

Tests

• The following tests are written for this issue:

  • TestClusterContext#testAssignmentStateMapLookup — verifies the new accessors return the correct instance-to-state map, and an empty map for a missing resource / missing partition / empty assignment.
  • TestClusterContext#testAssignmentStateMapMemoization — verifies a partition's underlying getReplicaMap is resolved only once across repeated lookups (proving memoization).
  • TestPartitionMovementConstraint — updated to stub the new cached accessors; existing assertions and expected scores are unchanged.

• Local code review completed

• The following is the result of the "mvn test" command on the appropriate module:

  Maven is not available in my local environment, so mvn test was not run locally; relying on the project CI to execute the helix-core test suite. The change is confined to helix-core and is covered by the unit tests listed above.

Changes that Break Backward Compatibility (Optional)

None. getStateMap was a package-private helper on the package-private abstract class AbstractPartitionMovementConstraint with no external callers; the new ClusterContext accessors are purely additive. No public API or behavior changes.

Documentation (Optional)

Not applicable — internal performance optimization with no new user-facing functionality.

Commits

• My commit uses an imperative subject separated from the body by a blank line, wrapped at 72 columns, and explains what and why rather than how.

Code Quality

• The diff follows the existing Helix formatting conventions in the touched files.

🤖 Generated with GitHub Copilot CLI

[kabragaurav]

Nit: -> should be -> As in above code I see -> has been used. Or for javadoc, let us make it uniform.

[kabragaurav]

Nit: "and best possible assignments above." If someone adds/removes/moves the L64 in future, "above" statement will be confusing and wrong.

[kabragaurav]

Same as above comment.

[kabragaurav]

Suggestion: Right now every lookup goes through computeIfAbsent lazily. Can we build the maps once and use getOrDefault instead?

Like:

// built once, never mutated after construction, so no concurrent map needed
private final Map<String, Map<String, Map<String, String>>> _baselineAssignmentStateMap;
private final Map<String, Map<String, Map<String, String>>> _bestPossibleAssignmentStateMap;

In the constructor, after  _baselineAssignment  /  _bestPossibleAssignment  are set:

_baselineAssignmentStateMap = buildAssignmentStateMap(getBaselineAssignment());
_bestPossibleAssignmentStateMap = buildAssignmentStateMap(getBestPossibleAssignment());

Helpers:

private static Map<String, Map<String, Map<String, String>>> buildAssignmentStateMap(
    Map<String, ResourceAssignment> assignment) {
  Map<String, Map<String, Map<String, String>>> result = new HashMap<>();
  for (Map.Entry<String, ResourceAssignment> entry : assignment.entrySet()) {
    ResourceAssignment resourceAssignment = entry.getValue();
    Map<String, Map<String, String>> byPartition = new HashMap<>();
    for (String partitionName : resourceAssignment.getMappedPartitions().stream()
        .map(Partition::getPartitionName).collect(Collectors.toList())) {
      byPartition.put(partitionName,
          resourceAssignment.getReplicaMap(new Partition(partitionName)));
    }
    result.put(entry.getKey(), byPartition);
  }
  return result;
}

And rewrite like:

public Map<String, String> getBaselineAssignmentStateMap(String resourceName, String partitionName) {
  return _baselineAssignmentStateMap.getOrDefault(resourceName, Collections.emptyMap())
      .getOrDefault(partitionName, Collections.emptyMap());
}

[sjainit]

I looked at building the maps eagerly in the constructor. The catch is that the  (resource, partition)  keys actually scored are only those in  toBeAssignedReplicas , which is computed later in  ConstraintBasedAlgorithm  and isn't available at  ClusterContext construction. So an eager build would materialize every partition of the full baseline + best-possible assignment — including the majority never scored in a partial rebalance, which is the hot path this PR targets. Lazy computeIfAbsent  gets that selectivity for free by populating on demand. I've kept lazy but hardened it: wrapped the returned view in unmodifiableMap  and added a 16-thread concurrent test to lock in the thread-safety invariant

Let me know if you want to discuss this!

[kabragaurav]

Got it. Thanks for clarification!

[kabragaurav]

The test loops 5 times on a single thread and asserts  getReplicaMap  was called once. That proves the cache works. Test should also cover many threads hitting the same key at once inside  parallelStream. So a future refactor (say a plain HashMap ) would pass.

Something like:

@Test
public void testAssignmentStateMapMemoizationIsThreadSafe() throws Exception {
  ResourceControllerDataProvider testCache = setupClusterDataCache();
  Set<AssignableReplica> assignmentSet = generateReplicas(testCache);
  String resourceName = _resourceNames.get(0);
  String partitionName = _partitionNames.get(0);
  Map<String, String> stateMap = ImmutableMap.of("instance0", "MASTER");

  ResourceAssignment bestPossibleAssignment = mock(ResourceAssignment.class);
  when(bestPossibleAssignment.getReplicaMap(new Partition(partitionName))).thenReturn(stateMap);

  ClusterContext context = new ClusterContext(assignmentSet, generateNodes(testCache),
      new HashMap<>(), ImmutableMap.of(resourceName, bestPossibleAssignment));

  int threads = 16;
  ExecutorService pool = Executors.newFixedThreadPool(threads);
  CountDownLatch start = new CountDownLatch(1);
  List<Future<Map<String, String>>> results = new ArrayList<>();
  for (int i = 0; i < threads; i++) {
    results.add(pool.submit(() -> {
      start.await();
      return context.getBestPossibleAssignmentStateMap(resourceName, partitionName);
    }));
  }
  start.countDown();
  for (Future<Map<String, String>> f : results) {
    Assert.assertEquals(f.get(), stateMap);
  }
  pool.shutdown();
  // still resolved exactly once despite concurrent first-access
  verify(bestPossibleAssignment, times(1)).getReplicaMap(new Partition(partitionName));
}

[sjainit]

made the tests cover concurrent scenarios

[kabragaurav]

Thanks!

[kabragaurav]

LGTM! Thanks for addressing comments.

[kabragaurav]

CI is failing, but that is due to other failure unrelated to this PR, as Copilot tells.