A68 random subsetting (part 1)

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

@@ -40,7 +40,7 @@ public void getClassesViaHardcoded_classesPresent() throws Exception {
   @Test
   public void stockProviders() {
     LoadBalancerRegistry defaultRegistry = LoadBalancerRegistry.getDefaultRegistry();
-    assertThat(defaultRegistry.providers()).hasSize(3);
+    assertThat(defaultRegistry.providers()).hasSize(4);
 
     LoadBalancerProvider pickFirst = defaultRegistry.getProvider("pick_first");
     assertThat(pickFirst).isInstanceOf(PickFirstLoadBalancerProvider.class);
@@ -56,6 +56,11 @@ public void stockProviders() {
     assertThat(outlierDetection.getClass().getName()).isEqualTo(
         "io.grpc.util.OutlierDetectionLoadBalancerProvider");
     assertThat(roundRobin.getPriority()).isEqualTo(5);
+
+    LoadBalancerProvider randomSubsetting = defaultRegistry.getProvider("random_subsetting");
+    assertThat(randomSubsetting.getClass().getName()).isEqualTo(
+        "io.grpc.util.RandomSubsettingLoadBalancerProvider");
+    assertThat(randomSubsetting.getPriority()).isEqualTo(5);
   }
 
   @Test

settings.gradle

@@ -93,6 +93,7 @@ include ":grpc-inprocess"
 include ":grpc-util"
 include ":grpc-opentelemetry"
 include ":grpc-context-override-opentelemetry"
+include ":grpc-third-party:zero-allocation-hashing"
 
 project(':grpc-api').projectDir = "$rootDir/api" as File
 project(':grpc-core').projectDir = "$rootDir/core" as File
@@ -130,6 +131,7 @@ project(':grpc-inprocess').projectDir = "$rootDir/inprocess" as File
 project(':grpc-util').projectDir = "$rootDir/util" as File
 project(':grpc-opentelemetry').projectDir = "$rootDir/opentelemetry" as File
 project(':grpc-context-override-opentelemetry').projectDir = "$rootDir/contextstorage" as File
+project(':grpc-third-party:zero-allocation-hashing').projectDir = "$rootDir/third-party/zero-allocation-hashing" as File
 
 if (settings.hasProperty('skipCodegen') && skipCodegen.toBoolean()) {
     println '*** Skipping the build of codegen and compilation of proto files because skipCodegen=true'

third-party/zero-allocation-hashing/BUILD.bazel

@@ -0,0 +1,26 @@
+load("@rules_java//java:defs.bzl", "java_binary", "java_library", "java_test")
+
+java_library(
+    name = "zero-allocation-hashing",
+    srcs = [
+        "src/main/java/io/grpc/tp/zah/XxHash64.java",
+    ],
+    deps = [
+        "@maven//:com_google_guava_guava",
+    ],
+    visibility = [
+        "//xds:__pkg__",
+        "//util:__pkg__",
+    ],
+)
+
+java_test(
+    name = "XxHash64Test",
+    size = "small",
+    srcs = ["src/test/java/io/grpc/tp/zah/XxHash64Test.java"],
+    deps = [
+        ":zero-allocation-hashing",
+        "@maven//:com_google_guava_guava",
+        "@maven//:junit_junit",
+    ],
+)

third-party/zero-allocation-hashing/build.gradle

@@ -0,0 +1,25 @@
+plugins {
+    id "java-library"
+}
+
+description = 'gRPC: Zero Allocation Hashing'
+
+dependencies {
+    implementation libraries.guava
+
+    testImplementation libraries.junit
+}
+
+tasks.named("jar").configure {
+    manifest {
+        attributes('Automatic-Module-Name': 'io.grpc.tp.zah')
+    }
+}
+
+tasks.named("checkstyleMain").configure {
+    enabled = false
+}
+
+tasks.named("checkstyleTest").configure {
+    enabled = false
+}

xds/third_party/zero-allocation-hashing/main/java/io/grpc/xds/XxHash64.java → third-party/zero-allocation-hashing/src/main/java/io/grpc/tp/zah/XxHash64.java

@@ -18,7 +18,7 @@
  * Modified by the gRPC Authors
  */
 
-package io.grpc.xds;
+package io.grpc.tp.zah;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
@@ -33,8 +33,8 @@
  * <a href="https://github.com/OpenHFT/Zero-Allocation-Hashing/blob/master/src/main/java/net/openhft/hashing/XxHash.java">
  * OpenHFT/Zero-Allocation-Hashing</a>.
  */
-final class XxHash64 {
-  static final XxHash64 INSTANCE = new XxHash64(0);
+final public class XxHash64 {
+  static public final XxHash64 INSTANCE = new XxHash64(0);
 
   // Primes if treated as unsigned
   private static final long P1 = -7046029288634856825L;
@@ -47,12 +47,12 @@ final class XxHash64 {
   private final long seed;
   private final long voidHash;
 
-  XxHash64(long seed) {
+  public XxHash64(long seed) {
     this.seed = seed;
     this.voidHash = finalize(seed + P5);
   }
 
-  long hashLong(long input) {
+  public long hashLong(long input) {
     input = byteOrder == ByteOrder.LITTLE_ENDIAN ? input : Long.reverseBytes(input);
     long hash = seed + P5 + 8;
     input *= P2;
@@ -63,15 +63,15 @@ long hashLong(long input) {
     return finalize(hash);
   }
 
-  long hashInt(int input) {
+  public long hashInt(int input) {
     input = byteOrder == ByteOrder.LITTLE_ENDIAN ? input : Integer.reverseBytes(input);
     long hash = seed + P5 + 4;
     hash ^= (input & 0xFFFFFFFFL) * P1;
     hash = Long.rotateLeft(hash, 23) * P2 + P3;
     return finalize(hash);
   }
 
-  long hashShort(short input) {
+  public long hashShort(short input) {
     input = byteOrder == ByteOrder.LITTLE_ENDIAN ? input : Short.reverseBytes(input);
     long hash = seed + P5 + 2;
     hash ^= (input & 0xFFL) * P5;
@@ -81,22 +81,22 @@ long hashShort(short input) {
     return finalize(hash);
   }
 
-  long hashChar(char input) {
+  public long hashChar(char input) {
     return hashShort((short) input);
   }
 
-  long hashByte(byte input) {
+  public long hashByte(byte input) {
     long hash = seed + P5 + 1;
     hash ^= (input & 0xFF) * P5;
     hash = Long.rotateLeft(hash, 11) * P1;
     return finalize(hash);
   }
 
-  long hashVoid() {
+  public long hashVoid() {
     return voidHash;
   }
 
-  long hashAsciiString(String input) {
+  public long hashAsciiString(String input) {
     ByteSupplier supplier = new AsciiStringByteSupplier(input);
     return hashBytes(supplier);
   }
@@ -106,7 +106,7 @@ long hashBytes(byte[] bytes) {
     return hashBytes(supplier);
   }
 
-  long hashBytes(byte[] bytes, int offset, int len) {
+  public long hashBytes(byte[] bytes, int offset, int len) {
     ByteSupplier supplier = new PlainByteSupplier(bytes, offset, len);
     return hashBytes(supplier);
   }

xds/third_party/zero-allocation-hashing/test/java/io/grpc/xds/XxHash64Test.java → third-party/zero-allocation-hashing/src/test/java/io/grpc/tp/zah/XxHash64Test.java

@@ -18,7 +18,7 @@
  * Modified by the gRPC Authors
  */
 
-package io.grpc.xds;
+package io.grpc.tp.zah;
 
 import static org.junit.Assert.assertEquals;
 

util/BUILD.bazel

@@ -13,6 +13,7 @@ java_library(
     deps = [
         "//api",
         "//core:internal",
+        "//third-party/zero-allocation-hashing",
         artifact("com.google.code.findbugs:jsr305"),
         artifact("com.google.errorprone:error_prone_annotations"),
         artifact("com.google.guava:guava"),

util/build.gradle

@@ -19,6 +19,7 @@ dependencies {
     api project(':grpc-api')
 
     implementation project(':grpc-core'),
+            project(':grpc-third-party:zero-allocation-hashing'),
             libraries.animalsniffer.annotations,
             libraries.guava
     testImplementation libraries.guava.testlib,

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

@@ -0,0 +1,158 @@
+/*
+ * Copyright 2025 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.util;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import io.grpc.EquivalentAddressGroup;
+import io.grpc.Internal;
+import io.grpc.LoadBalancer;
+import io.grpc.Status;
+import io.grpc.tp.zah.XxHash64;
+import java.security.SecureRandom;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+
+
+/**
+ * Wraps a child {@code LoadBalancer}, separating the total set of backends into smaller subsets for
+ * the child balancer to balance across.
+ *
+ * <p>This implements random subsetting gRFC:
+ * https://https://github.com/grpc/proposal/blob/master/A68-random-subsetting.md
+ */
+@Internal
+public final class RandomSubsettingLoadBalancer extends LoadBalancer {
+  private final GracefulSwitchLoadBalancer switchLb;
+
+  public RandomSubsettingLoadBalancer(Helper helper) {
+    switchLb = new GracefulSwitchLoadBalancer(checkNotNull(helper, "helper"));
+  }
+
+  @Override
+  public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+    RandomSubsettingLoadBalancerConfig config =
+        (RandomSubsettingLoadBalancerConfig)
+            resolvedAddresses.getLoadBalancingPolicyConfig();
+
+    ResolvedAddresses subsetAddresses = filterEndpoints(
+        resolvedAddresses, config.subsetSize, new SecureRandom().nextLong());
+
+    return switchLb.acceptResolvedAddresses(
+        subsetAddresses.toBuilder()
+            .setLoadBalancingPolicyConfig(config.childConfig)
+            .build());
+  }
+
+  // implements the subsetting algorithm, as described in A68:
+  // https://github.com/grpc/proposal/pull/423
+  private ResolvedAddresses filterEndpoints(
+      ResolvedAddresses resolvedAddresses, long subsetSize, long seed) {
+    // configured subset sizes in the range [Integer.MAX_VALUE, Long.MAX_VALUE] will always fall
+    // into this if statement due to collection indexing limitations in JVM
+    if (subsetSize >= resolvedAddresses.getAddresses().size()) {
+      return resolvedAddresses;
+    }
+
+    XxHash64 hashFunc = new XxHash64(seed);
+    ArrayList<EndpointWithHash> endpointWithHashList = new ArrayList<>();
+
+    for (EquivalentAddressGroup addressGroup : resolvedAddresses.getAddresses()) {
+      endpointWithHashList.add(
+          new EndpointWithHash(
+              addressGroup,
+              hashFunc.hashAsciiString(addressGroup.getAddresses().get(0).toString())));
+    }
+
+    Collections.sort(endpointWithHashList, new HashAddressComparator());
+
+    ArrayList<EquivalentAddressGroup> addressGroups = new ArrayList<>();
+
+    // for loop is executed for subset sizes in range [0, Integer.MAX_VALUE), therefore indexing
+    // variable is not going to overflow here
+    for (int idx = 0; idx < subsetSize; ++idx) {
+      addressGroups.add(endpointWithHashList.get(idx).addressGroup);
+    }
+
+    return resolvedAddresses.toBuilder().setAddresses(addressGroups).build();
+  }
+
+  @Override
+  public void handleNameResolutionError(Status error) {
+    switchLb.handleNameResolutionError(error);
+  }
+
+  @Override
+  public void shutdown() {
+    switchLb.shutdown();
+  }
+
+  private static final class EndpointWithHash {
+    public final EquivalentAddressGroup addressGroup;
+    public final long hash;
+
+    public EndpointWithHash(EquivalentAddressGroup addressGroup, long hash) {
+      this.addressGroup = addressGroup;
+      this.hash = hash;
+    }
+  }
+
+  private static final class HashAddressComparator implements Comparator<EndpointWithHash> {
+    @Override
+    public int compare(EndpointWithHash lhs, EndpointWithHash rhs) {
+      return Long.compare(lhs.hash, rhs.hash);
+    }
+  }
+
+  public static final class RandomSubsettingLoadBalancerConfig {
+    public final long subsetSize;
+    public final Object childConfig;
+
+    private RandomSubsettingLoadBalancerConfig(long subsetSize, Object childConfig) {
+      this.subsetSize = subsetSize;
+      this.childConfig = childConfig;
+    }
+
+    public static class Builder {
+      Long subsetSize;
+      Object childConfig;
+
+      public Builder setSubsetSize(Integer subsetSize) {
+        checkNotNull(subsetSize, "subsetSize");
+        // {@code Integer.toUnsignedLong(int)} is not part of Android API level 21, therefore doing
+        // it manually
+        Long subsetSizeAsLong = ((long) subsetSize) & 0xFFFFFFFFL;
+        checkArgument(subsetSizeAsLong > 0L, "Subset size must be greater than 0");
+        this.subsetSize = subsetSizeAsLong;
+        return this;
+      }
+
+      public Builder setChildConfig(Object childConfig) {
+        this.childConfig = checkNotNull(childConfig, "childConfig");
+        return this;
+      }
+
+      public RandomSubsettingLoadBalancerConfig build() {
+        return new RandomSubsettingLoadBalancerConfig(
+            checkNotNull(subsetSize, "subsetSize"),
+            checkNotNull(childConfig, "childConfig"));
+      }
+    }
+  }
+}