diff --git a/autosharding/build.gradle b/autosharding/build.gradle
new file mode 100644
index 00000000000..26ffdacdd6d
--- /dev/null
+++ b/autosharding/build.gradle
@@ -0,0 +1,63 @@
+plugins {
+ id "java-library"
+ id "maven-publish"
+ id "com.google.protobuf"
+ id "jacoco"
+ id "ru.vyarus.animalsniffer"
+}
+
+description = "gRPC: Autosharding Loadbalancing plugin"
+
+tasks.named("jar").configure {
+ manifest {
+ attributes('Automatic-Module-Name': 'io.grpc.autosharding')
+ }
+}
+
+dependencies {
+ implementation project(':grpc-util'),
+ project(':grpc-core'),
+ project(':grpc-protobuf'),
+ project(':grpc-stub'),
+ libraries.auto.value.annotations,
+ libraries.guava
+ annotationProcessor libraries.auto.value
+ testImplementation libraries.truth,
+ project(':grpc-grpclb'),
+ project(':grpc-inprocess'),
+ project(':grpc-testing'),
+ project(':grpc-testing-proto'),
+ testFixtures(project(':grpc-api')),
+ testFixtures(project(':grpc-core'))
+ signature (libraries.signature.java) {
+ artifact {
+ extension = "signature"
+ }
+ }
+}
+
+tasks.named("compileJava").configure {
+ it.options.compilerArgs += [
+ "-Xlint:-processing",
+ ]
+ appendToProperty(
+ it.options.errorprone.excludedPaths,
+ ".*/build/generated/sources/annotationProcessor/java/.*",
+ "|")
+}
+
+tasks.named("javadoc").configure {
+ failOnError = false
+ exclude 'io/grpc/autosharding/v1/**'
+ exclude 'io/grpc/autosharding/*Provider.java'
+ exclude 'io/grpc/autosharding/internal/**'
+ exclude 'io/grpc/autosharding/Internal*'
+}
+
+tasks.named("jacocoTestReport").configure {
+ classDirectories.from = sourceSets.main.output.collect {
+ fileTree(dir: it, exclude: ['**/io/grpc/autosharding/v1/**'])
+ }
+}
+
+configureProtoCompilation()
diff --git a/autosharding/src/generated/main/grpc/com/google/cloud/autosharding/v1main/DynamicShardingServiceGrpc.java b/autosharding/src/generated/main/grpc/com/google/cloud/autosharding/v1main/DynamicShardingServiceGrpc.java
new file mode 100644
index 00000000000..0105353640e
--- /dev/null
+++ b/autosharding/src/generated/main/grpc/com/google/cloud/autosharding/v1main/DynamicShardingServiceGrpc.java
@@ -0,0 +1,473 @@
+package com.google.cloud.autosharding.v1main;
+
+import static io.grpc.MethodDescriptor.generateFullMethodName;
+
+/**
+ *
+ * An auto-sharding service that assigns keys in an application's keyspace to
+ * abstract "endpoints", and which uses load information from clients to update
+ * that assignment over time. Assignments and load reports are scoped to an
+ * abstract "slicing target."
+ * Concrete examples of the concepts in this protocol:
+ * Endpoints:
+ * * Application servers
+ * * Pods in a Kubernetes cluster
+ * * Regions in a multi-regional service
+ * Keys:
+ * * User ids
+ * * Tenant ids in a multi-tenant system
+ * Load:
+ * * Request count
+ * * CPU cost of processing requests
+ *
+ */
+@io.grpc.stub.annotations.GrpcGenerated
+public final class DynamicShardingServiceGrpc {
+
+ private DynamicShardingServiceGrpc() {}
+
+ public static final java.lang.String SERVICE_NAME = "google.cloud.autosharding.v1main.DynamicShardingService";
+
+ // Static method descriptors that strictly reflect the proto.
+ private static volatile io.grpc.MethodDescriptor getWatchShardingAssignmentMethod;
+
+ @io.grpc.stub.annotations.RpcMethod(
+ fullMethodName = SERVICE_NAME + '/' + "WatchShardingAssignment",
+ requestType = com.google.cloud.autosharding.v1main.WatchShardingAssignmentRequest.class,
+ responseType = com.google.cloud.autosharding.v1main.WatchShardingAssignmentResponse.class,
+ methodType = io.grpc.MethodDescriptor.MethodType.BIDI_STREAMING)
+ public static io.grpc.MethodDescriptor getWatchShardingAssignmentMethod() {
+ io.grpc.MethodDescriptor getWatchShardingAssignmentMethod;
+ if ((getWatchShardingAssignmentMethod = DynamicShardingServiceGrpc.getWatchShardingAssignmentMethod) == null) {
+ synchronized (DynamicShardingServiceGrpc.class) {
+ if ((getWatchShardingAssignmentMethod = DynamicShardingServiceGrpc.getWatchShardingAssignmentMethod) == null) {
+ DynamicShardingServiceGrpc.getWatchShardingAssignmentMethod = getWatchShardingAssignmentMethod =
+ io.grpc.MethodDescriptor.newBuilder()
+ .setType(io.grpc.MethodDescriptor.MethodType.BIDI_STREAMING)
+ .setFullMethodName(generateFullMethodName(SERVICE_NAME, "WatchShardingAssignment"))
+ .setSampledToLocalTracing(true)
+ .setRequestMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
+ com.google.cloud.autosharding.v1main.WatchShardingAssignmentRequest.getDefaultInstance()))
+ .setResponseMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
+ com.google.cloud.autosharding.v1main.WatchShardingAssignmentResponse.getDefaultInstance()))
+ .setSchemaDescriptor(new DynamicShardingServiceMethodDescriptorSupplier("WatchShardingAssignment"))
+ .build();
+ }
+ }
+ }
+ return getWatchShardingAssignmentMethod;
+ }
+
+ /**
+ * Creates a new async stub that supports all call types for the service
+ */
+ public static DynamicShardingServiceStub newStub(io.grpc.Channel channel) {
+ io.grpc.stub.AbstractStub.StubFactory factory =
+ new io.grpc.stub.AbstractStub.StubFactory() {
+ @java.lang.Override
+ public DynamicShardingServiceStub newStub(io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ return new DynamicShardingServiceStub(channel, callOptions);
+ }
+ };
+ return DynamicShardingServiceStub.newStub(factory, channel);
+ }
+
+ /**
+ * Creates a new blocking-style stub that supports all types of calls on the service
+ */
+ public static DynamicShardingServiceBlockingV2Stub newBlockingV2Stub(
+ io.grpc.Channel channel) {
+ io.grpc.stub.AbstractStub.StubFactory factory =
+ new io.grpc.stub.AbstractStub.StubFactory() {
+ @java.lang.Override
+ public DynamicShardingServiceBlockingV2Stub newStub(io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ return new DynamicShardingServiceBlockingV2Stub(channel, callOptions);
+ }
+ };
+ return DynamicShardingServiceBlockingV2Stub.newStub(factory, channel);
+ }
+
+ /**
+ * Creates a new blocking-style stub that supports unary and streaming output calls on the service
+ */
+ public static DynamicShardingServiceBlockingStub newBlockingStub(
+ io.grpc.Channel channel) {
+ io.grpc.stub.AbstractStub.StubFactory factory =
+ new io.grpc.stub.AbstractStub.StubFactory() {
+ @java.lang.Override
+ public DynamicShardingServiceBlockingStub newStub(io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ return new DynamicShardingServiceBlockingStub(channel, callOptions);
+ }
+ };
+ return DynamicShardingServiceBlockingStub.newStub(factory, channel);
+ }
+
+ /**
+ * Creates a new ListenableFuture-style stub that supports unary calls on the service
+ */
+ public static DynamicShardingServiceFutureStub newFutureStub(
+ io.grpc.Channel channel) {
+ io.grpc.stub.AbstractStub.StubFactory factory =
+ new io.grpc.stub.AbstractStub.StubFactory() {
+ @java.lang.Override
+ public DynamicShardingServiceFutureStub newStub(io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ return new DynamicShardingServiceFutureStub(channel, callOptions);
+ }
+ };
+ return DynamicShardingServiceFutureStub.newStub(factory, channel);
+ }
+
+ /**
+ *
+ * An auto-sharding service that assigns keys in an application's keyspace to
+ * abstract "endpoints", and which uses load information from clients to update
+ * that assignment over time. Assignments and load reports are scoped to an
+ * abstract "slicing target."
+ * Concrete examples of the concepts in this protocol:
+ * Endpoints:
+ * * Application servers
+ * * Pods in a Kubernetes cluster
+ * * Regions in a multi-regional service
+ * Keys:
+ * * User ids
+ * * Tenant ids in a multi-tenant system
+ * Load:
+ * * Request count
+ * * CPU cost of processing requests
+ *
+ */
+ public interface AsyncService {
+
+ /**
+ *
+ * Opens a stream over which clients report load and the server delivers
+ * sharding assignments.
+ * A given client may use this stream just to receive assignments or also to
+ * report load, and it may opt in or out of reporting load at any time. (An
+ * implementation of the DynamicSharding service may impose restrictions on
+ * which clients are allowed to report load, and it may respond to clients
+ * improperly reporting load by ignoring their reports or terminating their
+ * streams with an error.)
+ * The client should keep this stream open at all times and reopen the stream
+ * after it closes, with backoff if the stream closed without delivering any
+ * data.
+ *
+ */
+ default io.grpc.stub.StreamObserver watchShardingAssignment(
+ io.grpc.stub.StreamObserver responseObserver) {
+ return io.grpc.stub.ServerCalls.asyncUnimplementedStreamingCall(getWatchShardingAssignmentMethod(), responseObserver);
+ }
+ }
+
+ /**
+ * Base class for the server implementation of the service DynamicShardingService.
+ *
+ * An auto-sharding service that assigns keys in an application's keyspace to
+ * abstract "endpoints", and which uses load information from clients to update
+ * that assignment over time. Assignments and load reports are scoped to an
+ * abstract "slicing target."
+ * Concrete examples of the concepts in this protocol:
+ * Endpoints:
+ * * Application servers
+ * * Pods in a Kubernetes cluster
+ * * Regions in a multi-regional service
+ * Keys:
+ * * User ids
+ * * Tenant ids in a multi-tenant system
+ * Load:
+ * * Request count
+ * * CPU cost of processing requests
+ *
+ */
+ public static abstract class DynamicShardingServiceImplBase
+ implements io.grpc.BindableService, AsyncService {
+
+ @java.lang.Override public final io.grpc.ServerServiceDefinition bindService() {
+ return DynamicShardingServiceGrpc.bindService(this);
+ }
+ }
+
+ /**
+ * A stub to allow clients to do asynchronous rpc calls to service DynamicShardingService.
+ *
+ * An auto-sharding service that assigns keys in an application's keyspace to
+ * abstract "endpoints", and which uses load information from clients to update
+ * that assignment over time. Assignments and load reports are scoped to an
+ * abstract "slicing target."
+ * Concrete examples of the concepts in this protocol:
+ * Endpoints:
+ * * Application servers
+ * * Pods in a Kubernetes cluster
+ * * Regions in a multi-regional service
+ * Keys:
+ * * User ids
+ * * Tenant ids in a multi-tenant system
+ * Load:
+ * * Request count
+ * * CPU cost of processing requests
+ *
+ */
+ public static final class DynamicShardingServiceStub
+ extends io.grpc.stub.AbstractAsyncStub {
+ private DynamicShardingServiceStub(
+ io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ super(channel, callOptions);
+ }
+
+ @java.lang.Override
+ protected DynamicShardingServiceStub build(
+ io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ return new DynamicShardingServiceStub(channel, callOptions);
+ }
+
+ /**
+ *
+ * Opens a stream over which clients report load and the server delivers
+ * sharding assignments.
+ * A given client may use this stream just to receive assignments or also to
+ * report load, and it may opt in or out of reporting load at any time. (An
+ * implementation of the DynamicSharding service may impose restrictions on
+ * which clients are allowed to report load, and it may respond to clients
+ * improperly reporting load by ignoring their reports or terminating their
+ * streams with an error.)
+ * The client should keep this stream open at all times and reopen the stream
+ * after it closes, with backoff if the stream closed without delivering any
+ * data.
+ *
+ */
+ public io.grpc.stub.StreamObserver watchShardingAssignment(
+ io.grpc.stub.StreamObserver responseObserver) {
+ return io.grpc.stub.ClientCalls.asyncBidiStreamingCall(
+ getChannel().newCall(getWatchShardingAssignmentMethod(), getCallOptions()), responseObserver);
+ }
+ }
+
+ /**
+ * A stub to allow clients to do synchronous rpc calls to service DynamicShardingService.
+ *
+ * An auto-sharding service that assigns keys in an application's keyspace to
+ * abstract "endpoints", and which uses load information from clients to update
+ * that assignment over time. Assignments and load reports are scoped to an
+ * abstract "slicing target."
+ * Concrete examples of the concepts in this protocol:
+ * Endpoints:
+ * * Application servers
+ * * Pods in a Kubernetes cluster
+ * * Regions in a multi-regional service
+ * Keys:
+ * * User ids
+ * * Tenant ids in a multi-tenant system
+ * Load:
+ * * Request count
+ * * CPU cost of processing requests
+ *
+ */
+ public static final class DynamicShardingServiceBlockingV2Stub
+ extends io.grpc.stub.AbstractBlockingStub {
+ private DynamicShardingServiceBlockingV2Stub(
+ io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ super(channel, callOptions);
+ }
+
+ @java.lang.Override
+ protected DynamicShardingServiceBlockingV2Stub build(
+ io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ return new DynamicShardingServiceBlockingV2Stub(channel, callOptions);
+ }
+
+ /**
+ *
+ * Opens a stream over which clients report load and the server delivers
+ * sharding assignments.
+ * A given client may use this stream just to receive assignments or also to
+ * report load, and it may opt in or out of reporting load at any time. (An
+ * implementation of the DynamicSharding service may impose restrictions on
+ * which clients are allowed to report load, and it may respond to clients
+ * improperly reporting load by ignoring their reports or terminating their
+ * streams with an error.)
+ * The client should keep this stream open at all times and reopen the stream
+ * after it closes, with backoff if the stream closed without delivering any
+ * data.
+ *
+ */
+ @io.grpc.ExperimentalApi("https://github.com/grpc/grpc-java/issues/10918")
+ public io.grpc.stub.BlockingClientCall
+ watchShardingAssignment() {
+ return io.grpc.stub.ClientCalls.blockingBidiStreamingCall(
+ getChannel(), getWatchShardingAssignmentMethod(), getCallOptions());
+ }
+ }
+
+ /**
+ * A stub to allow clients to do limited synchronous rpc calls to service DynamicShardingService.
+ *
+ * An auto-sharding service that assigns keys in an application's keyspace to
+ * abstract "endpoints", and which uses load information from clients to update
+ * that assignment over time. Assignments and load reports are scoped to an
+ * abstract "slicing target."
+ * Concrete examples of the concepts in this protocol:
+ * Endpoints:
+ * * Application servers
+ * * Pods in a Kubernetes cluster
+ * * Regions in a multi-regional service
+ * Keys:
+ * * User ids
+ * * Tenant ids in a multi-tenant system
+ * Load:
+ * * Request count
+ * * CPU cost of processing requests
+ *
+ */
+ public static final class DynamicShardingServiceBlockingStub
+ extends io.grpc.stub.AbstractBlockingStub {
+ private DynamicShardingServiceBlockingStub(
+ io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ super(channel, callOptions);
+ }
+
+ @java.lang.Override
+ protected DynamicShardingServiceBlockingStub build(
+ io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ return new DynamicShardingServiceBlockingStub(channel, callOptions);
+ }
+ }
+
+ /**
+ * A stub to allow clients to do ListenableFuture-style rpc calls to service DynamicShardingService.
+ *
+ * An auto-sharding service that assigns keys in an application's keyspace to
+ * abstract "endpoints", and which uses load information from clients to update
+ * that assignment over time. Assignments and load reports are scoped to an
+ * abstract "slicing target."
+ * Concrete examples of the concepts in this protocol:
+ * Endpoints:
+ * * Application servers
+ * * Pods in a Kubernetes cluster
+ * * Regions in a multi-regional service
+ * Keys:
+ * * User ids
+ * * Tenant ids in a multi-tenant system
+ * Load:
+ * * Request count
+ * * CPU cost of processing requests
+ *
+ */
+ public static final class DynamicShardingServiceFutureStub
+ extends io.grpc.stub.AbstractFutureStub {
+ private DynamicShardingServiceFutureStub(
+ io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ super(channel, callOptions);
+ }
+
+ @java.lang.Override
+ protected DynamicShardingServiceFutureStub build(
+ io.grpc.Channel channel, io.grpc.CallOptions callOptions) {
+ return new DynamicShardingServiceFutureStub(channel, callOptions);
+ }
+ }
+
+ private static final int METHODID_WATCH_SHARDING_ASSIGNMENT = 0;
+
+ private static final class MethodHandlers implements
+ io.grpc.stub.ServerCalls.UnaryMethod,
+ io.grpc.stub.ServerCalls.ServerStreamingMethod,
+ io.grpc.stub.ServerCalls.ClientStreamingMethod,
+ io.grpc.stub.ServerCalls.BidiStreamingMethod {
+ private final AsyncService serviceImpl;
+ private final int methodId;
+
+ MethodHandlers(AsyncService serviceImpl, int methodId) {
+ this.serviceImpl = serviceImpl;
+ this.methodId = methodId;
+ }
+
+ @java.lang.Override
+ @java.lang.SuppressWarnings("unchecked")
+ public void invoke(Req request, io.grpc.stub.StreamObserver responseObserver) {
+ switch (methodId) {
+ default:
+ throw new AssertionError();
+ }
+ }
+
+ @java.lang.Override
+ @java.lang.SuppressWarnings("unchecked")
+ public io.grpc.stub.StreamObserver invoke(
+ io.grpc.stub.StreamObserver responseObserver) {
+ switch (methodId) {
+ case METHODID_WATCH_SHARDING_ASSIGNMENT:
+ return (io.grpc.stub.StreamObserver) serviceImpl.watchShardingAssignment(
+ (io.grpc.stub.StreamObserver) responseObserver);
+ default:
+ throw new AssertionError();
+ }
+ }
+ }
+
+ public static final io.grpc.ServerServiceDefinition bindService(AsyncService service) {
+ return io.grpc.ServerServiceDefinition.builder(getServiceDescriptor())
+ .addMethod(
+ getWatchShardingAssignmentMethod(),
+ io.grpc.stub.ServerCalls.asyncBidiStreamingCall(
+ new MethodHandlers<
+ com.google.cloud.autosharding.v1main.WatchShardingAssignmentRequest,
+ com.google.cloud.autosharding.v1main.WatchShardingAssignmentResponse>(
+ service, METHODID_WATCH_SHARDING_ASSIGNMENT)))
+ .build();
+ }
+
+ private static abstract class DynamicShardingServiceBaseDescriptorSupplier
+ implements io.grpc.protobuf.ProtoFileDescriptorSupplier, io.grpc.protobuf.ProtoServiceDescriptorSupplier {
+ DynamicShardingServiceBaseDescriptorSupplier() {}
+
+ @java.lang.Override
+ public com.google.protobuf.Descriptors.FileDescriptor getFileDescriptor() {
+ return com.google.cloud.autosharding.v1main.DynamicShardingProto.getDescriptor();
+ }
+
+ @java.lang.Override
+ public com.google.protobuf.Descriptors.ServiceDescriptor getServiceDescriptor() {
+ return getFileDescriptor().findServiceByName("DynamicShardingService");
+ }
+ }
+
+ private static final class DynamicShardingServiceFileDescriptorSupplier
+ extends DynamicShardingServiceBaseDescriptorSupplier {
+ DynamicShardingServiceFileDescriptorSupplier() {}
+ }
+
+ private static final class DynamicShardingServiceMethodDescriptorSupplier
+ extends DynamicShardingServiceBaseDescriptorSupplier
+ implements io.grpc.protobuf.ProtoMethodDescriptorSupplier {
+ private final java.lang.String methodName;
+
+ DynamicShardingServiceMethodDescriptorSupplier(java.lang.String methodName) {
+ this.methodName = methodName;
+ }
+
+ @java.lang.Override
+ public com.google.protobuf.Descriptors.MethodDescriptor getMethodDescriptor() {
+ return getServiceDescriptor().findMethodByName(methodName);
+ }
+ }
+
+ private static volatile io.grpc.ServiceDescriptor serviceDescriptor;
+
+ public static io.grpc.ServiceDescriptor getServiceDescriptor() {
+ io.grpc.ServiceDescriptor result = serviceDescriptor;
+ if (result == null) {
+ synchronized (DynamicShardingServiceGrpc.class) {
+ result = serviceDescriptor;
+ if (result == null) {
+ serviceDescriptor = result = io.grpc.ServiceDescriptor.newBuilder(SERVICE_NAME)
+ .setSchemaDescriptor(new DynamicShardingServiceFileDescriptorSupplier())
+ .addMethod(getWatchShardingAssignmentMethod())
+ .build();
+ }
+ }
+ }
+ return result;
+ }
+}
diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java
new file mode 100644
index 00000000000..d211fdf2062
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java
@@ -0,0 +1,462 @@
+/*
+ * Copyright 2026 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.autosharding;
+
+import static io.grpc.ConnectivityState.CONNECTING;
+import static io.grpc.ConnectivityState.IDLE;
+import static io.grpc.ConnectivityState.READY;
+import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
+
+import com.google.cloud.autosharding.v1main.PerSliceEndpointState;
+import com.google.cloud.autosharding.v1main.SliceAssignment;
+import io.grpc.Attributes;
+import io.grpc.Channel;
+import io.grpc.ConnectivityState;
+import io.grpc.EquivalentAddressGroup;
+import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.LoadBalancerRegistry;
+import io.grpc.Status;
+import io.grpc.SynchronizationContext;
+import io.grpc.autosharding.AutoShardingLoadBalancerProvider.AutoShardingConfig;
+import io.grpc.autosharding.EndpointMap.EndpointHolder;
+import io.grpc.autosharding.SliceMap.SliceEntry;
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+public final class AutoShardingLoadBalancer extends LoadBalancer {
+ private static final Logger logger =
+ Logger.getLogger(AutoShardingLoadBalancer.class.getName());
+
+ /**
+ * ChannelFactory is used to create a Channel to the sharding service.
+ */
+ public interface ChannelFactory {
+ final class ChannelHolder implements AutoCloseable {
+ private final Channel channel;
+ private final Runnable releaseCallback;
+
+ public ChannelHolder(Channel channel, Runnable releaseCallback) {
+ this.channel = channel;
+ this.releaseCallback = releaseCallback;
+ }
+
+ public Channel getChannel() {
+ return channel;
+ }
+
+ @Override
+ public void close() {
+ if (releaseCallback != null) {
+ releaseCallback.run();
+ }
+ }
+ }
+
+ ChannelHolder createChannel(String target);
+ }
+
+ public static final Attributes.Key CHANNEL_FACTORY_KEY =
+ Attributes.Key.create("io.grpc.autosharding.AutoShardingLoadBalancer.CHANNEL_FACTORY");
+
+ public static final Attributes.Key LOCALITY_KEY =
+ Attributes.Key.create("io.grpc.autosharding.AutoShardingLoadBalancer.LOCALITY");
+
+ private final Helper helper;
+ private final SynchronizationContext syncContext;
+ private final LoadBalancerProvider pickFirstProvider;
+
+ private ChannelFactory.ChannelHolder shardingChannelHolder;
+ private Channel shardingChannel;
+ private ShardingClient shardingClient;
+ private String currentChannelFactoryKey;
+ private String currentSlicingTarget;
+
+ private boolean fallbackEnabled = false;
+ private String sliceKeyHeaderName = "";
+ private long initialAssignmentTimeoutNanos = TimeUnit.SECONDS.toNanos(60);
+
+ // Endpoint map: hostname -> EndpointHolder
+ private final EndpointMap endpointMap = new EndpointMap();
+
+ private SliceMap currentSliceMap;
+ private List latestSliceAssignments;
+ private List latestEndpointsProto;
+ private long latestGeneration = 0;
+
+ private SynchronizationContext.ScheduledHandle fallbackTimer;
+ private boolean fallbackTimerFired = false;
+
+ public AutoShardingLoadBalancer(Helper helper) {
+ this.helper = helper;
+ this.syncContext = helper.getSynchronizationContext();
+ this.pickFirstProvider =
+ LoadBalancerRegistry.getDefaultRegistry().getProvider("pick_first");
+ }
+
+ @Override
+ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+ AutoShardingConfig config =
+ (AutoShardingConfig) resolvedAddresses.getLoadBalancingPolicyConfig();
+ if (config == null) {
+ return Status.INVALID_ARGUMENT.withDescription("Missing AutoShardingConfig");
+ }
+
+ this.fallbackEnabled = config.enableFallback;
+ this.sliceKeyHeaderName = config.sliceKeyHeaderName;
+ if (config.initialAssignmentTimeoutNanos != null) {
+ this.initialAssignmentTimeoutNanos = config.initialAssignmentTimeoutNanos;
+ }
+
+ // Connect to sharding service if channelFactoryKey or slicingTarget changed
+ if (shardingClient == null
+ || !config.channelFactoryKey.equals(currentChannelFactoryKey)
+ || !config.slicingTarget.equals(currentSlicingTarget)) {
+ initShardingClient(
+ resolvedAddresses.getAttributes(),
+ config.channelFactoryKey,
+ config.slicingTarget);
+ }
+
+ // Process endpoints from Name Resolver
+ List addresses = resolvedAddresses.getAddresses();
+ if (addresses.isEmpty()) {
+ endpointMap.shutdownAll();
+ currentSliceMap = null;
+ helper.updateBalancingState(
+ TRANSIENT_FAILURE,
+ new FixedResultPicker(PickResult.withError(
+ Status.UNAVAILABLE.withDescription(
+ "NameResolver returned empty list of endpoints"))));
+ return Status.OK;
+ }
+
+ // Group addresses by hostname
+ Map> groupedByHostname = new LinkedHashMap<>();
+ for (EquivalentAddressGroup eag : addresses) {
+ String hostname = getHostname(eag);
+ List eags = groupedByHostname.get(hostname);
+ if (eags == null) {
+ eags = new ArrayList<>();
+ groupedByHostname.put(hostname, eags);
+ }
+ eags.add(eag);
+ }
+
+ Set newHostnames = groupedByHostname.keySet();
+ int index = 0;
+ for (Map.Entry> entry : groupedByHostname.entrySet()) {
+ String hostname = entry.getKey();
+ List eags = entry.getValue();
+
+ EndpointHolder holder = endpointMap.get(hostname);
+ if (holder == null) {
+ holder = new EndpointHolder(
+ index, helper, pickFirstProvider, this::updateAggregatedState);
+ endpointMap.put(hostname, holder);
+ } else {
+ holder.index = index;
+ }
+ holder.updateAddresses(eags, resolvedAddresses.getAttributes());
+ index++;
+ }
+
+ // Remove obsolete endpoints
+ List toRemove = new ArrayList<>();
+ for (String oldHost : endpointMap.keySet()) {
+ if (!newHostnames.contains(oldHost)) {
+ toRemove.add(oldHost);
+ }
+ }
+ for (String host : toRemove) {
+ EndpointHolder removed = endpointMap.remove(host);
+ if (removed != null) {
+ removed.shutdown();
+ }
+ }
+
+ // Re-index remaining endpoints so indices form contiguous 0..N-1
+ endpointMap.reindex();
+
+ // Build slice map if assignment received or if timer has fired
+ if (latestSliceAssignments != null || fallbackTimerFired) {
+ rebuildSliceMap();
+ }
+ updateAggregatedState();
+ return Status.OK;
+ }
+
+ private void closeShardingChannel() {
+ if (shardingClient != null) {
+ shardingClient.stop();
+ shardingClient = null;
+ }
+ if (shardingChannelHolder != null) {
+ try {
+ shardingChannelHolder.close();
+ } catch (Exception e) {
+ logger.log(Level.WARNING, "Error closing sharding channel", e);
+ }
+ shardingChannelHolder = null;
+ }
+ shardingChannel = null;
+ }
+
+ private void initShardingClient(
+ Attributes attributes, String channelFactoryKey, String slicingTarget) {
+ closeShardingChannel();
+ if (fallbackTimer != null) {
+ fallbackTimer.cancel();
+ fallbackTimer = null;
+ }
+
+ currentChannelFactoryKey = channelFactoryKey;
+ currentSlicingTarget = slicingTarget;
+
+ ChannelFactory factory = attributes.get(CHANNEL_FACTORY_KEY);
+ if (factory == null) {
+ logger.log(
+ Level.WARNING,
+ "No ChannelFactory attribute provided to AutoShardingLoadBalancer");
+ return;
+ }
+
+ String locality = attributes.get(LOCALITY_KEY);
+ if (locality == null) {
+ locality = "";
+ }
+ String actualTarget = slicingTarget.replace("%s", locality);
+
+ shardingChannelHolder = factory.createChannel(channelFactoryKey);
+ if (shardingChannelHolder == null || shardingChannelHolder.getChannel() == null) {
+ logger.log(
+ Level.WARNING,
+ "ChannelFactory returned null channel for target {0}",
+ channelFactoryKey);
+ return;
+ }
+ shardingChannel = shardingChannelHolder.getChannel();
+ shardingClient =
+ new ShardingClient(
+ shardingChannel,
+ actualTarget,
+ latestGeneration,
+ syncContext,
+ helper.getScheduledExecutorService(),
+ new ShardingCallback());
+ shardingClient.start();
+
+ // Start fallback-at-startup timer
+ fallbackTimerFired = false;
+ fallbackTimer = syncContext.schedule(
+ this::onFallbackTimerExpired,
+ initialAssignmentTimeoutNanos,
+ TimeUnit.NANOSECONDS,
+ helper.getScheduledExecutorService());
+ }
+
+ private void onFallbackTimerExpired() {
+ fallbackTimerFired = true;
+ fallbackTimer = null;
+ logger.log(Level.WARNING, "Initial assignment timeout expired. Entering fallback mode.");
+ latestSliceAssignments = null;
+ latestEndpointsProto = null;
+ latestGeneration = 0;
+ rebuildSliceMap();
+ updateAggregatedState();
+ }
+
+ private void rebuildSliceMap() {
+ if (endpointMap.isEmpty()) {
+ currentSliceMap = null;
+ return;
+ }
+
+ // Populate fallback_pool deterministically sorted by endpoint index
+ List fallbackPool = new ArrayList<>();
+ for (int i = 0; i < endpointMap.size(); i++) {
+ fallbackPool.add(i);
+ }
+
+ // If no assignment received yet (startup case), return early with empty slices
+ if (latestSliceAssignments == null || latestEndpointsProto == null) {
+ currentSliceMap = new SliceMap(Collections.emptyList(), fallbackPool, 0);
+ return;
+ }
+
+ List sliceEntries = new ArrayList<>();
+ for (SliceAssignment protoSlice : latestSliceAssignments) {
+ List sliceEndpoints = new ArrayList<>();
+ for (PerSliceEndpointState perSliceEp : protoSlice.getEndpointsList()) {
+ int protoEpIdx = perSliceEp.getEndpointIndex();
+ if (protoEpIdx >= 0 && protoEpIdx < latestEndpointsProto.size()) {
+ String hostname = latestEndpointsProto.get(protoEpIdx).getEndpoint();
+ EndpointHolder holder = endpointMap.get(hostname);
+ if (holder != null) {
+ sliceEndpoints.add(holder.index);
+ }
+ }
+ }
+ sliceEntries.add(
+ new SliceEntry(
+ protoSlice.getSlice().getStartKeyInclusive().toByteArray(), sliceEndpoints));
+ }
+
+ currentSliceMap = new SliceMap(sliceEntries, fallbackPool, latestGeneration);
+ }
+
+ private void updateAggregatedState() {
+ if (endpointMap.isEmpty()) {
+ helper.updateBalancingState(
+ TRANSIENT_FAILURE,
+ new FixedResultPicker(PickResult.withError(
+ Status.UNAVAILABLE.withDescription(
+ "NameResolver returned empty list of endpoints"))));
+ return;
+ }
+
+ int readyCount = 0;
+ int tfCount = 0;
+ int connectingCount = 0;
+ int idleCount = 0;
+
+ EndpointHolder firstIdle = null;
+
+ for (EndpointHolder holder : endpointMap.values()) {
+ ConnectivityState state = holder.state;
+ if (state == READY) {
+ readyCount++;
+ } else if (state == TRANSIENT_FAILURE) {
+ tfCount++;
+ } else if (state == CONNECTING) {
+ connectingCount++;
+ } else if (state == IDLE) {
+ idleCount++;
+ if (firstIdle == null) {
+ firstIdle = holder;
+ }
+ }
+ }
+
+ ConnectivityState aggregated;
+ int total = endpointMap.size();
+
+ if (readyCount > 0) {
+ aggregated = READY;
+ } else if (tfCount >= 2) {
+ aggregated = TRANSIENT_FAILURE;
+ } else if (connectingCount > 0) {
+ aggregated = CONNECTING;
+ } else if (tfCount == 1 && total > 1) {
+ aggregated = CONNECTING;
+ } else if (idleCount > 0) {
+ aggregated = IDLE;
+ } else {
+ aggregated = TRANSIENT_FAILURE;
+ }
+
+ // gRFC A119 heuristic: ensure at least one IDLE endpoint starts connecting
+ // if aggregated state is CONNECTING or TRANSIENT_FAILURE and none are CONNECTING
+ if ((aggregated == CONNECTING || aggregated == TRANSIENT_FAILURE)
+ && connectingCount == 0 && firstIdle != null) {
+ firstIdle.requestConnection();
+ }
+
+ SubchannelPicker picker;
+ if (currentSliceMap != null) {
+ List pickerEndpoints =
+ new ArrayList<>(Collections.nCopies(endpointMap.size(), null));
+ for (EndpointHolder holder : endpointMap.values()) {
+ pickerEndpoints.set(
+ holder.index,
+ new PickerEndpoint(
+ holder.state,
+ holder.picker,
+ () -> syncContext.execute(holder::requestConnection)));
+ }
+ picker = new AutoShardingPicker(
+ currentSliceMap, pickerEndpoints, fallbackEnabled, sliceKeyHeaderName);
+ } else {
+ picker = new FixedResultPicker(
+ PickResult.withNoResult(
+ "autosharding_assignment_pending", "Waiting for initial sharding assignment"));
+ }
+
+ helper.updateBalancingState(aggregated, picker);
+ }
+
+ @Override
+ public void handleNameResolutionError(Status error) {
+ helper.updateBalancingState(
+ TRANSIENT_FAILURE,
+ new FixedResultPicker(PickResult.withError(error)));
+ }
+
+ @Override
+ public void shutdown() {
+ closeShardingChannel();
+ if (fallbackTimer != null) {
+ fallbackTimer.cancel();
+ fallbackTimer = null;
+ }
+ endpointMap.shutdownAll();
+ }
+
+ private static String getHostname(EquivalentAddressGroup eag) {
+ String hostname = eag.getAttributes().get(EquivalentAddressGroup.ATTR_AUTHORITY_OVERRIDE);
+ if (hostname != null && !hostname.isEmpty()) {
+ return hostname;
+ }
+ SocketAddress address = eag.getAddresses().get(0);
+ if (address instanceof InetSocketAddress) {
+ return ((InetSocketAddress) address).getHostString();
+ }
+ return address.toString();
+ }
+
+ private final class ShardingCallback implements ShardingClient.Callback {
+ @Override
+ public void onAssignmentReceived(
+ List sliceAssignments,
+ List endpoints,
+ long generation) {
+ if (fallbackTimer != null) {
+ fallbackTimer.cancel();
+ fallbackTimer = null;
+ }
+ latestSliceAssignments = sliceAssignments;
+ latestEndpointsProto = endpoints;
+ latestGeneration = generation;
+ rebuildSliceMap();
+ updateAggregatedState();
+ }
+
+ @Override
+ public void onError(Throwable t) {
+ logger.log(Level.WARNING, "ShardingClient stream error", t);
+ }
+ }
+}
diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerProvider.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerProvider.java
new file mode 100644
index 00000000000..d27d79116cb
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerProvider.java
@@ -0,0 +1,162 @@
+/*
+ * Copyright 2026 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.autosharding;
+
+import com.google.common.base.MoreObjects;
+import io.grpc.Internal;
+import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.NameResolver.ConfigOrError;
+import io.grpc.Status;
+import io.grpc.internal.JsonUtil;
+import java.util.Map;
+import java.util.Objects;
+
+@Internal
+public final class AutoShardingLoadBalancerProvider extends LoadBalancerProvider {
+ static final String POLICY_NAME = "autosharding_experimental";
+
+ @Override
+ public boolean isAvailable() {
+ return true;
+ }
+
+ @Override
+ public int getPriority() {
+ return 5;
+ }
+
+ @Override
+ public String getPolicyName() {
+ return POLICY_NAME;
+ }
+
+ @Override
+ public LoadBalancer newLoadBalancer(LoadBalancer.Helper helper) {
+ return new AutoShardingLoadBalancer(helper);
+ }
+
+ @Override
+ public ConfigOrError parseLoadBalancingPolicyConfig(
+ Map rawLoadBalancingPolicyConfig) {
+ try {
+ String channelFactoryKey =
+ JsonUtil.getString(rawLoadBalancingPolicyConfig, "channelFactoryKey");
+ if (channelFactoryKey == null || channelFactoryKey.isEmpty()) {
+ return ConfigOrError.fromError(
+ Status.INVALID_ARGUMENT.withDescription(
+ "Missing required field 'channelFactoryKey' in autosharding config"));
+ }
+
+ String slicingTarget =
+ JsonUtil.getString(rawLoadBalancingPolicyConfig, "slicingTarget");
+ if (slicingTarget == null || slicingTarget.isEmpty()) {
+ return ConfigOrError.fromError(
+ Status.INVALID_ARGUMENT.withDescription(
+ "Missing required field 'slicingTarget' in autosharding config"));
+ }
+
+ String sliceKeyHeaderName =
+ JsonUtil.getString(rawLoadBalancingPolicyConfig, "sliceKeyHeaderName");
+ if (sliceKeyHeaderName == null || sliceKeyHeaderName.isEmpty()) {
+ return ConfigOrError.fromError(
+ Status.INVALID_ARGUMENT.withDescription(
+ "Missing required field 'sliceKeyHeaderName' in autosharding config"));
+ }
+
+ Boolean enableFallback =
+ JsonUtil.getBoolean(rawLoadBalancingPolicyConfig, "enableFallback");
+ if (enableFallback == null) {
+ enableFallback = false;
+ }
+
+ Long initialAssignmentTimeoutNanos =
+ JsonUtil.getStringAsDuration(
+ rawLoadBalancingPolicyConfig, "initialAssignmentTimeout");
+
+ return ConfigOrError.fromConfig(
+ new AutoShardingConfig(
+ channelFactoryKey,
+ slicingTarget,
+ sliceKeyHeaderName,
+ enableFallback,
+ initialAssignmentTimeoutNanos));
+ } catch (RuntimeException e) {
+ return ConfigOrError.fromError(
+ Status.INVALID_ARGUMENT.withDescription(
+ "Failed to parse autosharding config: " + e.getMessage()).withCause(e));
+ }
+ }
+
+ public static final class AutoShardingConfig {
+ final String channelFactoryKey;
+ final String slicingTarget;
+ final String sliceKeyHeaderName;
+ final boolean enableFallback;
+ final Long initialAssignmentTimeoutNanos;
+
+ public AutoShardingConfig(
+ String channelFactoryKey,
+ String slicingTarget,
+ String sliceKeyHeaderName,
+ boolean enableFallback,
+ Long initialAssignmentTimeoutNanos) {
+ this.channelFactoryKey = channelFactoryKey;
+ this.slicingTarget = slicingTarget;
+ this.sliceKeyHeaderName = sliceKeyHeaderName;
+ this.enableFallback = enableFallback;
+ this.initialAssignmentTimeoutNanos = initialAssignmentTimeoutNanos;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ AutoShardingConfig that = (AutoShardingConfig) o;
+ return enableFallback == that.enableFallback
+ && Objects.equals(channelFactoryKey, that.channelFactoryKey)
+ && Objects.equals(slicingTarget, that.slicingTarget)
+ && Objects.equals(sliceKeyHeaderName, that.sliceKeyHeaderName)
+ && Objects.equals(initialAssignmentTimeoutNanos, that.initialAssignmentTimeoutNanos);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(
+ channelFactoryKey,
+ slicingTarget,
+ sliceKeyHeaderName,
+ enableFallback,
+ initialAssignmentTimeoutNanos);
+ }
+
+ @Override
+ public String toString() {
+ return MoreObjects.toStringHelper(this)
+ .add("channelFactoryKey", channelFactoryKey)
+ .add("slicingTarget", slicingTarget)
+ .add("sliceKeyHeaderName", sliceKeyHeaderName)
+ .add("enableFallback", enableFallback)
+ .add("initialAssignmentTimeoutNanos", initialAssignmentTimeoutNanos)
+ .toString();
+ }
+ }
+}
diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java
new file mode 100644
index 00000000000..59bfe36a00c
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java
@@ -0,0 +1,156 @@
+/*
+ * Copyright 2026 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.autosharding;
+
+import io.grpc.ConnectivityState;
+import io.grpc.InternalMetadata;
+import io.grpc.LoadBalancer.PickResult;
+import io.grpc.LoadBalancer.PickSubchannelArgs;
+import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+
+final class AutoShardingPicker extends SubchannelPicker {
+ private static final byte[] EMPTY_BYTES = new byte[0];
+
+ private static final InternalMetadata.TrustedAsciiMarshaller RAW_ASCII_MARSHALLER =
+ new InternalMetadata.TrustedAsciiMarshaller() {
+ @Override
+ public byte[] toAsciiString(byte[] value) {
+ return value;
+ }
+
+ @Override
+ public byte[] parseAsciiString(byte[] serialized) {
+ return serialized;
+ }
+ };
+
+ private final SliceMap sliceMap;
+ private final List endpoints;
+ private final boolean[] sliceInFallback;
+ private final boolean fallbackEnabled;
+ private final Metadata.Key sliceKeyHeader;
+
+ AutoShardingPicker(
+ SliceMap sliceMap,
+ List endpoints,
+ boolean fallbackEnabled,
+ String sliceKeyHeaderName) {
+ this.sliceMap = sliceMap;
+ this.endpoints = Collections.unmodifiableList(new ArrayList<>(endpoints));
+ this.fallbackEnabled = fallbackEnabled;
+
+ if (sliceKeyHeaderName == null || sliceKeyHeaderName.isEmpty()) {
+ this.sliceKeyHeader = null;
+ } else if (sliceKeyHeaderName.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
+ this.sliceKeyHeader = Metadata.Key.of(sliceKeyHeaderName, Metadata.BINARY_BYTE_MARSHALLER);
+ } else {
+ this.sliceKeyHeader = InternalMetadata.keyOf(sliceKeyHeaderName, RAW_ASCII_MARSHALLER);
+ }
+
+ this.sliceInFallback = new boolean[sliceMap.getSlices().size()];
+ for (int i = 0; i < sliceInFallback.length; i++) {
+ this.sliceInFallback[i] = isPoolInFallback(sliceMap.getSlices().get(i).endpoints);
+ }
+ }
+
+ private boolean isPoolInFallback(List indices) {
+ if (indices.isEmpty()) {
+ return true;
+ }
+ for (int idx : indices) {
+ if (endpoints.get(idx).state != ConnectivityState.TRANSIENT_FAILURE) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ @Override
+ public PickResult pickSubchannel(PickSubchannelArgs args) {
+ byte[] key = extractKeyBytes(args.getHeaders());
+ Integer sliceIdx = sliceMap.lookup(key);
+
+ if (sliceIdx == null) {
+ if (fallbackEnabled) {
+ return pickFromEndpointIndices(sliceMap.getFallbackPool(), args);
+ } else {
+ return PickResult.withError(
+ Status.UNAVAILABLE.withDescription(
+ "No sharding assignment available and fallback disabled"));
+ }
+ }
+
+ if (sliceInFallback[sliceIdx] && fallbackEnabled) {
+ return pickFromEndpointIndices(sliceMap.getFallbackPool(), args);
+ }
+
+ SliceMap.SliceEntry sliceEntry = sliceMap.getSlices().get(sliceIdx);
+ return pickFromEndpointIndices(sliceEntry.endpoints, args);
+ }
+
+ private PickResult pickFromEndpointIndices(
+ List indices, PickSubchannelArgs args) {
+ if (indices.isEmpty()) {
+ return PickResult.withError(
+ Status.UNAVAILABLE.withDescription("No valid endpoints in slice and fallback disabled"));
+ }
+
+ int size = indices.size();
+ int firstIndex = ThreadLocalRandom.current().nextInt(size);
+ boolean requestedConnection = false;
+ boolean foundConnecting = false;
+
+ for (int i = 0; i < size; i++) {
+ int epIdx = indices.get((firstIndex + i) % size);
+ PickerEndpoint endpoint = endpoints.get(epIdx);
+
+ if (endpoint.state == ConnectivityState.READY) {
+ return endpoint.picker.pickSubchannel(args);
+ }
+
+ if (endpoint.state == ConnectivityState.CONNECTING) {
+ foundConnecting = true;
+ } else if (!requestedConnection && endpoint.state == ConnectivityState.IDLE) {
+ if (endpoint.requestConnection != null) {
+ endpoint.requestConnection.run();
+ }
+ requestedConnection = true;
+ }
+ }
+
+ if (requestedConnection || foundConnecting) {
+ return PickResult.withNoResult("connecting", "Waiting for endpoint connection");
+ }
+
+ int firstEpIdx = indices.get(firstIndex);
+ return endpoints.get(firstEpIdx).picker.pickSubchannel(args);
+ }
+
+ private byte[] extractKeyBytes(Metadata headers) {
+ if (sliceKeyHeader != null) {
+ byte[] val = headers.get(sliceKeyHeader);
+ return val != null ? val : EMPTY_BYTES;
+ }
+ return EMPTY_BYTES;
+ }
+}
diff --git a/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java
new file mode 100644
index 00000000000..8b55bed53ab
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java
@@ -0,0 +1,148 @@
+/*
+ * Copyright 2026 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.autosharding;
+
+import static io.grpc.ConnectivityState.IDLE;
+
+import io.grpc.Attributes;
+import io.grpc.ConnectivityState;
+import io.grpc.EquivalentAddressGroup;
+import io.grpc.LoadBalancer.FixedResultPicker;
+import io.grpc.LoadBalancer.Helper;
+import io.grpc.LoadBalancer.PickResult;
+import io.grpc.LoadBalancer.ResolvedAddresses;
+import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.util.ForwardingLoadBalancerHelper;
+import java.util.Collection;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Manages the mapping from endpoint hostname to {@link EndpointHolder} and their lifecycle.
+ */
+final class EndpointMap {
+ private final Map map = new LinkedHashMap<>();
+
+ EndpointHolder get(String hostname) {
+ return map.get(hostname);
+ }
+
+ void put(String hostname, EndpointHolder holder) {
+ map.put(hostname, holder);
+ }
+
+ EndpointHolder remove(String hostname) {
+ return map.remove(hostname);
+ }
+
+ Collection values() {
+ return map.values();
+ }
+
+ Set keySet() {
+ return map.keySet();
+ }
+
+ int size() {
+ return map.size();
+ }
+
+ boolean isEmpty() {
+ return map.isEmpty();
+ }
+
+ void clear() {
+ map.clear();
+ }
+
+ void shutdownAll() {
+ for (EndpointHolder holder : map.values()) {
+ holder.shutdown();
+ }
+ map.clear();
+ }
+
+ void reindex() {
+ int nextIdx = 0;
+ for (EndpointHolder holder : map.values()) {
+ holder.index = nextIdx++;
+ }
+ }
+
+ /**
+ * Holds the state and lazy child load balancer for a single endpoint.
+ */
+ static final class EndpointHolder {
+ int index;
+ final LazyChildLoadBalancer childLb;
+ ConnectivityState state = IDLE;
+ SubchannelPicker picker = new FixedResultPicker(PickResult.withNoResult());
+
+ EndpointHolder(
+ int index,
+ Helper helper,
+ LoadBalancerProvider pickFirstProvider,
+ Runnable stateUpdateCallback) {
+ this.index = index;
+ this.childLb = new LazyChildLoadBalancer(
+ new ChildHelper(helper, stateUpdateCallback), pickFirstProvider);
+ }
+
+ void updateAddresses(List eags, Attributes attributes) {
+ ResolvedAddresses childAddresses = ResolvedAddresses.newBuilder()
+ .setAddresses(eags)
+ .setAttributes(attributes)
+ .build();
+ childLb.acceptResolvedAddresses(childAddresses);
+ }
+
+ void requestConnection() {
+ childLb.requestConnection();
+ }
+
+ void shutdown() {
+ childLb.shutdown();
+ }
+
+ private final class ChildHelper extends ForwardingLoadBalancerHelper {
+ private final Helper delegateHelper;
+ private final Runnable stateUpdateCallback;
+
+ ChildHelper(Helper delegateHelper, Runnable stateUpdateCallback) {
+ this.delegateHelper = delegateHelper;
+ this.stateUpdateCallback = stateUpdateCallback;
+ }
+
+ @Override
+ protected Helper delegate() {
+ return delegateHelper;
+ }
+
+ @Override
+ public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
+ state = newState;
+ picker = newPicker;
+ if (stateUpdateCallback != null) {
+ stateUpdateCallback.run();
+ }
+ }
+ }
+ }
+}
diff --git a/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java
new file mode 100644
index 00000000000..618dbebdd70
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2026 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.autosharding;
+
+import io.grpc.ConnectivityState;
+import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.Status;
+
+/**
+ * A wrapper LoadBalancer that lazily creates and delegates to a child LoadBalancer
+ * (e.g. pick_first) only when {@link #requestConnection()} is explicitly called.
+ * Until then, it remains in IDLE state.
+ */
+final class LazyChildLoadBalancer extends LoadBalancer {
+ private final Helper helper;
+ private final LoadBalancerProvider delegateProvider;
+ private LoadBalancer delegate;
+ private ResolvedAddresses lastResolvedAddresses;
+ private boolean connectionRequested = false;
+
+ LazyChildLoadBalancer(Helper helper, LoadBalancerProvider delegateProvider) {
+ this.helper = helper;
+ this.delegateProvider = delegateProvider;
+ }
+
+ @Override
+ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+ lastResolvedAddresses = resolvedAddresses;
+ if (connectionRequested) {
+ if (delegate == null) {
+ delegate = delegateProvider.newLoadBalancer(helper);
+ }
+ return delegate.acceptResolvedAddresses(resolvedAddresses);
+ } else {
+ // Report IDLE state until connection is explicitly requested
+ helper.updateBalancingState(
+ ConnectivityState.IDLE,
+ new FixedResultPicker(PickResult.withNoResult()));
+ return Status.OK;
+ }
+ }
+
+ @Override
+ public void handleNameResolutionError(Status error) {
+ if (delegate != null) {
+ delegate.handleNameResolutionError(error);
+ } else {
+ helper.updateBalancingState(
+ ConnectivityState.TRANSIENT_FAILURE,
+ new FixedResultPicker(PickResult.withError(error)));
+ }
+ }
+
+ @Override
+ public void requestConnection() {
+ connectionRequested = true;
+ if (delegate == null && lastResolvedAddresses != null) {
+ delegate = delegateProvider.newLoadBalancer(helper);
+ delegate.acceptResolvedAddresses(lastResolvedAddresses);
+ } else if (delegate != null) {
+ delegate.requestConnection();
+ }
+ }
+
+ @Override
+ public void shutdown() {
+ if (delegate != null) {
+ delegate.shutdown();
+ delegate = null;
+ }
+ }
+
+ // Visible for testing
+ boolean isConnectionRequested() {
+ return connectionRequested;
+ }
+}
diff --git a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java
new file mode 100644
index 00000000000..54310a1bdb6
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2026 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.autosharding;
+
+import io.grpc.ConnectivityState;
+import io.grpc.LoadBalancer.SubchannelPicker;
+
+/**
+ * Immutable snapshot of endpoint state used by the AutoShardingPicker.
+ */
+final class PickerEndpoint {
+ final ConnectivityState state;
+ final SubchannelPicker picker;
+ final Runnable requestConnection;
+
+ PickerEndpoint(
+ ConnectivityState state, SubchannelPicker picker, Runnable requestConnection) {
+ this.state = state;
+ this.picker = picker;
+ this.requestConnection = requestConnection;
+ }
+}
diff --git a/autosharding/src/main/java/io/grpc/autosharding/ShardingClient.java b/autosharding/src/main/java/io/grpc/autosharding/ShardingClient.java
new file mode 100644
index 00000000000..2d5ce69e506
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/ShardingClient.java
@@ -0,0 +1,281 @@
+/*
+ * Copyright 2026 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.autosharding;
+
+import com.google.cloud.autosharding.v1main.AssignmentChunk;
+import com.google.cloud.autosharding.v1main.AssignmentMetadata;
+import com.google.cloud.autosharding.v1main.DynamicShardingServiceGrpc;
+import com.google.cloud.autosharding.v1main.EndpointState;
+import com.google.cloud.autosharding.v1main.InitialClientConfig;
+import com.google.cloud.autosharding.v1main.PerSliceEndpointState;
+import com.google.cloud.autosharding.v1main.SliceAssignment;
+import com.google.cloud.autosharding.v1main.WatchShardingAssignmentRequest;
+import com.google.cloud.autosharding.v1main.WatchShardingAssignmentResponse;
+import com.google.protobuf.ByteString;
+import io.grpc.Channel;
+import io.grpc.Context;
+import io.grpc.Status;
+import io.grpc.SynchronizationContext;
+import io.grpc.internal.BackoffPolicy;
+import io.grpc.internal.ExponentialBackoffPolicy;
+import io.grpc.stub.StreamObserver;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+final class ShardingClient {
+ private static final Logger logger = Logger.getLogger(ShardingClient.class.getName());
+
+ interface Callback {
+ void onAssignmentReceived(
+ List sliceAssignments,
+ List endpoints,
+ long generation);
+
+ void onError(Throwable t);
+ }
+
+ private final Channel channel;
+ private final String target;
+ private final String clientUuid;
+ private final SynchronizationContext syncContext;
+ private final ScheduledExecutorService scheduledExecutorService;
+ private final BackoffPolicy.Provider backoffPolicyProvider;
+ private final Callback callback;
+
+ private Context.CancellableContext cancellableContext;
+ private StreamObserver requestStream;
+ private long currentGeneration = 0;
+ private BackoffPolicy backoffPolicy;
+ private SynchronizationContext.ScheduledHandle retryTimer;
+ private boolean stopped = false;
+
+ // State for assembling chunks
+ private final List currentSliceAssignments = new ArrayList<>();
+ private final List currentEndpoints = new ArrayList<>();
+
+ ShardingClient(
+ Channel channel,
+ String target,
+ long initialGeneration,
+ SynchronizationContext syncContext,
+ ScheduledExecutorService scheduledExecutorService,
+ Callback callback) {
+ this(
+ channel,
+ target,
+ initialGeneration,
+ syncContext,
+ scheduledExecutorService,
+ new ExponentialBackoffPolicy.Provider(),
+ callback);
+ }
+
+ ShardingClient(
+ Channel channel,
+ String target,
+ long initialGeneration,
+ SynchronizationContext syncContext,
+ ScheduledExecutorService scheduledExecutorService,
+ BackoffPolicy.Provider backoffPolicyProvider,
+ Callback callback) {
+ this.channel = channel;
+ this.target = target;
+ this.clientUuid = UUID.randomUUID().toString();
+ this.currentGeneration = initialGeneration;
+ this.syncContext = syncContext;
+ this.scheduledExecutorService = scheduledExecutorService;
+ this.backoffPolicyProvider = backoffPolicyProvider;
+ this.callback = callback;
+ }
+
+ void start() {
+ if (stopped) {
+ return;
+ }
+ closeStream();
+
+ cancellableContext = Context.current().withCancellation();
+ cancellableContext.run(() -> {
+ DynamicShardingServiceGrpc.DynamicShardingServiceStub stub =
+ DynamicShardingServiceGrpc.newStub(channel).withWaitForReady();
+
+ requestStream = stub.watchShardingAssignment(
+ new StreamObserver() {
+ @Override
+ public void onNext(WatchShardingAssignmentResponse response) {
+ syncContext.execute(() -> handleResponse(response));
+ }
+
+ @Override
+ public void onError(Throwable t) {
+ syncContext.execute(() -> handleError(t));
+ }
+
+ @Override
+ public void onCompleted() {
+ syncContext.execute(() -> handleError(
+ Status.UNAVAILABLE.withDescription("Server closed stream").asRuntimeException()));
+ }
+ });
+
+ InitialClientConfig initConfig = InitialClientConfig.newBuilder()
+ .setTarget(target)
+ .setClientUuid(clientUuid)
+ .setCurrentGeneration(currentGeneration)
+ .build();
+
+ requestStream.onNext(WatchShardingAssignmentRequest.newBuilder().setInit(initConfig).build());
+ });
+ }
+
+ private void handleResponse(WatchShardingAssignmentResponse response) {
+ if (response.hasChunk()) {
+ AssignmentChunk chunk = response.getChunk();
+ currentSliceAssignments.addAll(chunk.getSliceAssignmentsList());
+ currentEndpoints.addAll(chunk.getEndpointsList());
+ } else if (response.hasMetadata()) {
+ List assembledSlices = new ArrayList<>(currentSliceAssignments);
+ List assembledEndpoints = new ArrayList<>(currentEndpoints);
+
+ currentSliceAssignments.clear();
+ currentEndpoints.clear();
+
+ // Validate assignment per gRFC A119
+ if (!validateAssignment(assembledSlices, assembledEndpoints)) {
+ logger.log(
+ Level.WARNING, "Assignment validation failed. Terminating stream to reconnect.");
+ handleError(
+ Status.INTERNAL.withDescription(
+ "Assignment validation failed: invalid key ranges or endpoint indices")
+ .asRuntimeException());
+ return;
+ }
+
+ // Reset backoff state upon receiving a good logical assignment from the server
+ backoffPolicy = null;
+ AssignmentMetadata metadata = response.getMetadata();
+ currentGeneration = metadata.getGeneration();
+
+ callback.onAssignmentReceived(assembledSlices, assembledEndpoints, currentGeneration);
+ } else if (response.hasConfig()) {
+ // Ignore for now as per gRFC
+ } else {
+ logger.log(Level.WARNING, "Received unknown or empty response from sharding service");
+ }
+ }
+
+ private static boolean validateAssignment(
+ List slices, List endpoints) {
+ if (slices.isEmpty()) {
+ return false;
+ }
+ int totalEndpoints = endpoints.size();
+
+ // 1. Ensure all endpoint indices are valid
+ for (SliceAssignment sa : slices) {
+ for (PerSliceEndpointState pse : sa.getEndpointsList()) {
+ int epIdx = pse.getEndpointIndex();
+ if (epIdx < 0 || epIdx >= totalEndpoints) {
+ return false;
+ }
+ }
+ }
+
+ // 2. Ensure no gaps in key ranges and covers full range ["" .. ""]
+ List sorted = new ArrayList<>(slices);
+ sorted.sort(
+ Comparator.comparing(
+ sa -> sa.getSlice().getStartKeyInclusive(),
+ ByteString.unsignedLexicographicalComparator()));
+
+ // First slice must start at empty ByteString (start of keyspace)
+ if (!sorted.get(0).getSlice().getStartKeyInclusive().isEmpty()) {
+ return false;
+ }
+
+ for (int i = 0; i < sorted.size() - 1; i++) {
+ ByteString currentEnd = sorted.get(i).getSlice().getEndKeyExclusive();
+ ByteString nextStart = sorted.get(i + 1).getSlice().getStartKeyInclusive();
+ // If end_key is unset (empty), this slice extends to the largest allowed key.
+ // It cannot have subsequent slices.
+ if (currentEnd.isEmpty()) {
+ return false;
+ }
+ if (!currentEnd.equals(nextStart)) {
+ return false; // Gap or overlap
+ }
+ }
+
+ // Last slice's end_key must be empty (sentinel indicating end of keyspace)
+ if (!sorted.get(sorted.size() - 1).getSlice().getEndKeyExclusive().isEmpty()) {
+ return false;
+ }
+
+ return true;
+ }
+
+ private void closeStream() {
+ if (cancellableContext != null) {
+ cancellableContext.cancel(
+ Status.CANCELLED.withDescription("Stream closed by client").asRuntimeException());
+ cancellableContext = null;
+ }
+ requestStream = null;
+ }
+
+ private void handleError(Throwable t) {
+ currentSliceAssignments.clear();
+ currentEndpoints.clear();
+ closeStream();
+ callback.onError(t);
+ if (!stopped) {
+ scheduleReconnect();
+ }
+ }
+
+ private void scheduleReconnect() {
+ if (backoffPolicy == null) {
+ backoffPolicy = backoffPolicyProvider.get();
+ }
+ long delayNanos = backoffPolicy.nextBackoffNanos();
+ logger.log(Level.INFO, "ShardingClient stream disconnected. Retrying in {0} ns", delayNanos);
+ if (retryTimer != null) {
+ retryTimer.cancel();
+ }
+ retryTimer =
+ syncContext.schedule(
+ this::start,
+ delayNanos,
+ TimeUnit.NANOSECONDS,
+ scheduledExecutorService);
+ }
+
+ void stop() {
+ stopped = true;
+ if (retryTimer != null) {
+ retryTimer.cancel();
+ retryTimer = null;
+ }
+ closeStream();
+ }
+}
diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java
new file mode 100644
index 00000000000..5444c252556
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright 2026 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.autosharding;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import javax.annotation.Nullable;
+
+final class SliceMap {
+
+ static final class SliceEntry {
+ final byte[] startKey;
+ final List endpoints;
+
+ SliceEntry(byte[] startKey, List endpoints) {
+ this.startKey = startKey;
+ this.endpoints = Collections.unmodifiableList(new ArrayList<>(endpoints));
+ }
+ }
+
+ private static final byte[] EMPTY_BYTES = new byte[0];
+
+ private final List slices;
+ private final List fallbackPool;
+ private final long generation;
+
+ SliceMap(List slices, List fallbackPool, long generation) {
+ List sortedSlices = new ArrayList<>(slices);
+ sortedSlices.sort((e1, e2) -> compareUnsigned(e1.startKey, e2.startKey));
+ this.slices = Collections.unmodifiableList(sortedSlices);
+ this.fallbackPool = Collections.unmodifiableList(new ArrayList<>(fallbackPool));
+ this.generation = generation;
+ }
+
+ /**
+ * Looks up the matching slice index for the given key.
+ * Returns null if slices is empty (e.g. startup/fallback case where there are no assignments).
+ */
+ @Nullable
+ Integer lookup(@Nullable byte[] key) {
+ if (slices.isEmpty()) {
+ return null;
+ }
+ byte[] searchKey = key != null ? key : EMPTY_BYTES;
+ int low = 0;
+ int high = slices.size() - 1;
+
+ while (low <= high) {
+ int mid = (low + high) >>> 1;
+ int cmp = compareUnsigned(slices.get(mid).startKey, searchKey);
+
+ if (cmp < 0) {
+ low = mid + 1;
+ } else if (cmp > 0) {
+ high = mid - 1;
+ } else {
+ return mid; // Exact match on startKey
+ }
+ }
+
+ if (low == 0) {
+ // Key is smaller than first slice's startKey
+ return null;
+ }
+ return low - 1;
+ }
+
+ private static int compareUnsigned(byte[] a, byte[] b) {
+ int minLength = Math.min(a.length, b.length);
+ for (int i = 0; i < minLength; i++) {
+ int result = (a[i] & 0xFF) - (b[i] & 0xFF);
+ if (result != 0) {
+ return result;
+ }
+ }
+ return a.length - b.length;
+ }
+
+ List getSlices() {
+ return slices;
+ }
+
+ List getFallbackPool() {
+ return fallbackPool;
+ }
+
+ long getGeneration() {
+ return generation;
+ }
+}
diff --git a/autosharding/src/main/proto/google/cloud/autosharding/v1main/dynamic_sharding.proto b/autosharding/src/main/proto/google/cloud/autosharding/v1main/dynamic_sharding.proto
new file mode 100644
index 00000000000..443e63f1306
--- /dev/null
+++ b/autosharding/src/main/proto/google/cloud/autosharding/v1main/dynamic_sharding.proto
@@ -0,0 +1,287 @@
+syntax = "proto3";
+
+package google.cloud.autosharding.v1main;
+
+import "google/api/field_behavior.proto";
+import "google/protobuf/duration.proto";
+import "google/protobuf/timestamp.proto";
+
+option java_package = "com.google.cloud.autosharding.v1main";
+option java_multiple_files = true;
+option java_outer_classname = "DynamicShardingProto";
+
+// An auto-sharding service that assigns keys in an application's keyspace to
+// abstract "endpoints", and which uses load information from clients to update
+// that assignment over time. Assignments and load reports are scoped to an
+// abstract "slicing target."
+//
+// Concrete examples of the concepts in this protocol:
+//
+// Endpoints:
+// * Application servers
+// * Pods in a Kubernetes cluster
+// * Regions in a multi-regional service
+//
+// Keys:
+// * User ids
+// * Tenant ids in a multi-tenant system
+//
+// Load:
+// * Request count
+// * CPU cost of processing requests
+service DynamicShardingService {
+ // Opens a stream over which clients report load and the server delivers
+ // sharding assignments.
+ //
+ // A given client may use this stream just to receive assignments or also to
+ // report load, and it may opt in or out of reporting load at any time. (An
+ // implementation of the DynamicSharding service may impose restrictions on
+ // which clients are allowed to report load, and it may respond to clients
+ // improperly reporting load by ignoring their reports or terminating their
+ // streams with an error.)
+ //
+ // The client should keep this stream open at all times and reopen the stream
+ // after it closes, with backoff if the stream closed without delivering any
+ // data.
+ rpc WatchShardingAssignment(stream WatchShardingAssignmentRequest)
+ returns (stream WatchShardingAssignmentResponse);
+}
+
+// In each instance of this message, the client MUST set exactly one field.
+message WatchShardingAssignmentRequest {
+ // Initial client-side configuration necessary to set up the stream. This
+ // field MUST be set in the first message sent, and it MUST NOT be set in
+ // subsequent messages.
+ InitialClientConfig init = 1 [(google.api.field_behavior) = OPTIONAL];
+
+ // An aggregated snapshot of load data collected by this client.
+ //
+ // The client MUST include all load data it has collected since the last
+ // report in this report (irrespective of whether the last report was sent
+ // on this or a prior physical stream), and it MUST NOT include that data in
+ // any subsequent report.
+ //
+ // This field MUST NOT be set in the first message sent.
+ LoadReport load_report = 2 [(google.api.field_behavior) = OPTIONAL];
+}
+
+// In each instance of this message, the server will set exactly one field.
+message WatchShardingAssignmentResponse {
+ // Configuration dictating how the client should aggregate and send load
+ // reports.
+ //
+ // The client MUST wait to send load reports until it has received the first
+ // version of this config from the server (on this physical stream).
+ LoadReportingConfig config = 1;
+
+ // A chunk of a logical assignment.
+ //
+ // The server delivers an assignment as a sequence of one or more
+ // `AssignmentChunk` messages followed by a single `AssignmentMetadata`
+ // message. Clients MUST wait for an `AssignmentMetadata` message and combine
+ // the prior chunks to form a logical assignment before using it.
+ AssignmentChunk chunk = 2;
+
+ // Metadata associated with the most recent chunked assignment. Acts as a
+ // trailer indicating that the server has finished delivering all chunks for
+ // that assignment.
+ AssignmentMetadata metadata = 3;
+}
+
+// Initial client-side configuration necessary to set up the stream.
+message InitialClientConfig {
+ // The slicing target to which this client's load reports and assignments
+ // apply.
+ //
+ // If this target is unknown to the server, the server will terminate the
+ // stream with a NOT_FOUND error.
+ //
+ // If this client does not have permission to connect to this target, the
+ // server will terminate the stream with a PERMISSION_DENIED error.
+ string target = 1 [(google.api.field_behavior) = REQUIRED];
+
+ // A UUID generated by this client upon instantiation. The client MUST reuse
+ // the same UUID across DynamicSharding stream reconnections. It is
+ // permissible but not expected for the client to use the same UUID across
+ // process restarts.
+ string client_uuid = 2 [(google.api.field_behavior) = REQUIRED];
+
+ // The latest assignment generation that this client has received. If set, the
+ // server will only send newer assignments.
+ int64 current_generation = 3 [(google.api.field_behavior) = OPTIONAL];
+}
+
+// Server-specified configuration dictating how the client should report load.
+message LoadReportingConfig {
+ // The maximum fraction of total load in a `Load` message that appears in a
+ // single `RangeLoad` message, if the range is not a single-key range. When
+ // aggregating load data into a load report, the client MUST honor the latest
+ // received value of this field as of when it started aggregating. (See the
+ // `RangeLoad` documentation for a description of how the client should
+ // aggregate load data into load reports.)
+ //
+ // This field will always be in the range [0, 1].
+ double load_quantum_frac = 1;
+
+ // The maximum length of slice keys in load reports. When aggregating load
+ // data into a load report, the client MUST honor the latest received value
+ // of this field as of when it started aggregating. (See the `RangeLoad`
+ // documentation for a description of how the client should aggregate load
+ // data into load reports.)
+ int32 max_load_balancing_resolution_bytes = 2;
+
+ // The interval at which this client should send `LoadReport` messages. After
+ // it sends a given load report to the server, the client MUST wait for
+ // roughly the interval specified by the latest received value of this field
+ // before it sends its next load report.
+ google.protobuf.Duration load_reporting_interval = 3;
+
+ // The lower bound on a load report's `start_time_inclusive`, relative to the
+ // client's current time when it starts computing the load report. Clients
+ // should discard load data before this threshold.
+ google.protobuf.Duration load_report_start_threshold = 4;
+}
+
+// A load report consisting of per-metric load snapshots.
+message LoadReport {
+ // Snapshots of load in the keyspace associated with one or more metrics.
+ //
+ // A key in this map is a named load metric to which the load applies. This
+ // metric MUST be a "cost" metric, i.e., one where the total load for a given
+ // application key over a time interval is the sum of the individual load
+ // values recorded for that key over the time interval, and where the number
+ // of endpoints assigned a given application key does not affect the total
+ // load for that key.
+ map metric_load = 1 [(google.api.field_behavior) = REQUIRED];
+
+ // The inclusive start timestamp of the time range to which this load report
+ // applies.
+ google.protobuf.Timestamp start_time_inclusive = 2
+ [(google.api.field_behavior) = REQUIRED];
+
+ // The exclusive end timestamp of the time range to which this load report
+ // applies.
+ //
+ // This end time MUST be roughly equal to the client's current time, i.e.,
+ // the client must send this load report to the server immediately after
+ // computing it.
+ //
+ // To ensure good load-balancing behavior, the clock used to set these
+ // timestamps MUST be synchronized to within ~one second of the clocks of
+ // other clients and the DynamicSharding service. Implementations of the
+ // DynamicSharding service should have some mechanism for informing client
+ // owners when they detect significantly delayed client load reports (which
+ // suggests an actionable issue related to either network delays or clock
+ // skew).
+ google.protobuf.Timestamp end_time_exclusive = 3
+ [(google.api.field_behavior) = REQUIRED];
+}
+
+// A per-key-range load snapshot.
+message Load {
+ // An ordered set of per-key-range load. Each key range's exclusive end is
+ // implicitly the start of the next range in the ordered set.
+ //
+ // Entries here MUST have lexicographically increasing start keys, and two
+ // entries MUST NOT have the same start key.
+ //
+ // Start keys in these entries MUST be at most
+ // `max_load_balancing_resolution_bytes` bytes (using the latest received
+ // value as of when the client started aggregating this load report).
+ //
+ // Clients MUST aggregate observed load values for a given key during this
+ // report's time interval into a single value for the key by summing the
+ // values, and clients MUST aggregate per-key load values into a single load
+ // value for a key range by summing the values.
+ //
+ // A single range MUST contain at most `load_quantum_frac` of the total
+ // load across all ranges, unless it is a single-key range. (A range is a
+ // single-key range if the next range's start key is the immediate
+ // lexicographic successor of this range's start key in the space of
+ // bytestrings of length `max_load_balancing_resolution_bytes`.)
+ //
+ // Subject to the above constraints, clients MUST aggregate this load data
+ // into as few key ranges as possible.
+ //
+ // Clients under memory pressure may report "inexact" load according to a
+ // strategy specified by the DynamicSharding service owner, e.g., by
+ // sampling observed load data rather than retaining each data point.
+ //
+ // Implementation sketch for aggregating load into a load report: maintain a
+ // hash map from key to total load observed since the last report. When it is
+ // time to send a new report, truncate the keys to
+ // `max_load_balancing_resolution_bytes`, then sort the load values
+ // lexicographically by key. Walk the keyspace starting from the min key,
+ // maintaining a current "open range" and its total load. If adding the
+ // current key to the current open range would push its total load over
+ // `load_quantum_frac`, then "close" the range and start a new open range.
+ // Otherwise, add the current key to the current open range.
+ repeated RangeLoad key_load = 1;
+}
+
+// Load associated with a key range (whose end is implicitly defined by the
+// next range's start).
+message RangeLoad {
+ // The key range's start.
+ bytes start_key_inclusive = 1;
+
+ // The aggregate load associated with this key range observed by this client
+ // during the time interval in this load report.
+ double load = 2;
+}
+
+// Metadata associated with an assignment.
+message AssignmentMetadata {
+ // The assignment's unique and monotonically-increasing version number.
+ int64 generation = 1;
+}
+
+// A portion of an assignment of slices to endpoints. The individual fields'
+// documentation describes how to reassemble a set of chunks into a complete
+// assignment.
+message AssignmentChunk {
+ // A chunk of the set of assignments for individual slices.
+ //
+ // Clients MUST reassemble the complete `slice_assignments` list by
+ // concatenating the contents of this field across chunks (in any order).
+ repeated SliceAssignment slice_assignments = 1;
+
+ // A chunk of the list of all endpoints present in this assignment.
+ //
+ // Clients MUST reassemble the complete `endpoints` list by concatenating the
+ // contents of this field across chunks, in chunk order.
+ repeated EndpointState endpoints = 2;
+}
+
+// An endpoint and any associated metadata.
+message EndpointState {
+ // The endpoint's name.
+ string endpoint = 1;
+}
+
+// The assignment for an individual slice.
+message SliceAssignment {
+ // The slice to which endpoints are assigned.
+ Slice slice = 1;
+
+ // The endpoints assigned to this slice.
+ repeated PerSliceEndpointState endpoints = 2;
+}
+
+// A single slice, i.e., a continuous range of keys.
+message Slice {
+ // The slice's start. Will be at most 512 bytes.
+ bytes start_key_inclusive = 1;
+
+ // The slice's end. Will be at most 512 bytes. If this field is unset, this
+ // slice contains the largest allowed key (the 512-byte string 0xFFFF...).
+ bytes end_key_exclusive = 2;
+}
+
+// An endpoint assigned to a slice, along with any associated metadata about
+// the endpoint in the context of that slice.
+message PerSliceEndpointState {
+ // An endpoint, specified as an index into the complete assignment's
+ // reassembled `endpoints` list.
+ int32 endpoint_index = 1;
+}
diff --git a/autosharding/src/main/resources/META-INF/services/io.grpc.LoadBalancerProvider b/autosharding/src/main/resources/META-INF/services/io.grpc.LoadBalancerProvider
new file mode 100644
index 00000000000..a977316d164
--- /dev/null
+++ b/autosharding/src/main/resources/META-INF/services/io.grpc.LoadBalancerProvider
@@ -0,0 +1 @@
+io.grpc.autosharding.AutoShardingLoadBalancerProvider
diff --git a/settings.gradle b/settings.gradle
index abd30a63956..0322a9a1fbd 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -95,6 +95,7 @@ include ":grpc-inprocess"
include ":grpc-util"
include ":grpc-opentelemetry"
include ":grpc-context-override-opentelemetry"
+include ":grpc-autosharding"
project(':grpc-api').projectDir = "$rootDir/api" as File
project(':grpc-core').projectDir = "$rootDir/core" as File
@@ -132,6 +133,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-autosharding').projectDir = "$rootDir/autosharding" as File
if (settings.hasProperty('skipCodegen') && skipCodegen.toBoolean()) {
println '*** Skipping the build of codegen and compilation of proto files because skipCodegen=true'