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