[router] Add weight-aware Helix group selection to balance multi-key RCU

internal/venice-common/src/main/java/com/linkedin/venice/ConfigKeys.java

@@ -1480,6 +1480,14 @@ private ConfigKeys() {
   public static final String ROUTER_HELIX_ASSISTED_ROUTING_GROUP_SELECTION_STRATEGY =
       "router.helix.assisted.routing.group.selection.strategy";
 
+  /**
+   * Whether Helix-assisted-routing group selection weights each request by its key count (estimated RCU)
+   * instead of counting every request as 1, so variable-size multi-get traffic is balanced by keys/RCU
+   * across Helix groups. Only affects {@literal HelixGroupLeastLoadedStrategy}. Default: false.
+   */
+  public static final String ROUTER_HELIX_ASSISTED_ROUTING_GROUP_SELECTION_WEIGHT_AWARE_ENABLED =
+      "router.helix.assisted.routing.group.selection.weight.aware.enabled";
+
   public static final String ROUTER_PER_STORE_ROUTER_QUOTA_BUFFER = "router.per.store.router.quota.buffer";
 
   /**

services/venice-router/src/main/java/com/linkedin/venice/router/VeniceRouterConfig.java

@@ -38,6 +38,7 @@
 import static com.linkedin.venice.ConfigKeys.ROUTER_FULL_PENDING_QUEUE_SERVER_OOR_MS;
 import static com.linkedin.venice.ConfigKeys.ROUTER_HEART_BEAT_ENABLED;
 import static com.linkedin.venice.ConfigKeys.ROUTER_HELIX_ASSISTED_ROUTING_GROUP_SELECTION_STRATEGY;
+import static com.linkedin.venice.ConfigKeys.ROUTER_HELIX_ASSISTED_ROUTING_GROUP_SELECTION_WEIGHT_AWARE_ENABLED;
 import static com.linkedin.venice.ConfigKeys.ROUTER_HTTP2_HEADER_TABLE_SIZE;
 import static com.linkedin.venice.ConfigKeys.ROUTER_HTTP2_INBOUND_ENABLED;
 import static com.linkedin.venice.ConfigKeys.ROUTER_HTTP2_INITIAL_WINDOW_SIZE;
@@ -211,6 +212,7 @@ public class VeniceRouterConfig implements RouterRetryConfig {
   private final int ioThreadCountInPoolMode;
   private final VeniceMultiKeyRoutingStrategy multiKeyRoutingStrategy;
   private final HelixGroupSelectionStrategyEnum helixGroupSelectionStrategy;
+  private final boolean helixGroupSelectionWeightAwareEnabled;
   private final String systemSchemaClusterName;
   private final int maxConcurrentSslHandshakes;
   private final int resolveThreads;
@@ -404,6 +406,8 @@ public VeniceRouterConfig(VeniceProperties props) {
                 + helixGroupSelectionStrategyStr + ", and allowed values: "
                 + Arrays.toString(HelixGroupSelectionStrategyEnum.values()));
       }
+      helixGroupSelectionWeightAwareEnabled =
+          props.getBoolean(ROUTER_HELIX_ASSISTED_ROUTING_GROUP_SELECTION_WEIGHT_AWARE_ENABLED, false);
       systemSchemaClusterName = props.getString(SYSTEM_SCHEMA_CLUSTER_NAME, "");
       routerHeartBeatEnabled = props.getBoolean(ROUTER_HEART_BEAT_ENABLED, true);
       httpClient5PoolSize = props.getInt(ROUTER_HTTP_CLIENT5_POOL_SIZE, 1);
@@ -752,6 +756,10 @@ public HelixGroupSelectionStrategyEnum getHelixGroupSelectionStrategy() {
     return helixGroupSelectionStrategy;
   }
 
+  public boolean isHelixGroupSelectionWeightAwareEnabled() {
+    return helixGroupSelectionWeightAwareEnabled;
+  }
+
   public String getSystemSchemaClusterName() {
     return systemSchemaClusterName;
   }

services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceDelegateMode.java

@@ -109,6 +109,7 @@ public class VeniceDelegateMode extends ScatterGatherMode {
   private final RouterStats<AggRouterHttpRequestStats> routerStats;
   private final RouterStats<RouteHttpStats> perRouteStatsByType;
   private final boolean latencyBasedRoutingEnabled;
+  private final boolean helixGroupSelectionWeightAwareEnabled;
 
   private final RoutingComputationMode routingComputationMode;
   private final ThreadPoolExecutor parallelRoutingExecutor;
@@ -124,6 +125,7 @@ public VeniceDelegateMode(
     this.routeHttpRequestStats = routeHttpRequestStats;
     this.perRouteStatsByType = perRouteStatsByType;
     this.latencyBasedRoutingEnabled = config.isLatencyBasedRoutingEnabled();
+    this.helixGroupSelectionWeightAwareEnabled = config.isHelixGroupSelectionWeightAwareEnabled();
     this.multiKeyRoutingStrategy = config.getMultiKeyRoutingStrategy();
     switch (this.multiKeyRoutingStrategy) {
       case GROUP_BY_PRIMARY_HOST_ROUTING:
@@ -824,8 +826,14 @@ protected int getAssignedHelixGroupId(VenicePath venicePath) {
         /**
          * This function only needs to assign a group id to the original Router request, and all the retried requests
          * will share the same group id as the original Router request.
+         *
+         * When weight-aware group selection is enabled, the request contributes load proportional to its key count
+         * (an estimate of its RCU cost) so variable-size multi-key requests are balanced by keys/RCU across Helix
+         * groups rather than by raw request count. When disabled, every request weighs 1 (legacy behavior).
          */
-        venicePath.setHelixGroupId(helixGroupSelector.selectGroup(venicePath.getRequestId(), getHelixGroupNum()));
+        int weight = helixGroupSelectionWeightAwareEnabled ? venicePath.getPartitionKeys().size() : 1;
+        venicePath
+            .setHelixGroupId(helixGroupSelector.selectGroup(venicePath.getRequestId(), getHelixGroupNum(), weight));
       }
       return venicePath.getHelixGroupId();
     }

