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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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"); - } -} From a92010a5dbec9e2bc02228da4b2ea29cce221365 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Mon, 21 Sep 2026 22:07:16 +0530 Subject: [PATCH 29/33] implementation of autoshardinglb and autosharding client and work on comments. --- autosharding/build.gradle | 3 + .../java/io/grpc/autosharding/Assignment.java | 134 +++ .../grpc/autosharding/AssignmentParser.java | 197 ++++ .../autosharding/AutoShardingAttributes.java | 55 + .../AutoShardingLoadBalancer.java | 452 ++++++++ .../AutoShardingLoadBalancerConfig.java | 110 ++ .../grpc/autosharding/AutoshardingClient.java | 390 +++++++ .../io/grpc/autosharding/ChannelFactory.java | 61 ++ .../io/grpc/autosharding/EndpointMap.java | 434 +++++--- .../autosharding/AssignmentParserTest.java | 359 +++++++ .../AutoShardingLoadBalancerTest.java | 961 ++++++++++++++++++ .../autosharding/AutoshardingClientTest.java | 527 ++++++++++ .../io/grpc/autosharding/EndpointMapTest.java | 602 ++++++++--- 13 files changed, 3985 insertions(+), 300 deletions(-) create mode 100644 autosharding/src/main/java/io/grpc/autosharding/Assignment.java create mode 100644 autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java create mode 100644 autosharding/src/main/java/io/grpc/autosharding/AutoShardingAttributes.java create mode 100644 autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java create mode 100644 autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerConfig.java create mode 100644 autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java create mode 100644 autosharding/src/main/java/io/grpc/autosharding/ChannelFactory.java create mode 100644 autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java create mode 100644 autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerTest.java create mode 100644 autosharding/src/test/java/io/grpc/autosharding/AutoshardingClientTest.java diff --git a/autosharding/build.gradle b/autosharding/build.gradle index 8c88f6d0e64..94565ef205c 100644 --- a/autosharding/build.gradle +++ b/autosharding/build.gradle @@ -60,6 +60,9 @@ tasks.named("javadoc").configure { exclude 'io/grpc/autosharding/*Provider.java' exclude 'io/grpc/autosharding/internal/**' exclude 'io/grpc/autosharding/Internal*' + // @Internal types, published only so that the xDS integration can inject them. + exclude 'io/grpc/autosharding/AutoShardingAttributes.java' + exclude 'io/grpc/autosharding/ChannelFactory*' } tasks.named("jacocoTestReport").configure { diff --git a/autosharding/src/main/java/io/grpc/autosharding/Assignment.java b/autosharding/src/main/java/io/grpc/autosharding/Assignment.java new file mode 100644 index 00000000000..0ef9784ca18 --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/Assignment.java @@ -0,0 +1,134 @@ +/* + * 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.base.MoreObjects; +import com.google.common.collect.ImmutableList; +import com.google.common.io.BaseEncoding; +import java.util.List; +import javax.annotation.Nullable; +import javax.annotation.concurrent.Immutable; + +/** + * An immutable, validated, gap-free snapshot of a logical assignment received from the + * autosharding service. + * + *

    Instances are produced exclusively by {@link AssignmentParser}, which guarantees the + * following invariants (see gRFC A119, "Contract of the AutoshardingClient"): + *

      + *
    1. The {@link #getSlices()} list covers the entire keyspace, starting at the minimum + * possible key (the empty byte string) and ending at the maximum possible key + * (infinity, represented by a {@code null} {@link Slice#getEndKey()}).
      2. + *
    2. The slices are sorted in ascending lexicographical (unsigned) order by + * {@link Slice#getStartKey()}.
      3. + *
    3. The partitioning is contiguous and non-overlapping: for every index {@code i} in + * {@code [0, N-2]}, {@code slices[i].endKey} is exactly {@code slices[i + 1].startKey}.
      4. + *
    4. Key ranges not assigned by the autosharding server are present as slices with an + * empty {@link Slice#getEndpoints()} list.
      5. + *
    + */ +@Immutable +final class Assignment { + + /** + * A single contiguous key range and the endpoints assigned to it. + */ + @Immutable + @SuppressWarnings("Immutable") // Defensive copies are made; arrays are never mutated. + static final class Slice { + private final byte[] startKey; + @Nullable private final byte[] endKey; + private final ImmutableList endpoints; + + /** + * Constructs a {@link Slice}. + * + * @param startKey the inclusive start key of the range + * @param endKey the exclusive end key of the range, or {@code null} for the infinity + * sentinel covering the largest allowed key + * @param endpoints indices into {@link Assignment#getEndpointNames()} assigned to this range + */ + Slice(byte[] startKey, @Nullable byte[] endKey, List endpoints) { + this.startKey = checkNotNull(startKey, "startKey").clone(); + this.endKey = endKey == null ? null : endKey.clone(); + this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints")); + } + + byte[] getStartKey() { + return startKey; + } + + @Nullable + byte[] getEndKey() { + return endKey; + } + + ImmutableList getEndpoints() { + return endpoints; + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("startKey", BaseEncoding.base16().encode(startKey)) + .add("endKey", endKey == null ? "inf" : BaseEncoding.base16().encode(endKey)) + .add("endpoints", endpoints) + .toString(); + } + } + + private final ImmutableList slices; + private final ImmutableList endpointNames; + private final long generation; + + /** + * Constructs an {@link Assignment}. + * + * @param slices the validated, sorted, contiguous and gap-free list of key-range slices + * @param endpointNames the complete list of endpoint names, combined across all chunks in + * chunk order + * @param generation the generation number of this logical assignment + */ + Assignment(List slices, List endpointNames, long generation) { + this.slices = ImmutableList.copyOf(checkNotNull(slices, "slices")); + this.endpointNames = ImmutableList.copyOf(checkNotNull(endpointNames, "endpointNames")); + this.generation = generation; + } + + ImmutableList getSlices() { + return slices; + } + + ImmutableList getEndpointNames() { + return endpointNames; + } + + long getGeneration() { + return generation; + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("generation", generation) + .add("endpointNames", endpointNames) + .add("slices", slices) + .toString(); + } +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java b/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java new file mode 100644 index 00000000000..712be2efdd0 --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java @@ -0,0 +1,197 @@ +/* + * 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.cloud.autosharding.v1.AssignmentChunk; +import com.google.cloud.autosharding.v1.EndpointState; +import com.google.cloud.autosharding.v1.PerSliceEndpointState; +import com.google.cloud.autosharding.v1.SliceAssignment; +import com.google.common.collect.ImmutableList; +import com.google.common.io.BaseEncoding; +import com.google.common.primitives.UnsignedBytes; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import javax.annotation.Nullable; + +/** + * Combines the {@link AssignmentChunk} messages of a single logical assignment into a + * validated, sorted, contiguous and gap-free {@link Assignment}. + * + *

    Validation follows gRFC A119, "Handling assignments from the Autosharding server": + *

      + *
    • Every endpoint index referenced by a slice must be valid once the endpoint names from + * all chunks are combined in chunk order.
      • + *
    • A slice's {@code startKey} must not be greater than its {@code endKey}.
      • + *
    • Key ranges must not overlap.
      • + *
    + * + *

    Gaps in the key ranges returned by the server are not validation failures. They are + * explicitly filled with slices containing no endpoints, so that RPCs matching them either fall + * back (when fallback is enabled) or fail. + */ +final class AssignmentParser { + + /** + * Thrown when an assignment received from the autosharding server fails validation. The + * message is suitable for use as the {@code error_message} of an {@code AssignmentAck}. + */ + static final class ValidationException extends Exception { + private static final long serialVersionUID = 0L; + + ValidationException(String message) { + super(message); + } + } + + private static final Comparator UNSIGNED_BYTES_COMPARATOR = + UnsignedBytes.lexicographicalComparator(); + private static final byte[] EMPTY_BYTES = new byte[0]; + + private AssignmentParser() {} + + /** + * Parses and validates the buffered chunks of a single logical assignment. + * + * @param chunks the chunks received since the last {@code AssignmentMetadata}, in the order + * they were received + * @param generation the generation number from the terminating {@code AssignmentMetadata} + * @return a validated, gap-free {@link Assignment} covering the entire keyspace + * @throws ValidationException if the assignment is invalid + */ + static Assignment parse(List chunks, long generation) + throws ValidationException { + checkNotNull(chunks, "chunks"); + + ImmutableList endpointNames = combineEndpointNames(chunks); + List slices = combineSlices(chunks, endpointNames.size()); + + slices.sort( + (s1, s2) -> UNSIGNED_BYTES_COMPARATOR.compare(s1.getStartKey(), s2.getStartKey())); + checkNoOverlaps(slices); + + return new Assignment(fillGaps(slices), endpointNames, generation); + } + + /** + * Concatenates the endpoint names across all chunks, in chunk order. Slice endpoint indices + * are defined against this combined list. + */ + private static ImmutableList combineEndpointNames(List chunks) { + ImmutableList.Builder names = ImmutableList.builder(); + for (AssignmentChunk chunk : chunks) { + for (EndpointState endpointState : chunk.getEndpointsList()) { + names.add(endpointState.getEndpoint()); + } + } + return names.build(); + } + + /** + * Concatenates the slice assignments across all chunks, validating endpoint indices and key + * range ordering along the way. Slice assignments may appear in any order across chunks. + */ + private static List combineSlices( + List chunks, int endpointCount) throws ValidationException { + List slices = new ArrayList<>(); + for (AssignmentChunk chunk : chunks) { + for (SliceAssignment sliceAssignment : chunk.getSliceAssignmentsList()) { + com.google.cloud.autosharding.v1.Slice slice = sliceAssignment.getSlice(); + byte[] startKey = slice.getStartKey().toByteArray(); + byte[] endKey = slice.hasEndKey() ? slice.getEndKey().toByteArray() : null; + + if (endKey != null && UNSIGNED_BYTES_COMPARATOR.compare(startKey, endKey) > 0) { + throw new ValidationException( + String.format( + "Slice has start_key %s greater than end_key %s", + encode(startKey), encode(endKey))); + } + + List endpoints = new ArrayList<>(sliceAssignment.getEndpointsCount()); + for (PerSliceEndpointState perSlice : sliceAssignment.getEndpointsList()) { + int index = perSlice.getEndpointIndex(); + if (index < 0 || index >= endpointCount) { + throw new ValidationException( + String.format( + "Slice starting at %s references out-of-range endpoint index %s;" + + " assignment contains %s endpoints", + encode(startKey), index, endpointCount)); + } + endpoints.add(index); + } + slices.add(new Assignment.Slice(startKey, endKey, endpoints)); + } + } + return slices; + } + + /** + * Verifies that no two slices in the sorted list cover the same key. + */ + private static void checkNoOverlaps(List sorted) throws ValidationException { + for (int i = 0; i + 1 < sorted.size(); i++) { + Assignment.Slice current = sorted.get(i); + Assignment.Slice next = sorted.get(i + 1); + if (current.getEndKey() == null) { + throw new ValidationException( + String.format( + "Slice starting at %s extends to the end of the keyspace but overlaps the slice" + + " starting at %s", + encode(current.getStartKey()), encode(next.getStartKey()))); + } + if (UNSIGNED_BYTES_COMPARATOR.compare(current.getEndKey(), next.getStartKey()) > 0) { + throw new ValidationException( + String.format( + "Slice [%s, %s) overlaps the slice starting at %s", + encode(current.getStartKey()), + encode(current.getEndKey()), + encode(next.getStartKey()))); + } + } + } + + /** + * Returns a contiguous list of slices covering {@code ["", inf)}, inserting endpoint-less + * slices wherever the sorted input leaves a gap. + */ + private static List fillGaps(List sorted) { + List filled = new ArrayList<>(sorted.size() + 1); + // Exclusive upper bound of the key range covered so far; null once infinity is reached. + byte[] cursor = EMPTY_BYTES; + for (Assignment.Slice slice : sorted) { + if (cursor == null) { + // Unreachable: checkNoOverlaps() rejects any slice following an infinity-ended slice. + break; + } + if (UNSIGNED_BYTES_COMPARATOR.compare(cursor, slice.getStartKey()) < 0) { + filled.add(new Assignment.Slice(cursor, slice.getStartKey(), ImmutableList.of())); + } + filled.add(slice); + cursor = slice.getEndKey(); + } + if (cursor != null) { + filled.add(new Assignment.Slice(cursor, null, ImmutableList.of())); + } + return filled; + } + + private static String encode(@Nullable byte[] key) { + return key == null ? "inf" : BaseEncoding.base16().encode(key); + } +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingAttributes.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingAttributes.java new file mode 100644 index 00000000000..a4b35b3f6bf --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingAttributes.java @@ -0,0 +1,55 @@ +/* + * 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.Attributes; +import io.grpc.EquivalentAddressGroup; +import io.grpc.Internal; + +/** + * Attribute keys used to inject data into the {@code autosharding_experimental} LB policy. + * + *

    Both keys are set on the resolver result by whoever is driving the policy: the + * {@code cds_experimental} LB policy in xDS deployments, or the application in non-xDS ones. + * They are internal to gRPC and carry no compatibility guarantee; the supported public API for + * configuring this policy is added separately. + */ +@Internal +public final class AutoShardingAttributes { + + /** + * Hostname associated with an endpoint, as described in gRFC A81. + * + *

    When absent, {@link EndpointMap} falls back to the string form of the endpoint's first + * address, per gRFC A119. The hostname is what assignments from the sharding service name + * their endpoints by, so it must match what that service reports. + */ + @EquivalentAddressGroup.Attr + public static final Attributes.Key ATTR_ENDPOINT_HOSTNAME = + Attributes.Key.create("io.grpc.autosharding.endpointHostname"); + + /** + * The "Channel Factory" used to create a channel to the sharding service. + * + *

    Supplied alongside the LB policy configuration, which carries only the opaque key that + * the factory resolves into a channel. + */ + public static final Attributes.Key ATTR_CHANNEL_FACTORY = + Attributes.Key.create("io.grpc.autosharding.channelFactory"); + + private AutoShardingAttributes() {} +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java new file mode 100644 index 00000000000..065c76a2d0c --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java @@ -0,0 +1,452 @@ +/* + * 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.CONNECTING; +import static io.grpc.ConnectivityState.TRANSIENT_FAILURE; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Stopwatch; +import com.google.common.base.Supplier; +import com.google.common.collect.ImmutableList; +import io.grpc.Channel; +import io.grpc.ConnectivityState; +import io.grpc.EquivalentAddressGroup; +import io.grpc.LoadBalancer; +import io.grpc.LoadBalancerProvider; +import io.grpc.LoadBalancerRegistry; +import io.grpc.Metadata; +import io.grpc.Status; +import io.grpc.SynchronizationContext; +import io.grpc.SynchronizationContext.ScheduledHandle; +import io.grpc.internal.BackoffPolicy; +import io.grpc.internal.ExponentialBackoffPolicy; +import io.grpc.internal.GrpcUtil; +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.logging.Level; +import java.util.logging.Logger; +import javax.annotation.Nullable; + +/** + * The {@code autosharding_experimental} load balancing policy. + * + *

    This policy shards RPCs across endpoints by an application-defined key carried in a request + * header. The mapping from key ranges to endpoints comes from an external sharding service, which + * an {@link AutoshardingClient} streams assignments from. See gRFC A119. + * + *

    Moving parts

    + * + *
      + *
    • {@link EndpointMap} owns one lazily-created {@code pick_first} child per resolved + * endpoint and assigns each a dense index. + *
    • {@link AutoshardingClient} produces validated {@link Assignment}s, which name endpoints + * by hostname. + *
    • {@link SliceMap} is the join of the two: the assignment's key ranges with hostnames + * translated into endpoint indices. It is rebuilt whenever either input changes. + *
    • {@link AutoShardingPicker} performs the per-RPC lookup against a {@link SliceMap} and a + * snapshot of endpoint states. It is rebuilt on every child state update too, reusing the + * existing {@link SliceMap} because the endpoint indices did not move. + *
    + * + *

    Startup

    + * + *

    Creating a channel to the sharding service starts the initial assignment timer. Until the + * first assignment arrives or that timer fires, RPCs are queued. Once the timer fires without an + * assignment, RPCs either spread across every resolved endpoint or fail outright, depending on + * {@code enable_fallback}. An assignment carried over from a previous channel keeps being used + * while the timer runs, so a change of sharding service does not interrupt traffic. + * + *

    Threading model

    + * + *

    All state lives on the {@link SynchronizationContext}. + */ +final class AutoShardingLoadBalancer extends LoadBalancer { + private static final Logger logger = + Logger.getLogger(AutoShardingLoadBalancer.class.getName()); + + /** + * Published while waiting for the first assignment. The delay type is consumed by the + * name-resolution delay tracking in gRFC A121. + */ + private static final SubchannelPicker ASSIGNMENT_PENDING_PICKER = + new FixedResultPicker( + PickResult.withNoResult( + "autosharding_assignment_pending", "Waiting for initial sharding assignment")); + + private final Helper helper; + private final SynchronizationContext syncContext; + private final ScheduledExecutorService timeService; + private final LoadBalancerProvider childProvider; + private final BackoffPolicy.Provider backoffPolicyProvider; + private final Supplier stopwatchSupplier; + + /** Identifies this client to the sharding service; stable across stream restarts. */ + private final String clientUuid; + + private final EndpointMap endpointMap; + + @Nullable private AutoShardingLoadBalancerConfig config; + @Nullable private Metadata.Key keyHeader; + + /** The factory last seen in the resolver attributes. */ + @Nullable private ChannelFactory channelFactory; + + /** Channel borrowed from {@link #channelFactory}; must be given back when we are done. */ + @Nullable private Channel shardingChannel; + + @Nullable private AutoshardingClient client; + + /** Most recent assignment accepted from the sharding service, retained across reconnects. */ + @Nullable private Assignment assignment; + + /** Join of {@link #assignment} and {@link #endpointMap}; null only before the first update. */ + @Nullable private SliceMap sliceMap; + + @Nullable private ScheduledHandle initialAssignmentTimer; + + /** + * True from the moment a channel to the sharding service is created until either an assignment + * arrives on it or {@link #initialAssignmentTimer} fires. Combined with a null + * {@link #assignment} it means RPCs must be queued rather than failed. + */ + private boolean awaitingInitialAssignment; + + private boolean shutdown; + + AutoShardingLoadBalancer(Helper helper) { + this( + helper, + LoadBalancerRegistry.getDefaultRegistry().getProvider("pick_first"), + new ExponentialBackoffPolicy.Provider(), + GrpcUtil.STOPWATCH_SUPPLIER, + UUID.randomUUID().toString()); + } + + /** + * Constructs a load balancer with injectable collaborators. + * + * @param childProvider provides the per-endpoint child load balancer, {@code pick_first} in + * production. {@link EndpointMap} takes care of deferring its instantiation, so this must + * not be wrapped in a {@link io.grpc.util.LazyLoadBalancer.Factory} by the caller + */ + @VisibleForTesting + AutoShardingLoadBalancer( + Helper helper, + LoadBalancerProvider childProvider, + BackoffPolicy.Provider backoffPolicyProvider, + Supplier stopwatchSupplier, + String clientUuid) { + this.helper = checkNotNull(helper, "helper"); + this.syncContext = helper.getSynchronizationContext(); + this.timeService = helper.getScheduledExecutorService(); + this.childProvider = checkNotNull(childProvider, "childProvider"); + this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider"); + this.stopwatchSupplier = checkNotNull(stopwatchSupplier, "stopwatchSupplier"); + this.clientUuid = checkNotNull(clientUuid, "clientUuid"); + this.endpointMap = new EndpointMap(helper, this.childProvider, this::onChildStateUpdate); + } + + @Override + public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) { + if (shutdown) { + return Status.OK; + } + Object rawConfig = resolvedAddresses.getLoadBalancingPolicyConfig(); + if (!(rawConfig instanceof AutoShardingLoadBalancerConfig)) { + return failPermanently("autosharding: missing or malformed load balancing configuration"); + } + AutoShardingLoadBalancerConfig newConfig = (AutoShardingLoadBalancerConfig) rawConfig; + + ChannelFactory factory = + resolvedAddresses.getAttributes().get(AutoShardingAttributes.ATTR_CHANNEL_FACTORY); + if (factory == null) { + return failPermanently("autosharding: no channel factory supplied to the LB policy"); + } + + List endpoints = resolvedAddresses.getAddresses(); + if (endpoints.isEmpty()) { + // Tear the children down so that in-flight picks stop resolving to endpoints the resolver + // has retracted. The assignment is kept: it stays valid if the endpoints come back. + endpointMap.updateEndpoints(ImmutableList.of(), resolvedAddresses.getAttributes()); + config = newConfig; + return failPermanently("autosharding: name resolver returned no endpoints"); + } + + Status channelStatus = updateShardingServiceChannel(factory, newConfig); + if (!channelStatus.isOk()) { + return channelStatus; + } + + if (config == null || !config.keyHeaderName.equals(newConfig.keyHeaderName)) { + keyHeader = AutoShardingPicker.createKeyHeader(newConfig.keyHeaderName); + } + config = newConfig; + + endpointMap.updateEndpoints(endpoints, resolvedAddresses.getAttributes()); + + // Safe to call on every update: the client compares against what it already has and only + // restarts the stream when the channel or the target actually changed. + client.update(shardingChannel, resolveTarget(newConfig, endpoints)); + + rebuildSliceMapAndPublish(); + return Status.OK; + } + + @Override + public void handleNameResolutionError(Status error) { + if (shutdown) { + return; + } + // Endpoints we already have remain usable; only report the failure if we have nothing. + if (endpointMap.size() > 0) { + logger.log(Level.FINE, "Ignoring name resolution error, endpoints still known: {0}", error); + return; + } + helper.updateBalancingState( + TRANSIENT_FAILURE, + new FixedResultPicker( + PickResult.withError( + error.getCode() == Status.Code.OK + ? Status.UNAVAILABLE.withDescription("autosharding: name resolution failed") + : error))); + } + + @Override + public void requestConnection() { + endpointMap.maybeWakeUpIdleEndpoint(); + } + + @Override + public void shutdown() { + if (shutdown) { + return; + } + shutdown = true; + cancelInitialAssignmentTimer(); + if (client != null) { + client.shutdown(); + client = null; + } + if (shardingChannel != null) { + channelFactory.releaseChannel(shardingChannel); + shardingChannel = null; + } + + endpointMap.shutdown(); + } + + /** + * Creates a channel to the sharding service if this is the first configuration update, or if + * the key or the factory changed. Doing so also restarts the initial assignment timer, per + * gRFC A119. + */ + private Status updateShardingServiceChannel( + ChannelFactory factory, AutoShardingLoadBalancerConfig newConfig) { + boolean keyChanged = + config == null || !config.channelFactoryKey.equals(newConfig.channelFactoryKey); + if (shardingChannel != null && factory == channelFactory && !keyChanged) { + return Status.OK; + } + + Channel newChannel; + try { + newChannel = factory.createChannel(newConfig.channelFactoryKey); + } catch (RuntimeException e) { + logger.log(Level.WARNING, "Failed to create a channel to the sharding service", e); + return failPermanently( + "autosharding: channel factory rejected key '" + + newConfig.channelFactoryKey + + "': " + + e.getMessage()); + } + + // Release through the factory that produced it, which is not necessarily the new one. + if (shardingChannel != null) { + channelFactory.releaseChannel(shardingChannel); + } + shardingChannel = newChannel; + channelFactory = factory; + + if (client == null) { + client = + new AutoshardingClient( + clientUuid, + syncContext, + timeService, + backoffPolicyProvider, + stopwatchSupplier, + this::onAssignment); + } + startInitialAssignmentTimer(newConfig.initialAssignmentTimeoutNanos); + return Status.OK; + } + + /** + * Substitutes the optional {@code %s} token in the configured target with the locality of the + * resolved endpoints, or with the empty string when no locality is available. + */ + private static String resolveTarget( + AutoShardingLoadBalancerConfig config, List endpoints) { + if (!config.autoshardingTarget.contains("%s")) { + return config.autoshardingTarget; + } + String locality = + endpoints.get(0).getAttributes().get(EquivalentAddressGroup.ATTR_LOCALITY_NAME); + return config.autoshardingTarget.replace("%s", locality == null ? "" : locality); + } + + private void startInitialAssignmentTimer(long timeoutNanos) { + cancelInitialAssignmentTimer(); + awaitingInitialAssignment = true; + initialAssignmentTimer = + syncContext.schedule( + this::onInitialAssignmentTimeout, timeoutNanos, TimeUnit.NANOSECONDS, timeService); + } + + private void cancelInitialAssignmentTimer() { + if (initialAssignmentTimer != null) { + initialAssignmentTimer.cancel(); + initialAssignmentTimer = null; + } + awaitingInitialAssignment = false; + } + + /** + * Gives up on hearing from the sharding service. Any queued RPCs are retried against whatever + * the current configuration allows: the full endpoint set if fallback is enabled, otherwise a + * failing picker. + */ + private void onInitialAssignmentTimeout() { + logger.log( + Level.WARNING, + "Timed out waiting for the initial assignment from the sharding service; " + + "proceeding {0} fallback", + config != null && config.enableFallback ? "with" : "without"); + awaitingInitialAssignment = false; + initialAssignmentTimer = null; + rebuildSliceMapAndPublish(); + } + + /** Called by {@link AutoshardingClient} on the synchronization context. */ + private void onAssignment(Assignment newAssignment) { + if (shutdown) { + return; + } + assignment = newAssignment; + cancelInitialAssignmentTimer(); + rebuildSliceMapAndPublish(); + } + + /** + * Called by {@link EndpointMap} when a child reports a new state or picker. The endpoint set + * and the indices into it are unchanged, so the existing {@link SliceMap} still applies and + * only the picker needs rebuilding. + */ + private void onChildStateUpdate() { + if (shutdown) { + return; + } + publishPicker(); + } + + private void rebuildSliceMapAndPublish() { + sliceMap = buildSliceMap(); + publishPicker(); + } + + /** + * Joins the current assignment with the current endpoints, translating the assignment's + * hostnames into endpoint indices. Hostnames the resolver has not given us are dropped, which + * can leave a slice with no endpoints; the picker treats such a slice as being in fallback. + * + *

    Before any assignment has been received the result has no slices, so every lookup misses + * and the picker routes through the fallback pool or fails, according to configuration. + */ + private SliceMap buildSliceMap() { + int endpointCount = endpointMap.size(); + List fallbackPool = new ArrayList<>(endpointCount); + for (int i = 0; i < endpointCount; i++) { + fallbackPool.add(i); + } + if (assignment == null) { + return new SliceMap(ImmutableList.of(), fallbackPool, 0); + } + + ImmutableList endpointNames = assignment.getEndpointNames(); + List entries = new ArrayList<>(assignment.getSlices().size()); + for (Assignment.Slice slice : assignment.getSlices()) { + List indices = new ArrayList<>(slice.getEndpoints().size()); + for (int nameIndex : slice.getEndpoints()) { + int endpointIndex = endpointMap.indexOf(endpointNames.get(nameIndex)); + if (endpointIndex != -1) { + indices.add(endpointIndex); + } + } + entries.add(new SliceMap.SliceEntry(slice.getStartKey(), indices)); + } + return new SliceMap(entries, fallbackPool, assignment.getGeneration()); + } + + private void publishPicker() { + if (shutdown || config == null) { + return; + } + if (endpointMap.size() == 0) { + // acceptResolvedAddresses already reported TRANSIENT_FAILURE for this case. + return; + } + if (awaitingInitialAssignment && assignment == null) { + helper.updateBalancingState(CONNECTING, ASSIGNMENT_PENDING_PICKER); + return; + } + + ConnectivityState state = endpointMap.aggregateConnectivityState(); + helper.updateBalancingState( + state, + new AutoShardingPicker( + sliceMap, endpointMap.toPickerEndpoints(), config.enableFallback, keyHeader)); + + // Nothing else will drive progress: this policy only connects in response to picks, so a + // CONNECTING or TRANSIENT_FAILURE aggregate could otherwise stick with no attempt in flight. + // The woken endpoint reports CONNECTING synchronously, re-entering publishPicker() once to + // publish the fresher picker; that pass finds an endpoint CONNECTING and wakes no one else. + if (state == CONNECTING || state == TRANSIENT_FAILURE) { + endpointMap.maybeWakeUpIdleEndpoint(); + } + } + + @VisibleForTesting + EndpointMap getEndpointMap() { + return endpointMap; + } + + /** + * Reports TRANSIENT_FAILURE with a picker that fails every RPC, and returns the same error for + * {@link #acceptResolvedAddresses} to hand back to the channel. + */ + private Status failPermanently(String description) { + Status error = Status.UNAVAILABLE.withDescription(description); + helper.updateBalancingState(TRANSIENT_FAILURE, new FixedResultPicker(PickResult.withError( + error))); + return error; + } +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerConfig.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerConfig.java new file mode 100644 index 00000000000..9648e153fa6 --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerConfig.java @@ -0,0 +1,110 @@ +/* + * 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.base.MoreObjects; +import java.util.Objects; +import java.util.concurrent.TimeUnit; + +/** + * Configuration for the {@code autosharding_experimental} LB policy, as specified by + * {@code AutoshardingLbConfig} in gRFC A119. + * + *

    Parsing this out of service config JSON belongs to the LB policy provider, which is added + * along with the policy's public API. + */ +final class AutoShardingLoadBalancerConfig { + + /** Default for {@link #initialAssignmentTimeoutNanos} when the field is unset. */ + static final long DEFAULT_INITIAL_ASSIGNMENT_TIMEOUT_NANOS = TimeUnit.SECONDS.toNanos(60); + + /** Opaque key passed to the "Channel Factory" to reach the sharding service. */ + final String channelFactoryKey; + + /** + * Identifies the assignments this client should receive. + * + *

    May contain a single {@code %s} token, which the LB policy replaces with the locality + * before sending it to the sharding service, or with the empty string when no locality is + * available. + */ + final String autoshardingTarget; + + /** + * Name of the request header holding the application-defined sharding key. Empty means every + * RPC is treated as having an empty key. + */ + final String keyHeaderName; + + /** Whether RPCs may fall back to the full set of resolved endpoints. */ + final boolean enableFallback; + + /** How long to wait for the first assignment after creating a channel to the service. */ + final long initialAssignmentTimeoutNanos; + + AutoShardingLoadBalancerConfig( + String channelFactoryKey, + String autoshardingTarget, + String keyHeaderName, + boolean enableFallback, + long initialAssignmentTimeoutNanos) { + this.channelFactoryKey = checkNotNull(channelFactoryKey, "channelFactoryKey"); + this.autoshardingTarget = checkNotNull(autoshardingTarget, "autoshardingTarget"); + this.keyHeaderName = checkNotNull(keyHeaderName, "keyHeaderName"); + this.enableFallback = enableFallback; + this.initialAssignmentTimeoutNanos = initialAssignmentTimeoutNanos; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof AutoShardingLoadBalancerConfig)) { + return false; + } + AutoShardingLoadBalancerConfig that = (AutoShardingLoadBalancerConfig) o; + return enableFallback == that.enableFallback + && initialAssignmentTimeoutNanos == that.initialAssignmentTimeoutNanos + && channelFactoryKey.equals(that.channelFactoryKey) + && autoshardingTarget.equals(that.autoshardingTarget) + && keyHeaderName.equals(that.keyHeaderName); + } + + @Override + public int hashCode() { + return Objects.hash( + channelFactoryKey, + autoshardingTarget, + keyHeaderName, + enableFallback, + initialAssignmentTimeoutNanos); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("channelFactoryKey", channelFactoryKey) + .add("autoshardingTarget", autoshardingTarget) + .add("keyHeaderName", keyHeaderName) + .add("enableFallback", enableFallback) + .add("initialAssignmentTimeoutNanos", initialAssignmentTimeoutNanos) + .toString(); + } +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java b/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java new file mode 100644 index 00000000000..4a6970dc003 --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java @@ -0,0 +1,390 @@ +/* + * 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 com.google.common.base.Preconditions.checkState; + +import com.google.cloud.autosharding.v1.AssignmentAck; +import com.google.cloud.autosharding.v1.AssignmentChunk; +import com.google.cloud.autosharding.v1.AutoshardingServiceGrpc; +import com.google.cloud.autosharding.v1.InitialClientConfig; +import com.google.cloud.autosharding.v1.WatchShardingAssignmentRequest; +import com.google.cloud.autosharding.v1.WatchShardingAssignmentResponse; +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Stopwatch; +import com.google.common.base.Supplier; +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.Status; +import io.grpc.SynchronizationContext; +import io.grpc.SynchronizationContext.ScheduledHandle; +import io.grpc.internal.BackoffPolicy; +import io.grpc.stub.ClientCallStreamObserver; +import io.grpc.stub.ClientCalls; +import io.grpc.stub.ClientResponseObserver; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.logging.Level; +import java.util.logging.Logger; +import javax.annotation.Nullable; +import javax.annotation.concurrent.NotThreadSafe; + +/** + * Encapsulates all communication with an external autosharding service over the + * {@code WatchShardingAssignment} streaming protocol. + * + *

    This component owns the stream lifecycle, buffers and reassembles chunked assignments, + * validates them, acknowledges them, and hands validated {@link Assignment}s to the parent load + * balancer. See gRFC A119, "Communicating with the Autosharding service". + * + *

    Threading model: This class is not thread-safe. All public methods must be invoked from the + * {@link SynchronizationContext} supplied at construction, and all callbacks to the + * {@link AssignmentWatcher} are delivered on that same context. + */ +@NotThreadSafe +final class AutoshardingClient { + private static final Logger logger = Logger.getLogger(AutoshardingClient.class.getName()); + + /** + * Receives validated assignments from the autosharding service. + */ + interface AssignmentWatcher { + /** + * Called with a newly accepted assignment. Invoked on the {@link SynchronizationContext}. + */ + void onAssignment(Assignment assignment); + } + + private final SynchronizationContext syncContext; + private final ScheduledExecutorService timerService; + private final BackoffPolicy.Provider backoffPolicyProvider; + private final Stopwatch retryStopwatch; + private final AssignmentWatcher watcher; + private final String clientUuid; + + @Nullable private Channel channel; + @Nullable private String target; + + /** + * Generation of the most recent accepted assignment. Sent to the server so that it can skip + * resending an assignment the client already has. Reset to zero whenever the channel or the + * target changes, because the stored value is meaningless against a different sharding server + * or a different resource. + */ + private long latestGeneration; + + @Nullable private BackoffPolicy retryBackoffPolicy; + @Nullable private ScheduledHandle retryTimer; + @Nullable private AutoshardingStream stream; + private boolean shutdown; + + /** + * Constructs an {@link AutoshardingClient}. No stream is created until + * {@link #update(Channel, String)} supplies a channel and a target. + * + * @param clientUuid a UUID generated once by the parent load balancer and reused across all + * stream restarts + * @param syncContext the context on which all state is mutated and callbacks are delivered + * @param timerService used to schedule stream retries + * @param backoffPolicyProvider supplies the exponential backoff sequence for stream retries + * @param stopwatchSupplier supplies the stopwatch measuring time spent in a stream attempt + * @param watcher receives validated assignments + */ + AutoshardingClient( + String clientUuid, + SynchronizationContext syncContext, + ScheduledExecutorService timerService, + BackoffPolicy.Provider backoffPolicyProvider, + Supplier stopwatchSupplier, + AssignmentWatcher watcher) { + this.clientUuid = checkNotNull(clientUuid, "clientUuid"); + this.syncContext = checkNotNull(syncContext, "syncContext"); + this.timerService = checkNotNull(timerService, "timerService"); + this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider"); + this.retryStopwatch = checkNotNull(stopwatchSupplier, "stopwatchSupplier").get(); + this.watcher = checkNotNull(watcher, "watcher"); + } + + /** + * Applies a new channel and/or resolved autosharding target. + * + *

    If either changed, any existing stream is torn down, the stored generation number is + * discarded, and a new stream is started immediately. A stored generation number is only + * meaningful for the combination of sharding server and target that produced it; retaining it + * across a change could cause the server to withhold assignments indefinitely. + * + * @param channel the channel to the sharding service, created via the "Channel Factory" + * @param target the autosharding target, with any {@code %s} token already substituted + */ + void update(Channel channel, String target) { + syncContext.throwIfNotInThisSynchronizationContext(); + checkNotNull(channel, "channel"); + checkNotNull(target, "target"); + if (shutdown) { + return; + } + if (channel.equals(this.channel) && target.equals(this.target)) { + return; + } + this.channel = channel; + this.target = target; + this.latestGeneration = 0; + this.retryBackoffPolicy = null; + restartStream(); + } + + /** + * Cancels any in-flight stream and pending retry. The channel is not shut down, because it is + * owned by the parent load balancer. + */ + void shutdown() { + syncContext.throwIfNotInThisSynchronizationContext(); + if (shutdown) { + return; + } + shutdown = true; + cancelRetryTimer(); + if (stream != null) { + stream.close(Status.CANCELLED.withDescription("AutoshardingClient shutdown")); + stream = null; + } + } + + @VisibleForTesting + long getLatestGeneration() { + return latestGeneration; + } + + private void restartStream() { + cancelRetryTimer(); + if (stream != null) { + stream.close(Status.CANCELLED.withDescription("stream restarted")); + stream = null; + } + startStream(); + } + + private void startStream() { + if (shutdown || channel == null || target == null) { + return; + } + checkState(stream == null, "previous stream has not been cleared yet"); + retryStopwatch.reset().start(); + stream = new AutoshardingStream(); + stream.start(); + } + + private void cancelRetryTimer() { + if (retryTimer != null) { + if (retryTimer.isPending()) { + retryTimer.cancel(); + } + retryTimer = null; + } + } + + /** + * Schedules the next stream attempt. Per gRFC A119, backoff only applies to streams that closed + * without delivering a good logical assignment; the backoff sequence is reset as soon as one is + * received. + */ + private void scheduleRetry(boolean receivedGoodAssignment) { + if (shutdown) { + return; + } + if (receivedGoodAssignment || retryBackoffPolicy == null) { + retryBackoffPolicy = backoffPolicyProvider.get(); + } + // The backoff sequence bounds the interval between consecutive stream starts, so the actual + // delay is reduced by however long the previous attempt lasted. The retry always goes through + // the timer service, even when no delay remains, so that a channel failing calls synchronously + // cannot drive unbounded recursion between startStream() and handleStreamClosed(). + long delayNanos = + Math.max( + 0, + retryBackoffPolicy.nextBackoffNanos() - retryStopwatch.elapsed(TimeUnit.NANOSECONDS)); + retryTimer = + syncContext.schedule(this::startStream, delayNanos, TimeUnit.NANOSECONDS, timerService); + } + + /** + * A single {@code WatchShardingAssignment} stream. + */ + private final class AutoshardingStream + implements ClientResponseObserver< + WatchShardingAssignmentRequest, WatchShardingAssignmentResponse> { + + /** + * Chunks received since the last {@code AssignmentMetadata}. A chunk's slices reference + * endpoint indices into the list combined across all chunks, so chunks cannot be used until + * the assignment is terminated by an {@code AssignmentMetadata} message. + */ + private final List bufferedChunks = new ArrayList<>(); + + @Nullable private ClientCallStreamObserver requestStream; + private boolean receivedGoodAssignment; + private boolean closed; + + @Override + public void beforeStart( + ClientCallStreamObserver requestStream) { + this.requestStream = requestStream; + } + + void start() { + // wait_for_ready keeps the stream pending through transient connectivity failures instead + // of failing it, which recovers faster than applying backoff around stream creation. + ClientCalls.asyncBidiStreamingCall( + channel.newCall( + AutoshardingServiceGrpc.getWatchShardingAssignmentMethod(), + CallOptions.DEFAULT.withWaitForReady()), + this); + sendInitialClientConfig(); + } + + private void sendInitialClientConfig() { + WatchShardingAssignmentRequest request = + WatchShardingAssignmentRequest.newBuilder() + .setInitialClientConfig( + InitialClientConfig.newBuilder() + .setTarget(target) + .setClientUuid(clientUuid) + .setLatestGeneration(latestGeneration)) + .build(); + requestStream.onNext(request); + } + + @Override + public void onNext(WatchShardingAssignmentResponse response) { + syncContext.execute(() -> handleResponse(response)); + } + + @Override + public void onError(Throwable t) { + syncContext.execute(() -> handleStreamClosed(Status.fromThrowable(t))); + } + + @Override + public void onCompleted() { + syncContext.execute( + () -> + handleStreamClosed( + Status.UNAVAILABLE.withDescription("autosharding stream closed by server"))); + } + + private void handleResponse(WatchShardingAssignmentResponse response) { + if (closed) { + return; + } + if (response.hasChunk()) { + bufferedChunks.add(response.getChunk()); + } else if (response.hasMetadata()) { + handleAssignmentComplete(response.getMetadata().getGeneration()); + } + // LoadReportingConfig is intentionally ignored; load reporting is not yet supported. + } + + /** + * Reassembles, validates and acknowledges the buffered chunks terminated by an + * {@code AssignmentMetadata} message. + */ + private void handleAssignmentComplete(long generation) { + List chunks = new ArrayList<>(bufferedChunks); + bufferedChunks.clear(); + + // Generations are monotonically increasing. Anything we have already seen is stale, and is + // dropped without acknowledgement. + if (generation <= latestGeneration) { + logger.log( + Level.FINE, + "Dropping autosharding assignment with stale generation {0}; latest is {1}", + new Object[] {generation, latestGeneration}); + return; + } + + Assignment assignment; + try { + assignment = AssignmentParser.parse(chunks, generation); + } catch (AssignmentParser.ValidationException e) { + logger.log( + Level.WARNING, + "Rejecting autosharding assignment with generation {0}: {1}", + new Object[] {generation, e.getMessage()}); + sendAck(generation, false, e.getMessage()); + return; + } + + sendAck(generation, true, null); + latestGeneration = generation; + if (!receivedGoodAssignment) { + receivedGoodAssignment = true; + } + watcher.onAssignment(assignment); + } + + private void sendAck(long generation, boolean accepted, @Nullable String errorMessage) { + AssignmentAck.Builder ack = + AssignmentAck.newBuilder().setGeneration(generation).setAccepted(accepted); + if (errorMessage != null) { + ack.setErrorMessage(truncateErrorMessage(errorMessage)); + } + requestStream.onNext( + WatchShardingAssignmentRequest.newBuilder().setAssignmentAck(ack).build()); + } + + private void handleStreamClosed(Status status) { + if (closed) { + return; + } + closed = true; + logger.log( + Level.FINE, + "Autosharding stream closed with status {0}: {1}", + new Object[] {status.getCode(), status.getDescription()}); + bufferedChunks.clear(); + if (stream == this) { + stream = null; + scheduleRetry(receivedGoodAssignment); + } + } + + /** + * Cancels the stream without scheduling a retry. Used when the client is shutting down or + * when the configuration changed and a fresh stream is being created. + */ + void close(Status status) { + if (closed) { + return; + } + closed = true; + bufferedChunks.clear(); + if (requestStream != null) { + requestStream.cancel(status.getDescription(), status.getCause()); + } + } + } + + /** + * The {@code error_message} field must not exceed 512 characters. + */ + private static String truncateErrorMessage(String message) { + return message.length() <= 512 ? message : message.substring(0, 512); + } +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/ChannelFactory.java b/autosharding/src/main/java/io/grpc/autosharding/ChannelFactory.java new file mode 100644 index 00000000000..5170472e944 --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/ChannelFactory.java @@ -0,0 +1,61 @@ +/* + * 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.Channel; +import io.grpc.Internal; + +/** + * Creates channels to the autosharding service. + * + *

    The LB policy configuration carries only an opaque {@code channel_factory_key}; the factory + * is responsible for translating that key into a fully configured channel. Credentials and + * per-request metadata are deliberately kept out of the configuration so that a compromised + * control plane cannot escalate privileges, per gRFC A102. Implementations must therefore ensure + * the key uniquely encodes every parameter needed to create the channel. + * + *

    Channels are borrowed rather than owned: implementations may return the same underlying + * channel for repeated calls with the same key, so a caller must never shut one down directly + * and must instead hand it back with {@link #releaseChannel}. + * + *

    Injected into the LB policy through + * {@link AutoShardingAttributes#ATTR_CHANNEL_FACTORY}. In xDS deployments the + * {@code cds_experimental} LB policy supplies it; in non-xDS deployments the application does. + * + *

    See gRFC A119, "Creating a gRPC Channel to the Autosharding Service". + */ +@Internal +public interface ChannelFactory { + + /** + * Returns a channel to the sharding service identified by {@code channelFactoryKey}. + * + *

    The caller must pass the returned channel to {@link #releaseChannel} exactly once when it + * is done with it. + * + * @throws IllegalArgumentException if the key is not recognized + */ + Channel createChannel(String channelFactoryKey); + + /** + * Gives back a channel previously obtained from {@link #createChannel} on this same factory. + * + *

    This releases the caller's claim on the channel. Whether the channel is actually shut + * down is up to the implementation, since it may still be lent out elsewhere. + */ + void releaseChannel(Channel channel); +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java index 2edf86e9599..f61c457afb5 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java @@ -17,240 +17,398 @@ 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.CONNECTING; import static io.grpc.ConnectivityState.IDLE; +import static io.grpc.ConnectivityState.READY; +import static io.grpc.ConnectivityState.TRANSIENT_FAILURE; 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; 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.Status; +import io.grpc.SynchronizationContext; import io.grpc.util.ForwardingLoadBalancerHelper; import io.grpc.util.LazyLoadBalancer; -import java.util.Collection; +import java.util.ArrayList; +import java.util.HashMap; 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 java.util.logging.Level; +import java.util.logging.Logger; import javax.annotation.concurrent.NotThreadSafe; /** - * Manages the mapping from endpoint hostname to {@link EndpointHolder} and coordinates - * child load balancer lifecycle and connectivity state updates. + * Owns one lazily-created {@code pick_first} child load balancer per resolved endpoint, keyed by + * endpoint hostname, and tracks the connectivity state and picker most recently reported by each + * child. * - *

    Threading model: This class is not thread-safe. All methods must be invoked from the - * {@link io.grpc.SynchronizationContext} by the parent load balancer. + *

    Endpoint indices

    + * + *

    Endpoints are identified throughout the LB policy by a dense index in {@code [0, size)}. + * The index of an endpoint is simply its position in the list most recently passed to + * {@link #updateEndpoints}, after duplicate hostnames have been dropped. Indices are not stored + * anywhere; they are a property of the map's iteration order. This makes it impossible for + * indices handed out by {@link #indexOf} to disagree with the positions in the list returned by + * {@link #toPickerEndpoints}, as long as both are obtained without an intervening + * {@link #updateEndpoints} call. The LB policy relies on that pairing when it builds a + * {@link SliceMap} and an {@link AutoShardingPicker} from the same snapshot. + * + *

    Note that gRFC A119 derives the index from the position in the resolver's endpoint list + * before de-duplication, which can leave gaps when two endpoints share a hostname. We + * index after de-duplication instead, so the index is always a valid offset into + * {@link #toPickerEndpoints}. + * + *

    Lifecycle

    + * + *

    gRFC A119 describes building a brand new map on every resolver update. This class instead + * keeps one long-lived instance and rebuilds its contents in {@link #updateEndpoints}, which is + * the only method that changes the set of endpoints or their indices. Retaining the instance + * lets child load balancers — and therefore established connections — survive a resolver update + * that merely adds or removes unrelated endpoints. + * + *

    Threading model

    + * + *

    This class is not thread-safe. Every method must be called from the + * {@link SynchronizationContext} of the {@link Helper} supplied at construction. The sole + * exception is {@link PickerEndpoint#requestConnection}, reached from RPC threads through the + * snapshots returned by {@link #toPickerEndpoints}; it hops onto the synchronization context + * before touching any state here. */ @NotThreadSafe final class EndpointMap { - private final Map map = new LinkedHashMap<>(); + private static final Logger logger = Logger.getLogger(EndpointMap.class.getName()); - @Nullable - EndpointHolder get(String hostname) { - return map.get(checkNotNull(hostname, "hostname")); - } + private final Helper helper; + private final LoadBalancerProvider childProvider; + private final Runnable childStateListener; - void put(String hostname, EndpointHolder holder) { - map.put(checkNotNull(hostname, "hostname"), checkNotNull(holder, "holder")); - } + // The endpoints, in index order: an endpoint's index is its position here, never stored. + // Rebuilt wholesale by updateEndpoints. + private final List holders = new ArrayList<>(); - @Nullable - EndpointHolder remove(String hostname) { - return map.remove(checkNotNull(hostname, "hostname")); - } + // Hostname to its position in holders. Derived from holders and rebuilt with it; exists so + // that translating an assignment's hostnames into indices stays linear in the assignment + // size, rather than scanning the endpoints once per name. + private final Map indexByHostname = new HashMap<>(); - Collection values() { - return map.values(); - } + /** + * Set while {@link #updateEndpoints} is running. Children report a state synchronously from + * within that method, and forwarding those reports would make the LB policy publish a picker + * built from a half-rebuilt map. gRFC A119 avoids this by building a whole new map and + * swapping it in; rebuilding in place is what makes the flag necessary. + */ + private boolean rebuilding; - Set keySet() { - return map.keySet(); + /** + * Constructs an empty map. + * + * @param helper the parent LB policy's helper, used for its synchronization context and passed + * through to child load balancers + * @param childProvider provides the per-endpoint child load balancer, normally {@code + * pick_first}. It is wrapped in a {@link LazyLoadBalancer} here, so the child is not + * instantiated, and therefore does not start connecting, until a pick asks for it + * @param childStateListener run after a child reports a new connectivity state or picker. + * Invoked on the synchronization context, never during {@link #updateEndpoints} or after + * {@link #shutdown} + */ + EndpointMap(Helper helper, LoadBalancerProvider childProvider, Runnable childStateListener) { + this.helper = checkNotNull(helper, "helper"); + this.childProvider = checkNotNull(childProvider, "childProvider"); + this.childStateListener = checkNotNull(childStateListener, "childStateListener"); } - int size() { - return map.size(); + /** + * Replaces the set of endpoints, assigning each a new index. + * + *

    An endpoint whose hostname appears in both the old and the new set keeps its child load + * balancer, along with its connections and last reported state; only its addresses and index + * are refreshed. Endpoints that disappear have their child load balancers shut down. New + * endpoints start out IDLE with no child load balancer instantiated. + * + *

    If several endpoints resolve to the same hostname, the first one wins and the rest are + * dropped, as permitted by gRFC A119. + * + * @param endpoints the endpoints from the resolver, in the order the resolver supplied them + * @param attributes the resolver attributes, forwarded to every child load balancer + */ + void updateEndpoints(List endpoints, Attributes attributes) { + Map addressesByHostname = new LinkedHashMap<>(); + for (EquivalentAddressGroup endpoint : endpoints) { + String hostname = hostnameOf(endpoint); + if (addressesByHostname.putIfAbsent(hostname, endpoint) != null) { + logger.log(Level.FINE, "Dropping duplicate endpoint for hostname {0}", hostname); + } + } + + // Children of endpoints the resolver no longer reports are shut down and dropped. + Map survivors = new HashMap<>(); + for (EndpointHolder holder : holders) { + if (addressesByHostname.containsKey(holder.hostname)) { + survivors.put(holder.hostname, holder); + } else { + holder.shutdown(); + } + } + + rebuilding = true; + try { + holders.clear(); + indexByHostname.clear(); + for (Map.Entry entry : addressesByHostname.entrySet()) { + String hostname = entry.getKey(); + EndpointHolder holder = survivors.get(hostname); + if (holder == null) { + holder = new EndpointHolder(hostname); + } + indexByHostname.put(hostname, holders.size()); + holders.add(holder); + holder.updateAddresses(entry.getValue(), attributes); + } + } finally { + rebuilding = false; + } } - boolean isEmpty() { - return map.isEmpty(); + /** Returns the number of endpoints currently held. */ + int size() { + return holders.size(); } - void clear() { - map.clear(); + /** + * Returns the index of {@code hostname}, or {@code -1} if no endpoint with that hostname is + * currently held. Used to translate the hostnames in an {@link Assignment} into the indices + * that {@link SliceMap} and {@link AutoShardingPicker} work with. + */ + int indexOf(String hostname) { + Integer index = indexByHostname.get(hostname); + return index == null ? -1 : index; } /** - * Re-assigns contiguous 0-based index values across all current endpoint holders. + * Returns an immutable snapshot of the current endpoint states, where element {@code i} + * describes the endpoint with index {@code i}. + * + *

    The snapshot is safe to hand to a picker running on RPC threads: it captures the + * connectivity state and picker by value, and reaches back into this class only through + * {@link PickerEndpoint#requestConnection}. */ - void reindex() { - int nextIdx = 0; - for (EndpointHolder holder : map.values()) { - holder.setIndex(nextIdx++); + ImmutableList toPickerEndpoints() { + ImmutableList.Builder snapshot = + ImmutableList.builderWithExpectedSize(holders.size()); + for (EndpointHolder holder : holders) { + snapshot.add(holder.toPickerEndpoint()); } + return snapshot.build(); } /** - * Shuts down all child load balancers and clears the map. + * Returns the aggregated connectivity state to report for the channel, using the {@code + * ring_hash} rules from gRFC A42 that gRFC A119 adopts: + * + *

      + *
    1. at least one endpoint READY, report READY; + *
    2. two or more endpoints TRANSIENT_FAILURE, report TRANSIENT_FAILURE; + *
    3. at least one endpoint CONNECTING, report CONNECTING; + *
    4. exactly one endpoint TRANSIENT_FAILURE and more than one endpoint, report CONNECTING; + *
    5. at least one endpoint IDLE, report IDLE; + *
    6. otherwise report TRANSIENT_FAILURE. + *
    + * + *

    An empty map reports TRANSIENT_FAILURE, matching rule 6. */ - void shutdownAll() { - for (EndpointHolder holder : map.values()) { - holder.shutdown(); + ConnectivityState aggregateConnectivityState() { + int connecting = 0; + int idle = 0; + int transientFailure = 0; + for (EndpointHolder holder : holders) { + switch (holder.state) { + case READY: + return READY; + case CONNECTING: + connecting++; + break; + case IDLE: + idle++; + break; + case TRANSIENT_FAILURE: + transientFailure++; + break; + default: + break; + } + } + if (transientFailure >= 2) { + return TRANSIENT_FAILURE; } - map.clear(); + if (connecting > 0) { + return CONNECTING; + } + if (transientFailure == 1 && holders.size() > 1) { + return CONNECTING; + } + if (idle > 0) { + return IDLE; + } + return TRANSIENT_FAILURE; } /** - * Builds an immutable snapshot list of {@link PickerEndpoint}s placed strictly at their - * corresponding {@link EndpointHolder#getIndex()} positions. + * Starts connecting on one IDLE endpoint, unless some endpoint is already CONNECTING or none + * is IDLE. * - * @throws IllegalStateException if endpoint indices are not contiguous from 0 to N-1 + *

    Because this policy only connects in response to picks, an aggregated state of CONNECTING + * or TRANSIENT_FAILURE could otherwise persist with nothing in flight to resolve it. gRFC A119 + * therefore has the policy nudge a single endpoint after every child state update and resolver + * update. Which endpoint is chosen does not matter; this picks the lowest-indexed IDLE one. */ - ImmutableList toPickerEndpoints() { - int size = map.size(); - if (size == 0) { - return ImmutableList.of(); + void maybeWakeUpIdleEndpoint() { + EndpointHolder firstIdle = null; + for (EndpointHolder holder : holders) { + if (holder.state == CONNECTING) { + return; + } + if (firstIdle == null && holder.state == IDLE) { + firstIdle = holder; + } + } + if (firstIdle != null) { + firstIdle.requestConnection(); } - PickerEndpoint[] array = new PickerEndpoint[size]; - for (EndpointHolder holder : map.values()) { - 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(); + } + + /** Shuts down every child load balancer and empties the map. Idempotent. */ + void shutdown() { + for (EndpointHolder holder : holders) { + holder.shutdown(); } - return ImmutableList.copyOf(array); + holders.clear(); + indexByHostname.clear(); + } + + /** + * Returns the hostname identifying {@code endpoint}. Falls back to the endpoint's first + * address when the hostname attribute from gRFC A81 is absent, per gRFC A119. + */ + private static String hostnameOf(EquivalentAddressGroup endpoint) { + String hostname = endpoint.getAttributes().get(AutoShardingAttributes.ATTR_ENDPOINT_HOSTNAME); + return hostname != null ? hostname : endpoint.getAddresses().get(0).toString(); } @Override public String toString() { - return MoreObjects.toStringHelper(this) - .add("map", map) - .toString(); + return MoreObjects.toStringHelper(this).add("endpoints", holders).toString(); } /** - * Holds the connectivity state, picker, and lazy child load balancer for a single endpoint. + * The child load balancer for a single endpoint, together with the connectivity state and + * picker it most recently reported. */ - static final class EndpointHolder { - private int index; + private final class EndpointHolder { + private final String hostname; 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()); + private boolean childShutdown; - EndpointHolder( - int index, - Helper helper, - LoadBalancer.Factory pickFirstFactory, - @Nullable Runnable stateUpdateCallback) { - this.index = index; - this.helper = checkNotNull(helper, "helper"); - this.childLb = new LazyLoadBalancer( - new ChildHelper(helper, stateUpdateCallback), - checkNotNull(pickFirstFactory, "pickFirstFactory")); - } - - int getIndex() { - return index; - } - - void setIndex(int index) { - this.index = index; - } - - ConnectivityState getState() { - return state; - } - - SubchannelPicker getPicker() { - return picker; - } - - LazyLoadBalancer getChildLb() { - return childLb; + EndpointHolder(String hostname) { + this.hostname = hostname; + this.childLb = new LazyLoadBalancer(new ChildHelper(), childProvider); } + /** Captures the current state for use by a picker on RPC threads. */ PickerEndpoint toPickerEndpoint() { 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(EquivalentAddressGroup endpoint, Attributes attributes) { + Status status = + childLb.acceptResolvedAddresses( + ResolvedAddresses.newBuilder() + .setAddresses(ImmutableList.of(endpoint)) + .setAttributes(attributes) + .build()); + if (!status.isOk()) { + // pick_first only rejects an address list it cannot use at all, which should not happen + // for the single well-formed endpoint we pass. Report it rather than silently dropping + // it; the endpoint simply stays in whatever state it was already in. + logger.log( + Level.WARNING, + "Child load balancer for endpoint {0} rejected its addresses: {1}", + new Object[] {hostname, status}); } } - void updateAddresses(List eags, Attributes attributes) { - ResolvedAddresses childAddresses = ResolvedAddresses.newBuilder() - .setAddresses(ImmutableList.copyOf(checkNotNull(eags, "eags"))) - .setAttributes(checkNotNull(attributes, "attributes")) - .build(); - childLb.acceptResolvedAddresses(childAddresses); - } - + /** + * Starts connecting if this endpoint is IDLE. Must be called from the synchronization + * context. + */ void requestConnection() { + if (childShutdown || state != IDLE) { + return; + } childLb.requestConnection(); } + /** + * The {@link PickerEndpoint.ExitIdler} handed to pickers. Called from RPC threads, so it + * hops onto the synchronization context before doing anything. + * + *

    The state is re-checked there rather than here, which is what makes repeated calls + * harmless: a picker snapshot may be shared by many concurrent RPCs that all observe the + * same IDLE endpoint, and the snapshot may outlive the endpoint entirely if a resolver + * update removed it in the meantime. By the time the second and later tasks run, either the + * child has moved to CONNECTING or the holder has been shut down, and they return early. + */ + private void exitIdle() { + helper.getSynchronizationContext().execute(this::requestConnection); + } + void shutdown() { + if (childShutdown) { + return; + } + childShutdown = true; childLb.shutdown(); } @Override public String toString() { return MoreObjects.toStringHelper(this) - .add("index", index) + .add("hostname", hostname) .add("state", state) - .add("childLb", childLb) .toString(); } + /** + * Intercepts the child's balancing state so that it is recorded here instead of being + * published straight to the channel. The LB policy aggregates across all endpoints and + * publishes a single state and picker of its own. + */ 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; + return helper; } @Override public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) { - state = checkNotNull(newState, "newState"); - picker = checkNotNull(newPicker, "newPicker"); - if (stateUpdateCallback != null) { - stateUpdateCallback.run(); + if (childShutdown) { + return; + } + state = newState; + picker = newPicker; + if (!rebuilding) { + childStateListener.run(); } - } - - @Override - public String toString() { - return MoreObjects.toStringHelper(this) - .add("delegateHelper", delegateHelper) - .toString(); } } } diff --git a/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java b/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java new file mode 100644 index 00000000000..f8a35dd7a1b --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java @@ -0,0 +1,359 @@ +/* + * 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 com.google.cloud.autosharding.v1.AssignmentChunk; +import com.google.cloud.autosharding.v1.EndpointState; +import com.google.cloud.autosharding.v1.PerSliceEndpointState; +import com.google.cloud.autosharding.v1.SliceAssignment; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.ByteString; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.List; +import javax.annotation.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link AssignmentParser}. */ +@RunWith(JUnit4.class) +public class AssignmentParserTest { + + @Test + public void parse_singleChunkCoveringWholeKeyspace() throws Exception { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addEndpoints(endpoint("host-b")) + .addSliceAssignments(sliceAssignment("", "m", 0)) + .addSliceAssignments(sliceAssignment("m", null, 1)) + .build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 7); + + assertThat(assignment.getGeneration()).isEqualTo(7); + assertThat(assignment.getEndpointNames()).containsExactly("host-a", "host-b").inOrder(); + assertThat(assignment.getSlices()).hasSize(2); + assertSlice(assignment.getSlices().get(0), "", "m", 0); + assertSlice(assignment.getSlices().get(1), "m", null, 1); + } + + @Test + public void parse_endpointNamesCombinedInChunkOrder() throws Exception { + AssignmentChunk chunk1 = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addEndpoints(endpoint("host-b")) + .build(); + AssignmentChunk chunk2 = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-c")) + // Index 2 only resolves once chunk1's endpoints are prepended. + .addSliceAssignments(sliceAssignment("", null, 2)) + .build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk1, chunk2), 1); + + assertThat(assignment.getEndpointNames()) + .containsExactly("host-a", "host-b", "host-c") + .inOrder(); + assertSlice(assignment.getSlices().get(0), "", null, 2); + } + + @Test + public void parse_slicesAcrossChunksAreSorted() throws Exception { + AssignmentChunk chunk1 = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("m", null, 0)) + .build(); + AssignmentChunk chunk2 = + AssignmentChunk.newBuilder().addSliceAssignments(sliceAssignment("", "m", 0)).build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk1, chunk2), 1); + + assertThat(assignment.getSlices()).hasSize(2); + assertSlice(assignment.getSlices().get(0), "", "m", 0); + assertSlice(assignment.getSlices().get(1), "m", null, 0); + } + + @Test + public void parse_fillsLeadingGap() throws Exception { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("d", null, 0)) + .build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(assignment.getSlices()).hasSize(2); + assertSlice(assignment.getSlices().get(0), "", "d"); + assertSlice(assignment.getSlices().get(1), "d", null, 0); + } + + @Test + public void parse_fillsTrailingGap() throws Exception { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("", "d", 0)) + .build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(assignment.getSlices()).hasSize(2); + assertSlice(assignment.getSlices().get(0), "", "d", 0); + assertSlice(assignment.getSlices().get(1), "d", null); + } + + @Test + public void parse_fillsInteriorGap() throws Exception { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addEndpoints(endpoint("host-b")) + .addSliceAssignments(sliceAssignment("", "d", 0)) + .addSliceAssignments(sliceAssignment("m", null, 1)) + .build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(assignment.getSlices()).hasSize(3); + assertSlice(assignment.getSlices().get(0), "", "d", 0); + assertSlice(assignment.getSlices().get(1), "d", "m"); + assertSlice(assignment.getSlices().get(2), "m", null, 1); + } + + @Test + public void parse_noSlices_yieldsSingleEmptySliceCoveringKeyspace() throws Exception { + Assignment assignment = + AssignmentParser.parse(ImmutableList.of(AssignmentChunk.getDefaultInstance()), 3); + + assertThat(assignment.getSlices()).hasSize(1); + assertSlice(assignment.getSlices().get(0), "", null); + assertThat(assignment.getEndpointNames()).isEmpty(); + assertThat(assignment.getGeneration()).isEqualTo(3); + } + + @Test + public void parse_noChunks_yieldsSingleEmptySliceCoveringKeyspace() throws Exception { + Assignment assignment = AssignmentParser.parse(ImmutableList.of(), 1); + + assertThat(assignment.getSlices()).hasSize(1); + assertSlice(assignment.getSlices().get(0), "", null); + } + + @Test + public void parse_sliceWithNoEndpoints_isPreserved() throws Exception { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("", "d")) + .addSliceAssignments(sliceAssignment("d", null, 0)) + .build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(assignment.getSlices()).hasSize(2); + assertSlice(assignment.getSlices().get(0), "", "d"); + assertSlice(assignment.getSlices().get(1), "d", null, 0); + } + + @Test + public void parse_multipleEndpointsPerSlice() throws Exception { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addEndpoints(endpoint("host-b")) + .addSliceAssignments(sliceAssignment("", null, 0, 1)) + .build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertSlice(assignment.getSlices().get(0), "", null, 0, 1); + } + + @Test + public void parse_endpointIndexOutOfRange_throws() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("", null, 1)) + .build(); + + AssignmentParser.ValidationException e = + assertThrows( + AssignmentParser.ValidationException.class, + () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); + assertThat(e).hasMessageThat().contains("out-of-range endpoint index 1"); + } + + @Test + public void parse_negativeEndpointIndex_throws() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("", null, -1)) + .build(); + + AssignmentParser.ValidationException e = + assertThrows( + AssignmentParser.ValidationException.class, + () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); + assertThat(e).hasMessageThat().contains("out-of-range endpoint index -1"); + } + + @Test + public void parse_startKeyGreaterThanEndKey_throws() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder().addSliceAssignments(sliceAssignment("z", "a")).build(); + + AssignmentParser.ValidationException e = + assertThrows( + AssignmentParser.ValidationException.class, + () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); + assertThat(e).hasMessageThat().contains("greater than end_key"); + } + + @Test + public void parse_overlappingSlices_throws() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addSliceAssignments(sliceAssignment("a", "m")) + .addSliceAssignments(sliceAssignment("d", null)) + .build(); + + AssignmentParser.ValidationException e = + assertThrows( + AssignmentParser.ValidationException.class, + () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); + assertThat(e).hasMessageThat().contains("overlaps"); + } + + @Test + public void parse_duplicateStartKeys_throws() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addSliceAssignments(sliceAssignment("a", "m")) + .addSliceAssignments(sliceAssignment("a", "z")) + .build(); + + AssignmentParser.ValidationException e = + assertThrows( + AssignmentParser.ValidationException.class, + () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); + assertThat(e).hasMessageThat().contains("overlaps"); + } + + @Test + public void parse_sliceExtendingToInfinityFollowedByAnother_throws() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addSliceAssignments(sliceAssignment("a", null)) + .addSliceAssignments(sliceAssignment("m", null)) + .build(); + + AssignmentParser.ValidationException e = + assertThrows( + AssignmentParser.ValidationException.class, + () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); + assertThat(e).hasMessageThat().contains("extends to the end of the keyspace"); + } + + @Test + public void parse_unsignedByteOrderingIsUsed() throws Exception { + // 0x80 is negative as a signed byte but must sort after 0x01. + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addSliceAssignments( + SliceAssignment.newBuilder() + .setSlice( + com.google.cloud.autosharding.v1.Slice.newBuilder() + .setStartKey(ByteString.copyFrom(new byte[] {(byte) 0x80})))) + .addSliceAssignments( + SliceAssignment.newBuilder() + .setSlice( + com.google.cloud.autosharding.v1.Slice.newBuilder() + .setStartKey(ByteString.copyFrom(new byte[] {0x01})) + .setEndKey(ByteString.copyFrom(new byte[] {(byte) 0x80})))) + .build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + // Leading gap ["", 0x01) plus the two declared slices. + assertThat(assignment.getSlices()).hasSize(3); + assertThat(assignment.getSlices().get(1).getStartKey()).isEqualTo(new byte[] {0x01}); + assertThat(assignment.getSlices().get(2).getStartKey()).isEqualTo(new byte[] {(byte) 0x80}); + assertThat(assignment.getSlices().get(2).getEndKey()).isNull(); + } + + @Test + public void parse_resultingSlicesAreContiguous() throws Exception { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("b", "d", 0)) + .addSliceAssignments(sliceAssignment("k", "m", 0)) + .build(); + + Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + List slices = assignment.getSlices(); + assertThat(slices.get(0).getStartKey()).isEqualTo(new byte[0]); + for (int i = 0; i + 1 < slices.size(); i++) { + assertThat(slices.get(i).getEndKey()).isEqualTo(slices.get(i + 1).getStartKey()); + } + assertThat(slices.get(slices.size() - 1).getEndKey()).isNull(); + } + + private static EndpointState endpoint(String name) { + return EndpointState.newBuilder().setEndpoint(name).build(); + } + + private static SliceAssignment sliceAssignment( + String startKey, @Nullable String endKey, int... endpointIndices) { + com.google.cloud.autosharding.v1.Slice.Builder slice = + com.google.cloud.autosharding.v1.Slice.newBuilder() + .setStartKey(ByteString.copyFromUtf8(startKey)); + if (endKey != null) { + slice.setEndKey(ByteString.copyFromUtf8(endKey)); + } + SliceAssignment.Builder builder = SliceAssignment.newBuilder().setSlice(slice); + for (int index : endpointIndices) { + builder.addEndpoints(PerSliceEndpointState.newBuilder().setEndpointIndex(index)); + } + return builder.build(); + } + + private static void assertSlice( + Assignment.Slice slice, String startKey, @Nullable String endKey, int... endpoints) { + assertThat(slice.getStartKey()).isEqualTo(startKey.getBytes(StandardCharsets.UTF_8)); + if (endKey == null) { + assertThat(slice.getEndKey()).isNull(); + } else { + assertThat(slice.getEndKey()).isEqualTo(endKey.getBytes(StandardCharsets.UTF_8)); + } + assertThat(slice.getEndpoints()) + .containsExactlyElementsIn(Arrays.stream(endpoints).boxed().toArray()) + .inOrder(); + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerTest.java new file mode 100644 index 00000000000..9fb40741322 --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerTest.java @@ -0,0 +1,961 @@ +/* + * 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.CONNECTING; +import static io.grpc.ConnectivityState.IDLE; +import static io.grpc.ConnectivityState.READY; +import static io.grpc.ConnectivityState.TRANSIENT_FAILURE; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.cloud.autosharding.v1.AssignmentChunk; +import com.google.cloud.autosharding.v1.AssignmentMetadata; +import com.google.cloud.autosharding.v1.AutoshardingServiceGrpc; +import com.google.cloud.autosharding.v1.EndpointState; +import com.google.cloud.autosharding.v1.PerSliceEndpointState; +import com.google.cloud.autosharding.v1.SliceAssignment; +import com.google.cloud.autosharding.v1.WatchShardingAssignmentRequest; +import com.google.cloud.autosharding.v1.WatchShardingAssignmentResponse; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.ByteString; +import io.grpc.Attributes; +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ConnectivityState; +import io.grpc.EquivalentAddressGroup; +import io.grpc.LoadBalancer; +import io.grpc.LoadBalancer.Helper; +import io.grpc.LoadBalancer.PickDetailsConsumer; +import io.grpc.LoadBalancer.PickResult; +import io.grpc.LoadBalancer.ResolvedAddresses; +import io.grpc.LoadBalancer.Subchannel; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.LoadBalancerProvider; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Status; +import io.grpc.SynchronizationContext; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.internal.FakeClock; +import io.grpc.internal.PickSubchannelArgsImpl; +import io.grpc.stub.StreamObserver; +import io.grpc.testing.GrpcCleanupRule; +import io.grpc.testing.TestMethodDescriptors; +import java.net.SocketAddress; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import javax.annotation.Nullable; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Unit tests for {@link AutoShardingLoadBalancer}. + * + *

    These drive a real {@link AutoshardingClient} against an in-process fake sharding service, + * so the path from a served assignment through to a routed pick is covered end to end. + */ +@RunWith(JUnit4.class) +public class AutoShardingLoadBalancerTest { + private static final String CHANNEL_FACTORY_KEY = "shard-service-key"; + private static final String OTHER_CHANNEL_FACTORY_KEY = "other-shard-service-key"; + private static final String UNKNOWN_CHANNEL_FACTORY_KEY = "unknown-key"; + private static final String TARGET = "autosharding-target"; + private static final String KEY_HEADER = "x-shard-key"; + private static final long ASSIGNMENT_TIMEOUT_NANOS = TimeUnit.SECONDS.toNanos(10); + private static final long POLL_TIMEOUT_SECONDS = 5; + private static final MethodDescriptor METHOD = TestMethodDescriptors.voidMethod(); + + @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule(); + + private final SynchronizationContext syncContext = + new SynchronizationContext( + (t, e) -> { + throw new AssertionError(e); + }); + private final FakeClock fakeClock = new FakeClock(); + private final FakeAutoshardingService service = new FakeAutoshardingService(); + private final Helper helper = mock(Helper.class); + private final FakeChildProvider childProvider = new FakeChildProvider(); + private final FakeChannelFactory channelFactory = new FakeChannelFactory(); + + private Channel shardingChannel; + private AutoShardingLoadBalancer loadBalancer; + + @Nullable private ConnectivityState currentState; + @Nullable private SubchannelPicker currentPicker; + @Nullable private StreamObserver serverStream; + private int balancingStateUpdates; + + @Before + public void setUp() throws Exception { + String serverName = InProcessServerBuilder.generateName(); + grpcCleanup.register( + InProcessServerBuilder.forName(serverName) + .directExecutor() + .addService(service) + .build() + .start()); + shardingChannel = + grpcCleanup.register(InProcessChannelBuilder.forName(serverName).directExecutor().build()); + + when(helper.getSynchronizationContext()).thenReturn(syncContext); + when(helper.getScheduledExecutorService()).thenReturn(fakeClock.getScheduledExecutorService()); + doAnswer( + invocation -> { + currentState = invocation.getArgument(0); + currentPicker = invocation.getArgument(1); + balancingStateUpdates++; + return null; + }) + .when(helper) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + + loadBalancer = + new AutoShardingLoadBalancer( + helper, + childProvider, + () -> () -> TimeUnit.SECONDS.toNanos(1), + fakeClock.getStopwatchSupplier(), + "client-uuid"); + } + + @After + public void tearDown() { + // Must run before GrpcCleanupRule terminates the channel, otherwise the assignment client + // keeps retrying against a shutting-down channel. + syncContext.execute(loadBalancer::shutdown); + } + + // --------------------------------------------------------------------------------------------- + // Configuration handling + // --------------------------------------------------------------------------------------------- + + @Test + public void missingConfig_reportsTransientFailure() { + Status status = + acceptAddresses( + ResolvedAddresses.newBuilder() + .setAddresses(endpoints("a")) + .setAttributes(attributesWithChannelFactory()) + .build()); + + assertThat(status.getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(currentState).isEqualTo(TRANSIENT_FAILURE); + } + + @Test + public void missingChannelFactory_reportsTransientFailure() { + Status status = + acceptAddresses( + ResolvedAddresses.newBuilder() + .setAddresses(endpoints("a")) + .setAttributes(Attributes.EMPTY) + .setLoadBalancingPolicyConfig(config(CHANNEL_FACTORY_KEY, true)) + .build()); + + assertThat(status.getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(status.getDescription()).contains("channel factory"); + assertThat(currentState).isEqualTo(TRANSIENT_FAILURE); + } + + @Test + public void unknownChannelFactoryKey_reportsTransientFailure() { + Status status = deliverAddresses(config(UNKNOWN_CHANNEL_FACTORY_KEY, true), "a"); + + assertThat(status.getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(currentState).isEqualTo(TRANSIENT_FAILURE); + } + + @Test + public void noEndpoints_reportsTransientFailureAndFailsRpcs() { + Status status = deliverAddresses(config(CHANNEL_FACTORY_KEY, true)); + + assertThat(status.getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(currentState).isEqualTo(TRANSIENT_FAILURE); + assertThat(pick("anything").getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + } + + @Test + public void endpointsRetracted_thenRestored_resumesServing() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + deliverAssignment(1, slice("", "a")); + + deliverAddresses(config(CHANNEL_FACTORY_KEY, true)); + assertThat(currentState).isEqualTo(TRANSIENT_FAILURE); + + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + reportReady("a"); + + assertThat(currentState).isEqualTo(READY); + assertThat(pickedHost(pick("k"))).isEqualTo("a"); + } + + // --------------------------------------------------------------------------------------------- + // Channel to the sharding service + // --------------------------------------------------------------------------------------------- + + @Test + public void firstUpdate_createsChannelAndOpensStream() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + + assertThat(channelFactory.keys).containsExactly(CHANNEL_FACTORY_KEY); + WatchShardingAssignmentRequest request = takeRequest(); + assertThat(request.getInitialClientConfig().getTarget()).isEqualTo(TARGET); + assertThat(request.getInitialClientConfig().getClientUuid()).isEqualTo("client-uuid"); + } + + @Test + public void unchangedKey_reusesChannelAndStream() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + takeRequest(); + + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + + assertThat(channelFactory.keys).containsExactly(CHANNEL_FACTORY_KEY); + assertThat(service.streamCount.get()).isEqualTo(1); + } + + @Test + public void changedKey_createsNewChannelClosesOldAndRestartsStream() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + takeRequest(); + + deliverAddresses(config(OTHER_CHANNEL_FACTORY_KEY, true), "a"); + + assertThat(channelFactory.keys) + .containsExactly(CHANNEL_FACTORY_KEY, OTHER_CHANNEL_FACTORY_KEY) + .inOrder(); + assertThat(channelFactory.isReleased(0)).isTrue(); + assertThat(channelFactory.isReleased(1)).isFalse(); + assertThat(service.streamCount.get()).isEqualTo(2); + } + + @Test + public void changedTarget_restartsStreamWithoutNewChannel() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + takeRequest(); + + AutoShardingLoadBalancerConfig retargeted = + new AutoShardingLoadBalancerConfig( + CHANNEL_FACTORY_KEY, "other-target", KEY_HEADER, true, ASSIGNMENT_TIMEOUT_NANOS); + deliverAddresses(retargeted, "a"); + + assertThat(channelFactory.keys).containsExactly(CHANNEL_FACTORY_KEY); + assertThat(service.streamCount.get()).isEqualTo(2); + assertThat(takeRequest().getInitialClientConfig().getTarget()).isEqualTo("other-target"); + } + + @Test + public void targetWithLocalityToken_isSubstituted() throws Exception { + AutoShardingLoadBalancerConfig localityConfig = + new AutoShardingLoadBalancerConfig( + CHANNEL_FACTORY_KEY, "target/%s", KEY_HEADER, true, ASSIGNMENT_TIMEOUT_NANOS); + EquivalentAddressGroup endpoint = + new EquivalentAddressGroup( + new NamedAddress("addr-a"), + Attributes.newBuilder() + .set(AutoShardingAttributes.ATTR_ENDPOINT_HOSTNAME, "a") + .set(EquivalentAddressGroup.ATTR_LOCALITY_NAME, "us-central1-a") + .build()); + + acceptAddresses( + ResolvedAddresses.newBuilder() + .setAddresses(ImmutableList.of(endpoint)) + .setAttributes(attributesWithChannelFactory()) + .setLoadBalancingPolicyConfig(localityConfig) + .build()); + + assertThat(takeRequest().getInitialClientConfig().getTarget()) + .isEqualTo("target/us-central1-a"); + } + + @Test + public void targetWithLocalityToken_noLocality_substitutesEmptyString() throws Exception { + AutoShardingLoadBalancerConfig localityConfig = + new AutoShardingLoadBalancerConfig( + CHANNEL_FACTORY_KEY, "target/%s", KEY_HEADER, true, ASSIGNMENT_TIMEOUT_NANOS); + + deliverAddresses(localityConfig, "a"); + + assertThat(takeRequest().getInitialClientConfig().getTarget()).isEqualTo("target/"); + } + + // --------------------------------------------------------------------------------------------- + // Startup: queuing, timeout and fallback + // --------------------------------------------------------------------------------------------- + + @Test + public void beforeFirstAssignment_queuesRpcs() { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + + assertThat(currentState).isEqualTo(CONNECTING); + PickResult result = pick("anything"); + assertThat(result.getSubchannel()).isNull(); + assertThat(result.getStatus().isOk()).isTrue(); + } + + @Test + public void beforeFirstAssignment_childStateChangesDoNotUnqueueRpcs() { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + activate("a"); + reportReady("a"); + + // Still waiting on the sharding service, so RPCs stay queued rather than being routed + // anywhere arbitrary. + assertThat(currentState).isEqualTo(CONNECTING); + assertThat(pick("k").getSubchannel()).isNull(); + } + + @Test + public void initialAssignmentTimeout_fallbackEnabled_spreadsAcrossAllEndpoints() { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + reportReady("a"); + reportReady("b"); + + fakeClock.forwardNanos(ASSIGNMENT_TIMEOUT_NANOS); + + assertThat(currentState).isEqualTo(READY); + assertThat(pickedHost(pick("k"))).isAnyOf("a", "b"); + } + + @Test + public void initialAssignmentTimeout_fallbackDisabled_failsRpcs() { + deliverAddresses(config(CHANNEL_FACTORY_KEY, false), "a"); + reportReady("a"); + + fakeClock.forwardNanos(ASSIGNMENT_TIMEOUT_NANOS); + + PickResult result = pick("k"); + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()).contains("fallback disabled"); + } + + @Test + public void assignmentBeforeTimeout_cancelsTheTimer() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + assertThat(fakeClock.numPendingTasks()).isEqualTo(1); + + deliverAssignment(1, slice("", "a")); + + assertThat(fakeClock.numPendingTasks()).isEqualTo(0); + } + + @Test + public void newChannel_keepsServingPreviousAssignmentWhileTimerPending() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + deliverAssignment(1, slice("", "a"), slice("m", "b")); + reportReady("a"); + reportReady("b"); + assertThat(pickedHost(pick("z"))).isEqualTo("b"); + + // Switching sharding service must not interrupt traffic. + deliverAddresses(config(OTHER_CHANNEL_FACTORY_KEY, true), "a", "b"); + + assertThat(pickedHost(pick("z"))).isEqualTo("b"); + assertThat(pickedHost(pick("a"))).isEqualTo("a"); + } + + // --------------------------------------------------------------------------------------------- + // Routing on assignments + // --------------------------------------------------------------------------------------------- + + @Test + public void assignmentRoutesByKeyRange() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + deliverAssignment(1, slice("", "a"), slice("m", "b")); + reportReady("a"); + reportReady("b"); + + assertThat(pickedHost(pick("alpha"))).isEqualTo("a"); + assertThat(pickedHost(pick("zulu"))).isEqualTo("b"); + } + + @Test + public void assignmentNamingUnknownHostname_dropsIt() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, false), "a"); + // The sharding service still believes "ghost" is serving; the resolver disagrees. + deliverAssignment(1, slice("", "ghost")); + reportReady("a"); + + PickResult result = pick("k"); + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + } + + @Test + public void assignmentNamingUnknownHostname_fallbackEnabled_usesFallbackPool() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + deliverAssignment(1, slice("", "ghost")); + reportReady("a"); + + assertThat(pickedHost(pick("k"))).isEqualTo("a"); + } + + @Test + public void resolverUpdateAfterAssignment_rebuildsSliceMapWithNewIndices() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + deliverAssignment(1, slice("", "b")); + reportReady("a"); + reportReady("b"); + assertThat(pickedHost(pick("k"))).isEqualTo("b"); + + // "b" moves from index 1 to index 0. If the slice map were not rebuilt, the stale index + // would now route to "a". + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "b", "a"); + + assertThat(pickedHost(pick("k"))).isEqualTo("b"); + } + + @Test + public void staleGenerationAssignment_isIgnored() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + deliverAssignment(5, slice("", "a")); + reportReady("a"); + reportReady("b"); + assertThat(pickedHost(pick("k"))).isEqualTo("a"); + + pushAssignment(3, slice("", "b")); + + assertThat(pickedHost(pick("k"))).isEqualTo("a"); + } + + // --------------------------------------------------------------------------------------------- + // Child state updates and aggregated connectivity state + // --------------------------------------------------------------------------------------------- + + @Test + public void childStateUpdate_republishesPicker() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + deliverAssignment(1, slice("", "a")); + int updatesBefore = balancingStateUpdates; + + reportReady("a"); + + assertThat(balancingStateUpdates).isGreaterThan(updatesBefore); + assertThat(currentState).isEqualTo(READY); + assertThat(pickedHost(pick("k"))).isEqualTo("a"); + } + + @Test + public void allEndpointsIdle_reportsIdle() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + deliverAssignment(1, slice("", "a")); + + assertThat(currentState).isEqualTo(IDLE); + } + + @Test + public void twoEndpointsInTransientFailure_reportsTransientFailure() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + deliverAssignment(1, slice("", "a"), slice("m", "b")); + + reportTransientFailure("a"); + reportTransientFailure("b"); + + assertThat(currentState).isEqualTo(TRANSIENT_FAILURE); + } + + @Test + public void oneEndpointInTransientFailure_wakesUpAnIdleEndpoint() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b", "c"); + deliverAssignment(1, slice("", "a")); + + reportTransientFailure("a"); + + // Aggregated state is CONNECTING, and nothing was connecting, so exactly one IDLE endpoint + // is nudged so the policy can recover without needing a pick. + assertThat(currentState).isEqualTo(CONNECTING); + assertThat(childProvider.children).hasSize(2); + assertThat(childForHost("b").requestConnectionCount).isEqualTo(1); + } + + @Test + public void endpointAlreadyConnecting_noAdditionalWakeUp() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b", "c"); + deliverAssignment(1, slice("", "a")); + activate("b"); + + reportTransientFailure("a"); + + // "b" is already CONNECTING, so "c" is left alone. + assertThat(currentState).isEqualTo(CONNECTING); + assertThat(activatedHostnames()).containsExactly("a", "b"); + } + + @Test + public void requestConnection_wakesUpAnIdleEndpoint() { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + + syncContext.execute(loadBalancer::requestConnection); + + assertThat(childProvider.children).hasSize(1); + } + + @Test + public void picksOnIdleEndpointTriggerConnection() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + deliverAssignment(1, slice("", "a")); + assertThat(childProvider.children).isEmpty(); + + PickResult result = pick("k"); + + assertThat(result.getSubchannel()).isNull(); + assertThat(childProvider.children).hasSize(1); + } + + // --------------------------------------------------------------------------------------------- + // Name resolution errors and shutdown + // --------------------------------------------------------------------------------------------- + + @Test + public void nameResolutionError_withKnownEndpoints_keepsServing() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + deliverAssignment(1, slice("", "a")); + reportReady("a"); + + syncContext.execute( + () -> loadBalancer.handleNameResolutionError(Status.UNAVAILABLE.withDescription("boom"))); + + assertThat(currentState).isEqualTo(READY); + assertThat(pickedHost(pick("k"))).isEqualTo("a"); + } + + @Test + public void nameResolutionError_withNoEndpoints_reportsTransientFailure() { + syncContext.execute( + () -> loadBalancer.handleNameResolutionError(Status.UNAVAILABLE.withDescription("boom"))); + + assertThat(currentState).isEqualTo(TRANSIENT_FAILURE); + assertThat(pick("k").getStatus().getDescription()).contains("boom"); + } + + @Test + public void shutdown_closesChannelAndChildren() { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + activate("a"); + + syncContext.execute(loadBalancer::shutdown); + + assertThat(channelFactory.isReleased(0)).isTrue(); + assertThat(childProvider.children.get(0).shutdown).isTrue(); + } + + @Test + public void shutdown_isIdempotent() { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + + syncContext.execute(loadBalancer::shutdown); + syncContext.execute(loadBalancer::shutdown); + + assertThat(channelFactory.isReleased(0)).isTrue(); + } + + @Test + public void shutdown_cancelsInitialAssignmentTimer() { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + assertThat(fakeClock.numPendingTasks()).isEqualTo(1); + + syncContext.execute(loadBalancer::shutdown); + + assertThat(fakeClock.numPendingTasks()).isEqualTo(0); + } + + // --------------------------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------------------------- + + private AutoShardingLoadBalancerConfig config(String channelFactoryKey, boolean enableFallback) { + return new AutoShardingLoadBalancerConfig( + channelFactoryKey, TARGET, KEY_HEADER, enableFallback, ASSIGNMENT_TIMEOUT_NANOS); + } + + private Attributes attributesWithChannelFactory() { + return Attributes.newBuilder() + .set(AutoShardingAttributes.ATTR_CHANNEL_FACTORY, channelFactory) + .build(); + } + + private Status deliverAddresses(AutoShardingLoadBalancerConfig config, String... hostnames) { + return acceptAddresses( + ResolvedAddresses.newBuilder() + .setAddresses(endpoints(hostnames)) + .setAttributes(attributesWithChannelFactory()) + .setLoadBalancingPolicyConfig(config) + .build()); + } + + private Status acceptAddresses(ResolvedAddresses resolvedAddresses) { + AtomicReference status = new AtomicReference<>(); + syncContext.execute(() -> status.set(loadBalancer.acceptResolvedAddresses(resolvedAddresses))); + return status.get(); + } + + private static List endpoints(String... hostnames) { + List eags = new ArrayList<>(); + for (String hostname : hostnames) { + eags.add( + new EquivalentAddressGroup( + new NamedAddress("addr-" + hostname), + Attributes.newBuilder() + .set(AutoShardingAttributes.ATTR_ENDPOINT_HOSTNAME, hostname) + .build())); + } + return ImmutableList.copyOf(eags); + } + + /** Sends an assignment from the fake service and waits for the load balancer to apply it. */ + private void deliverAssignment(long generation, SliceSpec... slices) throws Exception { + pushAssignment(generation, slices); + } + + private void pushAssignment(long generation, SliceSpec... slices) throws Exception { + StreamObserver serverStream = currentServerStream(); + List endpointNames = new ArrayList<>(); + for (SliceSpec spec : slices) { + if (!endpointNames.contains(spec.hostname)) { + endpointNames.add(spec.hostname); + } + } + + AssignmentChunk.Builder chunk = AssignmentChunk.newBuilder(); + for (String name : endpointNames) { + chunk.addEndpoints(EndpointState.newBuilder().setEndpoint(name)); + } + for (int i = 0; i < slices.length; i++) { + SliceSpec spec = slices[i]; + String endKey = i + 1 < slices.length ? slices[i + 1].startKey : null; + chunk.addSliceAssignments( + sliceAssignment(spec.startKey, endKey, endpointNames.indexOf(spec.hostname))); + } + + serverStream.onNext(WatchShardingAssignmentResponse.newBuilder().setChunk(chunk).build()); + serverStream.onNext( + WatchShardingAssignmentResponse.newBuilder() + .setMetadata(AssignmentMetadata.newBuilder().setGeneration(generation)) + .build()); + } + + private static SliceSpec slice(String startKey, String hostname) { + return new SliceSpec(startKey, hostname); + } + + private static final class SliceSpec { + final String startKey; + final String hostname; + + SliceSpec(String startKey, String hostname) { + this.startKey = startKey; + this.hostname = hostname; + } + } + + private static SliceAssignment sliceAssignment( + String startKey, @Nullable String endKey, int endpointIndex) { + com.google.cloud.autosharding.v1.Slice.Builder slice = + com.google.cloud.autosharding.v1.Slice.newBuilder() + .setStartKey(ByteString.copyFromUtf8(startKey)); + if (endKey != null) { + slice.setEndKey(ByteString.copyFromUtf8(endKey)); + } + return SliceAssignment.newBuilder() + .setSlice(slice) + .addEndpoints(PerSliceEndpointState.newBuilder().setEndpointIndex(endpointIndex)) + .build(); + } + + private PickResult pick(String key) { + Metadata headers = new Metadata(); + headers.put(Metadata.Key.of(KEY_HEADER, Metadata.ASCII_STRING_MARSHALLER), key); + return currentPicker.pickSubchannel( + new PickSubchannelArgsImpl(METHOD, headers, CallOptions.DEFAULT, new PickDetailsConsumer() { + })); + } + + /** Returns the hostname of the endpoint the pick landed on. */ + private String pickedHost(PickResult result) { + Subchannel subchannel = result.getSubchannel(); + if (subchannel == null) { + throw new AssertionError("Pick did not select a subchannel: " + result); + } + for (FakeChild child : childProvider.children) { + if (child.subchannel == subchannel) { + return child.hostname; + } + } + throw new AssertionError("Pick returned an unrecognized subchannel"); + } + + /** + * Instantiates the child load balancer for {@code hostname} by asking its endpoint to connect, + * which is how the picker brings an endpoint out of IDLE at runtime. + */ + private void activate(String hostname) { + syncContext.execute( + () -> { + EndpointMap endpointMap = loadBalancer.getEndpointMap(); + int index = endpointMap.indexOf(hostname); + if (index == -1) { + throw new AssertionError("Unknown endpoint hostname " + hostname); + } + endpointMap.toPickerEndpoints().get(index).requestConnection(); + }); + } + + private void reportReady(String hostname) { + activate(hostname); + syncContext.execute(() -> childForHost(hostname).reportReady()); + } + + private void reportTransientFailure(String hostname) { + activate(hostname); + syncContext.execute(() -> childForHost(hostname).reportTransientFailure()); + } + + private FakeChild childForHost(String hostname) { + for (FakeChild child : childProvider.children) { + if (hostname.equals(child.hostname)) { + return child; + } + } + throw new AssertionError("No child load balancer for hostname " + hostname); + } + + private WatchShardingAssignmentRequest takeRequest() throws Exception { + WatchShardingAssignmentRequest request = + service.requests.poll(POLL_TIMEOUT_SECONDS, TimeUnit.SECONDS); + if (request == null) { + fail("timed out waiting for a request to the sharding service"); + } + return request; + } + + /** + * Returns the stream the client currently has open to the sharding service, picking up a newly + * opened one if there is any. Streams are created synchronously by the in-process transport, so + * a non-blocking poll is enough once the first one exists. + */ + private StreamObserver currentServerStream() throws Exception { + StreamObserver next = service.serverStreams.poll(); + if (next != null) { + serverStream = next; + } else if (serverStream == null) { + serverStream = service.serverStreams.poll(POLL_TIMEOUT_SECONDS, TimeUnit.SECONDS); + if (serverStream == null) { + fail("timed out waiting for a stream to the sharding service"); + } + } + return serverStream; + } + + /** Hostnames whose child load balancer has been instantiated, in creation order. */ + private List activatedHostnames() { + List result = new ArrayList<>(); + for (FakeChild child : childProvider.children) { + result.add(child.hostname); + } + return result; + } + + /** A {@link SocketAddress} with a predictable {@link #toString}. */ + private static final class NamedAddress extends SocketAddress { + private static final long serialVersionUID = 0L; + private final String name; + + NamedAddress(String name) { + this.name = name; + } + + @Override + public String toString() { + return name; + } + } + + /** + * Hands out in-process channels, each wrapped so that successive borrows are distinguishable + * even though they share one transport. + */ + private final class FakeChannelFactory implements ChannelFactory { + final List keys = new ArrayList<>(); + final List created = new ArrayList<>(); + final List released = new ArrayList<>(); + + @Override + public Channel createChannel(String channelFactoryKey) { + if (UNKNOWN_CHANNEL_FACTORY_KEY.equals(channelFactoryKey)) { + throw new IllegalArgumentException("unknown channel factory key"); + } + keys.add(channelFactoryKey); + Channel channel = new WrappedChannel(shardingChannel); + created.add(channel); + return channel; + } + + @Override + public void releaseChannel(Channel channel) { + released.add(channel); + } + + boolean isReleased(int index) { + Channel channel = created.get(index); + for (Channel released : this.released) { + if (released == channel) { + return true; + } + } + return false; + } + } + + /** Gives each handle a distinct channel identity over one shared transport. */ + private static final class WrappedChannel extends Channel { + private final Channel delegate; + + WrappedChannel(Channel delegate) { + this.delegate = delegate; + } + + @Override + public String authority() { + return delegate.authority(); + } + + @Override + public ClientCall newCall( + MethodDescriptor methodDescriptor, CallOptions callOptions) { + return delegate.newCall(methodDescriptor, callOptions); + } + } + + private static final class FakeAutoshardingService + extends AutoshardingServiceGrpc.AutoshardingServiceImplBase { + final BlockingQueue requests = new LinkedBlockingQueue<>(); + final BlockingQueue> serverStreams = + new LinkedBlockingQueue<>(); + final AtomicInteger streamCount = new AtomicInteger(); + + @Override + public StreamObserver watchShardingAssignment( + StreamObserver responseObserver) { + streamCount.incrementAndGet(); + serverStreams.add(responseObserver); + return new StreamObserver() { + @Override + public void onNext(WatchShardingAssignmentRequest request) { + requests.add(request); + } + + @Override + public void onError(Throwable t) {} + + @Override + public void onCompleted() {} + }; + } + } + + private static final class FakeChildProvider extends LoadBalancerProvider { + final List children = new ArrayList<>(); + + @Override + public boolean isAvailable() { + return true; + } + + @Override + public int getPriority() { + return 5; + } + + @Override + public String getPolicyName() { + return "fake_child"; + } + + @Override + public LoadBalancer newLoadBalancer(Helper childHelper) { + FakeChild child = new FakeChild(childHelper); + children.add(child); + return child; + } + } + + /** Stands in for {@code pick_first}, reporting CONNECTING as soon as it is asked to connect. */ + private static final class FakeChild extends LoadBalancer { + private final Helper helper; + final Subchannel subchannel = mock(Subchannel.class); + @Nullable String hostname; + int requestConnectionCount; + boolean shutdown; + + FakeChild(Helper helper) { + this.helper = helper; + } + + @Override + public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) { + hostname = + resolvedAddresses + .getAddresses() + .get(0) + .getAttributes() + .get(AutoShardingAttributes.ATTR_ENDPOINT_HOSTNAME); + return Status.OK; + } + + @Override + public void handleNameResolutionError(Status error) {} + + @Override + public void requestConnection() { + requestConnectionCount++; + helper.updateBalancingState(CONNECTING, new FixedResultPicker(PickResult.withNoResult())); + } + + @Override + public void shutdown() { + shutdown = true; + } + + void reportReady() { + helper.updateBalancingState(READY, new FixedResultPicker(PickResult.withSubchannel( + subchannel))); + } + + void reportTransientFailure() { + helper.updateBalancingState( + TRANSIENT_FAILURE, + new FixedResultPicker( + PickResult.withError(Status.UNAVAILABLE.withDescription("endpoint down")))); + } + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoshardingClientTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoshardingClientTest.java new file mode 100644 index 00000000000..ade4ca1240c --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoshardingClientTest.java @@ -0,0 +1,527 @@ +/* + * 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.fail; + +import com.google.cloud.autosharding.v1.AssignmentChunk; +import com.google.cloud.autosharding.v1.AssignmentMetadata; +import com.google.cloud.autosharding.v1.AutoshardingServiceGrpc; +import com.google.cloud.autosharding.v1.EndpointState; +import com.google.cloud.autosharding.v1.LoadReportingConfig; +import com.google.cloud.autosharding.v1.PerSliceEndpointState; +import com.google.cloud.autosharding.v1.SliceAssignment; +import com.google.cloud.autosharding.v1.WatchShardingAssignmentRequest; +import com.google.cloud.autosharding.v1.WatchShardingAssignmentResponse; +import com.google.protobuf.ByteString; +import io.grpc.Channel; +import io.grpc.Status; +import io.grpc.SynchronizationContext; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.internal.BackoffPolicy; +import io.grpc.internal.FakeClock; +import io.grpc.stub.StreamObserver; +import io.grpc.testing.GrpcCleanupRule; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Unit tests for {@link AutoshardingClient}. */ +@RunWith(JUnit4.class) +public class AutoshardingClientTest { + private static final String CLIENT_UUID = "client-uuid-1"; + private static final String TARGET = "autosharding-target"; + private static final String OTHER_TARGET = "other-autosharding-target"; + private static final long TIMEOUT_SECONDS = 5; + private static final long BACKOFF_NANOS = TimeUnit.SECONDS.toNanos(1); + + @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule(); + + private final SynchronizationContext syncContext = + new SynchronizationContext( + (t, e) -> { + throw new AssertionError(e); + }); + private final FakeClock fakeClock = new FakeClock(); + private final FakeAutoshardingService service = new FakeAutoshardingService(); + private final BlockingQueue assignments = new LinkedBlockingQueue<>(); + private final RecordingBackoffPolicyProvider backoffPolicyProvider = + new RecordingBackoffPolicyProvider(); + + private Channel channel; + private AutoshardingClient client; + + @Before + public void setUp() throws Exception { + String serverName = InProcessServerBuilder.generateName(); + grpcCleanup.register( + InProcessServerBuilder.forName(serverName) + .directExecutor() + .addService(service) + .build() + .start()); + channel = + grpcCleanup.register( + InProcessChannelBuilder.forName(serverName).directExecutor().build()); + client = + new AutoshardingClient( + CLIENT_UUID, + syncContext, + fakeClock.getScheduledExecutorService(), + backoffPolicyProvider, + fakeClock.getStopwatchSupplier(), + assignments::add); + } + + @After + public void tearDown() { + // Must happen before GrpcCleanupRule shuts the channel down, otherwise the client keeps + // retrying against a terminating channel. + syncContext.execute(client::shutdown); + } + + @Test + public void update_startsStreamAndSendsInitialClientConfig() throws Exception { + update(channel, TARGET); + + WatchShardingAssignmentRequest request = takeRequest(); + assertThat(request.hasInitialClientConfig()).isTrue(); + assertThat(request.getInitialClientConfig().getTarget()).isEqualTo(TARGET); + assertThat(request.getInitialClientConfig().getClientUuid()).isEqualTo(CLIENT_UUID); + assertThat(request.getInitialClientConfig().getLatestGeneration()).isEqualTo(0); + } + + @Test + public void update_unchanged_doesNotRestartStream() throws Exception { + update(channel, TARGET); + takeRequest(); + + update(channel, TARGET); + + assertThat(service.streamCount.get()).isEqualTo(1); + } + + @Test + public void chunksBufferedUntilMetadata_thenAssignmentDeliveredAndAcked() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0))); + assertThat(assignments).isEmpty(); + assertThat(service.requests).isEmpty(); + + serverStream.onNext(metadataResponse(5)); + + Assignment assignment = takeAssignment(); + assertThat(assignment.getGeneration()).isEqualTo(5); + assertThat(assignment.getEndpointNames()).containsExactly("host-a"); + assertThat(assignment.getSlices()).hasSize(1); + + WatchShardingAssignmentRequest ack = takeRequest(); + assertThat(ack.hasAssignmentAck()).isTrue(); + assertThat(ack.getAssignmentAck().getGeneration()).isEqualTo(5); + assertThat(ack.getAssignmentAck().getAccepted()).isTrue(); + assertThat(ack.getAssignmentAck().getErrorMessage()).isEmpty(); + } + + @Test + public void multipleChunks_combinedIntoOneLogicalAssignment() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + serverStream.onNext( + chunkResponse(AssignmentChunk.newBuilder().addEndpoints(endpoint("host-a")).build())); + serverStream.onNext( + chunkResponse( + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-b")) + .addSliceAssignments(sliceAssignment("", null, 1)) + .build())); + serverStream.onNext(metadataResponse(1)); + + Assignment assignment = takeAssignment(); + assertThat(assignment.getEndpointNames()).containsExactly("host-a", "host-b").inOrder(); + assertThat(assignment.getSlices().get(0).getEndpoints()).containsExactly(1); + } + + @Test + public void invalidAssignment_nackedAndNotDelivered() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + // Endpoint index 3 does not exist in the combined endpoint list. + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 3))); + serverStream.onNext(metadataResponse(5)); + + WatchShardingAssignmentRequest nack = takeRequest(); + assertThat(nack.hasAssignmentAck()).isTrue(); + assertThat(nack.getAssignmentAck().getGeneration()).isEqualTo(5); + assertThat(nack.getAssignmentAck().getAccepted()).isFalse(); + assertThat(nack.getAssignmentAck().getErrorMessage()) + .contains("out-of-range endpoint index 3"); + assertThat(assignments).isEmpty(); + assertThat(client.getLatestGeneration()).isEqualTo(0); + } + + @Test + public void invalidAssignment_doesNotLeakChunksIntoNextAssignment() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 3))); + serverStream.onNext(metadataResponse(5)); + takeRequest(); // NACK + + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-b", "", null, 0))); + serverStream.onNext(metadataResponse(6)); + + Assignment assignment = takeAssignment(); + assertThat(assignment.getEndpointNames()).containsExactly("host-b"); + } + + @Test + public void staleGeneration_droppedWithoutAck() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0))); + serverStream.onNext(metadataResponse(5)); + takeAssignment(); + takeRequest(); // ACK for generation 5 + + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-b", "", null, 0))); + serverStream.onNext(metadataResponse(5)); + + assertThat(assignments).isEmpty(); + assertThat(service.requests).isEmpty(); + assertThat(client.getLatestGeneration()).isEqualTo(5); + } + + @Test + public void olderGeneration_droppedWithoutAck() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0))); + serverStream.onNext(metadataResponse(5)); + takeAssignment(); + takeRequest(); + + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-b", "", null, 0))); + serverStream.onNext(metadataResponse(4)); + + assertThat(assignments).isEmpty(); + assertThat(service.requests).isEmpty(); + } + + @Test + public void loadReportingConfig_ignored() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + serverStream.onNext( + WatchShardingAssignmentResponse.newBuilder() + .setConfig(LoadReportingConfig.newBuilder().setLoadQuantumFraction(0.5)) + .build()); + + assertThat(assignments).isEmpty(); + assertThat(service.requests).isEmpty(); + } + + @Test + public void streamFailure_reconnectsAndSendsLatestGeneration() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0))); + serverStream.onNext(metadataResponse(9)); + takeAssignment(); + takeRequest(); // ACK + + serverStream.onError(Status.UNAVAILABLE.asRuntimeException()); + fireRetryTimer(); + + WatchShardingAssignmentRequest retryRequest = takeRequest(); + assertThat(retryRequest.hasInitialClientConfig()).isTrue(); + assertThat(retryRequest.getInitialClientConfig().getLatestGeneration()).isEqualTo(9); + assertThat(retryRequest.getInitialClientConfig().getClientUuid()).isEqualTo(CLIENT_UUID); + assertThat(service.streamCount.get()).isEqualTo(2); + } + + @Test + public void streamFailure_doesNotReconnectBeforeBackoffElapses() throws Exception { + update(channel, TARGET); + takeRequest(); + + takeServerStream().onError(Status.UNAVAILABLE.asRuntimeException()); + + assertThat(fakeClock.numPendingTasks()).isEqualTo(1); + fakeClock.forwardNanos(BACKOFF_NANOS - 1); + assertThat(service.streamCount.get()).isEqualTo(1); + + fakeClock.forwardNanos(1); + takeRequest(); + assertThat(service.streamCount.get()).isEqualTo(2); + } + + @Test + public void streamCompletedByServer_reconnects() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + serverStream.onCompleted(); + fireRetryTimer(); + + takeRequest(); + assertThat(service.streamCount.get()).isEqualTo(2); + } + + @Test + public void backoffSequence_onlyResetAfterGoodAssignment() throws Exception { + update(channel, TARGET); + takeRequest(); + + // First failure with no assignment received: a backoff sequence is created. + takeServerStream().onError(Status.UNAVAILABLE.asRuntimeException()); + fireRetryTimer(); + takeRequest(); + assertThat(backoffPolicyProvider.timesCalled).isEqualTo(1); + + // Second failure with no assignment received: the existing sequence continues. + takeServerStream().onError(Status.UNAVAILABLE.asRuntimeException()); + fireRetryTimer(); + takeRequest(); + assertThat(backoffPolicyProvider.timesCalled).isEqualTo(1); + + // A good assignment resets the sequence when the stream later fails. + StreamObserver serverStream = takeServerStream(); + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0))); + serverStream.onNext(metadataResponse(1)); + takeAssignment(); + takeRequest(); // ACK + serverStream.onError(Status.UNAVAILABLE.asRuntimeException()); + fireRetryTimer(); + takeRequest(); + assertThat(backoffPolicyProvider.timesCalled).isEqualTo(2); + } + + @Test + public void targetChange_restartsStreamAndResetsGeneration() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0))); + serverStream.onNext(metadataResponse(9)); + takeAssignment(); + takeRequest(); // ACK + assertThat(client.getLatestGeneration()).isEqualTo(9); + + update(channel, OTHER_TARGET); + + WatchShardingAssignmentRequest request = takeRequest(); + assertThat(request.hasInitialClientConfig()).isTrue(); + assertThat(request.getInitialClientConfig().getTarget()).isEqualTo(OTHER_TARGET); + assertThat(request.getInitialClientConfig().getLatestGeneration()).isEqualTo(0); + assertThat(client.getLatestGeneration()).isEqualTo(0); + assertThat(service.streamCount.get()).isEqualTo(2); + } + + @Test + public void channelChange_restartsStreamAndResetsGeneration() throws Exception { + update(channel, TARGET); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0))); + serverStream.onNext(metadataResponse(9)); + takeAssignment(); + takeRequest(); // ACK + + String otherServerName = InProcessServerBuilder.generateName(); + grpcCleanup.register( + InProcessServerBuilder.forName(otherServerName) + .directExecutor() + .addService(service) + .build() + .start()); + Channel otherChannel = + grpcCleanup.register( + InProcessChannelBuilder.forName(otherServerName).directExecutor().build()); + + update(otherChannel, TARGET); + + WatchShardingAssignmentRequest request = takeRequest(); + assertThat(request.getInitialClientConfig().getLatestGeneration()).isEqualTo(0); + assertThat(service.streamCount.get()).isEqualTo(2); + } + + @Test + public void shutdown_cancelsStreamAndStopsReconnecting() throws Exception { + update(channel, TARGET); + takeRequest(); + + syncContext.execute(client::shutdown); + + assertThat(service.streamCount.get()).isEqualTo(1); + assertThat(fakeClock.numPendingTasks()).isEqualTo(0); + } + + @Test + public void shutdown_isIdempotentAndIgnoresLaterUpdates() throws Exception { + update(channel, TARGET); + takeRequest(); + + syncContext.execute(client::shutdown); + syncContext.execute(client::shutdown); + update(channel, OTHER_TARGET); + + assertThat(service.streamCount.get()).isEqualTo(1); + } + + private void update(Channel channel, String target) { + syncContext.execute(() -> client.update(channel, target)); + } + + /** Asserts that a retry was scheduled and advances the clock so that it runs. */ + private void fireRetryTimer() { + assertThat(fakeClock.numPendingTasks()).isEqualTo(1); + fakeClock.forwardNanos(BACKOFF_NANOS); + } + + private WatchShardingAssignmentRequest takeRequest() throws Exception { + WatchShardingAssignmentRequest request = + service.requests.poll(TIMEOUT_SECONDS, TimeUnit.SECONDS); + if (request == null) { + fail("timed out waiting for a request from the autosharding client"); + } + return request; + } + + private StreamObserver takeServerStream() throws Exception { + StreamObserver stream = + service.serverStreams.poll(TIMEOUT_SECONDS, TimeUnit.SECONDS); + if (stream == null) { + fail("timed out waiting for the autosharding client to open a stream"); + } + return stream; + } + + private Assignment takeAssignment() throws Exception { + Assignment assignment = assignments.poll(TIMEOUT_SECONDS, TimeUnit.SECONDS); + if (assignment == null) { + fail("timed out waiting for an assignment"); + } + return assignment; + } + + private static WatchShardingAssignmentResponse chunkResponse(AssignmentChunk chunk) { + return WatchShardingAssignmentResponse.newBuilder().setChunk(chunk).build(); + } + + private static WatchShardingAssignmentResponse metadataResponse(long generation) { + return WatchShardingAssignmentResponse.newBuilder() + .setMetadata(AssignmentMetadata.newBuilder().setGeneration(generation)) + .build(); + } + + private static AssignmentChunk chunkWithEndpoint( + String endpointName, String startKey, @Nullable String endKey, int endpointIndex) { + return AssignmentChunk.newBuilder() + .addEndpoints(endpoint(endpointName)) + .addSliceAssignments(sliceAssignment(startKey, endKey, endpointIndex)) + .build(); + } + + private static EndpointState endpoint(String name) { + return EndpointState.newBuilder().setEndpoint(name).build(); + } + + private static SliceAssignment sliceAssignment( + String startKey, @Nullable String endKey, int... endpointIndices) { + com.google.cloud.autosharding.v1.Slice.Builder slice = + com.google.cloud.autosharding.v1.Slice.newBuilder() + .setStartKey(ByteString.copyFromUtf8(startKey)); + if (endKey != null) { + slice.setEndKey(ByteString.copyFromUtf8(endKey)); + } + SliceAssignment.Builder builder = SliceAssignment.newBuilder().setSlice(slice); + for (int index : endpointIndices) { + builder.addEndpoints(PerSliceEndpointState.newBuilder().setEndpointIndex(index)); + } + return builder.build(); + } + + private static final class FakeAutoshardingService + extends AutoshardingServiceGrpc.AutoshardingServiceImplBase { + final BlockingQueue requests = new LinkedBlockingQueue<>(); + final BlockingQueue> serverStreams = + new LinkedBlockingQueue<>(); + final AtomicInteger streamCount = new AtomicInteger(); + + @Override + public StreamObserver watchShardingAssignment( + StreamObserver responseObserver) { + streamCount.incrementAndGet(); + serverStreams.add(responseObserver); + return new StreamObserver() { + @Override + public void onNext(WatchShardingAssignmentRequest request) { + requests.add(request); + } + + @Override + public void onError(Throwable t) {} + + @Override + public void onCompleted() {} + }; + } + } + + /** + * Hands out backoff policies with a fixed, non-zero delay so that retries are driven explicitly + * by the fake clock. The number of policies handed out reflects how many times the backoff + * sequence was reset. + */ + private static final class RecordingBackoffPolicyProvider implements BackoffPolicy.Provider { + int timesCalled; + + @Override + public BackoffPolicy get() { + timesCalled++; + return () -> BACKOFF_NANOS; + } + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java index 5dd6f4c5332..666baae839a 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java @@ -17,264 +17,542 @@ package io.grpc.autosharding; import static com.google.common.truth.Truth.assertThat; +import static io.grpc.ConnectivityState.CONNECTING; import static io.grpc.ConnectivityState.IDLE; import static io.grpc.ConnectivityState.READY; import static io.grpc.ConnectivityState.TRANSIENT_FAILURE; -import 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; import com.google.common.collect.ImmutableList; import io.grpc.Attributes; +import io.grpc.ConnectivityState; import io.grpc.EquivalentAddressGroup; import io.grpc.LoadBalancer; 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.Status; 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 java.util.ArrayList; +import java.util.List; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -import org.mockito.ArgumentCaptor; +/** Unit tests for {@link EndpointMap}. */ @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); - }); + new SynchronizationContext( + (t, e) -> { + throw new AssertionError("Unhandled exception in syncContext", e); + }); + private final Helper helper = mock(Helper.class); + private final FakeChildProvider childProvider = new FakeChildProvider(); + private final List stateUpdates = new ArrayList<>(); 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(); + when(helper.getSynchronizationContext()).thenReturn(syncContext); + endpointMap = new EndpointMap(helper, childProvider, () -> stateUpdates.add(1)); } - private EndpointHolder createHolder(int index) { - return new EndpointHolder(index, mockHelper, mockProvider, stateChangeCount::incrementAndGet); + // --------------------------------------------------------------------------------------------- + // Endpoint set and indices + // --------------------------------------------------------------------------------------------- + + @Test + public void updateEndpoints_assignsDenseIndicesInResolverOrder() { + endpointMap.updateEndpoints(endpoints("a", "b", "c"), Attributes.EMPTY); + + assertThat(endpointMap.size()).isEqualTo(3); + assertThat(endpointMap.indexOf("a")).isEqualTo(0); + assertThat(endpointMap.indexOf("b")).isEqualTo(1); + assertThat(endpointMap.indexOf("c")).isEqualTo(2); + assertThat(endpointMap.toPickerEndpoints()).hasSize(3); } @Test - public void basicMapOperations() { - assertThat(endpointMap.isEmpty()).isTrue(); - assertThat(endpointMap.size()).isEqualTo(0); + public void indexOf_unknownHostname_returnsMinusOne() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); - EndpointHolder h1 = createHolder(0); - EndpointHolder h2 = createHolder(1); + assertThat(endpointMap.indexOf("nope")).isEqualTo(-1); + } - endpointMap.put("host1", h1); - endpointMap.put("host2", h2); + @Test + public void updateEndpoints_duplicateHostnames_keepsOneEntry() { + endpointMap.updateEndpoints(endpoints("a", "b", "a"), Attributes.EMPTY); - assertThat(endpointMap.isEmpty()).isFalse(); + // Indices stay dense so that they remain valid offsets into toPickerEndpoints(). 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.toPickerEndpoints()).hasSize(2); + assertThat(endpointMap.indexOf("a")).isEqualTo(0); + assertThat(endpointMap.indexOf("b")).isEqualTo(1); + } + + @Test + public void updateEndpoints_duplicateHostnames_firstEndpointSuppliesTheAddresses() { + EquivalentAddressGroup first = endpointWithHostname("first-addr", "a"); + EquivalentAddressGroup second = endpointWithHostname("second-addr", "a"); + + endpointMap.updateEndpoints(ImmutableList.of(first, second), Attributes.EMPTY); + activate(0); + + assertThat(childProvider.children).hasSize(1); + assertThat(childProvider.children.get(0).lastAddresses.getAddresses()).containsExactly(first); + } + + @Test + public void updateEndpoints_reordering_movesIndicesAndPickerEndpoints() { + endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY); + activate(0); + reportState(0, READY); + + endpointMap.updateEndpoints(endpoints("b", "a"), Attributes.EMPTY); + + assertThat(endpointMap.indexOf("a")).isEqualTo(1); + assertThat(endpointMap.indexOf("b")).isEqualTo(0); + // The state moved with the endpoint, not with the index. + assertThat(stateAt(1)).isEqualTo(READY); + assertThat(stateAt(0)).isEqualTo(IDLE); + } + + @Test + public void hostnameAttributeAbsent_fallsBackToFirstAddress() { + EquivalentAddressGroup eag = new EquivalentAddressGroup(new NamedAddress("1.2.3.4:80")); + + endpointMap.updateEndpoints(ImmutableList.of(eag), Attributes.EMPTY); + + assertThat(endpointMap.indexOf("1.2.3.4:80")).isEqualTo(0); + } + + // --------------------------------------------------------------------------------------------- + // Child lifecycle across resolver updates + // --------------------------------------------------------------------------------------------- + + @Test + public void updateEndpoints_survivingHostname_keepsChildAndState() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + activate(0); + reportState(0, READY); + assertThat(childProvider.children).hasSize(1); + + endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY); + + // No new child for "a", and its connectivity state survived the update. + assertThat(childProvider.children).hasSize(1); + assertThat(childProvider.children.get(0).shutdown).isFalse(); + assertThat(stateAt(0)).isEqualTo(READY); + } + + @Test + public void updateEndpoints_survivingHostname_forwardsNewAddresses() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + activate(0); + FakeChild child = childProvider.children.get(0); + int acceptsBefore = child.acceptCount; + + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + + assertThat(child.acceptCount).isGreaterThan(acceptsBefore); + } + + @Test + public void updateEndpoints_removedHostname_shutsDownChild() { + endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY); + activate(0); + activate(1); + + endpointMap.updateEndpoints(endpoints("b"), Attributes.EMPTY); + + assertThat(childProvider.children.get(0).shutdown).isTrue(); + assertThat(childProvider.children.get(1).shutdown).isFalse(); assertThat(endpointMap.size()).isEqualTo(1); - assertThat(endpointMap.get("host1")).isNull(); + assertThat(endpointMap.indexOf("a")).isEqualTo(-1); + } + + @Test + public void updateEndpoints_toEmpty_shutsDownEverything() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + activate(0); + + endpointMap.updateEndpoints(ImmutableList.of(), Attributes.EMPTY); + + assertThat(endpointMap.size()).isEqualTo(0); + assertThat(endpointMap.toPickerEndpoints()).isEmpty(); + assertThat(childProvider.children.get(0).shutdown).isTrue(); } @Test - public void nullChecks() { - EndpointHolder h = createHolder(0); + public void updateEndpoints_doesNotNotifyListenerWhileRebuilding() { + // New children publish their initial IDLE state from inside updateEndpoints(). Forwarding + // those would make the LB policy build a picker from a half-rebuilt map. + endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY); + + assertThat(stateUpdates).isEmpty(); + } - 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)); + @Test + public void childStateUpdate_notifiesListener() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + activate(0); + stateUpdates.clear(); - assertThrows( - NullPointerException.class, - () -> new EndpointHolder(0, null, mockProvider, null)); - assertThrows( - NullPointerException.class, - () -> new EndpointHolder(0, mockHelper, null, null)); + reportState(0, READY); - assertThrows( - NullPointerException.class, - () -> h.updateAddresses(null, Attributes.EMPTY)); - assertThrows( - NullPointerException.class, - () -> h.updateAddresses(Collections.emptyList(), null)); + assertThat(stateUpdates).hasSize(1); + assertThat(stateAt(0)).isEqualTo(READY); } @Test - public void reindex_updatesIndicesContiguously() { - EndpointHolder h0 = createHolder(0); - EndpointHolder h1 = createHolder(1); - EndpointHolder h2 = createHolder(2); + public void childStateUpdate_afterEndpointRemoved_isIgnored() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + activate(0); + FakeChild child = childProvider.children.get(0); + + endpointMap.updateEndpoints(ImmutableList.of(), Attributes.EMPTY); + stateUpdates.clear(); + child.report(READY, mock(SubchannelPicker.class)); - endpointMap.put("host0", h0); - endpointMap.put("host1", h1); - endpointMap.put("host2", h2); + assertThat(stateUpdates).isEmpty(); + } + + // --------------------------------------------------------------------------------------------- + // Connecting lazily + // --------------------------------------------------------------------------------------------- - // Remove middle element - endpointMap.remove("host1"); - assertThat(h0.getIndex()).isEqualTo(0); - assertThat(h2.getIndex()).isEqualTo(2); + @Test + public void endpointsStartIdleWithoutCreatingChildren() { + endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY); - endpointMap.reindex(); - assertThat(h0.getIndex()).isEqualTo(0); - assertThat(h2.getIndex()).isEqualTo(1); + assertThat(childProvider.children).isEmpty(); + assertThat(stateAt(0)).isEqualTo(IDLE); + assertThat(stateAt(1)).isEqualTo(IDLE); } @Test - public void endpointHolder_childHelperUpdatesStateAndTriggersCallback() { - EndpointHolder holder = createHolder(0); - assertThat(holder.getState()).isEqualTo(IDLE); + public void pickerEndpoint_requestConnection_createsChildAndConnects() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); - // Capture child helper passed to LazyChildLoadBalancer - ArgumentCaptor helperCaptor = ArgumentCaptor.forClass(Helper.class); - verify(mockProvider, org.mockito.Mockito.never()).newLoadBalancer(any()); + endpointMap.toPickerEndpoints().get(0).requestConnection(); - // Trigger connection to create child helper and delegate - holder.updateAddresses( - Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), - Attributes.EMPTY); - holder.requestConnection(); + assertThat(childProvider.children).hasSize(1); + assertThat(childProvider.children.get(0).requestConnectionCount).isEqualTo(1); + } - verify(mockProvider).newLoadBalancer(helperCaptor.capture()); - Helper childHelper = helperCaptor.getValue(); + @Test + public void pickerEndpoint_repeatedRequestConnection_connectsOnce() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + // A single snapshot is shared by every concurrent RPC, so the same stale IDLE endpoint can + // be asked to connect many times over. + PickerEndpoint stale = endpointMap.toPickerEndpoints().get(0); + + stale.requestConnection(); + stale.requestConnection(); + stale.requestConnection(); + + assertThat(childProvider.children).hasSize(1); + assertThat(childProvider.children.get(0).requestConnectionCount).isEqualTo(1); + } - // Reset counter before state update to verify callback fires on update - stateChangeCount.set(0); + @Test + public void pickerEndpoint_requestConnectionAfterEndpointRemoved_isNoOp() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + PickerEndpoint stale = endpointMap.toPickerEndpoints().get(0); - // Simulate child balancer updating state - SubchannelPicker testPicker = mock(SubchannelPicker.class); - childHelper.updateBalancingState(READY, testPicker); + endpointMap.updateEndpoints(endpoints("b"), Attributes.EMPTY); + stale.requestConnection(); - assertThat(holder.getState()).isEqualTo(READY); - assertThat(holder.getPicker()).isSameInstanceAs(testPicker); - assertThat(stateChangeCount.get()).isEqualTo(1); + assertThat(childProvider.children).isEmpty(); } @Test - public void toPickerEndpoints_buildsImmutableListMatchingHoldersByIndex() { - EndpointHolder h0 = createHolder(0); - EndpointHolder h1 = createHolder(1); + public void pickerEndpoint_requestConnectionAfterShutdown_isNoOp() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + PickerEndpoint stale = endpointMap.toPickerEndpoints().get(0); - ArgumentCaptor helperCaptor = ArgumentCaptor.forClass(Helper.class); + endpointMap.shutdown(); + stale.requestConnection(); - // Trigger connections so child helpers are passed to provider - h0.updateAddresses( - Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), - Attributes.EMPTY); - h0.requestConnection(); + assertThat(childProvider.children).isEmpty(); + } - h1.updateAddresses( - Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), - Attributes.EMPTY); - h1.requestConnection(); + // --------------------------------------------------------------------------------------------- + // Aggregated connectivity state (gRFC A42 rules) + // --------------------------------------------------------------------------------------------- - verify(mockProvider, times(2)).newLoadBalancer(helperCaptor.capture()); - Helper childHelper0 = helperCaptor.getAllValues().get(0); - Helper childHelper1 = helperCaptor.getAllValues().get(1); + @Test + public void aggregate_empty_isTransientFailure() { + assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(TRANSIENT_FAILURE); + } - SubchannelPicker picker0 = mock(SubchannelPicker.class); - SubchannelPicker picker1 = mock(SubchannelPicker.class); + @Test + public void aggregate_anyReady_isReady() { + setUpStates(TRANSIENT_FAILURE, TRANSIENT_FAILURE, READY); - childHelper0.updateBalancingState(READY, picker0); - childHelper1.updateBalancingState(TRANSIENT_FAILURE, picker1); + assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(READY); + } - // Insert in reverse index order to verify explicit index placement - endpointMap.put("host1", h1); - endpointMap.put("host0", h0); + @Test + public void aggregate_twoTransientFailures_isTransientFailure() { + setUpStates(TRANSIENT_FAILURE, TRANSIENT_FAILURE, IDLE); - ImmutableList pickerEndpoints = endpointMap.toPickerEndpoints(); - assertThat(pickerEndpoints).hasSize(2); - 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); + assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(TRANSIENT_FAILURE); } @Test - public void toPickerEndpoints_emptyMap_returnsEmptyList() { - assertThat(endpointMap.toPickerEndpoints()).isEmpty(); + public void aggregate_anyConnecting_isConnecting() { + setUpStates(IDLE, CONNECTING); + + assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(CONNECTING); } @Test - public void toPickerEndpoints_duplicateOrOutOfBoundsIndex_throwsIllegalStateException() { - EndpointHolder h0 = createHolder(0); - EndpointHolder h0Duplicate = createHolder(0); + public void aggregate_oneTransientFailureAmongMany_isConnecting() { + setUpStates(TRANSIENT_FAILURE, IDLE); - endpointMap.put("host0", h0); - endpointMap.put("host1", h0Duplicate); + assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(CONNECTING); + } - assertThrows(IllegalStateException.class, () -> endpointMap.toPickerEndpoints()); + @Test + public void aggregate_soleEndpointInTransientFailure_isTransientFailure() { + setUpStates(TRANSIENT_FAILURE); - endpointMap.clear(); - EndpointHolder hOutOfBounds = createHolder(5); - endpointMap.put("host0", hOutOfBounds); + assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(TRANSIENT_FAILURE); + } - assertThrows(IllegalStateException.class, () -> endpointMap.toPickerEndpoints()); + @Test + public void aggregate_allIdle_isIdle() { + endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY); + + assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(IDLE); } + // --------------------------------------------------------------------------------------------- + // Waking up an idle endpoint + // --------------------------------------------------------------------------------------------- + @Test - public void shutdownAll_cleansUpAllHoldersAndClearsMap() { - EndpointHolder h0 = createHolder(0); - EndpointHolder h1 = createHolder(1); + public void maybeWakeUpIdleEndpoint_connectsLowestIndexedIdleEndpoint() { + setUpStates(TRANSIENT_FAILURE, IDLE, IDLE); + + endpointMap.maybeWakeUpIdleEndpoint(); - endpointMap.put("host0", h0); - endpointMap.put("host1", h1); + assertThat(stateAt(1)).isEqualTo(CONNECTING); + assertThat(stateAt(2)).isEqualTo(IDLE); + } + + @Test + public void maybeWakeUpIdleEndpoint_somethingAlreadyConnecting_doesNothing() { + setUpStates(CONNECTING, IDLE); + int childrenBefore = childProvider.children.size(); - // Trigger connections so delegates exist - h0.updateAddresses( - Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), - Attributes.EMPTY); - h0.requestConnection(); + endpointMap.maybeWakeUpIdleEndpoint(); - endpointMap.shutdownAll(); - assertThat(endpointMap.isEmpty()).isTrue(); - verify(mockDelegate).shutdown(); + assertThat(childProvider.children).hasSize(childrenBefore); + assertThat(stateAt(1)).isEqualTo(IDLE); } @Test - public void toPickerEndpoint_requestConnection_wakesUpChildBalancerOnSyncContext() { - EndpointHolder holder = createHolder(0); - holder.updateAddresses( - Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), - Attributes.EMPTY); + public void maybeWakeUpIdleEndpoint_noIdleEndpoints_doesNothing() { + setUpStates(TRANSIENT_FAILURE, TRANSIENT_FAILURE); + int childrenBefore = childProvider.children.size(); - PickerEndpoint pickerEndpoint = holder.toPickerEndpoint(); - verify(mockProvider, org.mockito.Mockito.never()).newLoadBalancer(any()); + endpointMap.maybeWakeUpIdleEndpoint(); - // Trigger connection through PickerEndpoint (simulate AutoShardingPicker encountering IDLE) - pickerEndpoint.requestConnection(); + assertThat(childProvider.children).hasSize(childrenBefore); + } + + // --------------------------------------------------------------------------------------------- + // Shutdown + // --------------------------------------------------------------------------------------------- + + @Test + public void shutdown_shutsDownChildrenAndEmptiesMap() { + endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY); + activate(0); + activate(1); + + endpointMap.shutdown(); - verify(mockProvider).newLoadBalancer(any()); - verify(mockDelegate).acceptResolvedAddresses(any()); - verify(mockDelegate).requestConnection(); + assertThat(childProvider.children.get(0).shutdown).isTrue(); + assertThat(childProvider.children.get(1).shutdown).isTrue(); + assertThat(endpointMap.size()).isEqualTo(0); + assertThat(endpointMap.indexOf("a")).isEqualTo(-1); } @Test - public void toString_containsDebugFields() { - EndpointHolder h = createHolder(3); - endpointMap.put("host3", h); + public void shutdown_isIdempotent() { + endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY); + activate(0); + + endpointMap.shutdown(); + endpointMap.shutdown(); + + assertThat(endpointMap.size()).isEqualTo(0); + assertThat(childProvider.children.get(0).shutdown).isTrue(); + } + + // --------------------------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------------------------- + + /** Drives the endpoints at indices 0..n-1 into the given states. */ + private void setUpStates(ConnectivityState... states) { + String[] names = new String[states.length]; + for (int i = 0; i < states.length; i++) { + names[i] = "host" + i; + } + endpointMap.updateEndpoints(endpoints(names), Attributes.EMPTY); + for (int i = 0; i < states.length; i++) { + if (states[i] == IDLE) { + continue; + } + activate(i); + reportState(i, states[i]); + } + stateUpdates.clear(); + } + + /** Instantiates the child load balancer behind the endpoint at {@code index}. */ + private void activate(int index) { + endpointMap.toPickerEndpoints().get(index).requestConnection(); + } + + private void reportState(int index, ConnectivityState state) { + childForHost(hostnames.get(index)).report(state, mock(SubchannelPicker.class)); + } + + private ConnectivityState stateAt(int index) { + return endpointMap.toPickerEndpoints().get(index).getState(); + } + + /** Returns the child load balancer created for {@code hostname}. */ + private FakeChild childForHost(String hostname) { + for (FakeChild child : childProvider.children) { + if (child.lastAddresses == null) { + continue; + } + String childHostname = + child + .lastAddresses + .getAddresses() + .get(0) + .getAttributes() + .get(AutoShardingAttributes.ATTR_ENDPOINT_HOSTNAME); + if (hostname.equals(childHostname)) { + return child; + } + } + throw new AssertionError("No child load balancer created for hostname " + hostname); + } + + /** Hostnames of the endpoints most recently produced by {@link #endpoints}. */ + private final List hostnames = new ArrayList<>(); + + private List endpoints(String... hostnameArgs) { + hostnames.clear(); + List eags = new ArrayList<>(); + for (String hostname : hostnameArgs) { + hostnames.add(hostname); + eags.add(endpointWithHostname("addr-" + hostname, hostname)); + } + return ImmutableList.copyOf(eags); + } + + /** An endpoint at {@code addressName} advertising {@code hostname}. */ + private static EquivalentAddressGroup endpointWithHostname(String addressName, String hostname) { + return new EquivalentAddressGroup( + new NamedAddress(addressName), + Attributes.newBuilder() + .set(AutoShardingAttributes.ATTR_ENDPOINT_HOSTNAME, hostname) + .build()); + } + + /** A {@link SocketAddress} with a predictable {@link #toString}. */ + private static final class NamedAddress extends SocketAddress { + private static final long serialVersionUID = 0L; + private final String name; + + NamedAddress(String name) { + this.name = name; + } + + @Override + public String toString() { + return name; + } + } + + private static final class FakeChildProvider extends LoadBalancerProvider { + final List children = new ArrayList<>(); + + @Override + public boolean isAvailable() { + return true; + } + + @Override + public int getPriority() { + return 5; + } + + @Override + public String getPolicyName() { + return "fake_child"; + } + + @Override + public LoadBalancer newLoadBalancer(Helper helper) { + FakeChild child = new FakeChild(helper); + children.add(child); + return child; + } + } - assertThat(endpointMap.toString()).contains("host3"); - assertThat(h.toString()).contains("index=3"); - assertThat(h.toString()).contains("state=IDLE"); + /** Stands in for {@code pick_first}, including its move to CONNECTING when asked to connect. */ + private static final class FakeChild extends LoadBalancer { + private final Helper helper; + ResolvedAddresses lastAddresses; + int acceptCount; + int requestConnectionCount; + boolean shutdown; + + FakeChild(Helper helper) { + this.helper = helper; + } + + @Override + public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) { + lastAddresses = resolvedAddresses; + acceptCount++; + return Status.OK; + } + + @Override + public void handleNameResolutionError(Status error) {} + + @Override + public void requestConnection() { + requestConnectionCount++; + report(CONNECTING, new FixedResultPicker(PickResult.withNoResult())); + } + + @Override + public void shutdown() { + shutdown = true; + } + + void report(ConnectivityState state, SubchannelPicker picker) { + helper.updateBalancingState(state, picker); + } } } From cc2822a6bc2c2e67eba183d01f50bb454f8d1fae Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 22 Sep 2026 14:19:44 +0530 Subject: [PATCH 30/33] updates --- .../grpc/autosharding/AssignmentParser.java | 178 +++++++---- .../AutoShardingLoadBalancer.java | 123 ++++++-- .../grpc/autosharding/AutoshardingClient.java | 120 ++++--- .../autosharding/AssignmentParserTest.java | 292 ++++++++++++------ .../AutoShardingLoadBalancerTest.java | 136 ++++++-- .../autosharding/AutoshardingClientTest.java | 234 ++++++++------ 6 files changed, 740 insertions(+), 343 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java b/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java index 712be2efdd0..10120312650 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java @@ -31,35 +31,59 @@ import javax.annotation.Nullable; /** - * Combines the {@link AssignmentChunk} messages of a single logical assignment into a - * validated, sorted, contiguous and gap-free {@link Assignment}. + * Combines the {@link AssignmentChunk} messages of a single logical assignment into a sorted, + * contiguous and gap-free {@link Assignment}. + * + *

    Validation follows gRFC A119, "Handling assignments from the Autosharding server". A slice + * is usable only if all of the following hold: * - *

    Validation follows gRFC A119, "Handling assignments from the Autosharding server": *

      - *
    • Every endpoint index referenced by a slice must be valid once the endpoint names from - * all chunks are combined in chunk order.
      • - *
    • A slice's {@code startKey} must not be greater than its {@code endKey}.
      • - *
    • Key ranges must not overlap.
      • + *
    • its {@code startKey} is strictly less than its {@code endKey}, or it has no + * {@code endKey} and so runs to the end of the keyspace; + *
    • every endpoint index it references is valid once the endpoint names from all chunks are + * combined in chunk order; + *
    • its key range does not overlap a slice that was already kept. *
    * - *

    Gaps in the key ranges returned by the server are not validation failures. They are - * explicitly filled with slices containing no endpoints, so that RPCs matching them either fall + *

    A slice that fails any of these is dropped and treated as a gap rather than + * invalidating the whole assignment. Gaps, whether they came from the server or from a dropped + * slice, are filled with slices containing no endpoints, so that RPCs matching them either fall * back (when fallback is enabled) or fail. */ final class AssignmentParser { /** - * Thrown when an assignment received from the autosharding server fails validation. The - * message is suitable for use as the {@code error_message} of an {@code AssignmentAck}. + * The outcome of parsing one logical assignment. + * + *

    Maps onto the three non-stale rows of the outcome table in gRFC A119, "Handling + * assignments from the Autosharding server": + * + *

      + *
    • every slice usable: {@link #assignment} set, {@link #errorMessage} null; + *
    • some slices dropped but at least one kept: both set; + *
    • slices were received but none was usable: {@link #assignment} null, {@link + * #errorMessage} set. + *
    */ - static final class ValidationException extends Exception { - private static final long serialVersionUID = 0L; + static final class Result { + /** The assignment to hand to the LB policy, or null if no usable slice remained. */ + @Nullable final Assignment assignment; - ValidationException(String message) { - super(message); + /** + * Describes the slices that were dropped, suitable for the {@code error_message} of an + * {@code AssignmentAck}. Null when every slice was usable. + */ + @Nullable final String errorMessage; + + private Result(@Nullable Assignment assignment, @Nullable String errorMessage) { + this.assignment = assignment; + this.errorMessage = errorMessage; } } + /** At most this many dropped slices are named in {@link Result#errorMessage}. */ + private static final int MAX_REPORTED_PROBLEMS = 3; + private static final Comparator UNSIGNED_BYTES_COMPARATOR = UnsignedBytes.lexicographicalComparator(); private static final byte[] EMPTY_BYTES = new byte[0]; @@ -69,24 +93,31 @@ private AssignmentParser() {} /** * Parses and validates the buffered chunks of a single logical assignment. * + *

    An assignment carrying no slices at all is not an error: the server is saying that nothing + * is assigned, and the result is a single endpoint-less slice spanning the keyspace. Only an + * assignment whose slices were all rejected is unusable. + * * @param chunks the chunks received since the last {@code AssignmentMetadata}, in the order * they were received * @param generation the generation number from the terminating {@code AssignmentMetadata} - * @return a validated, gap-free {@link Assignment} covering the entire keyspace - * @throws ValidationException if the assignment is invalid */ - static Assignment parse(List chunks, long generation) - throws ValidationException { + static Result parse(List chunks, long generation) { checkNotNull(chunks, "chunks"); ImmutableList endpointNames = combineEndpointNames(chunks); - List slices = combineSlices(chunks, endpointNames.size()); + List dropped = new ArrayList<>(); + List slices = combineSlices(chunks, endpointNames.size(), dropped); slices.sort( (s1, s2) -> UNSIGNED_BYTES_COMPARATOR.compare(s1.getStartKey(), s2.getStartKey())); - checkNoOverlaps(slices); + slices = dropOverlaps(slices, dropped); - return new Assignment(fillGaps(slices), endpointNames, generation); + String errorMessage = dropped.isEmpty() ? null : describe(dropped); + if (slices.isEmpty() && !dropped.isEmpty()) { + return new Result(null, errorMessage); + } + return new Result( + new Assignment(fillGaps(slices), endpointNames, generation), errorMessage); } /** @@ -104,11 +135,14 @@ private static ImmutableList combineEndpointNames(List } /** - * Concatenates the slice assignments across all chunks, validating endpoint indices and key - * range ordering along the way. Slice assignments may appear in any order across chunks. + * Concatenates the slice assignments across all chunks, dropping any whose key range is + * inverted or whose endpoint indices are out of range. Slice assignments may appear in any + * order across chunks. + * + * @param dropped collects a description of each slice that was dropped */ private static List combineSlices( - List chunks, int endpointCount) throws ValidationException { + List chunks, int endpointCount, List dropped) { List slices = new ArrayList<>(); for (AssignmentChunk chunk : chunks) { for (SliceAssignment sliceAssignment : chunk.getSliceAssignmentsList()) { @@ -116,25 +150,44 @@ private static List combineSlices( byte[] startKey = slice.getStartKey().toByteArray(); byte[] endKey = slice.hasEndKey() ? slice.getEndKey().toByteArray() : null; - if (endKey != null && UNSIGNED_BYTES_COMPARATOR.compare(startKey, endKey) > 0) { - throw new ValidationException( - String.format( - "Slice has start_key %s greater than end_key %s", - encode(startKey), encode(endKey))); + if (endKey != null) { + int keyOrder = UNSIGNED_BYTES_COMPARATOR.compare(startKey, endKey); + if (keyOrder > 0) { + dropped.add( + String.format( + "slice has start_key %s greater than end_key %s", + encode(startKey), encode(endKey))); + continue; + } + if (keyOrder == 0) { + // Permitted by the gRFC's "start_key <= end_key" rule, but it covers no keys and + // would put two entries with the same start key in the SliceMap, which makes the + // picker's binary search ambiguous. Dropping it leaves no gap: its neighbours + // already meet at this key. + dropped.add( + String.format("slice [%s, %s) is empty", encode(startKey), encode(endKey))); + continue; + } } List endpoints = new ArrayList<>(sliceAssignment.getEndpointsCount()); + String indexProblem = null; for (PerSliceEndpointState perSlice : sliceAssignment.getEndpointsList()) { int index = perSlice.getEndpointIndex(); if (index < 0 || index >= endpointCount) { - throw new ValidationException( + indexProblem = String.format( - "Slice starting at %s references out-of-range endpoint index %s;" + "slice starting at %s references out-of-range endpoint index %s;" + " assignment contains %s endpoints", - encode(startKey), index, endpointCount)); + encode(startKey), index, endpointCount); + break; } endpoints.add(index); } + if (indexProblem != null) { + dropped.add(indexProblem); + continue; + } slices.add(new Assignment.Slice(startKey, endKey, endpoints)); } } @@ -142,28 +195,40 @@ private static List combineSlices( } /** - * Verifies that no two slices in the sorted list cover the same key. + * Returns the slices of {@code sorted} that do not overlap one another, preferring the slice + * with the lower {@code startKey} whenever two of them collide. + * + * @param sorted slices in ascending {@code startKey} order + * @param dropped collects a description of each slice that was dropped */ - private static void checkNoOverlaps(List sorted) throws ValidationException { - for (int i = 0; i + 1 < sorted.size(); i++) { - Assignment.Slice current = sorted.get(i); - Assignment.Slice next = sorted.get(i + 1); - if (current.getEndKey() == null) { - throw new ValidationException( - String.format( - "Slice starting at %s extends to the end of the keyspace but overlaps the slice" - + " starting at %s", - encode(current.getStartKey()), encode(next.getStartKey()))); - } - if (UNSIGNED_BYTES_COMPARATOR.compare(current.getEndKey(), next.getStartKey()) > 0) { - throw new ValidationException( + private static List dropOverlaps( + List sorted, List dropped) { + List kept = new ArrayList<>(sorted.size()); + Assignment.Slice previous = null; + for (Assignment.Slice slice : sorted) { + if (previous != null && overlaps(previous, slice)) { + dropped.add( String.format( - "Slice [%s, %s) overlaps the slice starting at %s", - encode(current.getStartKey()), - encode(current.getEndKey()), - encode(next.getStartKey()))); + "slice starting at %s overlaps the slice [%s, %s)", + encode(slice.getStartKey()), + encode(previous.getStartKey()), + encode(previous.getEndKey()))); + continue; } + kept.add(slice); + previous = slice; } + return kept; + } + + /** + * Returns whether {@code later}, which starts at or after {@code earlier}, shares any key with + * it. A slice with no end key runs to the end of the keyspace and so overlaps everything that + * follows it. + */ + private static boolean overlaps(Assignment.Slice earlier, Assignment.Slice later) { + return earlier.getEndKey() == null + || UNSIGNED_BYTES_COMPARATOR.compare(earlier.getEndKey(), later.getStartKey()) > 0; } /** @@ -176,7 +241,7 @@ private static List fillGaps(List sorted) { byte[] cursor = EMPTY_BYTES; for (Assignment.Slice slice : sorted) { if (cursor == null) { - // Unreachable: checkNoOverlaps() rejects any slice following an infinity-ended slice. + // Unreachable: dropOverlaps() discards any slice following an infinity-ended slice. break; } if (UNSIGNED_BYTES_COMPARATOR.compare(cursor, slice.getStartKey()) < 0) { @@ -191,6 +256,15 @@ private static List fillGaps(List sorted) { return filled; } + /** Summarizes the dropped slices, capped so that the ack stays a reasonable size. */ + private static String describe(List dropped) { + if (dropped.size() <= MAX_REPORTED_PROBLEMS) { + return String.join("; ", dropped); + } + return String.join("; ", dropped.subList(0, MAX_REPORTED_PROBLEMS)) + + String.format("; and %s more", dropped.size() - MAX_REPORTED_PROBLEMS); + } + private static String encode(@Nullable byte[] key) { return key == null ? "inf" : BaseEncoding.base16().encode(key); } diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java index 065c76a2d0c..726feb213bb 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java @@ -69,10 +69,11 @@ * *

    Startup

    * - *

    Creating a channel to the sharding service starts the initial assignment timer. Until the - * first assignment arrives or that timer fires, RPCs are queued. Once the timer fires without an + *

    Creating an {@link AutoshardingClient} starts the initial assignment timer. Until the first + * assignment arrives or that timer fires, RPCs are queued. Once the timer fires without an * assignment, RPCs either spread across every resolved endpoint or fail outright, depending on - * {@code enable_fallback}. An assignment carried over from a previous channel keeps being used + * {@code enable_fallback}. A new client is created whenever the channel factory key or the + * sharding target changes; an assignment carried over from the previous client keeps being used * while the timer runs, so a change of sharding service does not interrupt traffic. * *

    Threading model

    @@ -113,6 +114,13 @@ final class AutoShardingLoadBalancer extends LoadBalancer { /** Channel borrowed from {@link #channelFactory}; must be given back when we are done. */ @Nullable private Channel shardingChannel; + /** + * The {@code autosharding_target} the current {@link #client} was created with, after {@code %s} + * substitution. Tracked separately from the config because the substitution depends on the + * resolved endpoints, so the target can change while the config does not. + */ + @Nullable private String shardingTarget; + @Nullable private AutoshardingClient client; /** Most recent assignment accepted from the sharding service, retained across reconnects. */ @@ -124,8 +132,8 @@ final class AutoShardingLoadBalancer extends LoadBalancer { @Nullable private ScheduledHandle initialAssignmentTimer; /** - * True from the moment a channel to the sharding service is created until either an assignment - * arrives on it or {@link #initialAssignmentTimer} fires. Combined with a null + * True from the moment an {@link AutoshardingClient} is created until either an assignment + * arrives from it or {@link #initialAssignmentTimer} fires. Combined with a null * {@link #assignment} it means RPCs must be queued rather than failed. */ private boolean awaitingInitialAssignment; @@ -191,6 +199,7 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) { return failPermanently("autosharding: name resolver returned no endpoints"); } + Channel previousChannel = shardingChannel; Status channelStatus = updateShardingServiceChannel(factory, newConfig); if (!channelStatus.isOk()) { return channelStatus; @@ -203,9 +212,11 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) { endpointMap.updateEndpoints(endpoints, resolvedAddresses.getAttributes()); - // Safe to call on every update: the client compares against what it already has and only - // restarts the stream when the channel or the target actually changed. - client.update(shardingChannel, resolveTarget(newConfig, endpoints)); + // The target is resolved against the endpoints, so it can change even when the config did not. + maybeRecreateClient( + shardingChannel != previousChannel, + resolveTarget(newConfig, endpoints), + newConfig.initialAssignmentTimeoutNanos); rebuildSliceMapAndPublish(); return Status.OK; @@ -246,6 +257,7 @@ public void shutdown() { client.shutdown(); client = null; } + shardingTarget = null; if (shardingChannel != null) { channelFactory.releaseChannel(shardingChannel); shardingChannel = null; @@ -256,8 +268,9 @@ public void shutdown() { /** * Creates a channel to the sharding service if this is the first configuration update, or if - * the key or the factory changed. Doing so also restarts the initial assignment timer, per - * gRFC A119. + * the {@code channel_factory_key} or the factory itself changed. Leaves {@link #shardingChannel} + * untouched when nothing changed, which is how the caller detects that no new channel was + * needed. */ private Status updateShardingServiceChannel( ChannelFactory factory, AutoShardingLoadBalancerConfig newConfig) { @@ -285,19 +298,43 @@ private Status updateShardingServiceChannel( } shardingChannel = newChannel; channelFactory = factory; + return Status.OK; + } - if (client == null) { - client = - new AutoshardingClient( - clientUuid, - syncContext, - timeService, - backoffPolicyProvider, - stopwatchSupplier, - this::onAssignment); + /** + * Replaces the {@link AutoshardingClient} when there is none yet, or when the channel to the + * sharding service or the resolved target changed. + * + *

    gRFC A119 calls for a new client rather than an in-place update because the client's + * accepted-generation watermark is only meaningful against the server and the resource it was + * learned from; carrying it over could make a different server withhold assignments + * indefinitely. + * + *

    Creating a client also restarts the initial assignment timer, since the new one has to + * start from scratch. Any assignment carried over from the previous client keeps being served + * while that timer runs. + */ + private void maybeRecreateClient(boolean channelChanged, String newTarget, long timeoutNanos) { + if (client != null && !channelChanged && newTarget.equals(shardingTarget)) { + return; } - startInitialAssignmentTimer(newConfig.initialAssignmentTimeoutNanos); - return Status.OK; + if (client != null) { + client.shutdown(); + } + shardingTarget = newTarget; + client = + new AutoshardingClient( + clientUuid, + syncContext, + timeService, + backoffPolicyProvider, + stopwatchSupplier, + shardingChannel, + newTarget, + new AssignmentWatcherImpl()); + // Armed before the stream opens so that an assignment delivered right away cancels it. + startInitialAssignmentTimer(timeoutNanos); + client.start(); } /** @@ -346,14 +383,44 @@ private void onInitialAssignmentTimeout() { rebuildSliceMapAndPublish(); } - /** Called by {@link AutoshardingClient} on the synchronization context. */ - private void onAssignment(Assignment newAssignment) { - if (shutdown) { - return; + /** + * Receives assignments from the current {@link AutoshardingClient}. Both callbacks arrive on + * the synchronization context. + * + *

    A client that has been replaced cannot deliver anything, because {@link + * AutoshardingClient#shutdown()} closes its stream, so there is no need to check which client a + * callback came from. + */ + private final class AssignmentWatcherImpl implements AutoshardingClient.AssignmentWatcher { + @Override + public void onAssignment(Assignment newAssignment) { + if (shutdown) { + return; + } + assignment = newAssignment; + cancelInitialAssignmentTimer(); + rebuildSliceMapAndPublish(); + } + + @Override + public void onError(Status error) { + if (shutdown) { + return; + } + if (assignment != null) { + // An assignment we can still use is in hand; the service only failed to replace it. + // Mirrors handleNameResolutionError: stale data beats no data. + logger.log(Level.WARNING, "Keeping the current sharding assignment: {0}", error); + return; + } + logger.log( + Level.WARNING, + "The sharding service sent no usable assignment; proceeding {0} fallback: {1}", + new Object[] {config != null && config.enableFallback ? "with" : "without", error}); + // Stop queuing RPCs: there is nothing left to wait for on this generation. + cancelInitialAssignmentTimer(); + rebuildSliceMapAndPublish(); } - assignment = newAssignment; - cancelInitialAssignmentTimer(); - rebuildSliceMapAndPublish(); } /** diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java b/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java index 4a6970dc003..0bfdbbb9cad 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java @@ -70,6 +70,13 @@ interface AssignmentWatcher { * Called with a newly accepted assignment. Invoked on the {@link SynchronizationContext}. */ void onAssignment(Assignment assignment); + + /** + * Called when the sharding service sent an assignment that could not be used at all, meaning + * every slice in it failed validation. Any assignment already in use remains valid; this + * reports that it could not be replaced. Invoked on the {@link SynchronizationContext}. + */ + void onError(Status error); } private final SynchronizationContext syncContext; @@ -78,15 +85,17 @@ interface AssignmentWatcher { private final Stopwatch retryStopwatch; private final AssignmentWatcher watcher; private final String clientUuid; - - @Nullable private Channel channel; - @Nullable private String target; + private final Channel channel; + private final String target; /** * Generation of the most recent accepted assignment. Sent to the server so that it can skip - * resending an assignment the client already has. Reset to zero whenever the channel or the - * target changes, because the stored value is meaningless against a different sharding server - * or a different resource. + * resending an assignment the client already has. + * + *

    This is why the parent load balancer replaces the whole client when the channel or the + * target changes: the stored value is meaningless against a different sharding server or a + * different resource, and retaining it could cause the server to withhold assignments + * indefinitely. */ private long latestGeneration; @@ -96,8 +105,7 @@ interface AssignmentWatcher { private boolean shutdown; /** - * Constructs an {@link AutoshardingClient}. No stream is created until - * {@link #update(Channel, String)} supplies a channel and a target. + * Constructs an {@link AutoshardingClient}. No stream is created until {@link #start()}. * * @param clientUuid a UUID generated once by the parent load balancer and reused across all * stream restarts @@ -105,6 +113,9 @@ interface AssignmentWatcher { * @param timerService used to schedule stream retries * @param backoffPolicyProvider supplies the exponential backoff sequence for stream retries * @param stopwatchSupplier supplies the stopwatch measuring time spent in a stream attempt + * @param channel the channel to the sharding service, created via the "Channel Factory" and + * owned by the parent load balancer + * @param target the autosharding target, with any {@code %s} token already substituted * @param watcher receives validated assignments */ AutoshardingClient( @@ -113,41 +124,24 @@ interface AssignmentWatcher { ScheduledExecutorService timerService, BackoffPolicy.Provider backoffPolicyProvider, Supplier stopwatchSupplier, + Channel channel, + String target, AssignmentWatcher watcher) { this.clientUuid = checkNotNull(clientUuid, "clientUuid"); this.syncContext = checkNotNull(syncContext, "syncContext"); this.timerService = checkNotNull(timerService, "timerService"); this.backoffPolicyProvider = checkNotNull(backoffPolicyProvider, "backoffPolicyProvider"); this.retryStopwatch = checkNotNull(stopwatchSupplier, "stopwatchSupplier").get(); + this.channel = checkNotNull(channel, "channel"); + this.target = checkNotNull(target, "target"); this.watcher = checkNotNull(watcher, "watcher"); } - /** - * Applies a new channel and/or resolved autosharding target. - * - *

    If either changed, any existing stream is torn down, the stored generation number is - * discarded, and a new stream is started immediately. A stored generation number is only - * meaningful for the combination of sharding server and target that produced it; retaining it - * across a change could cause the server to withhold assignments indefinitely. - * - * @param channel the channel to the sharding service, created via the "Channel Factory" - * @param target the autosharding target, with any {@code %s} token already substituted - */ - void update(Channel channel, String target) { + /** Opens the {@code WatchShardingAssignment} stream. Call once, on the sync context. */ + void start() { syncContext.throwIfNotInThisSynchronizationContext(); - checkNotNull(channel, "channel"); - checkNotNull(target, "target"); - if (shutdown) { - return; - } - if (channel.equals(this.channel) && target.equals(this.target)) { - return; - } - this.channel = channel; - this.target = target; - this.latestGeneration = 0; - this.retryBackoffPolicy = null; - restartStream(); + checkState(stream == null, "already started"); + startStream(); } /** @@ -172,17 +166,8 @@ long getLatestGeneration() { return latestGeneration; } - private void restartStream() { - cancelRetryTimer(); - if (stream != null) { - stream.close(Status.CANCELLED.withDescription("stream restarted")); - stream = null; - } - startStream(); - } - private void startStream() { - if (shutdown || channel == null || target == null) { + if (shutdown) { return; } checkState(stream == null, "previous stream has not been cleared yet"); @@ -304,39 +289,52 @@ private void handleResponse(WatchShardingAssignmentResponse response) { /** * Reassembles, validates and acknowledges the buffered chunks terminated by an * {@code AssignmentMetadata} message. + * + *

    Implements the outcome table in gRFC A119, "Handling assignments from the Autosharding + * server": every assignment is acknowledged, and only the ones carrying at least one usable + * slice reach the load balancer. */ private void handleAssignmentComplete(long generation) { List chunks = new ArrayList<>(bufferedChunks); bufferedChunks.clear(); - // Generations are monotonically increasing. Anything we have already seen is stale, and is - // dropped without acknowledgement. + // Generations are monotonically increasing, so anything we have already accepted is stale. + // It is still acknowledged, so that the server does not wait on a reply that never comes. if (generation <= latestGeneration) { - logger.log( - Level.FINE, - "Dropping autosharding assignment with stale generation {0}; latest is {1}", - new Object[] {generation, latestGeneration}); + String error = + String.format( + "stale generation %s; %s has already been accepted", generation, latestGeneration); + logger.log(Level.FINE, "Dropping autosharding assignment: {0}", error); + sendAck(generation, false, error); return; } - Assignment assignment; - try { - assignment = AssignmentParser.parse(chunks, generation); - } catch (AssignmentParser.ValidationException e) { + AssignmentParser.Result result = AssignmentParser.parse(chunks, generation); + if (result.assignment == null) { logger.log( Level.WARNING, - "Rejecting autosharding assignment with generation {0}: {1}", - new Object[] {generation, e.getMessage()}); - sendAck(generation, false, e.getMessage()); + "Rejecting autosharding assignment with generation {0}, no usable slices: {1}", + new Object[] {generation, result.errorMessage}); + sendAck(generation, false, result.errorMessage); + watcher.onError( + Status.UNAVAILABLE.withDescription( + "autosharding: no usable slices in assignment with generation " + + generation + + ": " + + result.errorMessage)); return; } - sendAck(generation, true, null); - latestGeneration = generation; - if (!receivedGoodAssignment) { - receivedGoodAssignment = true; + if (result.errorMessage != null) { + logger.log( + Level.WARNING, + "Accepting autosharding assignment with generation {0} after dropping slices: {1}", + new Object[] {generation, result.errorMessage}); } - watcher.onAssignment(assignment); + sendAck(generation, true, result.errorMessage); + latestGeneration = generation; + receivedGoodAssignment = true; + watcher.onAssignment(result.assignment); } private void sendAck(long generation, boolean accepted, @Nullable String errorMessage) { diff --git a/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java b/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java index f8a35dd7a1b..2303b12c1bb 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java @@ -17,7 +17,6 @@ package io.grpc.autosharding; import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.assertThrows; import com.google.cloud.autosharding.v1.AssignmentChunk; import com.google.cloud.autosharding.v1.EndpointState; @@ -38,7 +37,7 @@ public class AssignmentParserTest { @Test - public void parse_singleChunkCoveringWholeKeyspace() throws Exception { + public void parse_singleChunkCoveringWholeKeyspace() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) @@ -47,7 +46,7 @@ public void parse_singleChunkCoveringWholeKeyspace() throws Exception { .addSliceAssignments(sliceAssignment("m", null, 1)) .build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 7); + Assignment assignment = parseFully(ImmutableList.of(chunk), 7); assertThat(assignment.getGeneration()).isEqualTo(7); assertThat(assignment.getEndpointNames()).containsExactly("host-a", "host-b").inOrder(); @@ -57,7 +56,7 @@ public void parse_singleChunkCoveringWholeKeyspace() throws Exception { } @Test - public void parse_endpointNamesCombinedInChunkOrder() throws Exception { + public void parse_endpointNamesCombinedInChunkOrder() { AssignmentChunk chunk1 = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) @@ -70,7 +69,7 @@ public void parse_endpointNamesCombinedInChunkOrder() throws Exception { .addSliceAssignments(sliceAssignment("", null, 2)) .build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk1, chunk2), 1); + Assignment assignment = parseFully(ImmutableList.of(chunk1, chunk2), 1); assertThat(assignment.getEndpointNames()) .containsExactly("host-a", "host-b", "host-c") @@ -79,7 +78,7 @@ public void parse_endpointNamesCombinedInChunkOrder() throws Exception { } @Test - public void parse_slicesAcrossChunksAreSorted() throws Exception { + public void parse_slicesAcrossChunksAreSorted() { AssignmentChunk chunk1 = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) @@ -88,7 +87,7 @@ public void parse_slicesAcrossChunksAreSorted() throws Exception { AssignmentChunk chunk2 = AssignmentChunk.newBuilder().addSliceAssignments(sliceAssignment("", "m", 0)).build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk1, chunk2), 1); + Assignment assignment = parseFully(ImmutableList.of(chunk1, chunk2), 1); assertThat(assignment.getSlices()).hasSize(2); assertSlice(assignment.getSlices().get(0), "", "m", 0); @@ -96,14 +95,14 @@ public void parse_slicesAcrossChunksAreSorted() throws Exception { } @Test - public void parse_fillsLeadingGap() throws Exception { + public void parse_fillsLeadingGap() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) .addSliceAssignments(sliceAssignment("d", null, 0)) .build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + Assignment assignment = parseFully(ImmutableList.of(chunk), 1); assertThat(assignment.getSlices()).hasSize(2); assertSlice(assignment.getSlices().get(0), "", "d"); @@ -111,14 +110,14 @@ public void parse_fillsLeadingGap() throws Exception { } @Test - public void parse_fillsTrailingGap() throws Exception { + public void parse_fillsTrailingGap() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) .addSliceAssignments(sliceAssignment("", "d", 0)) .build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + Assignment assignment = parseFully(ImmutableList.of(chunk), 1); assertThat(assignment.getSlices()).hasSize(2); assertSlice(assignment.getSlices().get(0), "", "d", 0); @@ -126,7 +125,7 @@ public void parse_fillsTrailingGap() throws Exception { } @Test - public void parse_fillsInteriorGap() throws Exception { + public void parse_fillsInteriorGap() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) @@ -135,7 +134,7 @@ public void parse_fillsInteriorGap() throws Exception { .addSliceAssignments(sliceAssignment("m", null, 1)) .build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + Assignment assignment = parseFully(ImmutableList.of(chunk), 1); assertThat(assignment.getSlices()).hasSize(3); assertSlice(assignment.getSlices().get(0), "", "d", 0); @@ -144,9 +143,9 @@ public void parse_fillsInteriorGap() throws Exception { } @Test - public void parse_noSlices_yieldsSingleEmptySliceCoveringKeyspace() throws Exception { + public void parse_noSlices_yieldsSingleEmptySliceCoveringKeyspace() { Assignment assignment = - AssignmentParser.parse(ImmutableList.of(AssignmentChunk.getDefaultInstance()), 3); + parseFully(ImmutableList.of(AssignmentChunk.getDefaultInstance()), 3); assertThat(assignment.getSlices()).hasSize(1); assertSlice(assignment.getSlices().get(0), "", null); @@ -155,15 +154,15 @@ public void parse_noSlices_yieldsSingleEmptySliceCoveringKeyspace() throws Excep } @Test - public void parse_noChunks_yieldsSingleEmptySliceCoveringKeyspace() throws Exception { - Assignment assignment = AssignmentParser.parse(ImmutableList.of(), 1); + public void parse_noChunks_yieldsSingleEmptySliceCoveringKeyspace() { + Assignment assignment = parseFully(ImmutableList.of(), 1); assertThat(assignment.getSlices()).hasSize(1); assertSlice(assignment.getSlices().get(0), "", null); } @Test - public void parse_sliceWithNoEndpoints_isPreserved() throws Exception { + public void parse_sliceWithNoEndpoints_isPreserved() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) @@ -171,7 +170,7 @@ public void parse_sliceWithNoEndpoints_isPreserved() throws Exception { .addSliceAssignments(sliceAssignment("d", null, 0)) .build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + Assignment assignment = parseFully(ImmutableList.of(chunk), 1); assertThat(assignment.getSlices()).hasSize(2); assertSlice(assignment.getSlices().get(0), "", "d"); @@ -179,7 +178,7 @@ public void parse_sliceWithNoEndpoints_isPreserved() throws Exception { } @Test - public void parse_multipleEndpointsPerSlice() throws Exception { + public void parse_multipleEndpointsPerSlice() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) @@ -187,142 +186,249 @@ public void parse_multipleEndpointsPerSlice() throws Exception { .addSliceAssignments(sliceAssignment("", null, 0, 1)) .build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + Assignment assignment = parseFully(ImmutableList.of(chunk), 1); assertSlice(assignment.getSlices().get(0), "", null, 0, 1); } @Test - public void parse_endpointIndexOutOfRange_throws() { + public void parse_unsignedByteOrderingIsUsed() { + // 0x80 is negative as a signed byte but must sort after 0x01. + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addSliceAssignments( + SliceAssignment.newBuilder() + .setSlice( + com.google.cloud.autosharding.v1.Slice.newBuilder() + .setStartKey(ByteString.copyFrom(new byte[] {(byte) 0x80})))) + .addSliceAssignments( + SliceAssignment.newBuilder() + .setSlice( + com.google.cloud.autosharding.v1.Slice.newBuilder() + .setStartKey(ByteString.copyFrom(new byte[] {0x01})) + .setEndKey(ByteString.copyFrom(new byte[] {(byte) 0x80})))) + .build(); + + Assignment assignment = parseFully(ImmutableList.of(chunk), 1); + + // Leading gap ["", 0x01) plus the two declared slices. + assertThat(assignment.getSlices()).hasSize(3); + assertThat(assignment.getSlices().get(1).getStartKey()).isEqualTo(new byte[] {0x01}); + assertThat(assignment.getSlices().get(2).getStartKey()).isEqualTo(new byte[] {(byte) 0x80}); + assertThat(assignment.getSlices().get(2).getEndKey()).isNull(); + } + + @Test + public void parse_resultingSlicesAreContiguous() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) - .addSliceAssignments(sliceAssignment("", null, 1)) + .addSliceAssignments(sliceAssignment("b", "d", 0)) + .addSliceAssignments(sliceAssignment("k", "m", 0)) + .build(); + + Assignment assignment = parseFully(ImmutableList.of(chunk), 1); + + List slices = assignment.getSlices(); + assertThat(slices.get(0).getStartKey()).isEqualTo(new byte[0]); + for (int i = 0; i + 1 < slices.size(); i++) { + assertThat(slices.get(i).getEndKey()).isEqualTo(slices.get(i + 1).getStartKey()); + } + assertThat(slices.get(slices.size() - 1).getEndKey()).isNull(); + } + + @Test + public void parse_endpointIndexOutOfRange_sliceBecomesGap() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("", "m", 0)) + .addSliceAssignments(sliceAssignment("m", null, 5)) .build(); - AssignmentParser.ValidationException e = - assertThrows( - AssignmentParser.ValidationException.class, - () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); - assertThat(e).hasMessageThat().contains("out-of-range endpoint index 1"); + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(result.errorMessage).contains("out-of-range endpoint index 5"); + assertThat(result.assignment.getSlices()).hasSize(2); + assertSlice(result.assignment.getSlices().get(0), "", "m", 0); + assertSlice(result.assignment.getSlices().get(1), "m", null); } @Test - public void parse_negativeEndpointIndex_throws() { + public void parse_negativeEndpointIndex_sliceBecomesGap() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) - .addSliceAssignments(sliceAssignment("", null, -1)) + .addSliceAssignments(sliceAssignment("", "m", 0)) + .addSliceAssignments(sliceAssignment("m", null, -1)) .build(); - AssignmentParser.ValidationException e = - assertThrows( - AssignmentParser.ValidationException.class, - () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); - assertThat(e).hasMessageThat().contains("out-of-range endpoint index -1"); + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(result.errorMessage).contains("out-of-range endpoint index -1"); + assertThat(result.assignment.getSlices()).hasSize(2); + assertSlice(result.assignment.getSlices().get(1), "m", null); } @Test - public void parse_startKeyGreaterThanEndKey_throws() { + public void parse_endpointIndexOutOfRange_dropsTheWholeSliceNotJustThatEndpoint() { AssignmentChunk chunk = - AssignmentChunk.newBuilder().addSliceAssignments(sliceAssignment("z", "a")).build(); + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("", "m", 0)) + // Index 0 is valid, but the slice as a whole is rejected because index 5 is not. + .addSliceAssignments(sliceAssignment("m", null, 0, 5)) + .build(); - AssignmentParser.ValidationException e = - assertThrows( - AssignmentParser.ValidationException.class, - () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); - assertThat(e).hasMessageThat().contains("greater than end_key"); + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertSlice(result.assignment.getSlices().get(1), "m", null); } @Test - public void parse_overlappingSlices_throws() { + public void parse_startKeyGreaterThanEndKey_sliceBecomesGap() { AssignmentChunk chunk = AssignmentChunk.newBuilder() - .addSliceAssignments(sliceAssignment("a", "m")) - .addSliceAssignments(sliceAssignment("d", null)) + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("", "m", 0)) + .addSliceAssignments(sliceAssignment("z", "n")) + .build(); + + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(result.errorMessage).contains("greater than end_key"); + assertThat(result.assignment.getSlices()).hasSize(2); + assertSlice(result.assignment.getSlices().get(0), "", "m", 0); + assertSlice(result.assignment.getSlices().get(1), "m", null); + } + + @Test + public void parse_zeroWidthSlice_isDropped() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addEndpoints(endpoint("host-b")) + .addSliceAssignments(sliceAssignment("", "m", 0)) + // start_key == end_key satisfies the gRFC's "start_key <= end_key", but the slice + // covers no keys and would collide with the next slice's start key. + .addSliceAssignments(sliceAssignment("m", "m", 0)) + .addSliceAssignments(sliceAssignment("m", null, 1)) .build(); - AssignmentParser.ValidationException e = - assertThrows( - AssignmentParser.ValidationException.class, - () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); - assertThat(e).hasMessageThat().contains("overlaps"); + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(result.errorMessage).contains("is empty"); + // No gap appears where it was: its neighbours already met at "m". + assertThat(result.assignment.getSlices()).hasSize(2); + assertSlice(result.assignment.getSlices().get(0), "", "m", 0); + assertSlice(result.assignment.getSlices().get(1), "m", null, 1); } + /** + * The picker looks a key up by binary search over start keys, so two slices sharing one would + * make the result depend on where the search happened to land. + */ @Test - public void parse_duplicateStartKeys_throws() { + public void parse_startKeysAreUnique() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("", "", 0)) + .addSliceAssignments(sliceAssignment("", "m", 0)) + .addSliceAssignments(sliceAssignment("m", "m", 0)) + .addSliceAssignments(sliceAssignment("m", null, 0)) + .build(); + + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + List slices = result.assignment.getSlices(); + for (int i = 0; i + 1 < slices.size(); i++) { + assertThat(slices.get(i).getStartKey()).isNotEqualTo(slices.get(i + 1).getStartKey()); + } + } + + @Test + public void parse_overlappingSlices_theLaterOneIsDropped() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addSliceAssignments(sliceAssignment("a", "m")) - .addSliceAssignments(sliceAssignment("a", "z")) + .addSliceAssignments(sliceAssignment("d", null)) .build(); - AssignmentParser.ValidationException e = - assertThrows( - AssignmentParser.ValidationException.class, - () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); - assertThat(e).hasMessageThat().contains("overlaps"); + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(result.errorMessage).contains("overlaps"); + // ["", "a") and ["m", inf) are gaps around the slice that survived. + assertThat(result.assignment.getSlices()).hasSize(3); + assertSlice(result.assignment.getSlices().get(0), "", "a"); + assertSlice(result.assignment.getSlices().get(1), "a", "m"); + assertSlice(result.assignment.getSlices().get(2), "m", null); } @Test - public void parse_sliceExtendingToInfinityFollowedByAnother_throws() { + public void parse_duplicateStartKeys_theSecondOneIsDropped() { AssignmentChunk chunk = AssignmentChunk.newBuilder() - .addSliceAssignments(sliceAssignment("a", null)) - .addSliceAssignments(sliceAssignment("m", null)) + .addSliceAssignments(sliceAssignment("a", "m")) + .addSliceAssignments(sliceAssignment("a", "z")) .build(); - AssignmentParser.ValidationException e = - assertThrows( - AssignmentParser.ValidationException.class, - () -> AssignmentParser.parse(ImmutableList.of(chunk), 1)); - assertThat(e).hasMessageThat().contains("extends to the end of the keyspace"); + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(result.errorMessage).contains("overlaps"); + assertThat(result.assignment.getSlices()).hasSize(3); + assertSlice(result.assignment.getSlices().get(1), "a", "m"); } @Test - public void parse_unsignedByteOrderingIsUsed() throws Exception { - // 0x80 is negative as a signed byte but must sort after 0x01. + public void parse_sliceExtendingToInfinityFollowedByAnother_theLaterOneIsDropped() { AssignmentChunk chunk = AssignmentChunk.newBuilder() - .addSliceAssignments( - SliceAssignment.newBuilder() - .setSlice( - com.google.cloud.autosharding.v1.Slice.newBuilder() - .setStartKey(ByteString.copyFrom(new byte[] {(byte) 0x80})))) - .addSliceAssignments( - SliceAssignment.newBuilder() - .setSlice( - com.google.cloud.autosharding.v1.Slice.newBuilder() - .setStartKey(ByteString.copyFrom(new byte[] {0x01})) - .setEndKey(ByteString.copyFrom(new byte[] {(byte) 0x80})))) + .addSliceAssignments(sliceAssignment("a", null)) + .addSliceAssignments(sliceAssignment("m", null)) .build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); - // Leading gap ["", 0x01) plus the two declared slices. - assertThat(assignment.getSlices()).hasSize(3); - assertThat(assignment.getSlices().get(1).getStartKey()).isEqualTo(new byte[] {0x01}); - assertThat(assignment.getSlices().get(2).getStartKey()).isEqualTo(new byte[] {(byte) 0x80}); - assertThat(assignment.getSlices().get(2).getEndKey()).isNull(); + assertThat(result.errorMessage).contains("overlaps"); + assertThat(result.assignment.getSlices()).hasSize(2); + assertSlice(result.assignment.getSlices().get(0), "", "a"); + assertSlice(result.assignment.getSlices().get(1), "a", null); } @Test - public void parse_resultingSlicesAreContiguous() throws Exception { + public void parse_everySliceInvalid_yieldsNoAssignment() { AssignmentChunk chunk = AssignmentChunk.newBuilder() .addEndpoints(endpoint("host-a")) - .addSliceAssignments(sliceAssignment("b", "d", 0)) - .addSliceAssignments(sliceAssignment("k", "m", 0)) + .addSliceAssignments(sliceAssignment("", null, 1)) .build(); - Assignment assignment = AssignmentParser.parse(ImmutableList.of(chunk), 1); + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); - List slices = assignment.getSlices(); - assertThat(slices.get(0).getStartKey()).isEqualTo(new byte[0]); - for (int i = 0; i + 1 < slices.size(); i++) { - assertThat(slices.get(i).getEndKey()).isEqualTo(slices.get(i + 1).getStartKey()); + assertThat(result.assignment).isNull(); + assertThat(result.errorMessage).contains("out-of-range endpoint index 1"); + } + + @Test + public void parse_errorMessageIsCappedAtThreeProblems() { + AssignmentChunk.Builder chunk = AssignmentChunk.newBuilder(); + for (String startKey : new String[] {"v", "w", "x", "y", "z"}) { + chunk.addSliceAssignments(sliceAssignment(startKey, "a")); } - assertThat(slices.get(slices.size() - 1).getEndKey()).isNull(); + + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk.build()), 1); + + assertThat(result.assignment).isNull(); + assertThat(result.errorMessage).contains("and 2 more"); + } + + /** Parses chunks that are expected to be usable in their entirety. */ + private static Assignment parseFully(List chunks, long generation) { + AssignmentParser.Result result = AssignmentParser.parse(chunks, generation); + assertThat(result.errorMessage).isNull(); + assertThat(result.assignment).isNotNull(); + return result.assignment; } private static EndpointState endpoint(String name) { diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerTest.java index 9fb40741322..e33308f67d2 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerTest.java @@ -265,10 +265,7 @@ public void changedTarget_restartsStreamWithoutNewChannel() throws Exception { deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); takeRequest(); - AutoShardingLoadBalancerConfig retargeted = - new AutoShardingLoadBalancerConfig( - CHANNEL_FACTORY_KEY, "other-target", KEY_HEADER, true, ASSIGNMENT_TIMEOUT_NANOS); - deliverAddresses(retargeted, "a"); + deliverAddresses(retargetedConfig("other-target"), "a"); assertThat(channelFactory.keys).containsExactly(CHANNEL_FACTORY_KEY); assertThat(service.streamCount.get()).isEqualTo(2); @@ -277,35 +274,44 @@ public void changedTarget_restartsStreamWithoutNewChannel() throws Exception { @Test public void targetWithLocalityToken_isSubstituted() throws Exception { - AutoShardingLoadBalancerConfig localityConfig = - new AutoShardingLoadBalancerConfig( - CHANNEL_FACTORY_KEY, "target/%s", KEY_HEADER, true, ASSIGNMENT_TIMEOUT_NANOS); - EquivalentAddressGroup endpoint = - new EquivalentAddressGroup( - new NamedAddress("addr-a"), - Attributes.newBuilder() - .set(AutoShardingAttributes.ATTR_ENDPOINT_HOSTNAME, "a") - .set(EquivalentAddressGroup.ATTR_LOCALITY_NAME, "us-central1-a") - .build()); + acceptAddresses( + ResolvedAddresses.newBuilder() + .setAddresses(ImmutableList.of(endpointInLocality("a", "us-central1-a"))) + .setAttributes(attributesWithChannelFactory()) + .setLoadBalancingPolicyConfig(retargetedConfig("target/%s")) + .build()); + assertThat(takeRequest().getInitialClientConfig().getTarget()) + .isEqualTo("target/us-central1-a"); + } + + @Test + public void changedLocality_createsANewClientEvenThoughTheConfigIsUnchanged() throws Exception { + AutoShardingLoadBalancerConfig localityConfig = retargetedConfig("target/%s"); acceptAddresses( ResolvedAddresses.newBuilder() - .setAddresses(ImmutableList.of(endpoint)) + .setAddresses(ImmutableList.of(endpointInLocality("a", "us-central1-a"))) .setAttributes(attributesWithChannelFactory()) .setLoadBalancingPolicyConfig(localityConfig) .build()); + takeRequest(); + acceptAddresses( + ResolvedAddresses.newBuilder() + .setAddresses(ImmutableList.of(endpointInLocality("a", "us-central1-b"))) + .setAttributes(attributesWithChannelFactory()) + .setLoadBalancingPolicyConfig(localityConfig) + .build()); + + assertThat(channelFactory.keys).containsExactly(CHANNEL_FACTORY_KEY); + assertThat(service.streamCount.get()).isEqualTo(2); assertThat(takeRequest().getInitialClientConfig().getTarget()) - .isEqualTo("target/us-central1-a"); + .isEqualTo("target/us-central1-b"); } @Test public void targetWithLocalityToken_noLocality_substitutesEmptyString() throws Exception { - AutoShardingLoadBalancerConfig localityConfig = - new AutoShardingLoadBalancerConfig( - CHANNEL_FACTORY_KEY, "target/%s", KEY_HEADER, true, ASSIGNMENT_TIMEOUT_NANOS); - - deliverAddresses(localityConfig, "a"); + deliverAddresses(retargetedConfig("target/%s"), "a"); assertThat(takeRequest().getInitialClientConfig().getTarget()).isEqualTo("target/"); } @@ -385,6 +391,59 @@ public void newChannel_keepsServingPreviousAssignmentWhileTimerPending() throws assertThat(pickedHost(pick("a"))).isEqualTo("a"); } + @Test + public void newClient_restartsTheInitialAssignmentTimer() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a"); + deliverAssignment(1, slice("", "a")); + assertThat(fakeClock.numPendingTasks()).isEqualTo(0); + + deliverAddresses(retargetedConfig("other-target"), "a"); + + // The replacement client has to learn an assignment from scratch, so it gets the full + // timeout rather than inheriting the exhausted one. + assertThat(fakeClock.numPendingTasks()).isEqualTo(1); + } + + @Test + public void unusableAssignment_beforeAnyAssignment_stopsQueuingAndFallsBack() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + reportReady("a"); + reportReady("b"); + assertThat(currentState).isEqualTo(CONNECTING); + + pushUnusableAssignment(1); + + assertThat(currentState).isEqualTo(READY); + assertThat(pickedHost(pick("k"))).isAnyOf("a", "b"); + assertThat(fakeClock.numPendingTasks()).isEqualTo(0); + } + + @Test + public void unusableAssignment_beforeAnyAssignment_fallbackDisabled_failsRpcs() + throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, false), "a"); + reportReady("a"); + + pushUnusableAssignment(1); + + PickResult result = pick("k"); + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()).contains("fallback disabled"); + } + + @Test + public void unusableAssignment_afterAGoodOne_keepsServingTheGoodOne() throws Exception { + deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a", "b"); + deliverAssignment(1, slice("", "a"), slice("m", "b")); + reportReady("a"); + reportReady("b"); + + pushUnusableAssignment(2); + + assertThat(pickedHost(pick("alpha"))).isEqualTo("a"); + assertThat(pickedHost(pick("zulu"))).isEqualTo("b"); + } + // --------------------------------------------------------------------------------------------- // Routing on assignments // --------------------------------------------------------------------------------------------- @@ -598,6 +657,12 @@ private AutoShardingLoadBalancerConfig config(String channelFactoryKey, boolean channelFactoryKey, TARGET, KEY_HEADER, enableFallback, ASSIGNMENT_TIMEOUT_NANOS); } + /** The default config with a different {@code autosharding_target}. */ + private AutoShardingLoadBalancerConfig retargetedConfig(String target) { + return new AutoShardingLoadBalancerConfig( + CHANNEL_FACTORY_KEY, target, KEY_HEADER, true, ASSIGNMENT_TIMEOUT_NANOS); + } + private Attributes attributesWithChannelFactory() { return Attributes.newBuilder() .set(AutoShardingAttributes.ATTR_CHANNEL_FACTORY, channelFactory) @@ -632,11 +697,40 @@ private static List endpoints(String... hostnames) { return ImmutableList.copyOf(eags); } + private static EquivalentAddressGroup endpointInLocality(String hostname, String locality) { + return new EquivalentAddressGroup( + new NamedAddress("addr-" + hostname), + Attributes.newBuilder() + .set(AutoShardingAttributes.ATTR_ENDPOINT_HOSTNAME, hostname) + .set(EquivalentAddressGroup.ATTR_LOCALITY_NAME, locality) + .build()); + } + /** Sends an assignment from the fake service and waits for the load balancer to apply it. */ private void deliverAssignment(long generation, SliceSpec... slices) throws Exception { pushAssignment(generation, slices); } + /** + * Sends an assignment whose only slice fails validation. Nothing usable remains, so the client + * reports an error to the load balancer instead of an assignment. + */ + private void pushUnusableAssignment(long generation) throws Exception { + StreamObserver serverStream = currentServerStream(); + serverStream.onNext( + WatchShardingAssignmentResponse.newBuilder() + .setChunk( + AssignmentChunk.newBuilder() + .addEndpoints(EndpointState.newBuilder().setEndpoint("a")) + // Index 7 is past the end of the endpoint list above. + .addSliceAssignments(sliceAssignment("", null, 7))) + .build()); + serverStream.onNext( + WatchShardingAssignmentResponse.newBuilder() + .setMetadata(AssignmentMetadata.newBuilder().setGeneration(generation)) + .build()); + } + private void pushAssignment(long generation, SliceSpec... slices) throws Exception { StreamObserver serverStream = currentServerStream(); List endpointNames = new ArrayList<>(); diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoshardingClientTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoshardingClientTest.java index ade4ca1240c..8c3084443b4 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AutoshardingClientTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoshardingClientTest.java @@ -38,6 +38,8 @@ import io.grpc.internal.FakeClock; import io.grpc.stub.StreamObserver; import io.grpc.testing.GrpcCleanupRule; +import java.util.ArrayList; +import java.util.List; import java.util.concurrent.BlockingQueue; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; @@ -69,44 +71,35 @@ public class AutoshardingClientTest { private final FakeClock fakeClock = new FakeClock(); private final FakeAutoshardingService service = new FakeAutoshardingService(); private final BlockingQueue assignments = new LinkedBlockingQueue<>(); + private final BlockingQueue errors = new LinkedBlockingQueue<>(); private final RecordingBackoffPolicyProvider backoffPolicyProvider = new RecordingBackoffPolicyProvider(); + private final List clients = new ArrayList<>(); private Channel channel; private AutoshardingClient client; @Before public void setUp() throws Exception { - String serverName = InProcessServerBuilder.generateName(); - grpcCleanup.register( - InProcessServerBuilder.forName(serverName) - .directExecutor() - .addService(service) - .build() - .start()); - channel = - grpcCleanup.register( - InProcessChannelBuilder.forName(serverName).directExecutor().build()); - client = - new AutoshardingClient( - CLIENT_UUID, - syncContext, - fakeClock.getScheduledExecutorService(), - backoffPolicyProvider, - fakeClock.getStopwatchSupplier(), - assignments::add); + channel = newChannelToFakeService(); + client = newClient(channel, TARGET); } @After public void tearDown() { - // Must happen before GrpcCleanupRule shuts the channel down, otherwise the client keeps + // Must happen before GrpcCleanupRule shuts the channels down, otherwise a client keeps // retrying against a terminating channel. - syncContext.execute(client::shutdown); + syncContext.execute( + () -> { + for (AutoshardingClient created : clients) { + created.shutdown(); + } + }); } @Test - public void update_startsStreamAndSendsInitialClientConfig() throws Exception { - update(channel, TARGET); + public void start_opensStreamAndSendsInitialClientConfig() throws Exception { + start(client); WatchShardingAssignmentRequest request = takeRequest(); assertThat(request.hasInitialClientConfig()).isTrue(); @@ -115,19 +108,9 @@ public void update_startsStreamAndSendsInitialClientConfig() throws Exception { assertThat(request.getInitialClientConfig().getLatestGeneration()).isEqualTo(0); } - @Test - public void update_unchanged_doesNotRestartStream() throws Exception { - update(channel, TARGET); - takeRequest(); - - update(channel, TARGET); - - assertThat(service.streamCount.get()).isEqualTo(1); - } - @Test public void chunksBufferedUntilMetadata_thenAssignmentDeliveredAndAcked() throws Exception { - update(channel, TARGET); + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); @@ -151,7 +134,7 @@ public void chunksBufferedUntilMetadata_thenAssignmentDeliveredAndAcked() throws @Test public void multipleChunks_combinedIntoOneLogicalAssignment() throws Exception { - update(channel, TARGET); + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); @@ -171,12 +154,13 @@ public void multipleChunks_combinedIntoOneLogicalAssignment() throws Exception { } @Test - public void invalidAssignment_nackedAndNotDelivered() throws Exception { - update(channel, TARGET); + public void noUsableSlices_nackedAndReportedAsAnError() throws Exception { + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); - // Endpoint index 3 does not exist in the combined endpoint list. + // Endpoint index 3 does not exist in the combined endpoint list, so the only slice is + // dropped and nothing usable remains. serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 3))); serverStream.onNext(metadataResponse(5)); @@ -186,19 +170,51 @@ public void invalidAssignment_nackedAndNotDelivered() throws Exception { assertThat(nack.getAssignmentAck().getAccepted()).isFalse(); assertThat(nack.getAssignmentAck().getErrorMessage()) .contains("out-of-range endpoint index 3"); + assertThat(takeError().getDescription()).contains("out-of-range endpoint index 3"); assertThat(assignments).isEmpty(); + // A rejected assignment must not advance the watermark, or the server would stop resending. assertThat(client.getLatestGeneration()).isEqualTo(0); } @Test - public void invalidAssignment_doesNotLeakChunksIntoNextAssignment() throws Exception { - update(channel, TARGET); + public void someSlicesDropped_ackedWithErrorMessageAndStillDelivered() throws Exception { + start(client); + takeRequest(); + StreamObserver serverStream = takeServerStream(); + + serverStream.onNext( + chunkResponse( + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addSliceAssignments(sliceAssignment("", "m", 0)) + .addSliceAssignments(sliceAssignment("m", null, 3)) + .build())); + serverStream.onNext(metadataResponse(5)); + + Assignment assignment = takeAssignment(); + assertThat(assignment.getSlices()).hasSize(2); + assertThat(assignment.getSlices().get(0).getEndpoints()).containsExactly(0); + // The dropped slice was turned into a gap rather than invalidating the assignment. + assertThat(assignment.getSlices().get(1).getEndpoints()).isEmpty(); + + WatchShardingAssignmentRequest ack = takeRequest(); + assertThat(ack.getAssignmentAck().getAccepted()).isTrue(); + assertThat(ack.getAssignmentAck().getErrorMessage()) + .contains("out-of-range endpoint index 3"); + assertThat(errors).isEmpty(); + assertThat(client.getLatestGeneration()).isEqualTo(5); + } + + @Test + public void rejectedAssignment_doesNotLeakChunksIntoTheNextOne() throws Exception { + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 3))); serverStream.onNext(metadataResponse(5)); takeRequest(); // NACK + takeError(); serverStream.onNext(chunkResponse(chunkWithEndpoint("host-b", "", null, 0))); serverStream.onNext(metadataResponse(6)); @@ -208,8 +224,8 @@ public void invalidAssignment_doesNotLeakChunksIntoNextAssignment() throws Excep } @Test - public void staleGeneration_droppedWithoutAck() throws Exception { - update(channel, TARGET); + public void staleGeneration_nackedAndNotDelivered() throws Exception { + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); @@ -221,14 +237,19 @@ public void staleGeneration_droppedWithoutAck() throws Exception { serverStream.onNext(chunkResponse(chunkWithEndpoint("host-b", "", null, 0))); serverStream.onNext(metadataResponse(5)); + WatchShardingAssignmentRequest nack = takeRequest(); + assertThat(nack.getAssignmentAck().getGeneration()).isEqualTo(5); + assertThat(nack.getAssignmentAck().getAccepted()).isFalse(); + assertThat(nack.getAssignmentAck().getErrorMessage()).contains("stale generation"); + // A stale assignment tells the LB policy nothing it does not already know. assertThat(assignments).isEmpty(); - assertThat(service.requests).isEmpty(); + assertThat(errors).isEmpty(); assertThat(client.getLatestGeneration()).isEqualTo(5); } @Test - public void olderGeneration_droppedWithoutAck() throws Exception { - update(channel, TARGET); + public void olderGeneration_nackedAndNotDelivered() throws Exception { + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); @@ -240,13 +261,16 @@ public void olderGeneration_droppedWithoutAck() throws Exception { serverStream.onNext(chunkResponse(chunkWithEndpoint("host-b", "", null, 0))); serverStream.onNext(metadataResponse(4)); + WatchShardingAssignmentRequest nack = takeRequest(); + assertThat(nack.getAssignmentAck().getGeneration()).isEqualTo(4); + assertThat(nack.getAssignmentAck().getAccepted()).isFalse(); assertThat(assignments).isEmpty(); - assertThat(service.requests).isEmpty(); + assertThat(client.getLatestGeneration()).isEqualTo(5); } @Test public void loadReportingConfig_ignored() throws Exception { - update(channel, TARGET); + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); @@ -261,7 +285,7 @@ public void loadReportingConfig_ignored() throws Exception { @Test public void streamFailure_reconnectsAndSendsLatestGeneration() throws Exception { - update(channel, TARGET); + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); @@ -282,7 +306,7 @@ public void streamFailure_reconnectsAndSendsLatestGeneration() throws Exception @Test public void streamFailure_doesNotReconnectBeforeBackoffElapses() throws Exception { - update(channel, TARGET); + start(client); takeRequest(); takeServerStream().onError(Status.UNAVAILABLE.asRuntimeException()); @@ -298,7 +322,7 @@ public void streamFailure_doesNotReconnectBeforeBackoffElapses() throws Exceptio @Test public void streamCompletedByServer_reconnects() throws Exception { - update(channel, TARGET); + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); @@ -311,7 +335,7 @@ public void streamCompletedByServer_reconnects() throws Exception { @Test public void backoffSequence_onlyResetAfterGoodAssignment() throws Exception { - update(channel, TARGET); + start(client); takeRequest(); // First failure with no assignment received: a backoff sequence is created. @@ -338,9 +362,14 @@ public void backoffSequence_onlyResetAfterGoodAssignment() throws Exception { assertThat(backoffPolicyProvider.timesCalled).isEqualTo(2); } + /** + * The LB policy answers a target or channel change by replacing the client rather than by + * updating it, so the accepted-generation watermark never crosses over to a different server or + * resource. See gRFC A119, "Communicating with the Autosharding service". + */ @Test - public void targetChange_restartsStreamAndResetsGeneration() throws Exception { - update(channel, TARGET); + public void newClient_startsFromGenerationZero() throws Exception { + start(client); takeRequest(); StreamObserver serverStream = takeServerStream(); serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0))); @@ -348,48 +377,22 @@ public void targetChange_restartsStreamAndResetsGeneration() throws Exception { takeAssignment(); takeRequest(); // ACK assertThat(client.getLatestGeneration()).isEqualTo(9); + syncContext.execute(client::shutdown); - update(channel, OTHER_TARGET); + AutoshardingClient replacement = newClient(newChannelToFakeService(), OTHER_TARGET); + start(replacement); WatchShardingAssignmentRequest request = takeRequest(); assertThat(request.hasInitialClientConfig()).isTrue(); assertThat(request.getInitialClientConfig().getTarget()).isEqualTo(OTHER_TARGET); assertThat(request.getInitialClientConfig().getLatestGeneration()).isEqualTo(0); - assertThat(client.getLatestGeneration()).isEqualTo(0); - assertThat(service.streamCount.get()).isEqualTo(2); - } - - @Test - public void channelChange_restartsStreamAndResetsGeneration() throws Exception { - update(channel, TARGET); - takeRequest(); - StreamObserver serverStream = takeServerStream(); - serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0))); - serverStream.onNext(metadataResponse(9)); - takeAssignment(); - takeRequest(); // ACK - - String otherServerName = InProcessServerBuilder.generateName(); - grpcCleanup.register( - InProcessServerBuilder.forName(otherServerName) - .directExecutor() - .addService(service) - .build() - .start()); - Channel otherChannel = - grpcCleanup.register( - InProcessChannelBuilder.forName(otherServerName).directExecutor().build()); - - update(otherChannel, TARGET); - - WatchShardingAssignmentRequest request = takeRequest(); - assertThat(request.getInitialClientConfig().getLatestGeneration()).isEqualTo(0); + assertThat(replacement.getLatestGeneration()).isEqualTo(0); assertThat(service.streamCount.get()).isEqualTo(2); } @Test public void shutdown_cancelsStreamAndStopsReconnecting() throws Exception { - update(channel, TARGET); + start(client); takeRequest(); syncContext.execute(client::shutdown); @@ -399,19 +402,54 @@ public void shutdown_cancelsStreamAndStopsReconnecting() throws Exception { } @Test - public void shutdown_isIdempotentAndIgnoresLaterUpdates() throws Exception { - update(channel, TARGET); + public void shutdown_isIdempotent() throws Exception { + start(client); takeRequest(); syncContext.execute(client::shutdown); syncContext.execute(client::shutdown); - update(channel, OTHER_TARGET); assertThat(service.streamCount.get()).isEqualTo(1); } - private void update(Channel channel, String target) { - syncContext.execute(() -> client.update(channel, target)); + @Test + public void shutdown_beforeStart_leavesNothingBehind() { + syncContext.execute(client::shutdown); + + assertThat(service.streamCount.get()).isEqualTo(0); + assertThat(fakeClock.numPendingTasks()).isEqualTo(0); + } + + private Channel newChannelToFakeService() throws Exception { + String serverName = InProcessServerBuilder.generateName(); + grpcCleanup.register( + InProcessServerBuilder.forName(serverName) + .directExecutor() + .addService(service) + .build() + .start()); + return grpcCleanup.register( + InProcessChannelBuilder.forName(serverName).directExecutor().build()); + } + + /** Creates a client and registers it for shutdown, without starting it. */ + private AutoshardingClient newClient(Channel channel, String target) { + AutoshardingClient created = + new AutoshardingClient( + CLIENT_UUID, + syncContext, + fakeClock.getScheduledExecutorService(), + backoffPolicyProvider, + fakeClock.getStopwatchSupplier(), + channel, + target, + new RecordingWatcher()); + clients.add(created); + return created; + } + + private void start(AutoshardingClient target) { + syncContext.execute(target::start); } /** Asserts that a retry was scheduled and advances the clock so that it runs. */ @@ -446,6 +484,14 @@ private Assignment takeAssignment() throws Exception { return assignment; } + private Status takeError() throws Exception { + Status error = errors.poll(TIMEOUT_SECONDS, TimeUnit.SECONDS); + if (error == null) { + fail("timed out waiting for an error"); + } + return error; + } + private static WatchShardingAssignmentResponse chunkResponse(AssignmentChunk chunk) { return WatchShardingAssignmentResponse.newBuilder().setChunk(chunk).build(); } @@ -483,6 +529,18 @@ private static SliceAssignment sliceAssignment( return builder.build(); } + private final class RecordingWatcher implements AutoshardingClient.AssignmentWatcher { + @Override + public void onAssignment(Assignment assignment) { + assignments.add(assignment); + } + + @Override + public void onError(Status error) { + errors.add(error); + } + } + private static final class FakeAutoshardingService extends AutoshardingServiceGrpc.AutoshardingServiceImplBase { final BlockingQueue requests = new LinkedBlockingQueue<>(); From 34a7cce671b4b8a001e46b716ce6ab0ed946cd92 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 22 Sep 2026 15:40:40 +0530 Subject: [PATCH 31/33] updates --- .../grpc/autosharding/AssignmentParser.java | 81 +++++++++++++--- .../AutoShardingLoadBalancer.java | 7 ++ .../grpc/autosharding/AutoshardingClient.java | 22 ++++- .../autosharding/AssignmentParserTest.java | 93 ++++++++++++++++++- 4 files changed, 186 insertions(+), 17 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java b/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java index 10120312650..5f29f306faa 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AssignmentParser.java @@ -81,8 +81,19 @@ private Result(@Nullable Assignment assignment, @Nullable String errorMessage) { } } - /** At most this many dropped slices are named in {@link Result#errorMessage}. */ - private static final int MAX_REPORTED_PROBLEMS = 3; + /** + * The limit on {@code AssignmentAck.error_message}, from {@code autosharding.proto}: "The + * length of this field MUST NOT exceed 512 characters". + */ + private static final int MAX_ERROR_MESSAGE_CHARS = 512; + + /** + * How much of a key to hex-encode into a description. Keys may be up to 512 bytes, and only + * the leading bytes are needed to tell one slice from another in a log. + */ + private static final int MAX_ENCODED_KEY_BYTES = 8; + + private static final String SEPARATOR = "; "; private static final Comparator UNSIGNED_BYTES_COMPARATOR = UnsignedBytes.lexicographicalComparator(); @@ -160,10 +171,18 @@ private static List combineSlices( continue; } if (keyOrder == 0) { - // Permitted by the gRFC's "start_key <= end_key" rule, but it covers no keys and - // would put two entries with the same start key in the SliceMap, which makes the - // picker's binary search ambiguous. Dropping it leaves no gap: its neighbours - // already meet at this key. + // Satisfies the gRFC's "start_key <= end_key", but end_key is exclusive, so + // [k, k) is the empty range rather than the single key k. A server wanting to + // assign one key sends [k, k+1), i.e. an end_key of k with a trailing 0x00. + // + // Keeping it would put a second entry with the same start key in the SliceMap, + // and the picker resolves a key by binary search over start keys, so a key equal + // to this one could resolve to either entry. Any endpoints on it are no loss: + // no key can fall in an empty range, so they were unreachable through it anyway, + // and the assignment's endpoint list is built from the chunks, not from slices. + // + // Nor can dropping it open a gap -- it covered nothing -- so the gap-filling pass + // below produces the same coverage with or without it. dropped.add( String.format("slice [%s, %s) is empty", encode(startKey), encode(endKey))); continue; @@ -256,16 +275,54 @@ private static List fillGaps(List sorted) { return filled; } - /** Summarizes the dropped slices, capped so that the ack stays a reasonable size. */ + /** + * Summarizes the dropped slices, reporting as many as the {@code error_message} budget of an + * {@code AssignmentAck} allows and naming the count of those left out. + * + *

    The result is sized to fit within {@link #MAX_ERROR_MESSAGE_CHARS} so that + * {@code AutoshardingClient}'s final truncation never has to cut a description in half. + */ private static String describe(List dropped) { - if (dropped.size() <= MAX_REPORTED_PROBLEMS) { - return String.join("; ", dropped); + StringBuilder message = new StringBuilder(); + int reported = 0; + for (String problem : dropped) { + int separator = reported == 0 ? 0 : SEPARATOR.length(); + // Leave room for the suffix that will be needed if this is where we stop. + int reserved = andMore(dropped.size() - reported - 1).length(); + if (message.length() + separator + problem.length() + reserved + > MAX_ERROR_MESSAGE_CHARS) { + break; + } + if (reported > 0) { + message.append(SEPARATOR); + } + message.append(problem); + reported++; + } + if (reported == 0) { + // Not reachable while every description is bounded, but a lone oversized one is better + // reported in part than not at all; AutoshardingClient trims it to the limit. + return dropped.get(0); } - return String.join("; ", dropped.subList(0, MAX_REPORTED_PROBLEMS)) - + String.format("; and %s more", dropped.size() - MAX_REPORTED_PROBLEMS); + return message + andMore(dropped.size() - reported); + } + + private static String andMore(int omitted) { + return omitted == 0 ? "" : String.format("; and %s more", omitted); } + /** + * Hex-encodes a key for a human-readable description, shortening it if it is long. The + * protocol allows keys of up to 512 bytes, which would fill the entire error message budget + * twice over. + */ private static String encode(@Nullable byte[] key) { - return key == null ? "inf" : BaseEncoding.base16().encode(key); + if (key == null) { + return "inf"; + } + if (key.length <= MAX_ENCODED_KEY_BYTES) { + return BaseEncoding.base16().encode(key); + } + return BaseEncoding.base16().encode(key, 0, MAX_ENCODED_KEY_BYTES) + "..."; } } diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java index 726feb213bb..9168762b430 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java @@ -340,6 +340,13 @@ private void maybeRecreateClient(boolean channelChanged, String newTarget, long /** * Substitutes the optional {@code %s} token in the configured target with the locality of the * resolved endpoints, or with the empty string when no locality is available. + * + *

    The locality is read from {@link EquivalentAddressGroup#ATTR_LOCALITY_NAME}, which is a + * plain {@code io.grpc} attribute rather than an xDS-specific one. That keeps a single code path + * for both deployments: under xDS the attribute is populated by the cluster resolver, and + * without xDS gRFC A119 makes it the user's responsibility to have their name resolver populate + * it if their target contains a {@code %s} token. All endpoints handed to one instance of this + * policy belong to the same locality, so the first one is representative. */ private static String resolveTarget( AutoShardingLoadBalancerConfig config, List endpoints) { diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java b/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java index 0bfdbbb9cad..88c0e7d3daa 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java @@ -62,6 +62,9 @@ final class AutoshardingClient { private static final Logger logger = Logger.getLogger(AutoshardingClient.class.getName()); + /** The limit {@code autosharding.proto} places on {@code AssignmentAck.error_message}. */ + private static final int MAX_ERROR_MESSAGE_CODE_POINTS = 512; + /** * Receives validated assignments from the autosharding service. */ @@ -380,9 +383,24 @@ void close(Status status) { } /** - * The {@code error_message} field must not exceed 512 characters. + * Enforces the limit {@code autosharding.proto} places on {@code AssignmentAck.error_message}: + * "The length of this field MUST NOT exceed 512 characters (Unicode code points, see + * https://google.aip.dev/210)". + * + *

    A backstop only. {@link AssignmentParser} already assembles its summary to fit, so this + * should never actually cut anything. */ private static String truncateErrorMessage(String message) { - return message.length() <= 512 ? message : message.substring(0, 512); + // A string never has more code points than chars, so this settles the common case without + // walking it. + if (message.length() <= MAX_ERROR_MESSAGE_CODE_POINTS) { + return message; + } + if (message.codePointCount(0, message.length()) <= MAX_ERROR_MESSAGE_CODE_POINTS) { + return message; + } + // Cutting on a code point boundary rather than a char boundary keeps a surrogate pair from + // being split into an unpaired surrogate, which does not survive UTF-8 encoding. + return message.substring(0, message.offsetByCodePoints(0, MAX_ERROR_MESSAGE_CODE_POINTS)); } } diff --git a/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java b/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java index 2303b12c1bb..b8da410ad71 100644 --- a/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java +++ b/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java @@ -324,6 +324,54 @@ public void parse_zeroWidthSlice_isDropped() { assertSlice(result.assignment.getSlices().get(1), "m", null, 1); } + /** + * An empty range is unroutable whatever it carries, and it need not sit next to another + * slice, so dropping it has to fall through to ordinary gap filling. + */ + @Test + public void parse_zeroWidthSlice_withEndpointsAndNoNeighbour_leavesNoHole() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + .addEndpoints(endpoint("host-b")) + .addSliceAssignments(sliceAssignment("", "a", 0)) + .addSliceAssignments(sliceAssignment("m", "m", 1)) + .addSliceAssignments(sliceAssignment("z", null, 0)) + .build(); + + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(result.errorMessage).contains("is empty"); + assertThat(result.assignment.getSlices()).hasSize(3); + assertSlice(result.assignment.getSlices().get(0), "", "a", 0); + // ["a", "z") is one gap, not two slices meeting at "m". + assertSlice(result.assignment.getSlices().get(1), "a", "z"); + assertSlice(result.assignment.getSlices().get(2), "z", null, 0); + } + + @Test + public void parse_singleKeySlice_isKept() { + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addEndpoints(endpoint("host-a")) + // How a server actually assigns exactly one key: end_key is the successor of + // start_key, not start_key itself. + .addSliceAssignments( + SliceAssignment.newBuilder() + .setSlice( + com.google.cloud.autosharding.v1.Slice.newBuilder() + .setStartKey(ByteString.copyFromUtf8("m")) + .setEndKey(ByteString.copyFrom(new byte[] {'m', 0}))) + .addEndpoints(PerSliceEndpointState.newBuilder().setEndpointIndex(0))) + .build(); + + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + assertThat(result.errorMessage).isNull(); + assertThat(result.assignment.getSlices()).hasSize(3); + assertThat(result.assignment.getSlices().get(1).getEndpoints()).containsExactly(0); + } + /** * The picker looks a key up by binary search over start keys, so two slices sharing one would * make the result depend on where the search happened to land. @@ -411,7 +459,24 @@ public void parse_everySliceInvalid_yieldsNoAssignment() { } @Test - public void parse_errorMessageIsCappedAtThreeProblems() { + public void parse_manyProblems_errorMessageStaysWithinTheAckBudget() { + AssignmentChunk.Builder chunk = AssignmentChunk.newBuilder(); + for (int i = 0; i < 40; i++) { + // Inverted key range, so every one of them is dropped. + chunk.addSliceAssignments(sliceAssignment("z" + i, "a")); + } + + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk.build()), 1); + + assertThat(result.assignment).isNull(); + // autosharding.proto: error_message "MUST NOT exceed 512 characters". + assertThat(result.errorMessage.length()).isAtMost(512); + assertThat(result.errorMessage).contains("greater than end_key"); + assertThat(result.errorMessage).containsMatch("; and \\d+ more$"); + } + + @Test + public void parse_fewProblems_allAreReported() { AssignmentChunk.Builder chunk = AssignmentChunk.newBuilder(); for (String startKey : new String[] {"v", "w", "x", "y", "z"}) { chunk.addSliceAssignments(sliceAssignment(startKey, "a")); @@ -419,8 +484,30 @@ public void parse_errorMessageIsCappedAtThreeProblems() { AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk.build()), 1); - assertThat(result.assignment).isNull(); - assertThat(result.errorMessage).contains("and 2 more"); + // Five short descriptions fit comfortably, so nothing is elided. + assertThat(result.errorMessage).doesNotContain("more"); + assertThat(result.errorMessage.split("; ")).hasLength(5); + } + + @Test + public void parse_longKeysAreShortenedInTheErrorMessage() { + byte[] longKey = new byte[512]; + Arrays.fill(longKey, (byte) 0xAB); + AssignmentChunk chunk = + AssignmentChunk.newBuilder() + .addSliceAssignments( + SliceAssignment.newBuilder() + .setSlice( + com.google.cloud.autosharding.v1.Slice.newBuilder() + .setStartKey(ByteString.copyFrom(longKey)) + .setEndKey(ByteString.copyFromUtf8("a")))) + .build(); + + AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1); + + // Hex-encoding 512 bytes in full would be 1024 characters on its own. + assertThat(result.errorMessage.length()).isAtMost(512); + assertThat(result.errorMessage).contains("..."); } /** Parses chunks that are expected to be usable in their entirety. */ From 6e3cd8d4e35285f3c941dc9ddcc8827bcfe1e475 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 22 Sep 2026 15:48:40 +0530 Subject: [PATCH 32/33] updates --- .../java/io/grpc/autosharding/Assignment.java | 4 +--- .../autosharding/AutoShardingLoadBalancer.java | 9 +++------ .../io/grpc/autosharding/AutoshardingClient.java | 12 +++--------- .../java/io/grpc/autosharding/EndpointMap.java | 7 ++----- .../io/grpc/autosharding/PickerEndpoint.java | 15 ++------------- .../main/java/io/grpc/autosharding/SliceMap.java | 16 +++------------- 6 files changed, 14 insertions(+), 49 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/Assignment.java b/autosharding/src/main/java/io/grpc/autosharding/Assignment.java index 0ef9784ca18..d243e061f97 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/Assignment.java +++ b/autosharding/src/main/java/io/grpc/autosharding/Assignment.java @@ -46,9 +46,7 @@ @Immutable final class Assignment { - /** - * A single contiguous key range and the endpoints assigned to it. - */ + /** A single contiguous key range and the endpoints assigned to it. */ @Immutable @SuppressWarnings("Immutable") // Defensive copies are made; arrays are never mutated. static final class Slice { diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java index 9168762b430..9a6eea0197d 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java @@ -341,12 +341,9 @@ private void maybeRecreateClient(boolean channelChanged, String newTarget, long * Substitutes the optional {@code %s} token in the configured target with the locality of the * resolved endpoints, or with the empty string when no locality is available. * - *

    The locality is read from {@link EquivalentAddressGroup#ATTR_LOCALITY_NAME}, which is a - * plain {@code io.grpc} attribute rather than an xDS-specific one. That keeps a single code path - * for both deployments: under xDS the attribute is populated by the cluster resolver, and - * without xDS gRFC A119 makes it the user's responsibility to have their name resolver populate - * it if their target contains a {@code %s} token. All endpoints handed to one instance of this - * policy belong to the same locality, so the first one is representative. + *

    {@link EquivalentAddressGroup#ATTR_LOCALITY_NAME} is preferred over an xDS-specific + * attribute so that a plain name resolver can supply it too, which is what gRFC A119 asks of + * non-xDS users. All endpoints reaching one instance of this policy share a locality. */ private static String resolveTarget( AutoShardingLoadBalancerConfig config, List endpoints) { diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java b/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java index 88c0e7d3daa..5ec2b9cb526 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java @@ -65,13 +65,9 @@ final class AutoshardingClient { /** The limit {@code autosharding.proto} places on {@code AssignmentAck.error_message}. */ private static final int MAX_ERROR_MESSAGE_CODE_POINTS = 512; - /** - * Receives validated assignments from the autosharding service. - */ + /** Receives validated assignments from the autosharding service. */ interface AssignmentWatcher { - /** - * Called with a newly accepted assignment. Invoked on the {@link SynchronizationContext}. - */ + /** Called with a newly accepted assignment. Invoked on the sync context. */ void onAssignment(Assignment assignment); /** @@ -212,9 +208,7 @@ private void scheduleRetry(boolean receivedGoodAssignment) { syncContext.schedule(this::startStream, delayNanos, TimeUnit.NANOSECONDS, timerService); } - /** - * A single {@code WatchShardingAssignment} stream. - */ + /** A single {@code WatchShardingAssignment} stream. */ private final class AutoshardingStream implements ClientResponseObserver< WatchShardingAssignmentRequest, WatchShardingAssignmentResponse> { diff --git a/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java index f61c457afb5..a83aa634f84 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java @@ -178,7 +178,7 @@ void updateEndpoints(List endpoints, Attributes attribut } } - /** Returns the number of endpoints currently held. */ + // Returns the number of endpoints currently held int size() { return holders.size(); } @@ -347,10 +347,7 @@ void updateAddresses(EquivalentAddressGroup endpoint, Attributes attributes) { } } - /** - * Starts connecting if this endpoint is IDLE. Must be called from the synchronization - * context. - */ + /** Starts connecting if this endpoint is IDLE. */ void requestConnection() { if (childShutdown || state != IDLE) { return; diff --git a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java index d85189d222e..12f0fb61ff2 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java +++ b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java @@ -23,14 +23,10 @@ import io.grpc.LoadBalancer.SubchannelPicker; import javax.annotation.Nullable; -/** - * Immutable snapshot of endpoint state used by {@link AutoShardingPicker}. - */ +/** 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. - */ + /** Callback interface to trigger connection attempts on an IDLE endpoint's child balancer. */ @FunctionalInterface interface ExitIdler { /** @@ -46,13 +42,6 @@ interface ExitIdler { private final SubchannelPicker picker; @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 exitIdler a callback to trigger an IDLE child balancer to start connecting - */ PickerEndpoint( ConnectivityState state, SubchannelPicker picker, diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java index b5e640492c2..ff7b07a8adb 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -28,7 +28,7 @@ /** * An immutable lookup structure mapping application routing keys to slice indices. * - *

    The assignment provider guarantees that the assignment is pre-validated, gap-free, + *

    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 ""}).
      • @@ -45,28 +45,18 @@ * {@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 { - /** - * Represents a single key-range slice mapping to endpoint indices in the picker. - */ + /** Represents a single key-range slice mapping to endpoint indices in the picker. */ static final class SliceEntry { 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 - */ + /** Constructs an entry starting at the inclusive key {@code startKey}. */ SliceEntry(byte[] startKey, List endpoints) { this.startKey = checkNotNull(startKey, "startKey").clone(); this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints")); From a957a802d1b68c7445b67d5341bc479f7dd39825 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Tue, 22 Sep 2026 15:51:31 +0530 Subject: [PATCH 33/33] updates --- .../io/grpc/autosharding/PickerEndpoint.java | 15 +++++++++++++-- .../main/java/io/grpc/autosharding/SliceMap.java | 16 +++++++++++++--- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java index 12f0fb61ff2..d85189d222e 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java +++ b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java @@ -23,10 +23,14 @@ import io.grpc.LoadBalancer.SubchannelPicker; import javax.annotation.Nullable; -/** Immutable snapshot of endpoint state used by {@link AutoShardingPicker}. */ +/** + * 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. */ + /** + * Callback interface to trigger connection attempts on an IDLE endpoint's child balancer. + */ @FunctionalInterface interface ExitIdler { /** @@ -42,6 +46,13 @@ interface ExitIdler { private final SubchannelPicker picker; @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 exitIdler a callback to trigger an IDLE child balancer to start connecting + */ PickerEndpoint( ConnectivityState state, SubchannelPicker picker, diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java index ff7b07a8adb..b5e640492c2 100644 --- a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -28,7 +28,7 @@ /** * An immutable lookup structure mapping application routing keys to slice indices. * - *

    The assignment provider guarantees that the assignment is pre-validated, gap-free, + *

    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 ""}).
      • @@ -45,18 +45,28 @@ * {@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 { - /** Represents a single key-range slice mapping to endpoint indices in the picker. */ + /** + * Represents a single key-range slice mapping to endpoint indices in the picker. + */ static final class SliceEntry { private final byte[] startKey; private final ImmutableList endpoints; - /** Constructs an entry starting at the inclusive key {@code startKey}. */ + /** + * 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").clone(); this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints"));