From 13e6d52ca05d9425903cc311372d81c8b2d7fbe5 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Wed, 26 Aug 2026 21:24:02 +0530 Subject: [PATCH 01/28] autosharding: Add module build configuration and protobuf definitions --- autosharding/build.gradle | 63 +++ .../v1/AutoshardingServiceGrpc.java | 473 ++++++++++++++++++ .../proto/autosharding/v1/autosharding.proto | 327 ++++++++++++ settings.gradle | 2 + 4 files changed, 865 insertions(+) create mode 100644 autosharding/build.gradle create mode 100644 autosharding/src/generated/main/grpc/com/google/cloud/autosharding/v1/AutoshardingServiceGrpc.java create mode 100644 autosharding/src/main/proto/autosharding/v1/autosharding.proto 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/v1/AutoshardingServiceGrpc.java b/autosharding/src/generated/main/grpc/com/google/cloud/autosharding/v1/AutoshardingServiceGrpc.java new file mode 100644 index 00000000000..908872d03da --- /dev/null +++ b/autosharding/src/generated/main/grpc/com/google/cloud/autosharding/v1/AutoshardingServiceGrpc.java @@ -0,0 +1,473 @@ +package com.google.cloud.autosharding.v1; + +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 AutoshardingServiceGrpc { + + private AutoshardingServiceGrpc() {} + + public static final java.lang.String SERVICE_NAME = "google.cloud.autosharding.v1.AutoshardingService"; + + // 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.v1.WatchShardingAssignmentRequest.class, + responseType = com.google.cloud.autosharding.v1.WatchShardingAssignmentResponse.class, + methodType = io.grpc.MethodDescriptor.MethodType.BIDI_STREAMING) + public static io.grpc.MethodDescriptor getWatchShardingAssignmentMethod() { + io.grpc.MethodDescriptor getWatchShardingAssignmentMethod; + if ((getWatchShardingAssignmentMethod = AutoshardingServiceGrpc.getWatchShardingAssignmentMethod) == null) { + synchronized (AutoshardingServiceGrpc.class) { + if ((getWatchShardingAssignmentMethod = AutoshardingServiceGrpc.getWatchShardingAssignmentMethod) == null) { + AutoshardingServiceGrpc.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.v1.WatchShardingAssignmentRequest.getDefaultInstance())) + .setResponseMarshaller(io.grpc.protobuf.ProtoUtils.marshaller( + com.google.cloud.autosharding.v1.WatchShardingAssignmentResponse.getDefaultInstance())) + .setSchemaDescriptor(new AutoshardingServiceMethodDescriptorSupplier("WatchShardingAssignment")) + .build(); + } + } + } + return getWatchShardingAssignmentMethod; + } + + /** + * Creates a new async stub that supports all call types for the service + */ + public static AutoshardingServiceStub newStub(io.grpc.Channel channel) { + io.grpc.stub.AbstractStub.StubFactory factory = + new io.grpc.stub.AbstractStub.StubFactory() { + @java.lang.Override + public AutoshardingServiceStub newStub(io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + return new AutoshardingServiceStub(channel, callOptions); + } + }; + return AutoshardingServiceStub.newStub(factory, channel); + } + + /** + * Creates a new blocking-style stub that supports all types of calls on the service + */ + public static AutoshardingServiceBlockingV2Stub newBlockingV2Stub( + io.grpc.Channel channel) { + io.grpc.stub.AbstractStub.StubFactory factory = + new io.grpc.stub.AbstractStub.StubFactory() { + @java.lang.Override + public AutoshardingServiceBlockingV2Stub newStub(io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + return new AutoshardingServiceBlockingV2Stub(channel, callOptions); + } + }; + return AutoshardingServiceBlockingV2Stub.newStub(factory, channel); + } + + /** + * Creates a new blocking-style stub that supports unary and streaming output calls on the service + */ + public static AutoshardingServiceBlockingStub newBlockingStub( + io.grpc.Channel channel) { + io.grpc.stub.AbstractStub.StubFactory factory = + new io.grpc.stub.AbstractStub.StubFactory() { + @java.lang.Override + public AutoshardingServiceBlockingStub newStub(io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + return new AutoshardingServiceBlockingStub(channel, callOptions); + } + }; + return AutoshardingServiceBlockingStub.newStub(factory, channel); + } + + /** + * Creates a new ListenableFuture-style stub that supports unary calls on the service + */ + public static AutoshardingServiceFutureStub newFutureStub( + io.grpc.Channel channel) { + io.grpc.stub.AbstractStub.StubFactory factory = + new io.grpc.stub.AbstractStub.StubFactory() { + @java.lang.Override + public AutoshardingServiceFutureStub newStub(io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + return new AutoshardingServiceFutureStub(channel, callOptions); + } + }; + return AutoshardingServiceFutureStub.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 Autosharding 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 AutoshardingService. + * 
+   * 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 AutoshardingServiceImplBase + implements io.grpc.BindableService, AsyncService { + + @java.lang.Override public final io.grpc.ServerServiceDefinition bindService() { + return AutoshardingServiceGrpc.bindService(this); + } + } + + /** + * A stub to allow clients to do asynchronous rpc calls to service AutoshardingService. + * 
+   * 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 AutoshardingServiceStub + extends io.grpc.stub.AbstractAsyncStub { + private AutoshardingServiceStub( + io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + super(channel, callOptions); + } + + @java.lang.Override + protected AutoshardingServiceStub build( + io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + return new AutoshardingServiceStub(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 Autosharding 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 AutoshardingService. + * 
+   * 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 AutoshardingServiceBlockingV2Stub + extends io.grpc.stub.AbstractBlockingStub { + private AutoshardingServiceBlockingV2Stub( + io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + super(channel, callOptions); + } + + @java.lang.Override + protected AutoshardingServiceBlockingV2Stub build( + io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + return new AutoshardingServiceBlockingV2Stub(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 Autosharding 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 AutoshardingService. + * 
+   * 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 AutoshardingServiceBlockingStub + extends io.grpc.stub.AbstractBlockingStub { + private AutoshardingServiceBlockingStub( + io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + super(channel, callOptions); + } + + @java.lang.Override + protected AutoshardingServiceBlockingStub build( + io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + return new AutoshardingServiceBlockingStub(channel, callOptions); + } + } + + /** + * A stub to allow clients to do ListenableFuture-style rpc calls to service AutoshardingService. + * 
+   * 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 AutoshardingServiceFutureStub + extends io.grpc.stub.AbstractFutureStub { + private AutoshardingServiceFutureStub( + io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + super(channel, callOptions); + } + + @java.lang.Override + protected AutoshardingServiceFutureStub build( + io.grpc.Channel channel, io.grpc.CallOptions callOptions) { + return new AutoshardingServiceFutureStub(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.v1.WatchShardingAssignmentRequest, + com.google.cloud.autosharding.v1.WatchShardingAssignmentResponse>( + service, METHODID_WATCH_SHARDING_ASSIGNMENT))) + .build(); + } + + private static abstract class AutoshardingServiceBaseDescriptorSupplier + implements io.grpc.protobuf.ProtoFileDescriptorSupplier, io.grpc.protobuf.ProtoServiceDescriptorSupplier { + AutoshardingServiceBaseDescriptorSupplier() {} + + @java.lang.Override + public com.google.protobuf.Descriptors.FileDescriptor getFileDescriptor() { + return com.google.cloud.autosharding.v1.AutoshardingProto.getDescriptor(); + } + + @java.lang.Override + public com.google.protobuf.Descriptors.ServiceDescriptor getServiceDescriptor() { + return getFileDescriptor().findServiceByName("AutoshardingService"); + } + } + + private static final class AutoshardingServiceFileDescriptorSupplier + extends AutoshardingServiceBaseDescriptorSupplier { + AutoshardingServiceFileDescriptorSupplier() {} + } + + private static final class AutoshardingServiceMethodDescriptorSupplier + extends AutoshardingServiceBaseDescriptorSupplier + implements io.grpc.protobuf.ProtoMethodDescriptorSupplier { + private final java.lang.String methodName; + + AutoshardingServiceMethodDescriptorSupplier(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 (AutoshardingServiceGrpc.class) { + result = serviceDescriptor; + if (result == null) { + serviceDescriptor = result = io.grpc.ServiceDescriptor.newBuilder(SERVICE_NAME) + .setSchemaDescriptor(new AutoshardingServiceFileDescriptorSupplier()) + .addMethod(getWatchShardingAssignmentMethod()) + .build(); + } + } + } + return result; + } +} diff --git a/autosharding/src/main/proto/autosharding/v1/autosharding.proto b/autosharding/src/main/proto/autosharding/v1/autosharding.proto new file mode 100644 index 00000000000..c14a4b8921a --- /dev/null +++ b/autosharding/src/main/proto/autosharding/v1/autosharding.proto @@ -0,0 +1,327 @@ +syntax = "proto3"; + +package google.cloud.autosharding.v1; + +import "google/api/field_behavior.proto"; +import "google/protobuf/duration.proto"; +import "google/protobuf/timestamp.proto"; + +option go_package = "github.com/GoogleCloudPlatform/autosharding/apiv1;autoshardingpb"; +option java_package = "com.google.cloud.autosharding.v1"; +option java_multiple_files = true; +option java_outer_classname = "AutoshardingProto"; + +// 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 AutoshardingService { + // 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 Autosharding 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 initial_client_config = 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]; + + // The client's acknowledgment of the last received assignment. + // + // The client MUST respond with this message to every received + // `AssignmentMetadata` message from the server. + AssignmentAck assignment_ack = 3 [(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 Autosharding 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. If unset (= zero), the server + // will send the latest assignment immediately. + int64 latest_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_fraction = 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`, 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_loads = 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 = 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 Autosharding service. Implementations of the + // Autosharding 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 = 3 + [(google.api.field_behavior) = REQUIRED]; +} + +// A per-key-range load snapshot. +message Load { + // An ordered set of per-key-range load. Each range's inclusive start is + // implicitly the exclusive end of the previous 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_fraction` 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` or if the start + // key is the largest allowed key. + // + // 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 Autosharding 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_fraction`, then "close" the range and start a new open range. + // Otherwise, add the current key to the current open range. + repeated RangeLoad range_loads = 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 (inclusive). + bytes start_key = 1; + + // The aggregate load associated with this key range observed by this client + // during the time interval in this load report. + double load = 2; +} + +// The client's acknowledgment of the last received assignment. +message AssignmentAck { + // The generation of the assignment this acknowledgment applies to. + int64 generation = 1 [(google.api.field_behavior) = REQUIRED]; + + // Whether the client accepted the assignment. + // + // If the client did not accept the assignment, it MUST set this field to + // false and continue using the previous assignment, and it MUST set the + // `error_message` field to explain the reason for rejection. + bool accepted = 2 [(google.api.field_behavior) = REQUIRED]; + + // A textual error message from the client. + // + // This field MUST be set if and only if the client detected issues with the + // assignment, explaining those issues. The client may set this field even if + // it accepted the assignment, in which case the issues are not blocking from + // using the assignment. + // + // Service owners may use this information for fixing problems. + // + // The length of this field MUST NOT exceed 512 characters (Unicode code + // points, see https://google.aip.dev/210). + string error_message = 3 [ + (google.api.field_behavior) = OPTIONAL + ]; +} + +// 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 (inclusive). Will be at most 512 bytes. + bytes start_key = 1; + + // The slice's end (exclusive). Will be at most 512 bytes. If this field is + // unset, this slice contains the largest allowed key (the 512-byte string + // 0xFFFF...). + optional bytes end_key = 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/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' From 04c6963061a175726921d9b933aa1e80f42daf84 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Wed, 26 Aug 2026 21:37:40 +0530 Subject: [PATCH 02/28] add import.sh --- autosharding/import.sh | 56 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 autosharding/import.sh diff --git a/autosharding/import.sh b/autosharding/import.sh new file mode 100644 index 00000000000..545146aa077 --- /dev/null +++ b/autosharding/import.sh @@ -0,0 +1,56 @@ +#!/bin/bash +# 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. + +# Update VERSION then execute this script + +set -e +VERSION=3a21f69fe5516fe340af29ef68e327381473ecc7 +DOWNLOAD_URL="https://github.com/GoogleCloudPlatform/autosharding/archive/${VERSION}.tar.gz" +DOWNLOAD_BASE_DIR="autosharding-${VERSION}" +SOURCE_PROTO_BASE_DIR="${DOWNLOAD_BASE_DIR}/proto" +TARGET_PROTO_BASE_DIR="src/main/proto" +# Sorted alphabetically. +FILES=( +autosharding/v1/autosharding.proto +) + +pushd "$(git rev-parse --show-toplevel)/autosharding" > /dev/null + +# put the repo in a tmp directory +tmpdir="$(mktemp -d)" +trap "rm -rf ${tmpdir}" EXIT +curl -Ls "${DOWNLOAD_URL}" | tar xz -C "${tmpdir}" + +rm -rf "${TARGET_PROTO_BASE_DIR}" +mkdir -p "${TARGET_PROTO_BASE_DIR}" +pushd "${TARGET_PROTO_BASE_DIR}" > /dev/null + +# copy proto files to project directory +TOTAL=${#FILES[@]} +COPIED=0 +for file in "${FILES[@]}" +do + mkdir -p "$(dirname "${file}")" + cp -p "${tmpdir}/${SOURCE_PROTO_BASE_DIR}/${file}" "${file}" && (( ++COPIED )) +done +popd > /dev/null + +popd > /dev/null + +echo "Imported ${COPIED} files." +if (( COPIED != TOTAL )); then + echo "Failed importing $(( TOTAL - COPIED )) files." 1>&2 + exit 1 +fi From d16170599e00828be02dc78c9fae9cae885703de Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 1 Sep 2026 12:45:38 +0530 Subject: [PATCH 03/28] autosharding: Move proto and import.sh to third_party/autosharding directory --- autosharding/build.gradle | 8 + autosharding/third_party/autosharding/LICENSE | 201 ++++++++++++++++++ .../{ => third_party/autosharding}/import.sh | 4 +- .../proto/autosharding/v1/autosharding.proto | 0 4 files changed, 212 insertions(+), 1 deletion(-) create mode 100644 autosharding/third_party/autosharding/LICENSE rename autosharding/{ => third_party/autosharding}/import.sh (91%) mode change 100644 => 100755 rename autosharding/{ => third_party/autosharding}/src/main/proto/autosharding/v1/autosharding.proto (100%) diff --git a/autosharding/build.gradle b/autosharding/build.gradle index 26ffdacdd6d..8c88f6d0e64 100644 --- a/autosharding/build.gradle +++ b/autosharding/build.gradle @@ -14,6 +14,14 @@ tasks.named("jar").configure { } } +sourceSets { + main { + proto { + srcDir 'third_party/autosharding/src/main/proto' + } + } +} + dependencies { implementation project(':grpc-util'), project(':grpc-core'), diff --git a/autosharding/third_party/autosharding/LICENSE b/autosharding/third_party/autosharding/LICENSE new file mode 100644 index 00000000000..7260c1bae0d --- /dev/null +++ b/autosharding/third_party/autosharding/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2017, Google Research. + + 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. diff --git a/autosharding/import.sh b/autosharding/third_party/autosharding/import.sh old mode 100644 new mode 100755 similarity index 91% rename from autosharding/import.sh rename to autosharding/third_party/autosharding/import.sh index 545146aa077..d9e8fb1c886 --- a/autosharding/import.sh +++ b/autosharding/third_party/autosharding/import.sh @@ -26,13 +26,15 @@ FILES=( autosharding/v1/autosharding.proto ) -pushd "$(git rev-parse --show-toplevel)/autosharding" > /dev/null +pushd `git rev-parse --show-toplevel`/autosharding/third_party/autosharding > /dev/null # put the repo in a tmp directory tmpdir="$(mktemp -d)" trap "rm -rf ${tmpdir}" EXIT curl -Ls "${DOWNLOAD_URL}" | tar xz -C "${tmpdir}" +cp -p "${tmpdir}/${DOWNLOAD_BASE_DIR}/LICENSE" LICENSE + rm -rf "${TARGET_PROTO_BASE_DIR}" mkdir -p "${TARGET_PROTO_BASE_DIR}" pushd "${TARGET_PROTO_BASE_DIR}" > /dev/null diff --git a/autosharding/src/main/proto/autosharding/v1/autosharding.proto b/autosharding/third_party/autosharding/src/main/proto/autosharding/v1/autosharding.proto similarity index 100% rename from autosharding/src/main/proto/autosharding/v1/autosharding.proto rename to autosharding/third_party/autosharding/src/main/proto/autosharding/v1/autosharding.proto From 74bdfd125900ced685457c8db9dcfd48da7745ef Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 1 Sep 2026 08:25:23 +0530 Subject: [PATCH 04/28] autosharding: Add SliceMap and AutoShardingPicker --- .../grpc/autosharding/AutoShardingPicker.java | 156 +++++++++++ .../io/grpc/autosharding/PickerEndpoint.java | 36 +++ .../java/io/grpc/autosharding/SliceMap.java | 105 +++++++ .../autosharding/AutoShardingPickerTest.java | 259 ++++++++++++++++++ .../io/grpc/autosharding/SliceMapTest.java | 139 ++++++++++ 5 files changed, 695 insertions(+) create mode 100644 autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java create mode 100644 autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java create mode 100644 autosharding/src/main/java/io/grpc/autosharding/SliceMap.java create mode 100644 autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java create mode 100644 autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java 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/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/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/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java new file mode 100644 index 00000000000..03cb935bdb3 --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -0,0 +1,259 @@ +/* + * 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 com.google.common.truth.Truth.assertThat; + +import io.grpc.CallOptions; +import io.grpc.ConnectivityState; +import io.grpc.LoadBalancer.PickDetailsConsumer; +import io.grpc.LoadBalancer.PickResult; +import io.grpc.LoadBalancer.PickSubchannelArgs; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Status; +import io.grpc.autosharding.SliceMap.SliceEntry; +import io.grpc.internal.PickSubchannelArgsImpl; +import io.grpc.testing.TestMethodDescriptors; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AutoShardingPickerTest { + + private static final MethodDescriptor METHOD = TestMethodDescriptors.voidMethod(); + private static final Runnable NOOP = new Runnable() { + @Override + public void run() {} + }; + private static final PickDetailsConsumer NOOP_CONSUMER = new PickDetailsConsumer() {}; + + private PickSubchannelArgs createArgs(Metadata headers) { + return new PickSubchannelArgsImpl(METHOD, headers, CallOptions.DEFAULT, NOOP_CONSUMER); + } + + private static class FakePicker extends SubchannelPicker { + private final PickResult result; + + FakePicker(PickResult result) { + this.result = result; + } + + @Override + public PickResult pickSubchannel(PickSubchannelArgs args) { + return result; + } + } + + @Test + public void pick_noSliceMap_fallbackEnabled_picksFromFallbackPool() { + PickResult readyResult = PickResult.withNoResult(); // using as token + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(readyResult), NOOP); + + SliceMap emptySliceMap = new SliceMap( + Collections.emptyList(), Collections.singletonList(0), 1L); + AutoShardingPicker picker = new AutoShardingPicker( + emptySliceMap, Collections.singletonList(ep0), true, "x-slice-key"); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("x-slice-key", Metadata.ASCII_STRING_MARSHALLER), "user123"); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(readyResult); + } + + @Test + public void pick_noSliceMap_fallbackDisabled_returnsUnavailableError() { + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP); + + SliceMap emptySliceMap = new SliceMap( + Collections.emptyList(), Collections.singletonList(0), 1L); + AutoShardingPicker picker = new AutoShardingPicker( + emptySliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + Metadata headers = new Metadata(); + PickResult result = picker.pickSubchannel(createArgs(headers)); + + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()) + .contains("No sharding assignment available and fallback disabled"); + } + + @Test + public void pick_sliceFound_readyEndpoint_returnsPickResult() { + PickResult expectedResult = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(expectedResult), NOOP); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("x-slice-key", Metadata.ASCII_STRING_MARSHALLER), "anyKey"); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(expectedResult); + } + + @Test + public void pick_sliceFound_idleEndpoint_triggersConnectionAndQueues() { + AtomicInteger connectCalls = new AtomicInteger(0); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.IDLE, + new FakePicker(PickResult.withNoResult()), + connectCalls::incrementAndGet); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + + assertThat(connectCalls.get()).isEqualTo(1); + assertThat(result.hasResult()).isFalse(); + } + + @Test + public void pick_sliceFound_connectingEndpoint_queuesPick() { + AtomicInteger connectCalls = new AtomicInteger(0); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.CONNECTING, + new FakePicker(PickResult.withNoResult()), + connectCalls::incrementAndGet); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + + assertThat(connectCalls.get()).isEqualTo(0); + assertThat(result.hasResult()).isFalse(); + } + + @Test + public void pick_sliceFound_allTransientFailure_fallbackEnabled_picksFromFallbackPool() { + PickResult fallbackReadyResult = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.TRANSIENT_FAILURE, + new FakePicker(PickResult.withError(Status.UNAVAILABLE.withDescription("ep0 down"))), + NOOP); + PickerEndpoint ep1 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP); + + // Slice 0 only has ep0 (which is down) + SliceEntry slice0 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + // Fallback pool has ep1 (which is ready) + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice0), Collections.singletonList(1), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Arrays.asList(ep0, ep1), true, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result).isSameInstanceAs(fallbackReadyResult); + } + + @Test + public void pick_sliceFound_allTransientFailure_fallbackDisabled_delegatesToEndpointPicker() { + Status epError = Status.UNAVAILABLE.withDescription("connection refused to ep0"); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.TRANSIENT_FAILURE, + new FakePicker(PickResult.withError(epError)), + NOOP); + + SliceEntry slice0 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice0), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result.getStatus()).isEqualTo(epError); + } + + @Test + public void pick_binaryHeader_extractedProperly() { + PickResult ready0 = PickResult.withNoResult(); + PickResult ready1 = PickResult.withNoResult(); + + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready0), NOOP); + PickerEndpoint ep1 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready1), NOOP); + + SliceEntry s0 = new SliceEntry(new byte[] {0x00}, Collections.singletonList(0)); + SliceEntry s1 = new SliceEntry(new byte[] {0x50}, Collections.singletonList(1)); + SliceMap sliceMap = new SliceMap(Arrays.asList(s0, s1), Arrays.asList(0, 1), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Arrays.asList(ep0, ep1), false, "slice-key-bin"); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("slice-key-bin", Metadata.BINARY_BYTE_MARSHALLER), + new byte[] {0x60}); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(ready1); + } + + @Test + public void pick_emptySliceEndpoints_fallbackDisabled_returnsUnavailable() { + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP); + + SliceEntry emptySlice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.emptyList()); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(emptySlice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()) + .contains("No valid endpoints in slice and fallback disabled"); + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java new file mode 100644 index 00000000000..2ec753cf4e9 --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java @@ -0,0 +1,139 @@ +/* + * 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 com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import io.grpc.autosharding.SliceMap.SliceEntry; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class SliceMapTest { + + @Test + public void lookup_emptySlices_returnsNull() { + SliceMap sliceMap = new SliceMap(Collections.emptyList(), Arrays.asList(0, 1), 1L); + assertThat(sliceMap.lookup(new byte[] {1, 2, 3})).isNull(); + assertThat(sliceMap.lookup(null)).isNull(); + assertThat(sliceMap.lookup(new byte[0])).isNull(); + } + + @Test + public void lookup_singleSlice() { + byte[] startKey = new byte[0]; // Covers ["" .. inf) + SliceEntry slice = new SliceEntry(startKey, Arrays.asList(0, 1)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Arrays.asList(0, 1), 10L); + + assertThat(sliceMap.lookup(new byte[0])).isEqualTo(0); + assertThat(sliceMap.lookup("foo".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup(null)).isEqualTo(0); + } + + @Test + public void lookup_multipleSlices() { + // Slices: ["" .. "m"), ["m" .. "t"), ["t" .. inf) + SliceEntry s1 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceEntry s2 = new SliceEntry( + "m".getBytes(StandardCharsets.UTF_8), Collections.singletonList(1)); + SliceEntry s3 = new SliceEntry( + "t".getBytes(StandardCharsets.UTF_8), Collections.singletonList(2)); + + SliceMap sliceMap = new SliceMap(Arrays.asList(s3, s1, s2), Arrays.asList(0, 1, 2), 5L); + + // Exact matches + assertThat(sliceMap.lookup("".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("m".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("t".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + + // In-between matches + assertThat(sliceMap.lookup("a".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("l".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("n".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("s".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("u".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + assertThat(sliceMap.lookup("zzz".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + } + + @Test + public void lookup_keySmallerThanFirstSlice_returnsNull() { + // Slice starts at "m" + SliceEntry s1 = new SliceEntry( + "m".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(s1), Collections.singletonList(0), 1L); + + assertThat(sliceMap.lookup("a".getBytes(StandardCharsets.UTF_8))).isNull(); + assertThat(sliceMap.lookup("".getBytes(StandardCharsets.UTF_8))).isNull(); + assertThat(sliceMap.lookup(null)).isNull(); + assertThat(sliceMap.lookup("m".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("z".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + } + + @Test + public void lookup_unsignedByteComparison() { + // Test that 0x80 is treated as greater than 0x7F (unsigned) + byte[] key1 = new byte[] {0x7F}; + byte[] key2 = new byte[] {(byte) 0x80}; + byte[] key3 = new byte[] {(byte) 0xFF}; + + SliceEntry s1 = new SliceEntry(new byte[0], Collections.singletonList(0)); + SliceEntry s2 = new SliceEntry(key1, Collections.singletonList(1)); + SliceEntry s3 = new SliceEntry(key2, Collections.singletonList(2)); + SliceEntry s4 = new SliceEntry(key3, Collections.singletonList(3)); + + SliceMap sliceMap = new SliceMap( + Arrays.asList(s4, s2, s1, s3), Arrays.asList(0, 1, 2, 3), 1L); + + assertThat(sliceMap.lookup(new byte[] {0x10})).isEqualTo(0); + assertThat(sliceMap.lookup(new byte[] {0x7F})).isEqualTo(1); + assertThat(sliceMap.lookup(new byte[] {(byte) 0x80})).isEqualTo(2); + assertThat(sliceMap.lookup(new byte[] {(byte) 0x90})).isEqualTo(2); + assertThat(sliceMap.lookup(new byte[] {(byte) 0xFF})).isEqualTo(3); + assertThat(sliceMap.lookup(new byte[] {(byte) 0xFF, 0x01})).isEqualTo(3); + } + + @Test + public void gettersAndImmutability() { + List slices = new ArrayList<>(); + slices.add(new SliceEntry(new byte[] {1}, Arrays.asList(0, 1))); + List fallback = new ArrayList<>(Arrays.asList(0, 1)); + + SliceMap sliceMap = new SliceMap(slices, fallback, 42L); + + assertThat(sliceMap.getGeneration()).isEqualTo(42L); + assertThat(sliceMap.getFallbackPool()).containsExactly(0, 1).inOrder(); + assertThat(sliceMap.getSlices()).hasSize(1); + assertThat(sliceMap.getSlices().get(0).endpoints).containsExactly(0, 1).inOrder(); + + // Verify immutability + assertThrows(UnsupportedOperationException.class, () -> sliceMap.getSlices().clear()); + assertThrows(UnsupportedOperationException.class, () -> sliceMap.getFallbackPool().clear()); + assertThrows( + UnsupportedOperationException.class, + () -> sliceMap.getSlices().get(0).endpoints.clear()); + } +} From e5b4b22be205a8aeb4d764ed24f9bb92106fd2fc Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Wed, 2 Sep 2026 18:42:58 +0530 Subject: [PATCH 05/28] grfc updated --- .../grpc/autosharding/AutoShardingPicker.java | 18 +++++++------- .../autosharding/AutoShardingPickerTest.java | 24 +++++++++++++++++++ 2 files changed, 33 insertions(+), 9 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index 59bfe36a00c..b7ab4e71cea 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -48,23 +48,23 @@ public byte[] parseAsciiString(byte[] serialized) { private final List endpoints; private final boolean[] sliceInFallback; private final boolean fallbackEnabled; - private final Metadata.Key sliceKeyHeader; + private final Metadata.Key keyHeader; AutoShardingPicker( SliceMap sliceMap, List endpoints, boolean fallbackEnabled, - String sliceKeyHeaderName) { + String keyHeaderName) { 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); + if (keyHeaderName == null || keyHeaderName.isEmpty()) { + this.keyHeader = null; + } else if (keyHeaderName.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { + this.keyHeader = Metadata.Key.of(keyHeaderName, Metadata.BINARY_BYTE_MARSHALLER); } else { - this.sliceKeyHeader = InternalMetadata.keyOf(sliceKeyHeaderName, RAW_ASCII_MARSHALLER); + this.keyHeader = InternalMetadata.keyOf(keyHeaderName, RAW_ASCII_MARSHALLER); } this.sliceInFallback = new boolean[sliceMap.getSlices().size()]; @@ -147,8 +147,8 @@ private PickResult pickFromEndpointIndices( } private byte[] extractKeyBytes(Metadata headers) { - if (sliceKeyHeader != null) { - byte[] val = headers.get(sliceKeyHeader); + if (keyHeader != null) { + byte[] val = headers.get(keyHeader); return val != null ? val : EMPTY_BYTES; } return EMPTY_BYTES; diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java index 03cb935bdb3..f1d29062da7 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -256,4 +256,28 @@ public void pick_emptySliceEndpoints_fallbackDisabled_returnsUnavailable() { assertThat(result.getStatus().getDescription()) .contains("No valid endpoints in slice and fallback disabled"); } + + @Test + public void pick_emptySliceEndpoints_fallbackEnabled_routesToFallbackPool() { + PickResult fallbackReadyResult = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP); + + // Gap slice with empty endpoints list + SliceEntry gapSlice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.emptyList()); + // Fallback pool has ep0 + SliceMap sliceMap = new SliceMap( + Collections.singletonList(gapSlice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), true, "x-key"); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("x-key", Metadata.ASCII_STRING_MARSHALLER), "anyKey"); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(fallbackReadyResult); + } } From 7dfdf6b390e51fcfbfdbb794ff362d6e75722d84 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Fri, 4 Sep 2026 08:26:22 +0530 Subject: [PATCH 06/28] autosharding: Return primitive int from SliceMap.lookup to eliminate boxing --- .../io/grpc/autosharding/AutoShardingPicker.java | 4 ++-- .../main/java/io/grpc/autosharding/SliceMap.java | 10 +++++----- .../java/io/grpc/autosharding/SliceMapTest.java | 16 ++++++++-------- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index b7ab4e71cea..c0e655b0bf2 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -88,9 +88,9 @@ private boolean isPoolInFallback(List indices) { @Override public PickResult pickSubchannel(PickSubchannelArgs args) { byte[] key = extractKeyBytes(args.getHeaders()); - Integer sliceIdx = sliceMap.lookup(key); + int sliceIdx = sliceMap.lookup(key); - if (sliceIdx == null) { + if (sliceIdx == -1) { if (fallbackEnabled) { return pickFromEndpointIndices(sliceMap.getFallbackPool(), args); } else { diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java index 5444c252556..d83cbdfa46f 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -49,12 +49,12 @@ static final class SliceEntry { /** * 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). + * Returns -1 if slices is empty (e.g. startup/fallback case where there are no assignments) + * or if the key is smaller than the first slice's startKey. */ - @Nullable - Integer lookup(@Nullable byte[] key) { + int lookup(@Nullable byte[] key) { if (slices.isEmpty()) { - return null; + return -1; } byte[] searchKey = key != null ? key : EMPTY_BYTES; int low = 0; @@ -75,7 +75,7 @@ Integer lookup(@Nullable byte[] key) { if (low == 0) { // Key is smaller than first slice's startKey - return null; + return -1; } return low - 1; } diff --git a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java index 2ec753cf4e9..61cef33ae64 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java @@ -33,11 +33,11 @@ public class SliceMapTest { @Test - public void lookup_emptySlices_returnsNull() { + public void lookup_emptySlices_returnsInvalidIndex() { SliceMap sliceMap = new SliceMap(Collections.emptyList(), Arrays.asList(0, 1), 1L); - assertThat(sliceMap.lookup(new byte[] {1, 2, 3})).isNull(); - assertThat(sliceMap.lookup(null)).isNull(); - assertThat(sliceMap.lookup(new byte[0])).isNull(); + assertThat(sliceMap.lookup(new byte[] {1, 2, 3})).isEqualTo(-1); + assertThat(sliceMap.lookup(null)).isEqualTo(-1); + assertThat(sliceMap.lookup(new byte[0])).isEqualTo(-1); } @Test @@ -79,16 +79,16 @@ public void lookup_multipleSlices() { } @Test - public void lookup_keySmallerThanFirstSlice_returnsNull() { + public void lookup_keySmallerThanFirstSlice_returnsInvalidIndex() { // Slice starts at "m" SliceEntry s1 = new SliceEntry( "m".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); SliceMap sliceMap = new SliceMap( Collections.singletonList(s1), Collections.singletonList(0), 1L); - assertThat(sliceMap.lookup("a".getBytes(StandardCharsets.UTF_8))).isNull(); - assertThat(sliceMap.lookup("".getBytes(StandardCharsets.UTF_8))).isNull(); - assertThat(sliceMap.lookup(null)).isNull(); + assertThat(sliceMap.lookup("a".getBytes(StandardCharsets.UTF_8))).isEqualTo(-1); + assertThat(sliceMap.lookup("".getBytes(StandardCharsets.UTF_8))).isEqualTo(-1); + assertThat(sliceMap.lookup(null)).isEqualTo(-1); assertThat(sliceMap.lookup("m".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); assertThat(sliceMap.lookup("z".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); } From 67abc201235a0c031e40fbc5e27f6202325f16db Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Fri, 4 Sep 2026 08:33:38 +0530 Subject: [PATCH 07/28] use UnsignedBytes.lexicographicalComparator() --- .../java/io/grpc/autosharding/SliceMap.java | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java index d83cbdfa46f..e895b6ba010 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -16,8 +16,10 @@ package io.grpc.autosharding; +import com.google.common.primitives.UnsignedBytes; import java.util.ArrayList; import java.util.Collections; +import java.util.Comparator; import java.util.List; import javax.annotation.Nullable; @@ -33,6 +35,8 @@ static final class SliceEntry { } } + private static final Comparator UNSIGNED_BYTES_COMPARATOR = + UnsignedBytes.lexicographicalComparator(); private static final byte[] EMPTY_BYTES = new byte[0]; private final List slices; @@ -41,7 +45,7 @@ static final class SliceEntry { SliceMap(List slices, List fallbackPool, long generation) { List sortedSlices = new ArrayList<>(slices); - sortedSlices.sort((e1, e2) -> compareUnsigned(e1.startKey, e2.startKey)); + sortedSlices.sort((e1, e2) -> UNSIGNED_BYTES_COMPARATOR.compare(e1.startKey, e2.startKey)); this.slices = Collections.unmodifiableList(sortedSlices); this.fallbackPool = Collections.unmodifiableList(new ArrayList<>(fallbackPool)); this.generation = generation; @@ -62,7 +66,7 @@ int lookup(@Nullable byte[] key) { while (low <= high) { int mid = (low + high) >>> 1; - int cmp = compareUnsigned(slices.get(mid).startKey, searchKey); + int cmp = UNSIGNED_BYTES_COMPARATOR.compare(slices.get(mid).startKey, searchKey); if (cmp < 0) { low = mid + 1; @@ -80,17 +84,6 @@ int lookup(@Nullable byte[] key) { 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; } From 789de22f425bcaf387dee7bb8ba1a8a655e2a47f Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Fri, 4 Sep 2026 08:43:33 +0530 Subject: [PATCH 08/28] use ImmutableList for endpoints --- .../grpc/autosharding/AutoShardingPicker.java | 7 +++--- .../java/io/grpc/autosharding/SliceMap.java | 24 ++++++++++--------- .../io/grpc/autosharding/SliceMapTest.java | 12 ++++------ 3 files changed, 21 insertions(+), 22 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index c0e655b0bf2..a35b21818bb 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -16,6 +16,7 @@ package io.grpc.autosharding; +import com.google.common.collect.ImmutableList; import io.grpc.ConnectivityState; import io.grpc.InternalMetadata; import io.grpc.LoadBalancer.PickResult; @@ -23,8 +24,6 @@ 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; @@ -45,7 +44,7 @@ public byte[] parseAsciiString(byte[] serialized) { }; private final SliceMap sliceMap; - private final List endpoints; + private final ImmutableList endpoints; private final boolean[] sliceInFallback; private final boolean fallbackEnabled; private final Metadata.Key keyHeader; @@ -56,7 +55,7 @@ public byte[] parseAsciiString(byte[] serialized) { boolean fallbackEnabled, String keyHeaderName) { this.sliceMap = sliceMap; - this.endpoints = Collections.unmodifiableList(new ArrayList<>(endpoints)); + this.endpoints = ImmutableList.copyOf(endpoints); this.fallbackEnabled = fallbackEnabled; if (keyHeaderName == null || keyHeaderName.isEmpty()) { diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java index e895b6ba010..f53ef80e659 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -16,9 +16,11 @@ package io.grpc.autosharding; +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.common.collect.ImmutableList; import com.google.common.primitives.UnsignedBytes; import java.util.ArrayList; -import java.util.Collections; import java.util.Comparator; import java.util.List; import javax.annotation.Nullable; @@ -27,11 +29,11 @@ final class SliceMap { static final class SliceEntry { final byte[] startKey; - final List endpoints; + final ImmutableList endpoints; SliceEntry(byte[] startKey, List endpoints) { - this.startKey = startKey; - this.endpoints = Collections.unmodifiableList(new ArrayList<>(endpoints)); + this.startKey = checkNotNull(startKey, "startKey"); + this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints")); } } @@ -39,15 +41,15 @@ static final class SliceEntry { UnsignedBytes.lexicographicalComparator(); private static final byte[] EMPTY_BYTES = new byte[0]; - private final List slices; - private final List fallbackPool; + private final ImmutableList slices; + private final ImmutableList fallbackPool; private final long generation; SliceMap(List slices, List fallbackPool, long generation) { - List sortedSlices = new ArrayList<>(slices); + List sortedSlices = new ArrayList<>(checkNotNull(slices, "slices")); sortedSlices.sort((e1, e2) -> UNSIGNED_BYTES_COMPARATOR.compare(e1.startKey, e2.startKey)); - this.slices = Collections.unmodifiableList(sortedSlices); - this.fallbackPool = Collections.unmodifiableList(new ArrayList<>(fallbackPool)); + this.slices = ImmutableList.copyOf(sortedSlices); + this.fallbackPool = ImmutableList.copyOf(checkNotNull(fallbackPool, "fallbackPool")); this.generation = generation; } @@ -84,11 +86,11 @@ int lookup(@Nullable byte[] key) { return low - 1; } - List getSlices() { + ImmutableList getSlices() { return slices; } - List getFallbackPool() { + ImmutableList getFallbackPool() { return fallbackPool; } diff --git a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java index 61cef33ae64..4bcc4a32582 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java @@ -17,7 +17,6 @@ package io.grpc.autosharding; import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.assertThrows; import io.grpc.autosharding.SliceMap.SliceEntry; import java.nio.charset.StandardCharsets; @@ -129,11 +128,10 @@ public void gettersAndImmutability() { assertThat(sliceMap.getSlices()).hasSize(1); assertThat(sliceMap.getSlices().get(0).endpoints).containsExactly(0, 1).inOrder(); - // Verify immutability - assertThrows(UnsupportedOperationException.class, () -> sliceMap.getSlices().clear()); - assertThrows(UnsupportedOperationException.class, () -> sliceMap.getFallbackPool().clear()); - assertThrows( - UnsupportedOperationException.class, - () -> sliceMap.getSlices().get(0).endpoints.clear()); + // Verify defensive copying: mutating input collections does not affect sliceMap + slices.clear(); + fallback.clear(); + assertThat(sliceMap.getSlices()).hasSize(1); + assertThat(sliceMap.getFallbackPool()).containsExactly(0, 1).inOrder(); } } From 426621f4a17f6e88367a78f1ccb7c71853cbbc22 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Fri, 4 Sep 2026 09:15:32 +0530 Subject: [PATCH 09/28] add javadoc --- .../java/io/grpc/autosharding/SliceMap.java | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java index f53ef80e659..7df8164d6ab 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -25,8 +25,21 @@ import java.util.List; import javax.annotation.Nullable; +/** + * An immutable lookup structure mapping application routing keys to slice indices. + * + *

As defined in gRFC A119, the assignment provider guarantees that the assignment + * is pre-validated, gap-free, non-overlapping, and covers the entire keyspace {@code ["" .. inf)}. + * Any gaps returned by the autosharding server are filled as slice entries with an empty + * endpoints list. Therefore, each {@link SliceEntry} only needs to store {@code startKey} + * because the exclusive end key of slice {@code i} is implicitly the inclusive start key of + * slice {@code i + 1}. + */ final class SliceMap { + /** + * Represents a single key-range slice mapping to endpoint indices in the picker. + */ static final class SliceEntry { final byte[] startKey; final ImmutableList endpoints; @@ -45,6 +58,13 @@ static final class SliceEntry { private final ImmutableList fallbackPool; private final long generation; + /** + * Constructs an immutable {@link SliceMap}. + * + * @param slices the pre-validated list of key-range slice entries + * @param fallbackPool the list of all available endpoint indices for fallback routing + * @param generation the snapshot generation number from the assignment + */ SliceMap(List slices, List fallbackPool, long generation) { List sortedSlices = new ArrayList<>(checkNotNull(slices, "slices")); sortedSlices.sort((e1, e2) -> UNSIGNED_BYTES_COMPARATOR.compare(e1.startKey, e2.startKey)); From 1b63d200a4304584917066a0a6c01a44314ee98d Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Fri, 4 Sep 2026 10:14:12 +0530 Subject: [PATCH 10/28] improvements --- .../grpc/autosharding/AutoShardingPicker.java | 34 ++++++++----- .../io/grpc/autosharding/PickerEndpoint.java | 49 ++++++++++++++++--- .../java/io/grpc/autosharding/SliceMap.java | 25 ++++++++-- .../autosharding/AutoShardingPickerTest.java | 15 ++++++ .../io/grpc/autosharding/SliceMapTest.java | 3 +- 5 files changed, 102 insertions(+), 24 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index a35b21818bb..6d480a85833 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -27,6 +27,12 @@ import java.util.List; import java.util.concurrent.ThreadLocalRandom; +/** + * Subchannel picker for the auto-sharding load balancing policy. + * + *

Routes RPCs to backend endpoints based on a request metadata header key, matching against + * an immutable {@link SliceMap}. + */ final class AutoShardingPicker extends SubchannelPicker { private static final byte[] EMPTY_BYTES = new byte[0]; @@ -49,6 +55,14 @@ public byte[] parseAsciiString(byte[] serialized) { private final boolean fallbackEnabled; private final Metadata.Key keyHeader; + /** + * Constructs an {@link AutoShardingPicker}. + * + * @param sliceMap the pre-built, immutable mapping from key ranges to endpoint indices + * @param endpoints the list of endpoint snapshots corresponding 1:1 to endpoint indices + * @param fallbackEnabled whether fallback routing to all resolved endpoints is enabled + * @param keyHeaderName the metadata header name used to extract the routing key + */ AutoShardingPicker( SliceMap sliceMap, List endpoints, @@ -68,7 +82,7 @@ public byte[] parseAsciiString(byte[] serialized) { this.sliceInFallback = new boolean[sliceMap.getSlices().size()]; for (int i = 0; i < sliceInFallback.length; i++) { - this.sliceInFallback[i] = isPoolInFallback(sliceMap.getSlices().get(i).endpoints); + this.sliceInFallback[i] = isPoolInFallback(sliceMap.getSlices().get(i).getEndpoints()); } } @@ -77,7 +91,7 @@ private boolean isPoolInFallback(List indices) { return true; } for (int idx : indices) { - if (endpoints.get(idx).state != ConnectivityState.TRANSIENT_FAILURE) { + if (endpoints.get(idx).getState() != ConnectivityState.TRANSIENT_FAILURE) { return false; } } @@ -104,7 +118,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) { } SliceMap.SliceEntry sliceEntry = sliceMap.getSlices().get(sliceIdx); - return pickFromEndpointIndices(sliceEntry.endpoints, args); + return pickFromEndpointIndices(sliceEntry.getEndpoints(), args); } private PickResult pickFromEndpointIndices( @@ -123,16 +137,14 @@ private PickResult pickFromEndpointIndices( int epIdx = indices.get((firstIndex + i) % size); PickerEndpoint endpoint = endpoints.get(epIdx); - if (endpoint.state == ConnectivityState.READY) { - return endpoint.picker.pickSubchannel(args); + if (endpoint.getState() == ConnectivityState.READY) { + return endpoint.getPicker().pickSubchannel(args); } - if (endpoint.state == ConnectivityState.CONNECTING) { + if (endpoint.getState() == ConnectivityState.CONNECTING) { foundConnecting = true; - } else if (!requestedConnection && endpoint.state == ConnectivityState.IDLE) { - if (endpoint.requestConnection != null) { - endpoint.requestConnection.run(); - } + } else if (!requestedConnection && endpoint.getState() == ConnectivityState.IDLE) { + endpoint.requestConnection(); requestedConnection = true; } } @@ -142,7 +154,7 @@ private PickResult pickFromEndpointIndices( } int firstEpIdx = indices.get(firstIndex); - return endpoints.get(firstEpIdx).picker.pickSubchannel(args); + return endpoints.get(firstEpIdx).getPicker().pickSubchannel(args); } private byte[] extractKeyBytes(Metadata headers) { diff --git a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java index 54310a1bdb6..4fc9d98a850 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java +++ b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java @@ -16,21 +16,56 @@ package io.grpc.autosharding; +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.common.base.MoreObjects; import io.grpc.ConnectivityState; import io.grpc.LoadBalancer.SubchannelPicker; +import javax.annotation.Nullable; /** - * Immutable snapshot of endpoint state used by the AutoShardingPicker. + * Immutable snapshot of endpoint state used by {@link AutoShardingPicker}. */ final class PickerEndpoint { - final ConnectivityState state; - final SubchannelPicker picker; - final Runnable requestConnection; + private final ConnectivityState state; + private final SubchannelPicker picker; + @Nullable private final Runnable requestConnection; + /** + * Constructs a {@link PickerEndpoint}. + * + * @param state the current connectivity state of the endpoint + * @param picker the latest subchannel picker for the endpoint + * @param requestConnection a callback to trigger a connection attempt on the child balancer + */ PickerEndpoint( - ConnectivityState state, SubchannelPicker picker, Runnable requestConnection) { - this.state = state; - this.picker = picker; + ConnectivityState state, + SubchannelPicker picker, + @Nullable Runnable requestConnection) { + this.state = checkNotNull(state, "state"); + this.picker = checkNotNull(picker, "picker"); this.requestConnection = requestConnection; } + + ConnectivityState getState() { + return state; + } + + SubchannelPicker getPicker() { + return picker; + } + + void requestConnection() { + if (requestConnection != null) { + requestConnection.run(); + } + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("state", state) + .add("picker", picker) + .toString(); + } } diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java index 7df8164d6ab..269c1b0ca7e 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -41,13 +41,27 @@ final class SliceMap { * Represents a single key-range slice mapping to endpoint indices in the picker. */ static final class SliceEntry { - final byte[] startKey; - final ImmutableList endpoints; - + private final byte[] startKey; + private final ImmutableList endpoints; + + /** + * Constructs a {@link SliceEntry}. + * + * @param startKey the inclusive start key of the slice + * @param endpoints the list of endpoint indices assigned to this slice + */ SliceEntry(byte[] startKey, List endpoints) { this.startKey = checkNotNull(startKey, "startKey"); this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints")); } + + byte[] getStartKey() { + return startKey; + } + + ImmutableList getEndpoints() { + return endpoints; + } } private static final Comparator UNSIGNED_BYTES_COMPARATOR = @@ -67,7 +81,8 @@ static final class SliceEntry { */ SliceMap(List slices, List fallbackPool, long generation) { List sortedSlices = new ArrayList<>(checkNotNull(slices, "slices")); - sortedSlices.sort((e1, e2) -> UNSIGNED_BYTES_COMPARATOR.compare(e1.startKey, e2.startKey)); + sortedSlices.sort( + (e1, e2) -> UNSIGNED_BYTES_COMPARATOR.compare(e1.getStartKey(), e2.getStartKey())); this.slices = ImmutableList.copyOf(sortedSlices); this.fallbackPool = ImmutableList.copyOf(checkNotNull(fallbackPool, "fallbackPool")); this.generation = generation; @@ -88,7 +103,7 @@ int lookup(@Nullable byte[] key) { while (low <= high) { int mid = (low + high) >>> 1; - int cmp = UNSIGNED_BYTES_COMPARATOR.compare(slices.get(mid).startKey, searchKey); + int cmp = UNSIGNED_BYTES_COMPARATOR.compare(slices.get(mid).getStartKey(), searchKey); if (cmp < 0) { low = mid + 1; diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java index f1d29062da7..30beff751ca 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -280,4 +280,19 @@ public void pick_emptySliceEndpoints_fallbackEnabled_routesToFallbackPool() { PickResult result = picker.pickSubchannel(createArgs(headers)); assertThat(result).isSameInstanceAs(fallbackReadyResult); } + + @Test + public void pickerEndpoint_gettersAndToString() { + FakePicker fakePicker = new FakePicker(PickResult.withNoResult()); + AtomicInteger count = new AtomicInteger(); + PickerEndpoint ep = new PickerEndpoint( + ConnectivityState.IDLE, fakePicker, count::incrementAndGet); + + assertThat(ep.getState()).isEqualTo(ConnectivityState.IDLE); + assertThat(ep.getPicker()).isSameInstanceAs(fakePicker); + assertThat(ep.toString()).contains("state=IDLE"); + + ep.requestConnection(); + assertThat(count.get()).isEqualTo(1); + } } diff --git a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java index 4bcc4a32582..2ad3d13998c 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java @@ -126,7 +126,8 @@ public void gettersAndImmutability() { assertThat(sliceMap.getGeneration()).isEqualTo(42L); assertThat(sliceMap.getFallbackPool()).containsExactly(0, 1).inOrder(); assertThat(sliceMap.getSlices()).hasSize(1); - assertThat(sliceMap.getSlices().get(0).endpoints).containsExactly(0, 1).inOrder(); + assertThat(sliceMap.getSlices().get(0).getStartKey()).isEqualTo(new byte[] {1}); + assertThat(sliceMap.getSlices().get(0).getEndpoints()).containsExactly(0, 1).inOrder(); // Verify defensive copying: mutating input collections does not affect sliceMap slices.clear(); From 07714334cc6c795b2094979c5d53b28b5bde0ee5 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Fri, 4 Sep 2026 11:10:58 +0530 Subject: [PATCH 11/28] have exitIdler functional interface --- .../io/grpc/autosharding/PickerEndpoint.java | 27 ++++++++++++++----- .../autosharding/AutoShardingPickerTest.java | 25 ++++++++--------- 2 files changed, 34 insertions(+), 18 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java index 4fc9d98a850..d85189d222e 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java +++ b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java @@ -27,24 +27,39 @@ * Immutable snapshot of endpoint state used by {@link AutoShardingPicker}. */ final class PickerEndpoint { + + /** + * Callback interface to trigger connection attempts on an IDLE endpoint's child balancer. + */ + @FunctionalInterface + interface ExitIdler { + /** + * Requests the child load balancer to exit IDLE and initiate a connection. + * + *

Implementations MUST be thread-safe, non-blocking, idempotent, and dispatch + * execution to the {@link io.grpc.SynchronizationContext}. + */ + void exitIdle(); + } + private final ConnectivityState state; private final SubchannelPicker picker; - @Nullable private final Runnable requestConnection; + @Nullable private final ExitIdler exitIdler; /** * Constructs a {@link PickerEndpoint}. * * @param state the current connectivity state of the endpoint * @param picker the latest subchannel picker for the endpoint - * @param requestConnection a callback to trigger a connection attempt on the child balancer + * @param exitIdler a callback to trigger an IDLE child balancer to start connecting */ PickerEndpoint( ConnectivityState state, SubchannelPicker picker, - @Nullable Runnable requestConnection) { + @Nullable ExitIdler exitIdler) { this.state = checkNotNull(state, "state"); this.picker = checkNotNull(picker, "picker"); - this.requestConnection = requestConnection; + this.exitIdler = exitIdler; } ConnectivityState getState() { @@ -56,8 +71,8 @@ SubchannelPicker getPicker() { } void requestConnection() { - if (requestConnection != null) { - requestConnection.run(); + if (exitIdler != null) { + exitIdler.exitIdle(); } } diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java index 30beff751ca..89f314241bd 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -27,6 +27,7 @@ import io.grpc.Metadata; import io.grpc.MethodDescriptor; import io.grpc.Status; +import io.grpc.autosharding.PickerEndpoint.ExitIdler; import io.grpc.autosharding.SliceMap.SliceEntry; import io.grpc.internal.PickSubchannelArgsImpl; import io.grpc.testing.TestMethodDescriptors; @@ -42,9 +43,9 @@ public class AutoShardingPickerTest { private static final MethodDescriptor METHOD = TestMethodDescriptors.voidMethod(); - private static final Runnable NOOP = new Runnable() { + private static final ExitIdler NOOP_EXIT_IDLER = new ExitIdler() { @Override - public void run() {} + public void exitIdle() {} }; private static final PickDetailsConsumer NOOP_CONSUMER = new PickDetailsConsumer() {}; @@ -69,7 +70,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) { public void pick_noSliceMap_fallbackEnabled_picksFromFallbackPool() { PickResult readyResult = PickResult.withNoResult(); // using as token PickerEndpoint ep0 = new PickerEndpoint( - ConnectivityState.READY, new FakePicker(readyResult), NOOP); + ConnectivityState.READY, new FakePicker(readyResult), NOOP_EXIT_IDLER); SliceMap emptySliceMap = new SliceMap( Collections.emptyList(), Collections.singletonList(0), 1L); @@ -87,7 +88,7 @@ public void pick_noSliceMap_fallbackEnabled_picksFromFallbackPool() { @Test public void pick_noSliceMap_fallbackDisabled_returnsUnavailableError() { PickerEndpoint ep0 = new PickerEndpoint( - ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP); + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP_EXIT_IDLER); SliceMap emptySliceMap = new SliceMap( Collections.emptyList(), Collections.singletonList(0), 1L); @@ -106,7 +107,7 @@ public void pick_noSliceMap_fallbackDisabled_returnsUnavailableError() { public void pick_sliceFound_readyEndpoint_returnsPickResult() { PickResult expectedResult = PickResult.withNoResult(); PickerEndpoint ep0 = new PickerEndpoint( - ConnectivityState.READY, new FakePicker(expectedResult), NOOP); + ConnectivityState.READY, new FakePicker(expectedResult), NOOP_EXIT_IDLER); SliceEntry slice = new SliceEntry( "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); @@ -174,9 +175,9 @@ public void pick_sliceFound_allTransientFailure_fallbackEnabled_picksFromFallbac PickerEndpoint ep0 = new PickerEndpoint( ConnectivityState.TRANSIENT_FAILURE, new FakePicker(PickResult.withError(Status.UNAVAILABLE.withDescription("ep0 down"))), - NOOP); + NOOP_EXIT_IDLER); PickerEndpoint ep1 = new PickerEndpoint( - ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP); + ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP_EXIT_IDLER); // Slice 0 only has ep0 (which is down) SliceEntry slice0 = new SliceEntry( @@ -198,7 +199,7 @@ public void pick_sliceFound_allTransientFailure_fallbackDisabled_delegatesToEndp PickerEndpoint ep0 = new PickerEndpoint( ConnectivityState.TRANSIENT_FAILURE, new FakePicker(PickResult.withError(epError)), - NOOP); + NOOP_EXIT_IDLER); SliceEntry slice0 = new SliceEntry( "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); @@ -218,9 +219,9 @@ public void pick_binaryHeader_extractedProperly() { PickResult ready1 = PickResult.withNoResult(); PickerEndpoint ep0 = new PickerEndpoint( - ConnectivityState.READY, new FakePicker(ready0), NOOP); + ConnectivityState.READY, new FakePicker(ready0), NOOP_EXIT_IDLER); PickerEndpoint ep1 = new PickerEndpoint( - ConnectivityState.READY, new FakePicker(ready1), NOOP); + ConnectivityState.READY, new FakePicker(ready1), NOOP_EXIT_IDLER); SliceEntry s0 = new SliceEntry(new byte[] {0x00}, Collections.singletonList(0)); SliceEntry s1 = new SliceEntry(new byte[] {0x50}, Collections.singletonList(1)); @@ -241,7 +242,7 @@ public void pick_binaryHeader_extractedProperly() { @Test public void pick_emptySliceEndpoints_fallbackDisabled_returnsUnavailable() { PickerEndpoint ep0 = new PickerEndpoint( - ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP); + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP_EXIT_IDLER); SliceEntry emptySlice = new SliceEntry( "".getBytes(StandardCharsets.UTF_8), Collections.emptyList()); @@ -261,7 +262,7 @@ public void pick_emptySliceEndpoints_fallbackDisabled_returnsUnavailable() { public void pick_emptySliceEndpoints_fallbackEnabled_routesToFallbackPool() { PickResult fallbackReadyResult = PickResult.withNoResult(); PickerEndpoint ep0 = new PickerEndpoint( - ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP); + ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP_EXIT_IDLER); // Gap slice with empty endpoints list SliceEntry gapSlice = new SliceEntry( From 1db80174418bc060db48705fda021de6c050cd25 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Fri, 4 Sep 2026 11:34:32 +0530 Subject: [PATCH 12/28] create Metadata.Key statically --- .../grpc/autosharding/AutoShardingPicker.java | 39 +++++++---- .../autosharding/AutoShardingPickerTest.java | 70 ++++++++++++++++--- 2 files changed, 86 insertions(+), 23 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index 6d480a85833..9f27de5c5a8 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -16,6 +16,8 @@ package io.grpc.autosharding; +import static com.google.common.base.Preconditions.checkNotNull; + import com.google.common.collect.ImmutableList; import io.grpc.ConnectivityState; import io.grpc.InternalMetadata; @@ -26,6 +28,7 @@ import io.grpc.Status; import java.util.List; import java.util.concurrent.ThreadLocalRandom; +import javax.annotation.Nullable; /** * Subchannel picker for the auto-sharding load balancing policy. @@ -53,7 +56,24 @@ public byte[] parseAsciiString(byte[] serialized) { private final ImmutableList endpoints; private final boolean[] sliceInFallback; private final boolean fallbackEnabled; - private final Metadata.Key keyHeader; + @Nullable private final Metadata.Key keyHeader; + + /** + * Pre-creates a {@link Metadata.Key} for the given key header name. + * + * @param keyHeaderName the metadata header name, or {@code null}/empty if no header routing + * @return the pre-computed {@link Metadata.Key}, or {@code null} if keyHeaderName is null/empty + */ + @Nullable + static Metadata.Key createKeyHeader(@Nullable String keyHeaderName) { + if (keyHeaderName == null || keyHeaderName.isEmpty()) { + return null; + } else if (keyHeaderName.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { + return Metadata.Key.of(keyHeaderName, Metadata.BINARY_BYTE_MARSHALLER); + } else { + return InternalMetadata.keyOf(keyHeaderName, RAW_ASCII_MARSHALLER); + } + } /** * Constructs an {@link AutoShardingPicker}. @@ -61,24 +81,17 @@ public byte[] parseAsciiString(byte[] serialized) { * @param sliceMap the pre-built, immutable mapping from key ranges to endpoint indices * @param endpoints the list of endpoint snapshots corresponding 1:1 to endpoint indices * @param fallbackEnabled whether fallback routing to all resolved endpoints is enabled - * @param keyHeaderName the metadata header name used to extract the routing key + * @param keyHeader the pre-parsed metadata header key used to extract the routing key */ AutoShardingPicker( SliceMap sliceMap, List endpoints, boolean fallbackEnabled, - String keyHeaderName) { - this.sliceMap = sliceMap; - this.endpoints = ImmutableList.copyOf(endpoints); + @Nullable Metadata.Key keyHeader) { + this.sliceMap = checkNotNull(sliceMap, "sliceMap"); + this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints")); this.fallbackEnabled = fallbackEnabled; - - if (keyHeaderName == null || keyHeaderName.isEmpty()) { - this.keyHeader = null; - } else if (keyHeaderName.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { - this.keyHeader = Metadata.Key.of(keyHeaderName, Metadata.BINARY_BYTE_MARSHALLER); - } else { - this.keyHeader = InternalMetadata.keyOf(keyHeaderName, RAW_ASCII_MARSHALLER); - } + this.keyHeader = keyHeader; this.sliceInFallback = new boolean[sliceMap.getSlices().size()]; for (int i = 0; i < sliceInFallback.length; i++) { diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java index 89f314241bd..e17e81807b7 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -75,7 +75,10 @@ public void pick_noSliceMap_fallbackEnabled_picksFromFallbackPool() { SliceMap emptySliceMap = new SliceMap( Collections.emptyList(), Collections.singletonList(0), 1L); AutoShardingPicker picker = new AutoShardingPicker( - emptySliceMap, Collections.singletonList(ep0), true, "x-slice-key"); + emptySliceMap, + Collections.singletonList(ep0), + true, + AutoShardingPicker.createKeyHeader("x-slice-key")); Metadata headers = new Metadata(); headers.put( @@ -93,7 +96,10 @@ public void pick_noSliceMap_fallbackDisabled_returnsUnavailableError() { SliceMap emptySliceMap = new SliceMap( Collections.emptyList(), Collections.singletonList(0), 1L); AutoShardingPicker picker = new AutoShardingPicker( - emptySliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + emptySliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); Metadata headers = new Metadata(); PickResult result = picker.pickSubchannel(createArgs(headers)); @@ -115,7 +121,10 @@ public void pick_sliceFound_readyEndpoint_returnsPickResult() { Collections.singletonList(slice), Collections.singletonList(0), 1L); AutoShardingPicker picker = new AutoShardingPicker( - sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); Metadata headers = new Metadata(); headers.put( @@ -139,7 +148,10 @@ public void pick_sliceFound_idleEndpoint_triggersConnectionAndQueues() { Collections.singletonList(slice), Collections.singletonList(0), 1L); AutoShardingPicker picker = new AutoShardingPicker( - sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); PickResult result = picker.pickSubchannel(createArgs(new Metadata())); @@ -161,7 +173,10 @@ public void pick_sliceFound_connectingEndpoint_queuesPick() { Collections.singletonList(slice), Collections.singletonList(0), 1L); AutoShardingPicker picker = new AutoShardingPicker( - sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); PickResult result = picker.pickSubchannel(createArgs(new Metadata())); @@ -187,7 +202,10 @@ public void pick_sliceFound_allTransientFailure_fallbackEnabled_picksFromFallbac Collections.singletonList(slice0), Collections.singletonList(1), 1L); AutoShardingPicker picker = new AutoShardingPicker( - sliceMap, Arrays.asList(ep0, ep1), true, "x-slice-key"); + sliceMap, + Arrays.asList(ep0, ep1), + true, + AutoShardingPicker.createKeyHeader("x-slice-key")); PickResult result = picker.pickSubchannel(createArgs(new Metadata())); assertThat(result).isSameInstanceAs(fallbackReadyResult); @@ -207,7 +225,10 @@ public void pick_sliceFound_allTransientFailure_fallbackDisabled_delegatesToEndp Collections.singletonList(slice0), Collections.singletonList(0), 1L); AutoShardingPicker picker = new AutoShardingPicker( - sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); PickResult result = picker.pickSubchannel(createArgs(new Metadata())); assertThat(result.getStatus()).isEqualTo(epError); @@ -228,7 +249,10 @@ public void pick_binaryHeader_extractedProperly() { SliceMap sliceMap = new SliceMap(Arrays.asList(s0, s1), Arrays.asList(0, 1), 1L); AutoShardingPicker picker = new AutoShardingPicker( - sliceMap, Arrays.asList(ep0, ep1), false, "slice-key-bin"); + sliceMap, + Arrays.asList(ep0, ep1), + false, + AutoShardingPicker.createKeyHeader("slice-key-bin")); Metadata headers = new Metadata(); headers.put( @@ -250,7 +274,10 @@ public void pick_emptySliceEndpoints_fallbackDisabled_returnsUnavailable() { Collections.singletonList(emptySlice), Collections.singletonList(0), 1L); AutoShardingPicker picker = new AutoShardingPicker( - sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); PickResult result = picker.pickSubchannel(createArgs(new Metadata())); assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); @@ -272,7 +299,10 @@ public void pick_emptySliceEndpoints_fallbackEnabled_routesToFallbackPool() { Collections.singletonList(gapSlice), Collections.singletonList(0), 1L); AutoShardingPicker picker = new AutoShardingPicker( - sliceMap, Collections.singletonList(ep0), true, "x-key"); + sliceMap, + Collections.singletonList(ep0), + true, + AutoShardingPicker.createKeyHeader("x-key")); Metadata headers = new Metadata(); headers.put( @@ -296,4 +326,24 @@ public void pickerEndpoint_gettersAndToString() { ep.requestConnection(); assertThat(count.get()).isEqualTo(1); } + + @Test + public void createKeyHeader_nullOrEmpty_returnsNull() { + assertThat(AutoShardingPicker.createKeyHeader(null)).isNull(); + assertThat(AutoShardingPicker.createKeyHeader("")).isNull(); + } + + @Test + public void createKeyHeader_asciiHeader() { + Metadata.Key key = AutoShardingPicker.createKeyHeader("x-slice-key"); + assertThat(key).isNotNull(); + assertThat(key.name()).isEqualTo("x-slice-key"); + } + + @Test + public void createKeyHeader_binaryHeader() { + Metadata.Key key = AutoShardingPicker.createKeyHeader("x-slice-key-bin"); + assertThat(key).isNotNull(); + assertThat(key.name()).isEqualTo("x-slice-key-bin"); + } } From 29957b6f044db0023f06c402465557bd04585037 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Fri, 4 Sep 2026 11:46:38 +0530 Subject: [PATCH 13/28] add FunctionalInterface ThreadSafeRandom --- .../grpc/autosharding/AutoShardingPicker.java | 25 ++++++++++++- .../autosharding/AutoShardingPickerTest.java | 35 +++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index 9f27de5c5a8..a7c43852b41 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -18,6 +18,7 @@ import static com.google.common.base.Preconditions.checkNotNull; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import io.grpc.ConnectivityState; import io.grpc.InternalMetadata; @@ -29,6 +30,7 @@ import java.util.List; import java.util.concurrent.ThreadLocalRandom; import javax.annotation.Nullable; +import javax.annotation.concurrent.ThreadSafe; /** * Subchannel picker for the auto-sharding load balancing policy. @@ -39,6 +41,15 @@ final class AutoShardingPicker extends SubchannelPicker { private static final byte[] EMPTY_BYTES = new byte[0]; + @ThreadSafe + @FunctionalInterface + interface ThreadSafeRandom { + int nextInt(int bound); + } + + private static final ThreadSafeRandom DEFAULT_RANDOM = + bound -> ThreadLocalRandom.current().nextInt(bound); + private static final InternalMetadata.TrustedAsciiMarshaller RAW_ASCII_MARSHALLER = new InternalMetadata.TrustedAsciiMarshaller() { @Override @@ -57,6 +68,7 @@ public byte[] parseAsciiString(byte[] serialized) { private final boolean[] sliceInFallback; private final boolean fallbackEnabled; @Nullable private final Metadata.Key keyHeader; + private final ThreadSafeRandom random; /** * Pre-creates a {@link Metadata.Key} for the given key header name. @@ -88,10 +100,21 @@ static Metadata.Key createKeyHeader(@Nullable String keyHeaderName) { List endpoints, boolean fallbackEnabled, @Nullable Metadata.Key keyHeader) { + this(sliceMap, endpoints, fallbackEnabled, keyHeader, DEFAULT_RANDOM); + } + + @VisibleForTesting + AutoShardingPicker( + SliceMap sliceMap, + List endpoints, + boolean fallbackEnabled, + @Nullable Metadata.Key keyHeader, + ThreadSafeRandom random) { this.sliceMap = checkNotNull(sliceMap, "sliceMap"); this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints")); this.fallbackEnabled = fallbackEnabled; this.keyHeader = keyHeader; + this.random = checkNotNull(random, "random"); this.sliceInFallback = new boolean[sliceMap.getSlices().size()]; for (int i = 0; i < sliceInFallback.length; i++) { @@ -142,7 +165,7 @@ private PickResult pickFromEndpointIndices( } int size = indices.size(); - int firstIndex = ThreadLocalRandom.current().nextInt(size); + int firstIndex = random.nextInt(size); boolean requestedConnection = false; boolean foundConnecting = false; diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java index e17e81807b7..7bced71cd94 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -346,4 +346,39 @@ public void createKeyHeader_binaryHeader() { assertThat(key).isNotNull(); assertThat(key.name()).isEqualTo("x-slice-key-bin"); } + + @Test + public void pick_deterministicRandom_selectsExpectedEndpoint() { + PickResult ready0 = PickResult.withNoResult(); + PickResult ready1 = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready0), NOOP_EXIT_IDLER); + PickerEndpoint ep1 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready1), NOOP_EXIT_IDLER); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Arrays.asList(0, 1)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Arrays.asList(0, 1), 1L); + + // Test picking index 0 + AutoShardingPicker picker0 = new AutoShardingPicker( + sliceMap, + Arrays.asList(ep0, ep1), + false, + AutoShardingPicker.createKeyHeader("x-key"), + bound -> 0); + PickResult result0 = picker0.pickSubchannel(createArgs(new Metadata())); + assertThat(result0).isSameInstanceAs(ready0); + + // Test picking index 1 + AutoShardingPicker picker1 = new AutoShardingPicker( + sliceMap, + Arrays.asList(ep0, ep1), + false, + AutoShardingPicker.createKeyHeader("x-key"), + bound -> 1); + PickResult result1 = picker1.pickSubchannel(createArgs(new Metadata())); + assertThat(result1).isSameInstanceAs(ready1); + } } From 4105fd822c863182f937b0741298e08134040742 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Fri, 4 Sep 2026 18:26:29 +0530 Subject: [PATCH 14/28] fast path --- .../grpc/autosharding/AutoShardingPicker.java | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index a7c43852b41..81eeb9b370c 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -116,9 +116,23 @@ static Metadata.Key createKeyHeader(@Nullable String keyHeaderName) { this.keyHeader = keyHeader; this.random = checkNotNull(random, "random"); + boolean hasTransientFailure = false; + for (int i = 0; i < this.endpoints.size(); i++) { + if (this.endpoints.get(i).getState() == ConnectivityState.TRANSIENT_FAILURE) { + hasTransientFailure = true; + break; + } + } + this.sliceInFallback = new boolean[sliceMap.getSlices().size()]; - for (int i = 0; i < sliceInFallback.length; i++) { - this.sliceInFallback[i] = isPoolInFallback(sliceMap.getSlices().get(i).getEndpoints()); + if (!hasTransientFailure) { + for (int i = 0; i < sliceInFallback.length; i++) { + this.sliceInFallback[i] = sliceMap.getSlices().get(i).getEndpoints().isEmpty(); + } + } else { + for (int i = 0; i < sliceInFallback.length; i++) { + this.sliceInFallback[i] = isPoolInFallback(sliceMap.getSlices().get(i).getEndpoints()); + } } } From be312603abc993480bfc3006149919b4f1707fe3 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Mon, 7 Sep 2026 12:28:10 +0530 Subject: [PATCH 15/28] javadoc --- .../io/grpc/autosharding/AutoShardingPicker.java | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index 81eeb9b370c..26d01ef086c 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -37,6 +37,21 @@ * *

Routes RPCs to backend endpoints based on a request metadata header key, matching against * an immutable {@link SliceMap}. + * + *

Lifecycle & Assignment Expectations: + * This picker is intended for active routing and post-timeout fallback handling: + *

    + *
  • Pending Initial Assignment: During startup before the first assignment is received + * (and while the initial assignment timer is running), the load balancer handles RPC + * queuing polymorphically using a buffering picker in + * {@link io.grpc.ConnectivityState#CONNECTING}. + *
  • Active Routing: Once a valid assignment is received from the control plane, + * this picker routes RPCs according to the populated {@link SliceMap}. + *
  • Post-Timeout Fallback: If the initial assignment timer expires without an + * assignment, this picker is used with an empty {@link SliceMap} to either route across + * the {@code fallbackPool} (if fallback is enabled) or fail RPCs with + * {@link io.grpc.Status#UNAVAILABLE} (if fallback is disabled). + *
*/ final class AutoShardingPicker extends SubchannelPicker { private static final byte[] EMPTY_BYTES = new byte[0]; From b99fa0ccf04b50aa0f9faaa0e43c5943ceb18b87 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Mon, 7 Sep 2026 12:28:46 +0530 Subject: [PATCH 16/28] javadoc --- .../java/io/grpc/autosharding/AutoShardingPicker.java | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index 26d01ef086c..24086df3b5d 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -38,16 +38,15 @@ *

Routes RPCs to backend endpoints based on a request metadata header key, matching against * an immutable {@link SliceMap}. * - *

Lifecycle & Assignment Expectations: - * This picker is intended for active routing and post-timeout fallback handling: + *

This picker is intended for active routing and post-timeout fallback handling: *

    - *
  • Pending Initial Assignment: During startup before the first assignment is received + *
  • Pending Initial Assignment: During startup before the first assignment is received * (and while the initial assignment timer is running), the load balancer handles RPC * queuing polymorphically using a buffering picker in * {@link io.grpc.ConnectivityState#CONNECTING}. - *
  • Active Routing: Once a valid assignment is received from the control plane, + *
  • Active Routing: Once a valid assignment is received from the control plane, * this picker routes RPCs according to the populated {@link SliceMap}. - *
  • Post-Timeout Fallback: If the initial assignment timer expires without an + *
  • Post-Timeout Fallback: If the initial assignment timer expires without an * assignment, this picker is used with an empty {@link SliceMap} to either route across * the {@code fallbackPool} (if fallback is enabled) or fail RPCs with * {@link io.grpc.Status#UNAVAILABLE} (if fallback is disabled). From 12765cffe16f0af0a2e2a28d3db3b37304651cb3 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Mon, 7 Sep 2026 13:01:14 +0530 Subject: [PATCH 17/28] javadoc and unit test --- .../java/io/grpc/autosharding/SliceMap.java | 37 ++++++++--- .../io/grpc/autosharding/SliceMapTest.java | 66 +++++++++++++++++++ 2 files changed, 94 insertions(+), 9 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java index 269c1b0ca7e..ec3af3f0af4 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -28,12 +28,29 @@ /** * An immutable lookup structure mapping application routing keys to slice indices. * - *

    As defined in gRFC A119, the assignment provider guarantees that the assignment - * is pre-validated, gap-free, non-overlapping, and covers the entire keyspace {@code ["" .. inf)}. - * Any gaps returned by the autosharding server are filled as slice entries with an empty - * endpoints list. Therefore, each {@link SliceEntry} only needs to store {@code startKey} - * because the exclusive end key of slice {@code i} is implicitly the inclusive start key of - * slice {@code i + 1}. + *

    The assignment provider guarantees that the assignment is pre-validated, gap-free, + * non-overlapping, and covers the entire keyspace {@code ["" .. inf)}. + *

      + *
    • The first slice's {@code startKey} is expected to be {@code new byte[0]} ({@code ""}).
      • + *
    • Unassigned key ranges (gaps) returned by the autosharding server are filled as slice + * entries with an empty {@code endpoints} list.
      • + *
    • Endpoint indices in {@link SliceEntry#getEndpoints()} and {@link #getFallbackPool()} + * are non-negative indices corresponding 1:1 to the endpoint snapshot list in + * {@link AutoShardingPicker}.
      • + *
    + * + *

    Behavior on Invalid or Edge-case Inputs: + *

      + *
    • Empty slices list: {@link #lookup(byte[])} returns {@code -1}, allowing + * {@link AutoShardingPicker} to fall back to the fallback pool or fail with UNAVAILABLE.
      • + *
    • Key smaller than first slice start key: {@link #lookup(byte[])} returns {@code -1} + * if the first slice's {@code startKey} is not {@code ""} and the key precedes it.
      • + *
    • Null key: Treated as an empty byte array ({@code new byte[0]}).
      • + *
    • Unsorted slices: The constructor automatically sorts slices lexicographically + * using unsigned byte comparison.
      • + *
    • Null constructor arguments: Throws {@link NullPointerException} if {@code slices}, + * {@code fallbackPool}, {@code startKey}, or {@code endpoints} is {@code null}.
      • + *
    */ final class SliceMap { @@ -89,9 +106,11 @@ ImmutableList getEndpoints() { } /** - * Looks up the matching slice index for the given key. - * Returns -1 if slices is empty (e.g. startup/fallback case where there are no assignments) - * or if the key is smaller than the first slice's startKey. + * Looks up the matching slice index for the given key using binary search. + * + * @param key the routing key to look up, or {@code null} to search with an empty byte array + * @return the 0-based slice index in {@link #getSlices()}, or {@code -1} if {@code slices} + * is empty or if the key is smaller than the first slice's {@code startKey} */ int lookup(@Nullable byte[] key) { if (slices.isEmpty()) { diff --git a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java index 2ad3d13998c..ed9c2a8c055 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java @@ -135,4 +135,70 @@ public void gettersAndImmutability() { assertThat(sliceMap.getSlices()).hasSize(1); assertThat(sliceMap.getFallbackPool()).containsExactly(0, 1).inOrder(); } + + @Test + public void lookup_nullKey_treatedAsEmptyBytes() { + SliceEntry s1 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceEntry s2 = new SliceEntry( + "m".getBytes(StandardCharsets.UTF_8), Collections.singletonList(1)); + SliceMap sliceMap = new SliceMap(Arrays.asList(s1, s2), Arrays.asList(0, 1), 1L); + + assertThat(sliceMap.lookup(null)).isEqualTo(0); + } + + @Test + public void constructor_unsortedSlices_sortedLexicographically() { + SliceEntry s1 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceEntry s2 = new SliceEntry( + "m".getBytes(StandardCharsets.UTF_8), Collections.singletonList(1)); + SliceEntry s3 = new SliceEntry( + "z".getBytes(StandardCharsets.UTF_8), Collections.singletonList(2)); + + // Pass in reverse order + SliceMap sliceMap = new SliceMap(Arrays.asList(s3, s1, s2), Arrays.asList(0, 1, 2), 1L); + + assertThat(sliceMap.getSlices().get(0).getStartKey()) + .isEqualTo("".getBytes(StandardCharsets.UTF_8)); + assertThat(sliceMap.getSlices().get(1).getStartKey()) + .isEqualTo("m".getBytes(StandardCharsets.UTF_8)); + assertThat(sliceMap.getSlices().get(2).getStartKey()) + .isEqualTo("z".getBytes(StandardCharsets.UTF_8)); + + assertThat(sliceMap.lookup("abc".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("mmm".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("zzz".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + } + + @Test + public void lookup_duplicateStartKeys_matchesOne() { + SliceEntry s1 = new SliceEntry( + "a".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceEntry s2 = new SliceEntry( + "a".getBytes(StandardCharsets.UTF_8), Collections.singletonList(1)); + SliceMap sliceMap = new SliceMap(Arrays.asList(s1, s2), Arrays.asList(0, 1), 1L); + + int idx = sliceMap.lookup("a".getBytes(StandardCharsets.UTF_8)); + assertThat(idx).isAnyOf(0, 1); + } + + @Test + public void constructor_nullInputs_throwsNullPointerException() { + org.junit.Assert.assertThrows( + NullPointerException.class, + () -> new SliceMap(null, Collections.singletonList(0), 1L)); + + org.junit.Assert.assertThrows( + NullPointerException.class, + () -> new SliceMap(Collections.emptyList(), null, 1L)); + + org.junit.Assert.assertThrows( + NullPointerException.class, + () -> new SliceEntry(null, Collections.singletonList(0))); + + org.junit.Assert.assertThrows( + NullPointerException.class, + () -> new SliceEntry(new byte[0], null)); + } } From 903ec889634b7646b8dddee6250a683bd33b7ec3 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Mon, 7 Sep 2026 21:02:37 +0530 Subject: [PATCH 18/28] context specific error --- .../grpc/autosharding/AutoShardingPicker.java | 13 ++++++++----- .../autosharding/AutoShardingPickerTest.java | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java index 24086df3b5d..842b88d3a52 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -169,7 +169,8 @@ public PickResult pickSubchannel(PickSubchannelArgs args) { if (sliceIdx == -1) { if (fallbackEnabled) { - return pickFromEndpointIndices(sliceMap.getFallbackPool(), args); + return pickFromEndpointIndices( + sliceMap.getFallbackPool(), args, "No endpoints available in fallback pool"); } else { return PickResult.withError( Status.UNAVAILABLE.withDescription( @@ -178,18 +179,20 @@ public PickResult pickSubchannel(PickSubchannelArgs args) { } if (sliceInFallback[sliceIdx] && fallbackEnabled) { - return pickFromEndpointIndices(sliceMap.getFallbackPool(), args); + return pickFromEndpointIndices( + sliceMap.getFallbackPool(), args, "No endpoints available in fallback pool"); } SliceMap.SliceEntry sliceEntry = sliceMap.getSlices().get(sliceIdx); - return pickFromEndpointIndices(sliceEntry.getEndpoints(), args); + return pickFromEndpointIndices( + sliceEntry.getEndpoints(), args, "No valid endpoints in slice and fallback disabled"); } private PickResult pickFromEndpointIndices( - List indices, PickSubchannelArgs args) { + List indices, PickSubchannelArgs args, String emptyErrorDescription) { if (indices.isEmpty()) { return PickResult.withError( - Status.UNAVAILABLE.withDescription("No valid endpoints in slice and fallback disabled")); + Status.UNAVAILABLE.withDescription(emptyErrorDescription)); } int size = indices.size(); diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java index 7bced71cd94..18b432f0232 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -312,6 +312,25 @@ public void pick_emptySliceEndpoints_fallbackEnabled_routesToFallbackPool() { assertThat(result).isSameInstanceAs(fallbackReadyResult); } + @Test + public void pick_emptyFallbackPool_fallbackEnabled_returnsUnavailable() { + SliceEntry emptySlice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.emptyList()); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(emptySlice), Collections.emptyList(), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Collections.emptyList(), + true, + AutoShardingPicker.createKeyHeader("x-key")); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()) + .contains("No endpoints available in fallback pool"); + } + @Test public void pickerEndpoint_gettersAndToString() { FakePicker fakePicker = new FakePicker(PickResult.withNoResult()); From bb9a9009f3f7e7820723cf4c41a5e580f0e37e80 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Mon, 7 Sep 2026 21:08:49 +0530 Subject: [PATCH 19/28] clone start key --- .../src/main/java/io/grpc/autosharding/SliceMap.java | 2 +- .../src/test/java/io/grpc/autosharding/SliceMapTest.java | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java index ec3af3f0af4..b5e640492c2 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -68,7 +68,7 @@ static final class SliceEntry { * @param endpoints the list of endpoint indices assigned to this slice */ SliceEntry(byte[] startKey, List endpoints) { - this.startKey = checkNotNull(startKey, "startKey"); + this.startKey = checkNotNull(startKey, "startKey").clone(); this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints")); } diff --git a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java index ed9c2a8c055..126c008d392 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java @@ -117,8 +117,9 @@ public void lookup_unsignedByteComparison() { @Test public void gettersAndImmutability() { + byte[] inputKey = new byte[] {1}; List slices = new ArrayList<>(); - slices.add(new SliceEntry(new byte[] {1}, Arrays.asList(0, 1))); + slices.add(new SliceEntry(inputKey, Arrays.asList(0, 1))); List fallback = new ArrayList<>(Arrays.asList(0, 1)); SliceMap sliceMap = new SliceMap(slices, fallback, 42L); @@ -129,10 +130,12 @@ public void gettersAndImmutability() { assertThat(sliceMap.getSlices().get(0).getStartKey()).isEqualTo(new byte[] {1}); assertThat(sliceMap.getSlices().get(0).getEndpoints()).containsExactly(0, 1).inOrder(); - // Verify defensive copying: mutating input collections does not affect sliceMap + // Verify defensive copying: mutating input collections and key array does not affect sliceMap + inputKey[0] = 99; slices.clear(); fallback.clear(); assertThat(sliceMap.getSlices()).hasSize(1); + assertThat(sliceMap.getSlices().get(0).getStartKey()).isEqualTo(new byte[] {1}); assertThat(sliceMap.getFallbackPool()).containsExactly(0, 1).inOrder(); } From c11754e0b71633937ea3d3f384776cabe0811040 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Mon, 7 Sep 2026 21:12:36 +0530 Subject: [PATCH 20/28] add some behavioural unit tests --- .../autosharding/AutoShardingPickerTest.java | 96 +++++++++++++++++++ 1 file changed, 96 insertions(+) diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java index 18b432f0232..eefc65959c9 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -35,6 +35,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Assert; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -400,4 +401,99 @@ public void pick_deterministicRandom_selectsExpectedEndpoint() { PickResult result1 = picker1.pickSubchannel(createArgs(new Metadata())); assertThat(result1).isSameInstanceAs(ready1); } + + @Test + public void pick_invalidSliceEndpointIndex_throwsIndexOutOfBoundsException() { + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP_EXIT_IDLER); + + // Slice references index 5, but only 1 endpoint (index 0) exists + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(5)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-key")); + + Assert.assertThrows( + IndexOutOfBoundsException.class, + () -> picker.pickSubchannel(createArgs(new Metadata()))); + } + + @Test + public void pick_negativeSliceEndpointIndex_throwsIndexOutOfBoundsException() { + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP_EXIT_IDLER); + + // Slice references negative index -1 + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(-1)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-key")); + + Assert.assertThrows( + IndexOutOfBoundsException.class, + () -> picker.pickSubchannel(createArgs(new Metadata()))); + } + + @Test + public void pick_invalidFallbackPoolIndex_throwsIndexOutOfBoundsException() { + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP_EXIT_IDLER); + + // Fallback pool references index 10, but only 1 endpoint exists + SliceMap emptySliceMap = new SliceMap( + Collections.emptyList(), Collections.singletonList(10), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + emptySliceMap, + Collections.singletonList(ep0), + true, + AutoShardingPicker.createKeyHeader("x-key")); + + Assert.assertThrows( + IndexOutOfBoundsException.class, + () -> picker.pickSubchannel(createArgs(new Metadata()))); + } + + @Test + public void constructor_nullInputs_throwsNullPointerException() { + SliceMap sliceMap = new SliceMap( + Collections.emptyList(), Collections.emptyList(), 1L); + + Assert.assertThrows( + NullPointerException.class, + () -> new AutoShardingPicker( + null, + Collections.emptyList(), + false, + null)); + + Assert.assertThrows( + NullPointerException.class, + () -> new AutoShardingPicker( + sliceMap, + null, + false, + null)); + + Assert.assertThrows( + NullPointerException.class, + () -> new AutoShardingPicker( + sliceMap, + Collections.emptyList(), + false, + null, + null)); + } } From c79cc595c4cdfe2006ad27488413dd88c7819cc1 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 8 Sep 2026 11:30:44 +0530 Subject: [PATCH 21/28] autosharding: implementation of EndpointMap and LazyChildLB --- .../io/grpc/autosharding/EndpointMap.java | 226 ++++++++++++++++ .../autosharding/LazyChildLoadBalancer.java | 163 +++++++++++ .../io/grpc/autosharding/EndpointMapTest.java | 212 +++++++++++++++ .../LazyChildLoadBalancerTest.java | 255 ++++++++++++++++++ 4 files changed, 856 insertions(+) create mode 100644 autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java create mode 100644 autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java create mode 100644 autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java create mode 100644 autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java 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..11c0ee3e3ad --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java @@ -0,0 +1,226 @@ +/* + * 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 com.google.common.base.Preconditions.checkNotNull; +import static io.grpc.ConnectivityState.IDLE; + +import com.google.common.base.MoreObjects; +import com.google.common.collect.ImmutableList; +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; +import javax.annotation.Nullable; +import javax.annotation.concurrent.NotThreadSafe; + +/** + * Manages the mapping from endpoint hostname to {@link EndpointHolder} and coordinates + * child load balancer lifecycle and connectivity state updates. + * + *

    Threading model: This class is not thread-safe. All methods must be invoked from the + * {@link io.grpc.SynchronizationContext} by the parent load balancer. + */ +@NotThreadSafe +final class EndpointMap { + private final Map map = new LinkedHashMap<>(); + + @Nullable + EndpointHolder get(String hostname) { + return map.get(checkNotNull(hostname, "hostname")); + } + + void put(String hostname, EndpointHolder holder) { + map.put(checkNotNull(hostname, "hostname"), checkNotNull(holder, "holder")); + } + + @Nullable + EndpointHolder remove(String hostname) { + return map.remove(checkNotNull(hostname, "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(); + } + + /** + * Re-assigns contiguous 0-based index values across all current endpoint holders. + */ + void reindex() { + int nextIdx = 0; + for (EndpointHolder holder : map.values()) { + holder.setIndex(nextIdx++); + } + } + + /** + * Shuts down all child load balancers and clears the map. + */ + void shutdownAll() { + for (EndpointHolder holder : map.values()) { + holder.shutdown(); + } + map.clear(); + } + + /** + * Builds an immutable snapshot list of {@link PickerEndpoint}s ordered by index. + */ + ImmutableList toPickerEndpoints() { + ImmutableList.Builder builder = + ImmutableList.builderWithExpectedSize(map.size()); + for (EndpointHolder holder : map.values()) { + builder.add(holder.toPickerEndpoint()); + } + return builder.build(); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("map", map) + .toString(); + } + + /** + * Holds the connectivity state, picker, and lazy child load balancer for a single endpoint. + */ + static final class EndpointHolder { + private int index; + private final LazyChildLoadBalancer childLb; + private ConnectivityState state = IDLE; + private SubchannelPicker picker = new FixedResultPicker(PickResult.withNoResult()); + + EndpointHolder( + int index, + Helper helper, + LoadBalancerProvider pickFirstProvider, + @Nullable Runnable stateUpdateCallback) { + this.index = index; + this.childLb = new LazyChildLoadBalancer( + new ChildHelper(checkNotNull(helper, "helper"), stateUpdateCallback), + checkNotNull(pickFirstProvider, "pickFirstProvider")); + } + + int getIndex() { + return index; + } + + void setIndex(int index) { + this.index = index; + } + + ConnectivityState getState() { + return state; + } + + SubchannelPicker getPicker() { + return picker; + } + + LazyChildLoadBalancer getChildLb() { + return childLb; + } + + PickerEndpoint toPickerEndpoint() { + return new PickerEndpoint(state, picker, childLb); + } + + void updateAddresses(List eags, Attributes attributes) { + ResolvedAddresses childAddresses = ResolvedAddresses.newBuilder() + .setAddresses(ImmutableList.copyOf(checkNotNull(eags, "eags"))) + .setAttributes(checkNotNull(attributes, "attributes")) + .build(); + childLb.acceptResolvedAddresses(childAddresses); + } + + void requestConnection() { + childLb.requestConnection(); + } + + void shutdown() { + childLb.shutdown(); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("index", index) + .add("state", state) + .add("childLb", childLb) + .toString(); + } + + private final class ChildHelper extends ForwardingLoadBalancerHelper { + private final Helper delegateHelper; + @Nullable private final Runnable stateUpdateCallback; + + ChildHelper(Helper delegateHelper, @Nullable Runnable stateUpdateCallback) { + this.delegateHelper = checkNotNull(delegateHelper, "delegateHelper"); + this.stateUpdateCallback = stateUpdateCallback; + } + + @Override + protected Helper delegate() { + return delegateHelper; + } + + @Override + public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) { + state = checkNotNull(newState, "newState"); + picker = checkNotNull(newPicker, "newPicker"); + if (stateUpdateCallback != null) { + stateUpdateCallback.run(); + } + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("delegateHelper", delegateHelper) + .toString(); + } + } + } +} 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..49ab0bad7a0 --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java @@ -0,0 +1,163 @@ +/* + * 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 com.google.common.base.Preconditions.checkNotNull; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.MoreObjects; +import io.grpc.ConnectivityState; +import io.grpc.LoadBalancer; +import io.grpc.LoadBalancerProvider; +import io.grpc.Status; +import java.util.concurrent.atomic.AtomicBoolean; +import javax.annotation.Nullable; +import javax.annotation.concurrent.ThreadSafe; + +/** + * A wrapper {@link LoadBalancer} that lazily creates and delegates to a child LoadBalancer + * (typically {@code pick_first}) only when a connection attempt is explicitly requested. + * + *

    Until a connection is requested, this balancer remains in the {@link ConnectivityState#IDLE} + * state upon receiving resolved addresses without eagerly connecting. It implements + * {@link PickerEndpoint.ExitIdler} to allow worker threads in {@link AutoShardingPicker} to + * trigger connection attempts safely and non-blockingly via + * {@link io.grpc.SynchronizationContext}. + * + *

    Threading model: {@link #exitIdle()} is thread-safe and may be called concurrently by + * application/worker threads during picker execution. All other {@link LoadBalancer} lifecycle + * methods (such as {@link #acceptResolvedAddresses}, {@link #requestConnection}, + * {@link #handleNameResolutionError}, and {@link #shutdown}) must be invoked from the + * {@link io.grpc.SynchronizationContext}. + */ +@ThreadSafe +final class LazyChildLoadBalancer extends LoadBalancer implements PickerEndpoint.ExitIdler { + private final Helper helper; + private final LoadBalancerProvider delegateProvider; + private final AtomicBoolean connectingScheduled = new AtomicBoolean(false); + + @Nullable private LoadBalancer delegate; + @Nullable private ResolvedAddresses lastResolvedAddresses; + private boolean connectionRequested = false; + private boolean shutdown = false; + + /** + * Constructs a {@link LazyChildLoadBalancer}. + * + * @param helper the parent load balancer helper + * @param delegateProvider provider used to instantiate the child load balancer (e.g. pick_first) + */ + LazyChildLoadBalancer(Helper helper, LoadBalancerProvider delegateProvider) { + this.helper = checkNotNull(helper, "helper"); + this.delegateProvider = checkNotNull(delegateProvider, "delegateProvider"); + } + + @Override + public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) { + checkNotNull(resolvedAddresses, "resolvedAddresses"); + if (shutdown) { + return Status.FAILED_PRECONDITION.withDescription("LoadBalancer is shutdown"); + } + 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) { + checkNotNull(error, "error"); + if (shutdown) { + return; + } + if (delegate != null) { + delegate.handleNameResolutionError(error); + } else { + helper.updateBalancingState( + ConnectivityState.TRANSIENT_FAILURE, + new FixedResultPicker(PickResult.withError(error))); + } + } + + @Override + public void requestConnection() { + if (shutdown) { + return; + } + connectionRequested = true; + if (delegate == null && lastResolvedAddresses != null) { + delegate = delegateProvider.newLoadBalancer(helper); + delegate.acceptResolvedAddresses(lastResolvedAddresses); + } else if (delegate != null) { + delegate.requestConnection(); + } + } + + /** + * Callback invoked from worker threads during {@link AutoShardingPicker#pickSubchannel}. + * + *

    Protects against thundering herds by using an {@link AtomicBoolean} guard to schedule + * {@link #requestConnection()} onto the {@link io.grpc.SynchronizationContext} at most once. + */ + @Override + public void exitIdle() { + if (shutdown) { + return; + } + if (connectingScheduled.compareAndSet(false, true)) { + helper.getSynchronizationContext().execute(this::requestConnection); + } + } + + @Override + public void shutdown() { + shutdown = true; + if (delegate != null) { + delegate.shutdown(); + delegate = null; + } + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("connectionRequested", connectionRequested) + .add("shutdown", shutdown) + .add("delegate", delegate) + .toString(); + } + + @VisibleForTesting + boolean isConnectionRequested() { + return connectionRequested; + } + + @VisibleForTesting + @Nullable + LoadBalancer getDelegate() { + return delegate; + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java new file mode 100644 index 00000000000..c39a6375f59 --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java @@ -0,0 +1,212 @@ +/* + * 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 com.google.common.truth.Truth.assertThat; +import static io.grpc.ConnectivityState.IDLE; +import static io.grpc.ConnectivityState.READY; +import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.common.collect.ImmutableList; +import io.grpc.Attributes; +import io.grpc.EquivalentAddressGroup; +import io.grpc.LoadBalancer; +import io.grpc.LoadBalancer.Helper; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.LoadBalancerProvider; +import io.grpc.SynchronizationContext; +import io.grpc.autosharding.EndpointMap.EndpointHolder; +import java.net.SocketAddress; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; + +@RunWith(JUnit4.class) +public class EndpointMapTest { + + private final Helper mockHelper = mock(Helper.class); + private final LoadBalancerProvider mockProvider = mock(LoadBalancerProvider.class); + private final LoadBalancer mockDelegate = mock(LoadBalancer.class); + private final SynchronizationContext syncContext = + new SynchronizationContext((t, e) -> { + throw new AssertionError("Unhandled exception in syncContext", e); + }); + + private EndpointMap endpointMap; + private final AtomicInteger stateChangeCount = new AtomicInteger(0); + + @Before + public void setUp() { + when(mockHelper.getSynchronizationContext()).thenReturn(syncContext); + when(mockProvider.newLoadBalancer(any())).thenReturn(mockDelegate); + endpointMap = new EndpointMap(); + } + + private EndpointHolder createHolder(int index) { + return new EndpointHolder(index, mockHelper, mockProvider, stateChangeCount::incrementAndGet); + } + + @Test + public void basicMapOperations() { + assertThat(endpointMap.isEmpty()).isTrue(); + assertThat(endpointMap.size()).isEqualTo(0); + + EndpointHolder h1 = createHolder(0); + EndpointHolder h2 = createHolder(1); + + endpointMap.put("host1", h1); + endpointMap.put("host2", h2); + + assertThat(endpointMap.isEmpty()).isFalse(); + assertThat(endpointMap.size()).isEqualTo(2); + assertThat(endpointMap.get("host1")).isSameInstanceAs(h1); + assertThat(endpointMap.get("host2")).isSameInstanceAs(h2); + assertThat(endpointMap.get("unknown")).isNull(); + assertThat(endpointMap.keySet()).containsExactly("host1", "host2").inOrder(); + assertThat(endpointMap.values()).containsExactly(h1, h2).inOrder(); + + EndpointHolder removed = endpointMap.remove("host1"); + assertThat(removed).isSameInstanceAs(h1); + assertThat(endpointMap.size()).isEqualTo(1); + assertThat(endpointMap.get("host1")).isNull(); + } + + @Test + public void nullChecks() { + EndpointHolder h = createHolder(0); + + assertThrows(NullPointerException.class, () -> endpointMap.get(null)); + assertThrows(NullPointerException.class, () -> endpointMap.put(null, h)); + assertThrows(NullPointerException.class, () -> endpointMap.put("host", null)); + assertThrows(NullPointerException.class, () -> endpointMap.remove(null)); + + assertThrows( + NullPointerException.class, + () -> new EndpointHolder(0, null, mockProvider, null)); + assertThrows( + NullPointerException.class, + () -> new EndpointHolder(0, mockHelper, null, null)); + + assertThrows( + NullPointerException.class, + () -> h.updateAddresses(null, Attributes.EMPTY)); + assertThrows( + NullPointerException.class, + () -> h.updateAddresses(Collections.emptyList(), null)); + } + + @Test + public void reindex_updatesIndicesContiguously() { + EndpointHolder h0 = createHolder(0); + EndpointHolder h1 = createHolder(1); + EndpointHolder h2 = createHolder(2); + + endpointMap.put("host0", h0); + endpointMap.put("host1", h1); + endpointMap.put("host2", h2); + + // Remove middle element + endpointMap.remove("host1"); + assertThat(h0.getIndex()).isEqualTo(0); + assertThat(h2.getIndex()).isEqualTo(2); + + endpointMap.reindex(); + assertThat(h0.getIndex()).isEqualTo(0); + assertThat(h2.getIndex()).isEqualTo(1); + } + + @Test + public void endpointHolder_childHelperUpdatesStateAndTriggersCallback() { + EndpointHolder holder = createHolder(0); + assertThat(holder.getState()).isEqualTo(IDLE); + + // Capture child helper passed to LazyChildLoadBalancer + ArgumentCaptor helperCaptor = ArgumentCaptor.forClass(Helper.class); + verify(mockProvider, org.mockito.Mockito.never()).newLoadBalancer(any()); + + // Trigger connection to create child helper and delegate + holder.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + holder.requestConnection(); + + verify(mockProvider).newLoadBalancer(helperCaptor.capture()); + Helper childHelper = helperCaptor.getValue(); + + // Reset counter before state update to verify callback fires on update + stateChangeCount.set(0); + + // Simulate child balancer updating state + SubchannelPicker testPicker = mock(SubchannelPicker.class); + childHelper.updateBalancingState(READY, testPicker); + + assertThat(holder.getState()).isEqualTo(READY); + assertThat(holder.getPicker()).isSameInstanceAs(testPicker); + assertThat(stateChangeCount.get()).isEqualTo(1); + } + + @Test + public void toPickerEndpoints_buildsImmutableListMatchingHolders() { + EndpointHolder h0 = createHolder(0); + EndpointHolder h1 = createHolder(1); + + endpointMap.put("host0", h0); + endpointMap.put("host1", h1); + + ImmutableList pickerEndpoints = endpointMap.toPickerEndpoints(); + assertThat(pickerEndpoints).hasSize(2); + assertThat(pickerEndpoints.get(0).getState()).isEqualTo(IDLE); + assertThat(pickerEndpoints.get(1).getState()).isEqualTo(IDLE); + } + + @Test + public void shutdownAll_cleansUpAllHoldersAndClearsMap() { + EndpointHolder h0 = createHolder(0); + EndpointHolder h1 = createHolder(1); + + endpointMap.put("host0", h0); + endpointMap.put("host1", h1); + + // Trigger connections so delegates exist + h0.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + h0.requestConnection(); + + endpointMap.shutdownAll(); + assertThat(endpointMap.isEmpty()).isTrue(); + verify(mockDelegate).shutdown(); + } + + @Test + public void toString_containsDebugFields() { + EndpointHolder h = createHolder(3); + endpointMap.put("host3", h); + + assertThat(endpointMap.toString()).contains("host3"); + assertThat(h.toString()).contains("index=3"); + assertThat(h.toString()).contains("state=IDLE"); + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java new file mode 100644 index 00000000000..b77ee964861 --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java @@ -0,0 +1,255 @@ +/* + * 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 com.google.common.truth.Truth.assertThat; +import static io.grpc.ConnectivityState.IDLE; +import static io.grpc.ConnectivityState.TRANSIENT_FAILURE; +import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.grpc.Attributes; +import io.grpc.EquivalentAddressGroup; +import io.grpc.LoadBalancer; +import io.grpc.LoadBalancer.Helper; +import io.grpc.LoadBalancer.PickSubchannelArgs; +import io.grpc.LoadBalancer.ResolvedAddresses; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.LoadBalancerProvider; +import io.grpc.Status; +import io.grpc.SynchronizationContext; +import java.net.SocketAddress; +import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; + +@RunWith(JUnit4.class) +public class LazyChildLoadBalancerTest { + + private final Helper mockHelper = mock(Helper.class); + private final LoadBalancerProvider mockProvider = mock(LoadBalancerProvider.class); + private final LoadBalancer mockDelegate = mock(LoadBalancer.class); + private final SynchronizationContext syncContext = + new SynchronizationContext((t, e) -> { + throw new AssertionError("Unhandled exception in syncContext", e); + }); + + private LazyChildLoadBalancer lazyLb; + private ResolvedAddresses resolvedAddresses; + + @Before + public void setUp() { + when(mockHelper.getSynchronizationContext()).thenReturn(syncContext); + when(mockProvider.newLoadBalancer(any())).thenReturn(mockDelegate); + when(mockDelegate.acceptResolvedAddresses(any())).thenReturn(Status.OK); + + lazyLb = new LazyChildLoadBalancer(mockHelper, mockProvider); + resolvedAddresses = ResolvedAddresses.newBuilder() + .setAddresses(Collections.singletonList( + new EquivalentAddressGroup(new SocketAddress() {}))) + .setAttributes(Attributes.EMPTY) + .build(); + } + + @Test + public void constructor_nullArguments_throwsNullPointerException() { + assertThrows( + NullPointerException.class, + () -> new LazyChildLoadBalancer(null, mockProvider)); + assertThrows( + NullPointerException.class, + () -> new LazyChildLoadBalancer(mockHelper, null)); + } + + @Test + public void initialResolution_reportsIdle_doesNotCreateChildPolicy() { + Status status = lazyLb.acceptResolvedAddresses(resolvedAddresses); + assertThat(status.isOk()).isTrue(); + + ArgumentCaptor pickerCaptor = + ArgumentCaptor.forClass(SubchannelPicker.class); + verify(mockHelper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); + assertThat(pickerCaptor.getValue().pickSubchannel(mock(PickSubchannelArgs.class)).hasResult()) + .isFalse(); + + verify(mockProvider, never()).newLoadBalancer(any()); + assertThat(lazyLb.getDelegate()).isNull(); + assertThat(lazyLb.isConnectionRequested()).isFalse(); + } + + @Test + public void requestConnection_createsChildPolicy_andForwardsAddresses() { + lazyLb.acceptResolvedAddresses(resolvedAddresses); + assertThat(lazyLb.getDelegate()).isNull(); + + lazyLb.requestConnection(); + + assertThat(lazyLb.isConnectionRequested()).isTrue(); + verify(mockProvider).newLoadBalancer(mockHelper); + verify(mockDelegate).acceptResolvedAddresses(resolvedAddresses); + assertThat(lazyLb.getDelegate()).isSameInstanceAs(mockDelegate); + } + + @Test + public void requestConnection_beforeResolvedAddresses_createsPolicyWhenAddressesArrive() { + lazyLb.requestConnection(); + assertThat(lazyLb.isConnectionRequested()).isTrue(); + verify(mockProvider, never()).newLoadBalancer(any()); + + lazyLb.acceptResolvedAddresses(resolvedAddresses); + verify(mockProvider).newLoadBalancer(mockHelper); + verify(mockDelegate).acceptResolvedAddresses(resolvedAddresses); + } + + @Test + public void requestConnection_whenAlreadyCreated_delegatesRequestConnection() { + lazyLb.acceptResolvedAddresses(resolvedAddresses); + lazyLb.requestConnection(); + verify(mockDelegate, never()).requestConnection(); + + lazyLb.requestConnection(); + verify(mockDelegate).requestConnection(); + } + + @Test + public void acceptResolvedAddresses_afterConnectionRequested_forwardsDirectly() { + lazyLb.acceptResolvedAddresses(resolvedAddresses); + lazyLb.requestConnection(); + + ResolvedAddresses newAddresses = ResolvedAddresses.newBuilder() + .setAddresses(Collections.singletonList( + new EquivalentAddressGroup(new SocketAddress() {}))) + .setAttributes(Attributes.EMPTY) + .build(); + + lazyLb.acceptResolvedAddresses(newAddresses); + verify(mockDelegate).acceptResolvedAddresses(newAddresses); + } + + @Test + public void acceptResolvedAddresses_null_throwsNullPointerException() { + assertThrows( + NullPointerException.class, + () -> lazyLb.acceptResolvedAddresses(null)); + } + + @Test + public void exitIdle_schedulesRequestConnectionOnSyncContext_once() throws Exception { + lazyLb.acceptResolvedAddresses(resolvedAddresses); + + int threadCount = 10; + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(threadCount); + + for (int i = 0; i < threadCount; i++) { + executor.execute(() -> { + try { + startLatch.await(); + lazyLb.exitIdle(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + doneLatch.countDown(); + } + }); + } + + startLatch.countDown(); + assertThat(doneLatch.await(5, TimeUnit.SECONDS)).isTrue(); + executor.shutdown(); + + // Verify child policy instantiated exactly once despite concurrent exitIdle calls + verify(mockProvider, times(1)).newLoadBalancer(mockHelper); + verify(mockDelegate, times(1)).acceptResolvedAddresses(resolvedAddresses); + } + + @Test + public void handleNameResolutionError_null_throwsNullPointerException() { + assertThrows( + NullPointerException.class, + () -> lazyLb.handleNameResolutionError(null)); + } + + @Test + public void handleNameResolutionError_beforeConnectionRequested_reportsTransientFailure() { + Status error = Status.UNAVAILABLE.withDescription("dns failed"); + lazyLb.handleNameResolutionError(error); + + ArgumentCaptor pickerCaptor = + ArgumentCaptor.forClass(SubchannelPicker.class); + verify(mockHelper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); + assertThat(pickerCaptor.getValue().pickSubchannel(mock(PickSubchannelArgs.class)).getStatus()) + .isEqualTo(error); + } + + @Test + public void handleNameResolutionError_afterConnectionRequested_forwardsToDelegate() { + lazyLb.acceptResolvedAddresses(resolvedAddresses); + lazyLb.requestConnection(); + + Status error = Status.UNAVAILABLE.withDescription("dns failed"); + lazyLb.handleNameResolutionError(error); + verify(mockDelegate).handleNameResolutionError(error); + } + + @Test + public void shutdown_cleansUpDelegate() { + lazyLb.acceptResolvedAddresses(resolvedAddresses); + lazyLb.requestConnection(); + assertThat(lazyLb.getDelegate()).isNotNull(); + + lazyLb.shutdown(); + verify(mockDelegate).shutdown(); + assertThat(lazyLb.getDelegate()).isNull(); + } + + @Test + public void operationsAfterShutdown_areNoOpsOrReturnError() { + lazyLb.acceptResolvedAddresses(resolvedAddresses); + lazyLb.shutdown(); + + Status status = lazyLb.acceptResolvedAddresses(resolvedAddresses); + assertThat(status.getCode()).isEqualTo(Status.Code.FAILED_PRECONDITION); + + // None of these should throw or create a child load balancer + lazyLb.requestConnection(); + lazyLb.exitIdle(); + lazyLb.handleNameResolutionError(Status.UNAVAILABLE); + verify(mockProvider, never()).newLoadBalancer(any()); + } + + @Test + public void toString_containsDebugFields() { + String str = lazyLb.toString(); + assertThat(str).contains("connectionRequested=false"); + assertThat(str).contains("shutdown=false"); + } +} From e35068500d6776432bda09913f5939d46bacf2e6 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 8 Sep 2026 14:44:23 +0530 Subject: [PATCH 22/28] call request conn in parent policy --- .../io/grpc/autosharding/LazyChildLoadBalancer.java | 13 ++++++++++--- .../autosharding/LazyChildLoadBalancerTest.java | 12 +++++++++--- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java index 49ab0bad7a0..698893824f1 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java +++ b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java @@ -53,7 +53,7 @@ final class LazyChildLoadBalancer extends LoadBalancer implements PickerEndpoint @Nullable private LoadBalancer delegate; @Nullable private ResolvedAddresses lastResolvedAddresses; private boolean connectionRequested = false; - private boolean shutdown = false; + private volatile boolean shutdown = false; /** * Constructs a {@link LazyChildLoadBalancer}. @@ -74,10 +74,16 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) { } lastResolvedAddresses = resolvedAddresses; if (connectionRequested) { + boolean newlyCreated = false; if (delegate == null) { delegate = delegateProvider.newLoadBalancer(helper); + newlyCreated = true; } - return delegate.acceptResolvedAddresses(resolvedAddresses); + Status status = delegate.acceptResolvedAddresses(resolvedAddresses); + if (newlyCreated && status.isOk()) { + delegate.requestConnection(); + } + return status; } else { // Report IDLE state until connection is explicitly requested helper.updateBalancingState( @@ -111,7 +117,8 @@ public void requestConnection() { if (delegate == null && lastResolvedAddresses != null) { delegate = delegateProvider.newLoadBalancer(helper); delegate.acceptResolvedAddresses(lastResolvedAddresses); - } else if (delegate != null) { + } + if (delegate != null) { delegate.requestConnection(); } } diff --git a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java index b77ee964861..f6841867768 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java @@ -105,7 +105,7 @@ public void initialResolution_reportsIdle_doesNotCreateChildPolicy() { } @Test - public void requestConnection_createsChildPolicy_andForwardsAddresses() { + public void requestConnection_createsChildPolicy_forwardsAddresses_andRequestsConnection() { lazyLb.acceptResolvedAddresses(resolvedAddresses); assertThat(lazyLb.getDelegate()).isNull(); @@ -114,6 +114,7 @@ public void requestConnection_createsChildPolicy_andForwardsAddresses() { assertThat(lazyLb.isConnectionRequested()).isTrue(); verify(mockProvider).newLoadBalancer(mockHelper); verify(mockDelegate).acceptResolvedAddresses(resolvedAddresses); + verify(mockDelegate).requestConnection(); assertThat(lazyLb.getDelegate()).isSameInstanceAs(mockDelegate); } @@ -126,22 +127,24 @@ public void requestConnection_beforeResolvedAddresses_createsPolicyWhenAddresses lazyLb.acceptResolvedAddresses(resolvedAddresses); verify(mockProvider).newLoadBalancer(mockHelper); verify(mockDelegate).acceptResolvedAddresses(resolvedAddresses); + verify(mockDelegate).requestConnection(); } @Test public void requestConnection_whenAlreadyCreated_delegatesRequestConnection() { lazyLb.acceptResolvedAddresses(resolvedAddresses); lazyLb.requestConnection(); - verify(mockDelegate, never()).requestConnection(); + verify(mockDelegate, times(1)).requestConnection(); lazyLb.requestConnection(); - verify(mockDelegate).requestConnection(); + verify(mockDelegate, times(2)).requestConnection(); } @Test public void acceptResolvedAddresses_afterConnectionRequested_forwardsDirectly() { lazyLb.acceptResolvedAddresses(resolvedAddresses); lazyLb.requestConnection(); + verify(mockDelegate, times(1)).requestConnection(); ResolvedAddresses newAddresses = ResolvedAddresses.newBuilder() .setAddresses(Collections.singletonList( @@ -151,6 +154,8 @@ public void acceptResolvedAddresses_afterConnectionRequested_forwardsDirectly() lazyLb.acceptResolvedAddresses(newAddresses); verify(mockDelegate).acceptResolvedAddresses(newAddresses); + // Should not request connection again on subsequent address update + verify(mockDelegate, times(1)).requestConnection(); } @Test @@ -189,6 +194,7 @@ public void exitIdle_schedulesRequestConnectionOnSyncContext_once() throws Excep // Verify child policy instantiated exactly once despite concurrent exitIdle calls verify(mockProvider, times(1)).newLoadBalancer(mockHelper); verify(mockDelegate, times(1)).acceptResolvedAddresses(resolvedAddresses); + verify(mockDelegate, times(1)).requestConnection(); } @Test From ca8fd50f2499a49107acd453f3eacd311d318fde Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 8 Sep 2026 15:00:06 +0530 Subject: [PATCH 23/28] use exact same indices in toPickerEndpoints --- .../io/grpc/autosharding/EndpointMap.java | 27 +++++++-- .../io/grpc/autosharding/EndpointMapTest.java | 58 +++++++++++++++++-- 2 files changed, 76 insertions(+), 9 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java index 11c0ee3e3ad..0a4388b526d 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java @@ -17,6 +17,7 @@ package io.grpc.autosharding; import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; import static io.grpc.ConnectivityState.IDLE; import com.google.common.base.MoreObjects; @@ -105,15 +106,31 @@ void shutdownAll() { } /** - * Builds an immutable snapshot list of {@link PickerEndpoint}s ordered by index. + * Builds an immutable snapshot list of {@link PickerEndpoint}s placed strictly at their + * corresponding {@link EndpointHolder#getIndex()} positions. + * + * @throws IllegalStateException if endpoint indices are not contiguous from 0 to N-1 */ ImmutableList toPickerEndpoints() { - ImmutableList.Builder builder = - ImmutableList.builderWithExpectedSize(map.size()); + int size = map.size(); + if (size == 0) { + return ImmutableList.of(); + } + PickerEndpoint[] array = new PickerEndpoint[size]; for (EndpointHolder holder : map.values()) { - builder.add(holder.toPickerEndpoint()); + int idx = holder.getIndex(); + checkState( + idx >= 0 && idx < size, + "Endpoint holder index %s is out of bounds for size %s", + idx, + size); + checkState( + array[idx] == null, + "Duplicate endpoint holder index %s detected", + idx); + array[idx] = holder.toPickerEndpoint(); } - return builder.build(); + return ImmutableList.copyOf(array); } @Override diff --git a/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java index c39a6375f59..8c9cd43bfdd 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java @@ -19,9 +19,11 @@ import static com.google.common.truth.Truth.assertThat; import static io.grpc.ConnectivityState.IDLE; import static io.grpc.ConnectivityState.READY; +import static io.grpc.ConnectivityState.TRANSIENT_FAILURE; import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -168,17 +170,65 @@ public void endpointHolder_childHelperUpdatesStateAndTriggersCallback() { } @Test - public void toPickerEndpoints_buildsImmutableListMatchingHolders() { + public void toPickerEndpoints_buildsImmutableListMatchingHoldersByIndex() { EndpointHolder h0 = createHolder(0); EndpointHolder h1 = createHolder(1); - endpointMap.put("host0", h0); + ArgumentCaptor helperCaptor = ArgumentCaptor.forClass(Helper.class); + + // Trigger connections so child helpers are passed to provider + h0.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + h0.requestConnection(); + + h1.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + h1.requestConnection(); + + verify(mockProvider, times(2)).newLoadBalancer(helperCaptor.capture()); + Helper childHelper0 = helperCaptor.getAllValues().get(0); + Helper childHelper1 = helperCaptor.getAllValues().get(1); + + SubchannelPicker picker0 = mock(SubchannelPicker.class); + SubchannelPicker picker1 = mock(SubchannelPicker.class); + + childHelper0.updateBalancingState(READY, picker0); + childHelper1.updateBalancingState(TRANSIENT_FAILURE, picker1); + + // Insert in reverse index order to verify explicit index placement endpointMap.put("host1", h1); + endpointMap.put("host0", h0); ImmutableList pickerEndpoints = endpointMap.toPickerEndpoints(); assertThat(pickerEndpoints).hasSize(2); - assertThat(pickerEndpoints.get(0).getState()).isEqualTo(IDLE); - assertThat(pickerEndpoints.get(1).getState()).isEqualTo(IDLE); + assertThat(pickerEndpoints.get(0).getState()).isEqualTo(READY); + assertThat(pickerEndpoints.get(0).getPicker()).isSameInstanceAs(picker0); + assertThat(pickerEndpoints.get(1).getState()).isEqualTo(TRANSIENT_FAILURE); + assertThat(pickerEndpoints.get(1).getPicker()).isSameInstanceAs(picker1); + } + + @Test + public void toPickerEndpoints_emptyMap_returnsEmptyList() { + assertThat(endpointMap.toPickerEndpoints()).isEmpty(); + } + + @Test + public void toPickerEndpoints_duplicateOrOutOfBoundsIndex_throwsIllegalStateException() { + EndpointHolder h0 = createHolder(0); + EndpointHolder h0Duplicate = createHolder(0); + + endpointMap.put("host0", h0); + endpointMap.put("host1", h0Duplicate); + + assertThrows(IllegalStateException.class, () -> endpointMap.toPickerEndpoints()); + + endpointMap.clear(); + EndpointHolder hOutOfBounds = createHolder(5); + endpointMap.put("host0", hOutOfBounds); + + assertThrows(IllegalStateException.class, () -> endpointMap.toPickerEndpoints()); } @Test From 3a3f66700238f290a8bd2e3b049621a5eb1a3c30 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 8 Sep 2026 15:03:42 +0530 Subject: [PATCH 24/28] clear resources in shutdown --- .../java/io/grpc/autosharding/LazyChildLoadBalancer.java | 7 +++++++ .../io/grpc/autosharding/LazyChildLoadBalancerTest.java | 4 +++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java index 698893824f1..ff23db0f3b6 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java +++ b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java @@ -142,6 +142,7 @@ public void exitIdle() { @Override public void shutdown() { shutdown = true; + lastResolvedAddresses = null; if (delegate != null) { delegate.shutdown(); delegate = null; @@ -167,4 +168,10 @@ boolean isConnectionRequested() { LoadBalancer getDelegate() { return delegate; } + + @VisibleForTesting + @Nullable + ResolvedAddresses getLastResolvedAddresses() { + return lastResolvedAddresses; + } } diff --git a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java index f6841867768..07e6803e9cb 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java @@ -227,14 +227,16 @@ public void handleNameResolutionError_afterConnectionRequested_forwardsToDelegat } @Test - public void shutdown_cleansUpDelegate() { + public void shutdown_cleansUpDelegateAndAddresses() { lazyLb.acceptResolvedAddresses(resolvedAddresses); + assertThat(lazyLb.getLastResolvedAddresses()).isNotNull(); lazyLb.requestConnection(); assertThat(lazyLb.getDelegate()).isNotNull(); lazyLb.shutdown(); verify(mockDelegate).shutdown(); assertThat(lazyLb.getDelegate()).isNull(); + assertThat(lazyLb.getLastResolvedAddresses()).isNull(); } @Test From bbf9a6f18412f4915bcbe12f5a730035a0ee531d Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 8 Sep 2026 15:17:56 +0530 Subject: [PATCH 25/28] reset connectingScheduled flag in exitIdle --- .../io/grpc/autosharding/LazyChildLoadBalancer.java | 5 ++++- .../grpc/autosharding/LazyChildLoadBalancerTest.java | 12 ++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java index ff23db0f3b6..1c2591238a8 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java +++ b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java @@ -135,7 +135,10 @@ public void exitIdle() { return; } if (connectingScheduled.compareAndSet(false, true)) { - helper.getSynchronizationContext().execute(this::requestConnection); + helper.getSynchronizationContext().execute(() -> { + connectingScheduled.set(false); + requestConnection(); + }); } } diff --git a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java index 07e6803e9cb..464a72c0aae 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java @@ -197,6 +197,18 @@ public void exitIdle_schedulesRequestConnectionOnSyncContext_once() throws Excep verify(mockDelegate, times(1)).requestConnection(); } + @Test + public void exitIdle_resetsFlagOnSyncContext_allowsSubsequentExitIdle() { + lazyLb.acceptResolvedAddresses(resolvedAddresses); + + lazyLb.exitIdle(); + verify(mockDelegate, times(1)).requestConnection(); + + // Subsequent exitIdle after syncContext execution should request connection again + lazyLb.exitIdle(); + verify(mockDelegate, times(2)).requestConnection(); + } + @Test public void handleNameResolutionError_null_throwsNullPointerException() { assertThrows( From f47816f07a3393fa90f22c97d3dab8790896d324 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 8 Sep 2026 16:04:26 +0530 Subject: [PATCH 26/28] update unit test --- .../java/io/grpc/autosharding/LazyChildLoadBalancerTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java index 464a72c0aae..b2d84459a8c 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java @@ -22,6 +22,7 @@ import static org.junit.Assert.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; @@ -194,7 +195,7 @@ public void exitIdle_schedulesRequestConnectionOnSyncContext_once() throws Excep // Verify child policy instantiated exactly once despite concurrent exitIdle calls verify(mockProvider, times(1)).newLoadBalancer(mockHelper); verify(mockDelegate, times(1)).acceptResolvedAddresses(resolvedAddresses); - verify(mockDelegate, times(1)).requestConnection(); + verify(mockDelegate, atLeastOnce()).requestConnection(); } @Test From fa764498cbc2e14e22d7b1a1c5c6422d5765fa8f Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Wed, 9 Sep 2026 14:20:04 +0530 Subject: [PATCH 27/28] util: refactor LazyLoadBalancer into util to use it in autosharding --- util/build.gradle | 1 + .../src/main/java/io/grpc/util}/LazyLoadBalancer.java | 7 ++++--- .../src/test/java/io/grpc/util}/LazyLoadBalancerTest.java | 4 ++-- xds/src/main/java/io/grpc/xds/RingHashLoadBalancer.java | 1 + 4 files changed, 8 insertions(+), 5 deletions(-) rename {xds/src/main/java/io/grpc/xds => util/src/main/java/io/grpc/util}/LazyLoadBalancer.java (97%) rename {xds/src/test/java/io/grpc/xds => util/src/test/java/io/grpc/util}/LazyLoadBalancerTest.java (97%) diff --git a/util/build.gradle b/util/build.gradle index 846b110b106..ce9c6557c50 100644 --- a/util/build.gradle +++ b/util/build.gradle @@ -56,6 +56,7 @@ animalsniffer { } tasks.named("javadoc").configure { + exclude 'io/grpc/util/LazyLoadBalancer*' exclude 'io/grpc/util/MultiChildLoadBalancer.java' exclude 'io/grpc/util/OutlierDetectionLoadBalancer*' exclude 'io/grpc/util/RandomSubsettingLoadBalancer*' diff --git a/xds/src/main/java/io/grpc/xds/LazyLoadBalancer.java b/util/src/main/java/io/grpc/util/LazyLoadBalancer.java similarity index 97% rename from xds/src/main/java/io/grpc/xds/LazyLoadBalancer.java rename to util/src/main/java/io/grpc/util/LazyLoadBalancer.java index 8dbf021775b..0b84db3d006 100644 --- a/xds/src/main/java/io/grpc/xds/LazyLoadBalancer.java +++ b/util/src/main/java/io/grpc/util/LazyLoadBalancer.java @@ -14,19 +14,20 @@ * limitations under the License. */ -package io.grpc.xds; +package io.grpc.util; import com.google.common.base.Preconditions; import io.grpc.ConnectivityState; +import io.grpc.Internal; import io.grpc.LoadBalancer; import io.grpc.Status; -import io.grpc.util.ForwardingLoadBalancer; /** * A load balancer that starts in IDLE instead of CONNECTING. Once it starts connecting, it * instantiates its delegate. */ -final class LazyLoadBalancer extends ForwardingLoadBalancer { +@Internal +public final class LazyLoadBalancer extends ForwardingLoadBalancer { private LoadBalancer delegate; public LazyLoadBalancer(Helper helper, LoadBalancer.Factory delegateFactory) { diff --git a/xds/src/test/java/io/grpc/xds/LazyLoadBalancerTest.java b/util/src/test/java/io/grpc/util/LazyLoadBalancerTest.java similarity index 97% rename from xds/src/test/java/io/grpc/xds/LazyLoadBalancerTest.java rename to util/src/test/java/io/grpc/util/LazyLoadBalancerTest.java index c79d048c9d3..403336ebe87 100644 --- a/xds/src/test/java/io/grpc/xds/LazyLoadBalancerTest.java +++ b/util/src/test/java/io/grpc/util/LazyLoadBalancerTest.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package io.grpc.xds; +package io.grpc.util; import static com.google.common.truth.Truth.assertThat; @@ -34,7 +34,7 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -/** Unit test for {@link io.grpc.xds.LazyLoadBalancer}. */ +/** Unit test for {@link LazyLoadBalancer}. */ @RunWith(JUnit4.class) public final class LazyLoadBalancerTest { private SynchronizationContext syncContext = diff --git a/xds/src/main/java/io/grpc/xds/RingHashLoadBalancer.java b/xds/src/main/java/io/grpc/xds/RingHashLoadBalancer.java index eb8ba235d82..322723c5ba1 100644 --- a/xds/src/main/java/io/grpc/xds/RingHashLoadBalancer.java +++ b/xds/src/main/java/io/grpc/xds/RingHashLoadBalancer.java @@ -39,6 +39,7 @@ import io.grpc.Metadata; import io.grpc.Status; import io.grpc.SynchronizationContext; +import io.grpc.util.LazyLoadBalancer; import io.grpc.util.MultiChildLoadBalancer; import io.grpc.xds.ThreadSafeRandom.ThreadSafeRandomImpl; import io.grpc.xds.client.XdsLogger; From d32c49560f9f47875e13a632f816568a2726c9c7 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Thu, 10 Sep 2026 16:00:14 +0530 Subject: [PATCH 28/28] refactor to use LazyLB --- .../io/grpc/autosharding/EndpointMap.java | 30 +- .../autosharding/LazyChildLoadBalancer.java | 180 ------------ .../io/grpc/autosharding/EndpointMapTest.java | 18 ++ .../LazyChildLoadBalancerTest.java | 276 ------------------ 4 files changed, 40 insertions(+), 464 deletions(-) delete mode 100644 autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java delete mode 100644 autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java diff --git a/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java index 0a4388b526d..2edf86e9599 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java @@ -25,18 +25,20 @@ import io.grpc.Attributes; import io.grpc.ConnectivityState; import io.grpc.EquivalentAddressGroup; +import io.grpc.LoadBalancer; 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 io.grpc.util.LazyLoadBalancer; import java.util.Collection; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import javax.annotation.Nullable; import javax.annotation.concurrent.NotThreadSafe; @@ -145,19 +147,22 @@ public String toString() { */ static final class EndpointHolder { private int index; - private final LazyChildLoadBalancer childLb; + private final LazyLoadBalancer childLb; + private final AtomicBoolean connectingScheduled = new AtomicBoolean(false); + private final Helper helper; private ConnectivityState state = IDLE; private SubchannelPicker picker = new FixedResultPicker(PickResult.withNoResult()); EndpointHolder( int index, Helper helper, - LoadBalancerProvider pickFirstProvider, + LoadBalancer.Factory pickFirstFactory, @Nullable Runnable stateUpdateCallback) { this.index = index; - this.childLb = new LazyChildLoadBalancer( - new ChildHelper(checkNotNull(helper, "helper"), stateUpdateCallback), - checkNotNull(pickFirstProvider, "pickFirstProvider")); + this.helper = checkNotNull(helper, "helper"); + this.childLb = new LazyLoadBalancer( + new ChildHelper(helper, stateUpdateCallback), + checkNotNull(pickFirstFactory, "pickFirstFactory")); } int getIndex() { @@ -176,12 +181,21 @@ SubchannelPicker getPicker() { return picker; } - LazyChildLoadBalancer getChildLb() { + LazyLoadBalancer getChildLb() { return childLb; } PickerEndpoint toPickerEndpoint() { - return new PickerEndpoint(state, picker, childLb); + return new PickerEndpoint(state, picker, this::exitIdle); + } + + private void exitIdle() { + if (connectingScheduled.compareAndSet(false, true)) { + helper.getSynchronizationContext().execute(() -> { + connectingScheduled.set(false); + childLb.requestConnection(); + }); + } } void updateAddresses(List eags, Attributes attributes) { diff --git a/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java deleted file mode 100644 index 1c2591238a8..00000000000 --- a/autosharding/src/main/java/io/grpc/autosharding/LazyChildLoadBalancer.java +++ /dev/null @@ -1,180 +0,0 @@ -/* - * 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 com.google.common.base.Preconditions.checkNotNull; - -import com.google.common.annotations.VisibleForTesting; -import com.google.common.base.MoreObjects; -import io.grpc.ConnectivityState; -import io.grpc.LoadBalancer; -import io.grpc.LoadBalancerProvider; -import io.grpc.Status; -import java.util.concurrent.atomic.AtomicBoolean; -import javax.annotation.Nullable; -import javax.annotation.concurrent.ThreadSafe; - -/** - * A wrapper {@link LoadBalancer} that lazily creates and delegates to a child LoadBalancer - * (typically {@code pick_first}) only when a connection attempt is explicitly requested. - * - *

    Until a connection is requested, this balancer remains in the {@link ConnectivityState#IDLE} - * state upon receiving resolved addresses without eagerly connecting. It implements - * {@link PickerEndpoint.ExitIdler} to allow worker threads in {@link AutoShardingPicker} to - * trigger connection attempts safely and non-blockingly via - * {@link io.grpc.SynchronizationContext}. - * - *

    Threading model: {@link #exitIdle()} is thread-safe and may be called concurrently by - * application/worker threads during picker execution. All other {@link LoadBalancer} lifecycle - * methods (such as {@link #acceptResolvedAddresses}, {@link #requestConnection}, - * {@link #handleNameResolutionError}, and {@link #shutdown}) must be invoked from the - * {@link io.grpc.SynchronizationContext}. - */ -@ThreadSafe -final class LazyChildLoadBalancer extends LoadBalancer implements PickerEndpoint.ExitIdler { - private final Helper helper; - private final LoadBalancerProvider delegateProvider; - private final AtomicBoolean connectingScheduled = new AtomicBoolean(false); - - @Nullable private LoadBalancer delegate; - @Nullable private ResolvedAddresses lastResolvedAddresses; - private boolean connectionRequested = false; - private volatile boolean shutdown = false; - - /** - * Constructs a {@link LazyChildLoadBalancer}. - * - * @param helper the parent load balancer helper - * @param delegateProvider provider used to instantiate the child load balancer (e.g. pick_first) - */ - LazyChildLoadBalancer(Helper helper, LoadBalancerProvider delegateProvider) { - this.helper = checkNotNull(helper, "helper"); - this.delegateProvider = checkNotNull(delegateProvider, "delegateProvider"); - } - - @Override - public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) { - checkNotNull(resolvedAddresses, "resolvedAddresses"); - if (shutdown) { - return Status.FAILED_PRECONDITION.withDescription("LoadBalancer is shutdown"); - } - lastResolvedAddresses = resolvedAddresses; - if (connectionRequested) { - boolean newlyCreated = false; - if (delegate == null) { - delegate = delegateProvider.newLoadBalancer(helper); - newlyCreated = true; - } - Status status = delegate.acceptResolvedAddresses(resolvedAddresses); - if (newlyCreated && status.isOk()) { - delegate.requestConnection(); - } - return status; - } 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) { - checkNotNull(error, "error"); - if (shutdown) { - return; - } - if (delegate != null) { - delegate.handleNameResolutionError(error); - } else { - helper.updateBalancingState( - ConnectivityState.TRANSIENT_FAILURE, - new FixedResultPicker(PickResult.withError(error))); - } - } - - @Override - public void requestConnection() { - if (shutdown) { - return; - } - connectionRequested = true; - if (delegate == null && lastResolvedAddresses != null) { - delegate = delegateProvider.newLoadBalancer(helper); - delegate.acceptResolvedAddresses(lastResolvedAddresses); - } - if (delegate != null) { - delegate.requestConnection(); - } - } - - /** - * Callback invoked from worker threads during {@link AutoShardingPicker#pickSubchannel}. - * - *

    Protects against thundering herds by using an {@link AtomicBoolean} guard to schedule - * {@link #requestConnection()} onto the {@link io.grpc.SynchronizationContext} at most once. - */ - @Override - public void exitIdle() { - if (shutdown) { - return; - } - if (connectingScheduled.compareAndSet(false, true)) { - helper.getSynchronizationContext().execute(() -> { - connectingScheduled.set(false); - requestConnection(); - }); - } - } - - @Override - public void shutdown() { - shutdown = true; - lastResolvedAddresses = null; - if (delegate != null) { - delegate.shutdown(); - delegate = null; - } - } - - @Override - public String toString() { - return MoreObjects.toStringHelper(this) - .add("connectionRequested", connectionRequested) - .add("shutdown", shutdown) - .add("delegate", delegate) - .toString(); - } - - @VisibleForTesting - boolean isConnectionRequested() { - return connectionRequested; - } - - @VisibleForTesting - @Nullable - LoadBalancer getDelegate() { - return delegate; - } - - @VisibleForTesting - @Nullable - ResolvedAddresses getLastResolvedAddresses() { - return lastResolvedAddresses; - } -} diff --git a/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java index 8c9cd43bfdd..5dd6f4c5332 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java @@ -250,6 +250,24 @@ public void shutdownAll_cleansUpAllHoldersAndClearsMap() { verify(mockDelegate).shutdown(); } + @Test + public void toPickerEndpoint_requestConnection_wakesUpChildBalancerOnSyncContext() { + EndpointHolder holder = createHolder(0); + holder.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + + PickerEndpoint pickerEndpoint = holder.toPickerEndpoint(); + verify(mockProvider, org.mockito.Mockito.never()).newLoadBalancer(any()); + + // Trigger connection through PickerEndpoint (simulate AutoShardingPicker encountering IDLE) + pickerEndpoint.requestConnection(); + + verify(mockProvider).newLoadBalancer(any()); + verify(mockDelegate).acceptResolvedAddresses(any()); + verify(mockDelegate).requestConnection(); + } + @Test public void toString_containsDebugFields() { EndpointHolder h = createHolder(3); diff --git a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java b/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java deleted file mode 100644 index b2d84459a8c..00000000000 --- a/autosharding/src/test/java/io/grpc/autosharding/LazyChildLoadBalancerTest.java +++ /dev/null @@ -1,276 +0,0 @@ -/* - * 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 com.google.common.truth.Truth.assertThat; -import static io.grpc.ConnectivityState.IDLE; -import static io.grpc.ConnectivityState.TRANSIENT_FAILURE; -import static org.junit.Assert.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import io.grpc.Attributes; -import io.grpc.EquivalentAddressGroup; -import io.grpc.LoadBalancer; -import io.grpc.LoadBalancer.Helper; -import io.grpc.LoadBalancer.PickSubchannelArgs; -import io.grpc.LoadBalancer.ResolvedAddresses; -import io.grpc.LoadBalancer.SubchannelPicker; -import io.grpc.LoadBalancerProvider; -import io.grpc.Status; -import io.grpc.SynchronizationContext; -import java.net.SocketAddress; -import java.util.Collections; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import org.junit.Before; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.ArgumentCaptor; - -@RunWith(JUnit4.class) -public class LazyChildLoadBalancerTest { - - private final Helper mockHelper = mock(Helper.class); - private final LoadBalancerProvider mockProvider = mock(LoadBalancerProvider.class); - private final LoadBalancer mockDelegate = mock(LoadBalancer.class); - private final SynchronizationContext syncContext = - new SynchronizationContext((t, e) -> { - throw new AssertionError("Unhandled exception in syncContext", e); - }); - - private LazyChildLoadBalancer lazyLb; - private ResolvedAddresses resolvedAddresses; - - @Before - public void setUp() { - when(mockHelper.getSynchronizationContext()).thenReturn(syncContext); - when(mockProvider.newLoadBalancer(any())).thenReturn(mockDelegate); - when(mockDelegate.acceptResolvedAddresses(any())).thenReturn(Status.OK); - - lazyLb = new LazyChildLoadBalancer(mockHelper, mockProvider); - resolvedAddresses = ResolvedAddresses.newBuilder() - .setAddresses(Collections.singletonList( - new EquivalentAddressGroup(new SocketAddress() {}))) - .setAttributes(Attributes.EMPTY) - .build(); - } - - @Test - public void constructor_nullArguments_throwsNullPointerException() { - assertThrows( - NullPointerException.class, - () -> new LazyChildLoadBalancer(null, mockProvider)); - assertThrows( - NullPointerException.class, - () -> new LazyChildLoadBalancer(mockHelper, null)); - } - - @Test - public void initialResolution_reportsIdle_doesNotCreateChildPolicy() { - Status status = lazyLb.acceptResolvedAddresses(resolvedAddresses); - assertThat(status.isOk()).isTrue(); - - ArgumentCaptor pickerCaptor = - ArgumentCaptor.forClass(SubchannelPicker.class); - verify(mockHelper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); - assertThat(pickerCaptor.getValue().pickSubchannel(mock(PickSubchannelArgs.class)).hasResult()) - .isFalse(); - - verify(mockProvider, never()).newLoadBalancer(any()); - assertThat(lazyLb.getDelegate()).isNull(); - assertThat(lazyLb.isConnectionRequested()).isFalse(); - } - - @Test - public void requestConnection_createsChildPolicy_forwardsAddresses_andRequestsConnection() { - lazyLb.acceptResolvedAddresses(resolvedAddresses); - assertThat(lazyLb.getDelegate()).isNull(); - - lazyLb.requestConnection(); - - assertThat(lazyLb.isConnectionRequested()).isTrue(); - verify(mockProvider).newLoadBalancer(mockHelper); - verify(mockDelegate).acceptResolvedAddresses(resolvedAddresses); - verify(mockDelegate).requestConnection(); - assertThat(lazyLb.getDelegate()).isSameInstanceAs(mockDelegate); - } - - @Test - public void requestConnection_beforeResolvedAddresses_createsPolicyWhenAddressesArrive() { - lazyLb.requestConnection(); - assertThat(lazyLb.isConnectionRequested()).isTrue(); - verify(mockProvider, never()).newLoadBalancer(any()); - - lazyLb.acceptResolvedAddresses(resolvedAddresses); - verify(mockProvider).newLoadBalancer(mockHelper); - verify(mockDelegate).acceptResolvedAddresses(resolvedAddresses); - verify(mockDelegate).requestConnection(); - } - - @Test - public void requestConnection_whenAlreadyCreated_delegatesRequestConnection() { - lazyLb.acceptResolvedAddresses(resolvedAddresses); - lazyLb.requestConnection(); - verify(mockDelegate, times(1)).requestConnection(); - - lazyLb.requestConnection(); - verify(mockDelegate, times(2)).requestConnection(); - } - - @Test - public void acceptResolvedAddresses_afterConnectionRequested_forwardsDirectly() { - lazyLb.acceptResolvedAddresses(resolvedAddresses); - lazyLb.requestConnection(); - verify(mockDelegate, times(1)).requestConnection(); - - ResolvedAddresses newAddresses = ResolvedAddresses.newBuilder() - .setAddresses(Collections.singletonList( - new EquivalentAddressGroup(new SocketAddress() {}))) - .setAttributes(Attributes.EMPTY) - .build(); - - lazyLb.acceptResolvedAddresses(newAddresses); - verify(mockDelegate).acceptResolvedAddresses(newAddresses); - // Should not request connection again on subsequent address update - verify(mockDelegate, times(1)).requestConnection(); - } - - @Test - public void acceptResolvedAddresses_null_throwsNullPointerException() { - assertThrows( - NullPointerException.class, - () -> lazyLb.acceptResolvedAddresses(null)); - } - - @Test - public void exitIdle_schedulesRequestConnectionOnSyncContext_once() throws Exception { - lazyLb.acceptResolvedAddresses(resolvedAddresses); - - int threadCount = 10; - ExecutorService executor = Executors.newFixedThreadPool(threadCount); - CountDownLatch startLatch = new CountDownLatch(1); - CountDownLatch doneLatch = new CountDownLatch(threadCount); - - for (int i = 0; i < threadCount; i++) { - executor.execute(() -> { - try { - startLatch.await(); - lazyLb.exitIdle(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } finally { - doneLatch.countDown(); - } - }); - } - - startLatch.countDown(); - assertThat(doneLatch.await(5, TimeUnit.SECONDS)).isTrue(); - executor.shutdown(); - - // Verify child policy instantiated exactly once despite concurrent exitIdle calls - verify(mockProvider, times(1)).newLoadBalancer(mockHelper); - verify(mockDelegate, times(1)).acceptResolvedAddresses(resolvedAddresses); - verify(mockDelegate, atLeastOnce()).requestConnection(); - } - - @Test - public void exitIdle_resetsFlagOnSyncContext_allowsSubsequentExitIdle() { - lazyLb.acceptResolvedAddresses(resolvedAddresses); - - lazyLb.exitIdle(); - verify(mockDelegate, times(1)).requestConnection(); - - // Subsequent exitIdle after syncContext execution should request connection again - lazyLb.exitIdle(); - verify(mockDelegate, times(2)).requestConnection(); - } - - @Test - public void handleNameResolutionError_null_throwsNullPointerException() { - assertThrows( - NullPointerException.class, - () -> lazyLb.handleNameResolutionError(null)); - } - - @Test - public void handleNameResolutionError_beforeConnectionRequested_reportsTransientFailure() { - Status error = Status.UNAVAILABLE.withDescription("dns failed"); - lazyLb.handleNameResolutionError(error); - - ArgumentCaptor pickerCaptor = - ArgumentCaptor.forClass(SubchannelPicker.class); - verify(mockHelper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - assertThat(pickerCaptor.getValue().pickSubchannel(mock(PickSubchannelArgs.class)).getStatus()) - .isEqualTo(error); - } - - @Test - public void handleNameResolutionError_afterConnectionRequested_forwardsToDelegate() { - lazyLb.acceptResolvedAddresses(resolvedAddresses); - lazyLb.requestConnection(); - - Status error = Status.UNAVAILABLE.withDescription("dns failed"); - lazyLb.handleNameResolutionError(error); - verify(mockDelegate).handleNameResolutionError(error); - } - - @Test - public void shutdown_cleansUpDelegateAndAddresses() { - lazyLb.acceptResolvedAddresses(resolvedAddresses); - assertThat(lazyLb.getLastResolvedAddresses()).isNotNull(); - lazyLb.requestConnection(); - assertThat(lazyLb.getDelegate()).isNotNull(); - - lazyLb.shutdown(); - verify(mockDelegate).shutdown(); - assertThat(lazyLb.getDelegate()).isNull(); - assertThat(lazyLb.getLastResolvedAddresses()).isNull(); - } - - @Test - public void operationsAfterShutdown_areNoOpsOrReturnError() { - lazyLb.acceptResolvedAddresses(resolvedAddresses); - lazyLb.shutdown(); - - Status status = lazyLb.acceptResolvedAddresses(resolvedAddresses); - assertThat(status.getCode()).isEqualTo(Status.Code.FAILED_PRECONDITION); - - // None of these should throw or create a child load balancer - lazyLb.requestConnection(); - lazyLb.exitIdle(); - lazyLb.handleNameResolutionError(Status.UNAVAILABLE); - verify(mockProvider, never()).newLoadBalancer(any()); - } - - @Test - public void toString_containsDebugFields() { - String str = lazyLb.toString(); - assertThat(str).contains("connectionRequested=false"); - assertThat(str).contains("shutdown=false"); - } -}