services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupLeastLoadedStrategy.java

@@ -3,7 +3,6 @@
 import com.linkedin.alpini.base.concurrency.TimeoutProcessor;
 import com.linkedin.venice.exceptions.VeniceException;
 import com.linkedin.venice.stats.routing.HelixGroupStats;
-import com.linkedin.venice.utils.Pair;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -27,7 +26,7 @@ public class HelixGroupLeastLoadedStrategy implements HelixGroupSelectionStrateg
   private final int[] counters = new int[MAX_ALLOWED_GROUP];
   private final TimeoutProcessor timeoutProcessor;
   private final long timeoutInMS;
-  private final Map<Long, Pair<Integer, TimeoutProcessor.TimeoutFuture>> requestTimeoutFutureMap = new HashMap<>();
+  private final Map<Long, RequestGroupAssignment> requestTimeoutFutureMap = new HashMap<>();
   private final HelixGroupStats helixGroupStats;
 
   public HelixGroupLeastLoadedStrategy(
@@ -40,11 +39,18 @@ public HelixGroupLeastLoadedStrategy(
   }
 
   @Override
-  public int selectGroup(long requestId, int groupCount) {
+  public int selectGroup(long requestId, int groupCount, int weight) {
     if (groupCount > MAX_ALLOWED_GROUP || groupCount <= 0) {
       throw new VeniceException(
           "The valid group num must fail into this range: [1, " + MAX_ALLOWED_GROUP + "], but received: " + groupCount);
     }
+    /**
+     * Each request contributes at least 1 unit of load to the assigned group's counter, so a burst of
+     * zero/negative-weight requests cannot all pile onto a single group. A larger weight (e.g. the request's
+     * key count / estimated RCU) makes the request contribute proportionally more load, so variable-size
+     * multi-key requests are balanced by keys rather than by raw request count.
+     */
+    int effectiveWeight = Math.max(1, weight);
     int smallestCounter = Integer.MAX_VALUE;
     double lowestAvgLatency = Double.MAX_VALUE;
     int leastLoadedGroup = 0;
@@ -85,14 +91,15 @@ public int selectGroup(long requestId, int groupCount) {
        */
       requestTimeoutFutureMap.put(
           requestId,
-          new Pair<>(
+          new RequestGroupAssignment(
               leastLoadedGroup,
+              effectiveWeight,
               timeoutProcessor.schedule(
                   () -> timeoutRequest(requestId, finalLeastLoadedGroup, false),
                   timeoutInMS,
                   TimeUnit.MILLISECONDS)));
 
-      ++counters[leastLoadedGroup];
+      counters[leastLoadedGroup] += effectiveWeight;
     }
     helixGroupStats.recordGroupPendingRequest(leastLoadedGroup, counters[leastLoadedGroup]);
 
@@ -118,26 +125,27 @@ private void timeoutRequest(long requestId, int groupId, boolean cancelTimeoutFu
       helixGroupStats.recordGroupResponseWaitingTime(groupId, timeoutInMS);
     }
     synchronized (this) {
-      Pair<Integer, TimeoutProcessor.TimeoutFuture> timeoutFuturePair = requestTimeoutFutureMap.get(requestId);
-      if (timeoutFuturePair == null) {
+      RequestGroupAssignment assignment = requestTimeoutFutureMap.get(requestId);
+      if (assignment == null) {
         /**
          * Request has already timed out or already finished.
          */
         return;
       }
-      if (groupId != timeoutFuturePair.getFirst()) {
+      if (groupId != assignment.groupId) {
         throw new VeniceException(
-            "Group id for request with id: " + requestId + " should be: " + timeoutFuturePair.getFirst()
-                + ", but received: " + groupId);
+            "Group id for request with id: " + requestId + " should be: " + assignment.groupId + ", but received: "
+                + groupId);
       }
-      if (--counters[groupId] < 0) {
+      counters[groupId] -= assignment.weight;
+      if (counters[groupId] < 0) {
         counters[groupId] = 0;
         throw new VeniceException(
             "The counter for group: " + groupId + " became negative, something wrong happened, will reset it to be 0.");
       }
       if (cancelTimeoutFuture) {
         // Cancel the timeout future
-        timeoutFuturePair.getSecond().cancel();
+        assignment.timeoutFuture.cancel();
       } else {
         LOGGER.info(
             "Request with id: {} has timed out with threshold: {}ms, and the counter of group: {} will be reset for this request",
@@ -154,4 +162,20 @@ public void finishRequest(long requestId, int groupId, double latency) {
     timeoutRequest(requestId, groupId, true);
     helixGroupStats.recordGroupResponseWaitingTime(groupId, latency);
   }
+
+  /**
+   * Holds the per-request group assignment: the selected group, the load {@code weight} that was added to that
+   * group's counter (so the same amount can be subtracted on completion/timeout), and the leak-guard timeout future.
+   */
+  private static class RequestGroupAssignment {
+    final int groupId;
+    final int weight;
+    final TimeoutProcessor.TimeoutFuture timeoutFuture;
+
+    RequestGroupAssignment(int groupId, int weight, TimeoutProcessor.TimeoutFuture timeoutFuture) {
+      this.groupId = groupId;
+      this.weight = weight;
+      this.timeoutFuture = timeoutFuture;
+    }
+  }
 }

services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupRoundRobinStrategy.java

@@ -5,7 +5,7 @@
  */
 public class HelixGroupRoundRobinStrategy implements HelixGroupSelectionStrategy {
   @Override
-  public int selectGroup(long requestId, int groupNum) {
+  public int selectGroup(long requestId, int groupNum, int weight) {
     int assignedGroupId = 0;
     if (groupNum > 0) {
       assignedGroupId = (int) (requestId % groupNum);

services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelectionStrategy.java

@@ -2,9 +2,21 @@
 
 public interface HelixGroupSelectionStrategy {
   /**
-   * Select a Helix Group for the current request.
+   * Select a Helix Group for the current request, weighting the request's load contribution to the assigned
+   * group by {@code weight} (for example, its key count / estimated RCU). A larger weight makes the assigned
+   * group appear more loaded to subsequent selections, so variable-size multi-key requests are balanced by
+   * keys/RCU rather than by raw request count. Implementations must decrement by the same weight when the
+   * request finishes (or times out).
    */
-  int selectGroup(long requestId, int groupCount);
+  int selectGroup(long requestId, int groupCount, int weight);
+
+  /**
+   * Select a Helix Group for the current request, weighting every request equally (weight = 1). Preserves the
+   * legacy request-count-based selection behavior.
+   */
+  default int selectGroup(long requestId, int groupCount) {
+    return selectGroup(requestId, groupCount, 1);
+  }
 
   /**
    * Notify the corresponding Helix Group that the request is completed, and the implementation will decide whether

services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelector.java

@@ -54,9 +54,9 @@ public int getGroupCount() {
   }
 
   @Override
-  public int selectGroup(long requestId, int groupNum) {
+  public int selectGroup(long requestId, int groupNum, int weight) {
     helixGroupStats.recordGroupNum(groupNum);
-    int assignedGroupId = selectionStrategy.selectGroup(requestId, groupNum);
+    int assignedGroupId = selectionStrategy.selectGroup(requestId, groupNum, weight);
     helixGroupStats.recordGroupRequest(assignedGroupId);
     return assignedGroupId;
   }

services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupLeastLoadedStrategy.java

@@ -52,4 +52,95 @@ public void testLatencyBasedGroupSelection() {
     strategy.finishRequest(2, 2, 1);
     Assert.assertEquals(strategy.selectGroup(3, groupNum), 2);
   }
+
+  /**
+   * A high-weight (many-key) request should make its assigned group appear proportionally more loaded, so
+   * subsequent lighter requests are steered to the other groups until they catch up. This is the core of the
+   * weight-aware balancing that prevents one Helix group from absorbing disproportionate RCU.
+   */
+  @Test
+  public void testWeightedGroupSelectionSteersAwayFromHeavyGroup() {
+    TimeoutProcessor timeoutProcessor = mock(TimeoutProcessor.class);
+    doReturn(mock(TimeoutProcessor.TimeoutFuture.class)).when(timeoutProcessor).schedule(any(), anyLong(), any());
+    HelixGroupLeastLoadedStrategy strategy =
+        new HelixGroupLeastLoadedStrategy(timeoutProcessor, 10000, mock(HelixGroupStats.class));
+    int groupNum = 3;
+    // A single 10-key request lands on group 0, contributing 10 units of load.
+    Assert.assertEquals(strategy.selectGroup(0, groupNum, 10), 0);
+    // The next 9 single-key requests must all avoid group 0, since the other groups together can only reach
+    // 9 units of load and each remains below group 0's 10.
+    for (int requestId = 1; requestId <= 9; requestId++) {
+      Assert.assertNotEquals(
+          strategy.selectGroup(requestId, groupNum, 1),
+          0,
+          "Group 0 holds 10 units of weighted load and should be avoided by lighter requests");
+    }
+  }
+
+  /**
+   * Finishing a request must subtract exactly the weight that was added when the group was selected, so a heavy
+   * request fully releases its load and the group becomes available again (decrement symmetry). Without storing
+   * the per-request weight, a weighted increment paired with a unit decrement would leak load and permanently
+   * skew selection.
+   */
+  @Test
+  public void testWeightedDecrementSymmetry() {
+    TimeoutProcessor timeoutProcessor = mock(TimeoutProcessor.class);
+    doReturn(mock(TimeoutProcessor.TimeoutFuture.class)).when(timeoutProcessor).schedule(any(), anyLong(), any());
+    HelixGroupLeastLoadedStrategy strategy =
+        new HelixGroupLeastLoadedStrategy(timeoutProcessor, 10000, mock(HelixGroupStats.class));
+    int groupNum = 3;
+    // Heavy request on group 0, plus one unit request on each of the other two groups.
+    Assert.assertEquals(strategy.selectGroup(0, groupNum, 10), 0);
+    Assert.assertEquals(strategy.selectGroup(1, groupNum, 1), 1);
+    Assert.assertEquals(strategy.selectGroup(2, groupNum, 1), 2);
+    // Release the heavy request: its 10 units must be fully subtracted, leaving group 0 empty (0 units) while
+    // groups 1 and 2 still hold 1 unit each.
+    strategy.finishRequest(0, 0, 1);
+    // The now-empty group 0 is the least loaded and must be selected next.
+    Assert.assertEquals(strategy.selectGroup(100, groupNum, 1), 0);
+  }
+
+  /**
+   * A zero or negative weight must be clamped to 1 so that a flood of zero-weight requests cannot all pile onto
+   * a single group without moving its counter. Each request contributes at least one unit of load.
+   */
+  @Test
+  public void testNonPositiveWeightClampedToOne() {
+    TimeoutProcessor timeoutProcessor = mock(TimeoutProcessor.class);
+    doReturn(mock(TimeoutProcessor.TimeoutFuture.class)).when(timeoutProcessor).schedule(any(), anyLong(), any());
+    HelixGroupLeastLoadedStrategy strategy =
+        new HelixGroupLeastLoadedStrategy(timeoutProcessor, 10000, mock(HelixGroupStats.class));
+    int groupNum = 3;
+    // A zero-weight request on group 0 must still add 1 unit of load.
+    Assert.assertEquals(strategy.selectGroup(0, groupNum, 0), 0);
+    // Because group 0 now holds 1 unit (not 0), a subsequent request that starts scanning at group 0 must skip
+    // it in favor of an empty group. If the zero weight had NOT been clamped, group 0 would tie at 0 and win.
+    Assert.assertEquals(strategy.selectGroup(3, groupNum, 1), 1);
+    // Finishing the clamped request subtracts exactly 1 and must not drive the counter negative or throw.
+    strategy.finishRequest(0, 0, 1);
+    strategy.finishRequest(3, 1, 1);
+    // Both groups are empty again; the next request starting at group 0 selects group 0.
+    Assert.assertEquals(strategy.selectGroup(6, groupNum, 1), 0);
+  }
+
+  /**
+   * The weight-aware overload must remain backward compatible: the two-argument {@code selectGroup} default and
+   * an explicit weight of 1 must produce identical selection behavior (every request counts as exactly 1).
+   */
+  @Test
+  public void testDefaultOverloadEquivalentToWeightOne() {
+    TimeoutProcessor timeoutProcessor = mock(TimeoutProcessor.class);
+    doReturn(mock(TimeoutProcessor.TimeoutFuture.class)).when(timeoutProcessor).schedule(any(), anyLong(), any());
+    HelixGroupLeastLoadedStrategy strategy =
+        new HelixGroupLeastLoadedStrategy(timeoutProcessor, 10000, mock(HelixGroupStats.class));
+    int groupNum = 3;
+    // Two-arg default overload (weight defaults to 1).
+    Assert.assertEquals(strategy.selectGroup(0, groupNum), 0);
+    // Explicit weight of 1 on the next group.
+    Assert.assertEquals(strategy.selectGroup(1, groupNum, 1), 1);
+    Assert.assertEquals(strategy.selectGroup(2, groupNum), 2);
+    // With all groups holding exactly 1 unit, the next request starting at group 0 picks group 0.
+    Assert.assertEquals(strategy.selectGroup(3, groupNum, 1), 0);
+  }
 }