Skip to content

Commit 5332787

Browse files
Zgoda91Zgoda91
committed
Implemented random_subsetting LB policy
1 parent 8d349df commit 5332787

File tree

8 files changed

+714
-1
lines changed

8 files changed

+714
-1
lines changed

‎api/src/test/java/io/grpc/LoadBalancerRegistryTest.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ public void getClassesViaHardcoded_classesPresent() throws Exception {
4040
@Test
4141
public void stockProviders() {
4242
LoadBalancerRegistry defaultRegistry = LoadBalancerRegistry.getDefaultRegistry();
43-
assertThat(defaultRegistry.providers()).hasSize(3);
43+
assertThat(defaultRegistry.providers()).hasSize(4);
4444

4545
LoadBalancerProvider pickFirst = defaultRegistry.getProvider("pick_first");
4646
assertThat(pickFirst).isInstanceOf(PickFirstLoadBalancerProvider.class);
@@ -56,6 +56,11 @@ public void stockProviders() {
5656
assertThat(outlierDetection.getClass().getName()).isEqualTo(
5757
"io.grpc.util.OutlierDetectionLoadBalancerProvider");
5858
assertThat(roundRobin.getPriority()).isEqualTo(5);
59+
60+
LoadBalancerProvider randomSubsetting = defaultRegistry.getProvider("random_subsetting");
61+
assertThat(randomSubsetting.getClass().getName()).isEqualTo(
62+
"io.grpc.util.RandomSubsettingLoadBalancerProvider");
63+
assertThat(randomSubsetting.getPriority()).isEqualTo(5);
5964
}
6065

6166
@Test

‎util/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ java_library(
1313
deps = [
1414
"//api",
1515
"//core:internal",
16+
"//third-party/zero-allocation-hashing",
1617
artifact("com.google.code.findbugs:jsr305"),
1718
artifact("com.google.errorprone:error_prone_annotations"),
1819
artifact("com.google.guava:guava"),

‎util/build.gradle‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ dependencies {
1919
api project(':grpc-api')
2020

2121
implementation project(':grpc-core'),
22+
project(':grpc-third-party:zero-allocation-hashing'),
2223
libraries.animalsniffer.annotations,
2324
libraries.guava
2425
testImplementation libraries.guava.testlib,

‎util/src/main/java/io/grpc/util/RandomSubsettingLoadBalancer.java‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
/*
2+
* Copyright 2025 The gRPC Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.grpc.util;
18+
19+
import static com.google.common.base.Preconditions.checkArgument;
20+
import static com.google.common.base.Preconditions.checkNotNull;
21+
22+
import com.google.common.annotations.VisibleForTesting;
23+
import io.grpc.EquivalentAddressGroup;
24+
import io.grpc.Internal;
25+
import io.grpc.LoadBalancer;
26+
import io.grpc.Status;
27+
import io.grpc.tp.zah.XxHash64;
28+
import java.security.SecureRandom;
29+
import java.util.ArrayList;
30+
import java.util.Collections;
31+
import java.util.Comparator;
32+
33+
34+
/**
35+
* Wraps a child {@code LoadBalancer}, separating the total set of backends into smaller subsets for
36+
* the child balancer to balance across.
37+
*
38+
* <p>This implements random subsetting gRFC:
39+
* https://https://github.com/grpc/proposal/blob/master/A68-random-subsetting.md
40+
*/
41+
@Internal
42+
public final class RandomSubsettingLoadBalancer extends LoadBalancer {
43+
private final GracefulSwitchLoadBalancer switchLb;
44+
45+
public RandomSubsettingLoadBalancer(Helper helper) {
46+
switchLb = new GracefulSwitchLoadBalancer(checkNotNull(helper, "helper"));
47+
}
48+
49+
@Override
50+
public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
51+
RandomSubsettingLoadBalancerConfig config =
52+
(RandomSubsettingLoadBalancerConfig)
53+
resolvedAddresses.getLoadBalancingPolicyConfig();
54+
55+
ResolvedAddresses subsetAddresses = filterEndpoints(
56+
resolvedAddresses, config.subsetSize, new SecureRandom().nextLong());
57+
58+
return switchLb.acceptResolvedAddresses(
59+
subsetAddresses.toBuilder()
60+
.setLoadBalancingPolicyConfig(config.childConfig)
61+
.build());
62+
}
63+
64+
// implements the subsetting algorithm, as described in A68:
65+
// https://github.com/grpc/proposal/pull/423
66+
private ResolvedAddresses filterEndpoints(
67+
ResolvedAddresses resolvedAddresses, long subsetSize, long seed) {
68+
// configured subset sizes in the range [Integer.MAX_VALUE, Long.MAX_VALUE] will always fall
69+
// into this if statement due to collection indexing limitations in JVM
70+
if (subsetSize >= resolvedAddresses.getAddresses().size()) {
71+
return resolvedAddresses;
72+
}
73+
74+
XxHash64 hashFunc = new XxHash64(seed);
75+
ArrayList<EndpointWithHash> endpointWithHashList = new ArrayList<>();
76+
77+
for (EquivalentAddressGroup addressGroup : resolvedAddresses.getAddresses()) {
78+
endpointWithHashList.add(
79+
new EndpointWithHash(
80+
addressGroup,
81+
hashFunc.hashAsciiString(addressGroup.getAddresses().get(0).toString())));
82+
}
83+
84+
Collections.sort(endpointWithHashList, new HashAddressComparator());
85+
86+
ArrayList<EquivalentAddressGroup> addressGroups = new ArrayList<>();
87+
88+
// for loop is executed for subset sizes in range [0, Integer.MAX_VALUE), therefore indexing
89+
// variable is not going to overflow here
90+
for (int idx = 0; idx < subsetSize; ++idx) {
91+
addressGroups.add(endpointWithHashList.get(idx).addressGroup);
92+
}
93+
94+
return resolvedAddresses.toBuilder().setAddresses(addressGroups).build();
95+
}
96+
97+
@Override
98+
public void handleNameResolutionError(Status error) {
99+
switchLb.handleNameResolutionError(error);
100+
}
101+
102+
@Override
103+
public void shutdown() {
104+
switchLb.shutdown();
105+
}
106+
107+
private static final class EndpointWithHash {
108+
public final EquivalentAddressGroup addressGroup;
109+
public final long hash;
110+
111+
public EndpointWithHash(EquivalentAddressGroup addressGroup, long hash) {
112+
this.addressGroup = addressGroup;
113+
this.hash = hash;
114+
}
115+
}
116+
117+
@VisibleForTesting
118+
static class HashAddressComparator implements Comparator<EndpointWithHash> {
119+
@Override
120+
public int compare(EndpointWithHash lhs, EndpointWithHash rhs) {
121+
return Long.compare(lhs.hash, rhs.hash);
122+
}
123+
}
124+
125+
public static final class RandomSubsettingLoadBalancerConfig {
126+
public final long subsetSize;
127+
public final Object childConfig;
128+
129+
private RandomSubsettingLoadBalancerConfig(long subsetSize, Object childConfig) {
130+
this.subsetSize = subsetSize;
131+
this.childConfig = childConfig;
132+
}
133+
134+
public static class Builder {
135+
Long subsetSize;
136+
Object childConfig;
137+
138+
public Builder setSubsetSize(Integer subsetSize) {
139+
checkNotNull(subsetSize, "subsetSize");
140+
// {@code Integer.toUnsignedLong(int)} is not part of Android API level 21, therefore doing
141+
// it manually
142+
Long subsetSizeAsLong = ((long) subsetSize) & 0xFFFFFFFFL;
143+
checkArgument(subsetSizeAsLong > 0L, "Subset size must be greater than 0");
144+
this.subsetSize = subsetSizeAsLong;
145+
return this;
146+
}
147+
148+
public Builder setChildConfig(Object childConfig) {
149+
this.childConfig = checkNotNull(childConfig, "childConfig");
150+
return this;
151+
}
152+
153+
public RandomSubsettingLoadBalancerConfig build() {
154+
return new RandomSubsettingLoadBalancerConfig(
155+
checkNotNull(subsetSize, "subsetSize"),
156+
checkNotNull(childConfig, "childConfig"));
157+
}
158+
}
159+
}
160+
}

‎util/src/main/java/io/grpc/util/RandomSubsettingLoadBalancerProvider.java‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
/*
2+
* Copyright 2025 The gRPC Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.grpc.util;
18+
19+
import io.grpc.Internal;
20+
import io.grpc.LoadBalancer;
21+
import io.grpc.LoadBalancerProvider;
22+
import io.grpc.NameResolver.ConfigOrError;
23+
import io.grpc.Status;
24+
import io.grpc.internal.JsonUtil;
25+
import java.util.Map;
26+
27+
@Internal
28+
public final class RandomSubsettingLoadBalancerProvider extends LoadBalancerProvider {
29+
private static final String POLICY_NAME = "random_subsetting";
30+
31+
@Override
32+
public LoadBalancer newLoadBalancer(LoadBalancer.Helper helper) {
33+
return new RandomSubsettingLoadBalancer(helper);
34+
}
35+
36+
@Override
37+
public boolean isAvailable() {
38+
return true;
39+
}
40+
41+
@Override
42+
public int getPriority() {
43+
return 5;
44+
}
45+
46+
@Override
47+
public String getPolicyName() {
48+
return POLICY_NAME;
49+
}
50+
51+
@Override
52+
public ConfigOrError parseLoadBalancingPolicyConfig(Map<String, ?> rawConfig) {
53+
try {
54+
return parseLoadBalancingPolicyConfigInternal(rawConfig);
55+
} catch (RuntimeException e) {
56+
return ConfigOrError.fromError(
57+
Status.UNAVAILABLE
58+
.withCause(e)
59+
.withDescription("Failed parsing configuration for " + getPolicyName()));
60+
}
61+
}
62+
63+
private ConfigOrError parseLoadBalancingPolicyConfigInternal(Map<String, ?> rawConfig) {
64+
Integer subsetSize = JsonUtil.getNumberAsInteger(rawConfig, "subsetSize");
65+
if (subsetSize == null) {
66+
return ConfigOrError.fromError(
67+
Status.INTERNAL.withDescription(
68+
"Subset size missing in " + getPolicyName() + ", LB policy config=" + rawConfig));
69+
}
70+
71+
ConfigOrError childConfig = GracefulSwitchLoadBalancer.parseLoadBalancingPolicyConfig(
72+
JsonUtil.getListOfObjects(rawConfig, "childPolicy"));
73+
if (childConfig.getError() != null) {
74+
return ConfigOrError.fromError(Status.INTERNAL
75+
.withDescription(
76+
"Failed to parse child in " + getPolicyName() + ", LB policy config=" + rawConfig)
77+
.withCause(childConfig.getError().asRuntimeException()));
78+
}
79+
80+
return ConfigOrError.fromConfig(
81+
new RandomSubsettingLoadBalancer.RandomSubsettingLoadBalancerConfig.Builder()
82+
.setSubsetSize(subsetSize)
83+
.setChildConfig(childConfig.getConfig())
84+
.build());
85+
}
86+
}

‎util/src/main/resources/META-INF/services/io.grpc.LoadBalancerProvider‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
io.grpc.util.SecretRoundRobinLoadBalancerProvider$Provider
22
io.grpc.util.OutlierDetectionLoadBalancerProvider
3+
io.grpc.util.RandomSubsettingLoadBalancerProvider

0 commit comments

Comments
 (0)