diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/ConfigKeys.java b/internal/venice-common/src/main/java/com/linkedin/venice/ConfigKeys.java index 12feb3d7bc3..5a50d430bb9 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/ConfigKeys.java +++ b/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"; /** diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/VeniceRouterConfig.java b/services/venice-router/src/main/java/com/linkedin/venice/router/VeniceRouterConfig.java index 064677474dd..5e291f4ce77 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/VeniceRouterConfig.java +++ b/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; } diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceDelegateMode.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceDelegateMode.java index ee10178f72a..efc511c2784 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceDelegateMode.java +++ b/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 routerStats; private final RouterStats 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(); } diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupLeastLoadedStrategy.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupLeastLoadedStrategy.java index 4c64cccd2e6..2f69cdf0b59 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupLeastLoadedStrategy.java +++ b/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> requestTimeoutFutureMap = new HashMap<>(); + private final Map 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 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; + } + } } diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupRoundRobinStrategy.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupRoundRobinStrategy.java index 59cb7547ea1..284646edbec 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupRoundRobinStrategy.java +++ b/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); diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelectionStrategy.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelectionStrategy.java index c513e9c604f..21b525ab15d 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelectionStrategy.java +++ b/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 diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelector.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelector.java index 1831c57d7d5..17fc2837649 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelector.java +++ b/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; } diff --git a/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupLeastLoadedStrategy.java b/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupLeastLoadedStrategy.java index 827f9f434c5..017ea05a22c 100644 --- a/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupLeastLoadedStrategy.java +++ b/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); + } }