Validation follows gRFC A119, "Handling assignments from the Autosharding server". A slice
+ * is usable only if all of the following hold:
+ *
+ *
Maps onto the three non-stale rows of the outcome table in gRFC A119, "Handling
+ * assignments from the Autosharding server":
+ *
+ *
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.
+ *
+ * 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}
+ */
+ static Result parse(List chunks, long generation) {
+ checkNotNull(chunks, "chunks");
+
+ ImmutableList endpointNames = combineEndpointNames(chunks);
+ List dropped = new ArrayList<>();
+ List slices = combineSlices(chunks, endpointNames.size(), dropped);
+
+ slices.sort(
+ (s1, s2) -> UNSIGNED_BYTES_COMPARATOR.compare(s1.getStartKey(), s2.getStartKey()));
+ slices = dropOverlaps(slices, dropped);
+
+ 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);
+ }
+
+ /**
+ * 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, 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, List dropped) {
+ 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) {
+ 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) {
+ // 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. Dropping this cannot open a gap, because it covered nothing.
+ 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) {
+ indexProblem =
+ String.format(
+ "slice starting at %s references out-of-range endpoint index %s;"
+ + " assignment contains %s endpoints",
+ encode(startKey), index, endpointCount);
+ break;
+ }
+ endpoints.add(index);
+ }
+ if (indexProblem != null) {
+ dropped.add(indexProblem);
+ continue;
+ }
+ slices.add(new Assignment.Slice(startKey, endKey, endpoints));
+ }
+ }
+ return slices;
+ }
+
+ /**
+ * Returns the slices of {@code sorted} that overlap no other slice.
+ *
+ * When slices overlap, every one of them is dropped. The server has told us two different
+ * things about the same key and there is no basis for preferring either, so the keys they cover
+ * become a gap.
+ *
+ *
Overlap is transitive here in the sense that matters: a slice is dropped when it overlaps
+ * any other slice, even one it only reaches through a third. {@code ["a", "z")}, {@code ["b",
+ * "c")} and {@code ["d", "e")} all go, because the first overlaps the other two.
+ *
+ * @param sorted slices in ascending {@code startKey} order
+ * @param dropped collects a description of each slice that was dropped
+ */
+ private static List dropOverlaps(
+ List sorted, List dropped) {
+ List kept = new ArrayList<>(sorted.size());
+ int index = 0;
+ while (index < sorted.size()) {
+ // Extend a run of slices for as long as the keys covered so far reach into the next one.
+ // Every slice that joins overlaps some earlier member, and the first two overlap directly,
+ // so a run longer than one slice consists entirely of slices that overlap something.
+ byte[] runEndKey = sorted.get(index).getEndKey();
+ int end = index + 1;
+ while (end < sorted.size() && reaches(runEndKey, sorted.get(end).getStartKey())) {
+ byte[] endKey = sorted.get(end).getEndKey();
+ if (endKey == null || UNSIGNED_BYTES_COMPARATOR.compare(endKey, runEndKey) > 0) {
+ runEndKey = endKey;
+ }
+ end++;
+ }
+
+ if (end - index == 1) {
+ kept.add(sorted.get(index));
+ } else {
+ for (Assignment.Slice slice : sorted.subList(index, end)) {
+ dropped.add(
+ String.format(
+ "slice [%s, %s) overlaps another slice",
+ encode(slice.getStartKey()), encode(slice.getEndKey())));
+ }
+ }
+ index = end;
+ }
+ return kept;
+ }
+
+ /**
+ * Returns whether a range ending at {@code endKey} covers {@code startKey}, which is known not
+ * to precede it. A null {@code endKey} runs to the end of the keyspace and so covers everything.
+ */
+ private static boolean reaches(@Nullable byte[] endKey, byte[] startKey) {
+ return endKey == null || UNSIGNED_BYTES_COMPARATOR.compare(endKey, startKey) > 0;
+ }
+
+ /**
+ * 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: an infinity-ended slice overlaps anything after it, so dropOverlaps()
+ // drops the whole run; a kept one is always last.
+ 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;
+ }
+
+ /**
+ * 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) {
+ 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 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) {
+ 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/AutoShardingAttributes.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingAttributes.java
new file mode 100644
index 00000000000..f5668aff304
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingAttributes.java
@@ -0,0 +1,72 @@
+/*
+ * 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");
+
+ /**
+ * Locality this instance of the LB policy is balancing within, substituted for the {@code %s}
+ * token in {@code autosharding_target}.
+ *
+ * A resolver-state attribute, not a per-endpoint one, because it describes the policy
+ * instance rather than any single endpoint. gRFC A119 only defines it for the mode where this
+ * policy sits under a locality picker and therefore sees one locality; when the policy handles
+ * locality picking itself it sees endpoints from every locality, and the absence of this
+ * attribute is what makes the {@code %s} token correctly resolve to the empty string.
+ *
+ *
In xDS deployments the xDS integration populates this from the locality name that
+ * {@code weighted_target_experimental} publishes. Otherwise the application's name resolver is
+ * responsible for setting it, and need only do so if its target contains a {@code %s} token.
+ */
+ public static final Attributes.Key ATTR_LOCALITY =
+ Attributes.Key.create("io.grpc.autosharding.locality");
+
+ 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..c4730f7eb93
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancer.java
@@ -0,0 +1,542 @@
+/*
+ * 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.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.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 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}. 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
+ *
+ * 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;
+
+ /**
+ * 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. */
+ @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 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;
+
+ 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");
+ }
+
+ // gRFC A119 gives the configuration and the endpoints separate handling rules, and an empty
+ // endpoint set only speaks to the latter. Everything below that is driven by comparing the
+ // new configuration against the old one therefore has to run first: storing the new
+ // configuration without acting on it would destroy the comparison, and the change would then
+ // be lost for good, since the following update would no longer look like a change at all.
+ List endpoints = resolvedAddresses.getAddresses();
+
+ Channel previousChannel = shardingChannel;
+ 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;
+
+ // When the set is empty this tears the children down, so that in-flight picks stop resolving
+ // to endpoints the resolver has retracted.
+ endpointMap.updateEndpoints(endpoints, resolvedAddresses.getAttributes());
+
+ // The locality arrives in the resolver attributes, so the target can change even when the
+ // config did not.
+ maybeRecreateClient(
+ shardingChannel != previousChannel,
+ resolveTarget(newConfig, resolvedAddresses.getAttributes()),
+ newConfig.initialAssignmentTimeoutNanos);
+
+ if (endpoints.isEmpty()) {
+ // Any assignment is kept: it stays valid if the endpoints come back. Until they do,
+ // publishPicker() leaves this failure in place rather than publishing over it.
+ return failPermanently("autosharding: name resolver returned no endpoints");
+ }
+
+ rebuildSliceMapAndPublish();
+ return Status.OK;
+ }
+
+ @Override
+ public void handleNameResolutionError(Status error) {
+ if (shutdown) {
+ return;
+ }
+ // Endpoints from an earlier resolution stay usable, and the error is
+ // only reported when we are not already serving with them. Reporting it in that case is what
+ // makes a broken resolver visible; otherwise RPCs would fail with whatever the stale
+ // endpoints happen to be failing with, which names the wrong cause.
+ // The one addition is the initial assignment wait: RPCs queue until the timer
+ // fires, so a failed refresh must not turn that queue into failures.
+ boolean queueingForInitialAssignment = awaitingInitialAssignment && assignment == null;
+ if (endpointMap.size() > 0
+ && (queueingForInitialAssignment
+ || endpointMap.aggregateConnectivityState() == ConnectivityState.READY)) {
+ logger.log(Level.FINE, "Ignoring name resolution error, still serving: {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;
+ }
+ shardingTarget = 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 {@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) {
+ 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;
+ return Status.OK;
+ }
+
+ /**
+ * Replaces the {@link AutoshardingClient} when there is none yet, or when the channel to the
+ * sharding service or the resolved target changed.
+ *
+ * What gRFC A119 requires is a new channel and a new stream on it; whether the existing
+ * client is handed the new channel or a new client is built around it is left open. We replace
+ * the client because its 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;
+ }
+ 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();
+ }
+
+ /**
+ * Substitutes the optional {@code %s} token in the configured target with the locality this
+ * policy instance is balancing within, or with the empty string when no locality is available.
+ *
+ *
The locality is a property of the policy instance, so it is read from the resolver
+ * attributes rather than from any endpoint. Reading it from an endpoint would be wrong in the
+ * mode where this policy does its own locality picking: it is then handed endpoints from every
+ * locality, and picking one of them would make the target depend on resolver ordering. gRFC
+ * A119 says the token is not meant to be used in that mode, and leaving
+ * {@link AutoShardingAttributes#ATTR_LOCALITY} unset there resolves it to the empty string.
+ */
+ private static String resolveTarget(
+ AutoShardingLoadBalancerConfig config, Attributes resolverAttributes) {
+ if (!config.autoshardingTarget.contains("%s")) {
+ return config.autoshardingTarget;
+ }
+ String locality = resolverAttributes.get(AutoShardingAttributes.ATTR_LOCALITY);
+ 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();
+ }
+
+ /**
+ * 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();
+ }
+ }
+
+ /**
+ * 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..21892c4ccc3
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerConfig.java
@@ -0,0 +1,117 @@
+/*
+ * 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.checkArgument;
+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.
+ *
+ * {@link AutoShardingLoadBalancerProvider} parses this out of service config JSON and reports
+ * a bad configuration as a {@link io.grpc.NameResolver.ConfigOrError}, so the channel gets a
+ * useful message instead of an exception. The constructor here re-checks the same constraints as
+ * a backstop that cannot be bypassed.
+ */
+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. Never empty. */
+ 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");
+ checkArgument(!keyHeaderName.isEmpty(), "keyHeaderName is empty");
+ checkArgument(
+ initialAssignmentTimeoutNanos >= 0,
+ "initialAssignmentTimeoutNanos is negative: %s",
+ initialAssignmentTimeoutNanos);
+ 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/AutoShardingLoadBalancerProvider.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerProvider.java
new file mode 100644
index 00000000000..a9059333c4e
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingLoadBalancerProvider.java
@@ -0,0 +1,128 @@
+/*
+ * 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.Internal;
+import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.Metadata;
+import io.grpc.NameResolver.ConfigOrError;
+import io.grpc.Status;
+import io.grpc.internal.JsonUtil;
+import java.util.Map;
+
+/**
+ * Provider for the {@code autosharding_experimental} balancing policy.
+ *
+ *
Registering this on the classpath is what makes the policy reachable by name, both from a
+ * service config and from the xDS integration.
+ */
+@Internal
+public final class AutoShardingLoadBalancerProvider extends LoadBalancerProvider {
+ private static final String POLICY_NAME = "autosharding_experimental";
+
+ @Override
+ public LoadBalancer newLoadBalancer(LoadBalancer.Helper helper) {
+ return new AutoShardingLoadBalancer(helper);
+ }
+
+ @Override
+ public boolean isAvailable() {
+ return true;
+ }
+
+ @Override
+ public int getPriority() {
+ return 5;
+ }
+
+ @Override
+ public String getPolicyName() {
+ return POLICY_NAME;
+ }
+
+ @Override
+ public ConfigOrError parseLoadBalancingPolicyConfig(Map rawConfig) {
+ try {
+ return parseLoadBalancingPolicyConfigInternal(rawConfig);
+ } catch (RuntimeException e) {
+ return ConfigOrError.fromError(
+ Status.UNAVAILABLE
+ .withCause(e)
+ .withDescription("Failed parsing configuration for " + getPolicyName()));
+ }
+ }
+
+ /**
+ * Translates {@code AutoshardingLbConfig} from its JSON form, as specified by gRFC A119.
+ *
+ * The configuration is a proto3 message, so an absent string field and an empty one are the
+ * same value and no field is required. What gets rejected here is therefore a judgment call
+ * rather than something the gRFC spells out, and the bar is deliberately high: a value is only
+ * turned away when it has no coherent reading, or when honouring it would go wrong quietly.
+ * Anything the sharding service or the "Channel Factory" is the authority on is left to fail
+ * visibly at runtime instead of being guessed at here.
+ */
+ private ConfigOrError parseLoadBalancingPolicyConfigInternal(Map rawConfig) {
+ String channelFactoryKey = JsonUtil.getString(rawConfig, "channelFactoryKey");
+ if (channelFactoryKey == null) {
+ channelFactoryKey = "";
+ }
+
+ String autoshardingTarget = JsonUtil.getString(rawConfig, "autoshardingTarget");
+ if (autoshardingTarget == null) {
+ autoshardingTarget = "";
+ }
+
+ String keyHeaderName = JsonUtil.getString(rawConfig, "keyHeaderName");
+ if (keyHeaderName == null || keyHeaderName.isEmpty()) {
+ return error("'keyHeaderName' is required, LB policy config=" + rawConfig);
+ }
+ try {
+ // Rejects names that are not valid header names. Doing it here means a typo surfaces as a
+ // channel error instead of an exception thrown on the synchronization context later.
+ Metadata.Key unused = AutoShardingPicker.createKeyHeader(keyHeaderName);
+ } catch (IllegalArgumentException e) {
+ return error("'keyHeaderName' is not a valid header name: " + keyHeaderName);
+ }
+
+ Boolean enableFallback = JsonUtil.getBoolean(rawConfig, "enableFallback");
+
+ Long initialAssignmentTimeoutNanos =
+ JsonUtil.getStringAsDuration(rawConfig, "initialAssignmentTimeout");
+ if (initialAssignmentTimeoutNanos == null) {
+ initialAssignmentTimeoutNanos =
+ AutoShardingLoadBalancerConfig.DEFAULT_INITIAL_ASSIGNMENT_TIMEOUT_NANOS;
+ } else if (initialAssignmentTimeoutNanos < 0) {
+ return error(
+ "'initialAssignmentTimeout' must not be negative, LB policy config=" + rawConfig);
+ }
+
+ return ConfigOrError.fromConfig(
+ new AutoShardingLoadBalancerConfig(
+ channelFactoryKey,
+ autoshardingTarget,
+ keyHeaderName,
+ enableFallback != null && enableFallback,
+ initialAssignmentTimeoutNanos));
+ }
+
+ private static ConfigOrError error(String description) {
+ return ConfigOrError.fromError(
+ Status.UNAVAILABLE.withDescription("autosharding: " + description));
+ }
+}
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..55274333562
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/AutoshardingClient.java
@@ -0,0 +1,412 @@
+/*
+ * 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());
+
+ /** 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. */
+ interface AssignmentWatcher {
+ /** Called with a newly accepted assignment. Invoked on the sync context. */
+ 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;
+ private final ScheduledExecutorService timerService;
+ private final BackoffPolicy.Provider backoffPolicyProvider;
+ private final Stopwatch retryStopwatch;
+ private final AssignmentWatcher watcher;
+ private final String clientUuid;
+ 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.
+ *
+ *
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;
+
+ @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 #start()}.
+ *
+ * @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 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(
+ String clientUuid,
+ SynchronizationContext syncContext,
+ 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");
+ }
+
+ /** Opens the {@code WatchShardingAssignment} stream. Call once, on the sync context. */
+ void start() {
+ syncContext.throwIfNotInThisSynchronizationContext();
+ checkState(stream == null, "already started");
+ startStream();
+ }
+
+ /**
+ * 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 startStream() {
+ if (shutdown) {
+ 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<>();
+
+ /**
+ * The sending half of the stream. {@link ClientCalls} invokes {@link #beforeStart} before it
+ * starts the call and before it returns, so this is set before {@link #start} can send and
+ * before this stream is reachable by anything else.
+ */
+ private ClientCallStreamObserver requestStream;
+
+ /**
+ * Whether this stream delivered an assignment the load balancer could use, which is the only
+ * thing gRFC A119 resets the retry backoff on. An assignment dropped as stale, or rejected
+ * for having no usable slice, leaves the balancer with nothing newer than it already had, so
+ * the following attempt still backs off. A server that honours {@code latest_generation} does
+ * not resend an already-accepted assignment after a reconnect anyway.
+ */
+ 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);
+ // The call is started inside the above, so the config cannot go out from beforeStart().
+ 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.
+ *
+ * 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, 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) {
+ 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;
+ }
+
+ AssignmentParser.Result result = AssignmentParser.parse(chunks, generation);
+ if (result.assignment == null) {
+ logger.log(
+ Level.WARNING,
+ "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;
+ }
+
+ if (result.errorMessage != null) {
+ logger.log(
+ Level.WARNING,
+ "Accepting autosharding assignment with generation {0} after dropping slices: {1}",
+ new Object[] {generation, result.errorMessage});
+ }
+ sendAck(generation, true, result.errorMessage);
+ latestGeneration = generation;
+ receivedGoodAssignment = true;
+ watcher.onAssignment(result.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();
+ requestStream.cancel(status.getDescription(), status.getCause());
+ }
+ }
+
+ /**
+ * 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) {
+ // 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/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
new file mode 100644
index 00000000000..5d7597c9d17
--- /dev/null
+++ b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java
@@ -0,0 +1,431 @@
+/*
+ * 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.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.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.ArrayList;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.concurrent.NotThreadSafe;
+
+/**
+ * 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.
+ *
+ *
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 says the policy "must create a new {@code EndpointMap} whenever it receives
+ * endpoints from the Name Resolver". 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.
+ *
+ *
The difference is mechanical, not observable. Taken literally the gRFC's pseudocode builds
+ * fresh endpoint states with no child load balancer carried over, which would drop every
+ * connection on every resolver update; the C++ implementation accordingly builds a new map but
+ * moves surviving endpoints into it. Retaining the instance achieves the same thing and 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 static final Logger logger = Logger.getLogger(EndpointMap.class.getName());
+
+ private final Helper helper;
+ private final LoadBalancerProvider childProvider;
+ private final Runnable childStateListener;
+
+ // 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<>();
+
+ // 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<>();
+
+ /**
+ * Set while children are being handed their new addresses in {@link #updateEndpoints}. A child
+ * usually reports a state synchronously from that call, and forwarding every one would have
+ * the LB policy publish a picker per endpoint for a single resolver update. The state and
+ * picker are still recorded; only the notification is skipped, and the caller publishes once
+ * afterwards.
+ *
+ * Not specific to this policy: {@code MultiChildLoadBalancer} in {@code io.grpc.util}, which
+ * backs {@code round_robin}, {@code ring_hash}, {@code weighted_target} and the rest, carries
+ * the same flag as {@code resolvingAddresses} for the same reason.
+ */
+ private boolean rebuilding;
+
+ /**
+ * 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");
+ }
+
+ /**
+ * 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.
+ *
+ *
Children handed new addresses here often report a connectivity state before this method
+ * returns. Those reports are recorded but not forwarded to the {@code childStateListener}, so
+ * that one resolver update produces one picker rather than one per endpoint. The
+ * caller must therefore publish a picker itself once this returns, or the channel is
+ * left holding a picker built from the previous endpoint set.
+ *
+ * @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();
+ }
+ }
+
+ // Install the whole endpoint set and its indices before touching any child. A child given
+ // addresses in the second pass can call back in synchronously, and everything it can reach
+ // -- size(), indexOf(), toPickerEndpoints() -- has to already agree on the new set.
+ holders.clear();
+ indexByHostname.clear();
+ for (String hostname : addressesByHostname.keySet()) {
+ EndpointHolder survivor = survivors.get(hostname);
+ indexByHostname.put(hostname, holders.size());
+ holders.add(survivor != null ? survivor : new EndpointHolder(hostname));
+ }
+
+ rebuilding = true;
+ try {
+ for (EndpointHolder holder : holders) {
+ holder.updateAddresses(addressesByHostname.get(holder.hostname), attributes);
+ }
+ } finally {
+ rebuilding = false;
+ }
+ }
+
+ // Returns the number of endpoints currently held
+ int size() {
+ return holders.size();
+ }
+
+ /**
+ * 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;
+ }
+
+ /**
+ * 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}.
+ */
+ ImmutableList toPickerEndpoints() {
+ ImmutableList.Builder snapshot =
+ ImmutableList.builderWithExpectedSize(holders.size());
+ for (EndpointHolder holder : holders) {
+ snapshot.add(holder.toPickerEndpoint());
+ }
+ return snapshot.build();
+ }
+
+ /**
+ * Returns the aggregated connectivity state to report for the channel, using the {@code
+ * ring_hash} rules from gRFC A42 that gRFC A119 adopts:
+ *
+ *
+ * - at least one endpoint READY, report READY;
+ *
- two or more endpoints TRANSIENT_FAILURE, report TRANSIENT_FAILURE;
+ *
- at least one endpoint CONNECTING, report CONNECTING;
+ *
- exactly one endpoint TRANSIENT_FAILURE and more than one endpoint, report CONNECTING;
+ *
- at least one endpoint IDLE, report IDLE;
+ *
- otherwise report TRANSIENT_FAILURE.
+ *
+ *
+ * An empty map reports TRANSIENT_FAILURE, matching rule 6.
+ */
+ 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;
+ }
+ if (connecting > 0) {
+ return CONNECTING;
+ }
+ if (transientFailure == 1 && holders.size() > 1) {
+ return CONNECTING;
+ }
+ if (idle > 0) {
+ return IDLE;
+ }
+ return TRANSIENT_FAILURE;
+ }
+
+ /**
+ * Starts connecting on one IDLE endpoint, unless some endpoint is already CONNECTING or none
+ * is IDLE.
+ *
+ *
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.
+ */
+ 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();
+ }
+ }
+
+ /** Shuts down every child load balancer and empties the map. Idempotent. */
+ void shutdown() {
+ for (EndpointHolder holder : holders) {
+ holder.shutdown();
+ }
+ 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("endpoints", holders).toString();
+ }
+
+ /**
+ * The child load balancer for a single endpoint, together with the connectivity state and
+ * picker it most recently reported.
+ */
+ private final class EndpointHolder {
+ private final String hostname;
+ private final LazyLoadBalancer childLb;
+ private ConnectivityState state = IDLE;
+ private SubchannelPicker picker = new FixedResultPicker(PickResult.withNoResult());
+ private boolean childShutdown;
+
+ 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);
+ }
+
+ 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});
+ }
+ }
+
+ /** Starts connecting if this endpoint is IDLE. */
+ 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("hostname", hostname)
+ .add("state", state)
+ .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 {
+ @Override
+ protected Helper delegate() {
+ return helper;
+ }
+
+ @Override
+ public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
+ if (childShutdown) {
+ return;
+ }
+ state = newState;
+ picker = newPicker;
+ if (!rebuilding) {
+ childStateListener.run();
+ }
+ }
+ }
+ }
+}
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/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..4f5739170af
--- /dev/null
+++ b/autosharding/src/test/java/io/grpc/autosharding/AssignmentParserTest.java
@@ -0,0 +1,583 @@
+/*
+ * 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 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() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addEndpoints(endpoint("host-a"))
+ .addEndpoints(endpoint("host-b"))
+ .addSliceAssignments(sliceAssignment("", "m", 0))
+ .addSliceAssignments(sliceAssignment("m", null, 1))
+ .build();
+
+ Assignment assignment = parseFully(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() {
+ 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 = parseFully(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() {
+ 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 = parseFully(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() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addEndpoints(endpoint("host-a"))
+ .addSliceAssignments(sliceAssignment("d", null, 0))
+ .build();
+
+ Assignment assignment = parseFully(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() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addEndpoints(endpoint("host-a"))
+ .addSliceAssignments(sliceAssignment("", "d", 0))
+ .build();
+
+ Assignment assignment = parseFully(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() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addEndpoints(endpoint("host-a"))
+ .addEndpoints(endpoint("host-b"))
+ .addSliceAssignments(sliceAssignment("", "d", 0))
+ .addSliceAssignments(sliceAssignment("m", null, 1))
+ .build();
+
+ Assignment assignment = parseFully(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() {
+ Assignment assignment =
+ parseFully(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() {
+ Assignment assignment = parseFully(ImmutableList.of(), 1);
+
+ assertThat(assignment.getSlices()).hasSize(1);
+ assertSlice(assignment.getSlices().get(0), "", null);
+ }
+
+ @Test
+ public void parse_sliceWithNoEndpoints_isPreserved() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addEndpoints(endpoint("host-a"))
+ .addSliceAssignments(sliceAssignment("", "d"))
+ .addSliceAssignments(sliceAssignment("d", null, 0))
+ .build();
+
+ Assignment assignment = parseFully(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() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addEndpoints(endpoint("host-a"))
+ .addEndpoints(endpoint("host-b"))
+ .addSliceAssignments(sliceAssignment("", null, 0, 1))
+ .build();
+
+ Assignment assignment = parseFully(ImmutableList.of(chunk), 1);
+
+ assertSlice(assignment.getSlices().get(0), "", null, 0, 1);
+ }
+
+ @Test
+ 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("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.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_sliceBecomesGap() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addEndpoints(endpoint("host-a"))
+ .addSliceAssignments(sliceAssignment("", "m", 0))
+ .addSliceAssignments(sliceAssignment("m", null, -1))
+ .build();
+
+ 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_endpointIndexOutOfRange_dropsTheWholeSliceNotJustThatEndpoint() {
+ AssignmentChunk chunk =
+ 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.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1);
+
+ assertSlice(result.assignment.getSlices().get(1), "m", null);
+ }
+
+ @Test
+ public void parse_startKeyGreaterThanEndKey_sliceBecomesGap() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .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.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);
+ }
+
+ /**
+ * 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.
+ */
+ @Test
+ 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_bothAreDropped() {
+ // There is no way to tell which of the two the server meant, so neither is used and the keys
+ // they covered become a gap.
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addSliceAssignments(sliceAssignment("a", "m"))
+ .addSliceAssignments(sliceAssignment("d", null))
+ .build();
+
+ AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1);
+
+ assertThat(result.errorMessage).contains("overlaps");
+ assertThat(result.assignment).isNull();
+ }
+
+ @Test
+ public void parse_overlappingSlices_aSliceClearOfThemSurvives() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addSliceAssignments(sliceAssignment("a", "m"))
+ .addSliceAssignments(sliceAssignment("d", "p"))
+ .addSliceAssignments(sliceAssignment("p", "z"))
+ .build();
+
+ AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1);
+
+ assertThat(result.errorMessage).contains("overlaps");
+ // ["p", "z") abuts the overlap without entering it, so only it is kept.
+ assertThat(result.assignment.getSlices()).hasSize(3);
+ assertSlice(result.assignment.getSlices().get(0), "", "p");
+ assertSlice(result.assignment.getSlices().get(1), "p", "z");
+ assertSlice(result.assignment.getSlices().get(2), "z", null);
+ assertThat(result.assignment.getSlices().get(0).getEndpoints()).isEmpty();
+ }
+
+ @Test
+ public void parse_slicesOverlappingOnlyThroughAThird_areAllDropped() {
+ // ["b", "c") and ["d", "e") are disjoint, but both collide with ["a", "z"), so all three go.
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addSliceAssignments(sliceAssignment("a", "z"))
+ .addSliceAssignments(sliceAssignment("b", "c"))
+ .addSliceAssignments(sliceAssignment("d", "e"))
+ .build();
+
+ AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1);
+
+ assertThat(result.errorMessage).contains("overlaps");
+ assertThat(result.assignment).isNull();
+ }
+
+ @Test
+ public void parse_duplicateStartKeys_bothAreDropped() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addSliceAssignments(sliceAssignment("a", "m"))
+ .addSliceAssignments(sliceAssignment("a", "z"))
+ .build();
+
+ AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1);
+
+ assertThat(result.errorMessage).contains("overlaps");
+ assertThat(result.assignment).isNull();
+ }
+
+ @Test
+ public void parse_sliceExtendingToInfinityFollowedByAnother_bothAreDropped() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addSliceAssignments(sliceAssignment("a", null))
+ .addSliceAssignments(sliceAssignment("m", null))
+ .build();
+
+ AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1);
+
+ assertThat(result.errorMessage).contains("overlaps");
+ assertThat(result.assignment).isNull();
+ }
+
+ @Test
+ public void parse_everySliceInvalid_yieldsNoAssignment() {
+ AssignmentChunk chunk =
+ AssignmentChunk.newBuilder()
+ .addEndpoints(endpoint("host-a"))
+ .addSliceAssignments(sliceAssignment("", null, 1))
+ .build();
+
+ AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk), 1);
+
+ assertThat(result.assignment).isNull();
+ assertThat(result.errorMessage).contains("out-of-range endpoint index 1");
+ }
+
+ @Test
+ 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"));
+ }
+
+ AssignmentParser.Result result = AssignmentParser.parse(ImmutableList.of(chunk.build()), 1);
+
+ // 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. */
+ 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) {
+ 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/AutoShardingLoadBalancerProviderTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerProviderTest.java
new file mode 100644
index 00000000000..d59ef90df15
--- /dev/null
+++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerProviderTest.java
@@ -0,0 +1,188 @@
+/*
+ * 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 static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import io.grpc.InternalServiceProviders;
+import io.grpc.LoadBalancer.Helper;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.NameResolver.ConfigOrError;
+import io.grpc.Status;
+import io.grpc.SynchronizationContext;
+import io.grpc.internal.JsonParser;
+import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Unit tests for {@link AutoShardingLoadBalancerProvider}. */
+@RunWith(JUnit4.class)
+public class AutoShardingLoadBalancerProviderTest {
+
+ private final SynchronizationContext syncContext =
+ new SynchronizationContext((t, e) -> {
+ throw new AssertionError(e);
+ });
+
+ private final AutoShardingLoadBalancerProvider provider =
+ new AutoShardingLoadBalancerProvider();
+
+ @Test
+ public void provided() {
+ for (LoadBalancerProvider current :
+ InternalServiceProviders.getCandidatesViaServiceLoader(
+ LoadBalancerProvider.class, getClass().getClassLoader())) {
+ if (current instanceof AutoShardingLoadBalancerProvider) {
+ return;
+ }
+ }
+ fail("AutoShardingLoadBalancerProvider not registered");
+ }
+
+ @Test
+ public void providerProperties() {
+ assertThat(provider.getPolicyName()).isEqualTo("autosharding_experimental");
+ assertThat(provider.isAvailable()).isTrue();
+ assertThat(provider.getPriority()).isEqualTo(5);
+ }
+
+ @Test
+ public void providesLoadBalancer() {
+ Helper helper = mock(Helper.class);
+ when(helper.getSynchronizationContext()).thenReturn(syncContext);
+ assertThat(provider.newLoadBalancer(helper)).isInstanceOf(AutoShardingLoadBalancer.class);
+ }
+
+ @Test
+ public void parse_allFieldsPresent() throws IOException {
+ AutoShardingLoadBalancerConfig config =
+ parseSuccessfully(
+ "{"
+ + "\"channelFactoryKey\": \"shard-service\","
+ + "\"autoshardingTarget\": \"my-service/%s\","
+ + "\"keyHeaderName\": \"x-shard-key\","
+ + "\"enableFallback\": true,"
+ + "\"initialAssignmentTimeout\": \"5.5s\""
+ + "}");
+
+ assertThat(config.channelFactoryKey).isEqualTo("shard-service");
+ assertThat(config.autoshardingTarget).isEqualTo("my-service/%s");
+ assertThat(config.keyHeaderName).isEqualTo("x-shard-key");
+ assertThat(config.enableFallback).isTrue();
+ assertThat(config.initialAssignmentTimeoutNanos)
+ .isEqualTo(TimeUnit.MILLISECONDS.toNanos(5500));
+ }
+
+ @Test
+ public void parse_optionalFieldsAbsent_useGrfcDefaults() throws IOException {
+ AutoShardingLoadBalancerConfig config =
+ parseSuccessfully("{\"keyHeaderName\": \"x-shard-key\"}");
+
+ assertThat(config.channelFactoryKey).isEmpty();
+ assertThat(config.autoshardingTarget).isEmpty();
+ assertThat(config.enableFallback).isFalse();
+ assertThat(config.initialAssignmentTimeoutNanos)
+ .isEqualTo(AutoShardingLoadBalancerConfig.DEFAULT_INITIAL_ASSIGNMENT_TIMEOUT_NANOS);
+ assertThat(config.initialAssignmentTimeoutNanos).isEqualTo(TimeUnit.SECONDS.toNanos(60));
+ }
+
+ @Test
+ public void parse_binaryKeyHeaderName() throws IOException {
+ AutoShardingLoadBalancerConfig config =
+ parseSuccessfully("{\"keyHeaderName\": \"x-shard-key-bin\"}");
+
+ assertThat(config.keyHeaderName).isEqualTo("x-shard-key-bin");
+ }
+
+ @Test
+ public void parse_missingKeyHeaderName_isRejected() throws IOException {
+ assertThat(parseError("{}")).contains("'keyHeaderName' is required");
+ }
+
+ @Test
+ public void parse_emptyKeyHeaderName_isRejected() throws IOException {
+ // An empty name would leave every RPC with the empty key, pinning the channel to one shard.
+ assertThat(parseError("{\"keyHeaderName\": \"\"}")).contains("'keyHeaderName' is required");
+ }
+
+ @Test
+ public void parse_malformedKeyHeaderName_isRejected() throws IOException {
+ // Reported here rather than thrown out of Metadata.Key on the synchronization context.
+ assertThat(parseError("{\"keyHeaderName\": \"not a header\"}"))
+ .contains("'keyHeaderName' is not a valid header name");
+ }
+
+ @Test
+ public void parse_zeroInitialAssignmentTimeout_isAccepted() throws IOException {
+ // A coherent request: skip the wait, start in fallback, and upgrade once the first assignment
+ // lands. The gRFC does not forbid it, so the parser does not either.
+ AutoShardingLoadBalancerConfig config =
+ parseSuccessfully(
+ "{\"keyHeaderName\": \"x-shard-key\", \"initialAssignmentTimeout\": \"0s\"}");
+
+ assertThat(config.initialAssignmentTimeoutNanos).isEqualTo(0);
+ }
+
+ @Test
+ public void parse_negativeInitialAssignmentTimeout_isRejected() throws IOException {
+ assertThat(
+ parseError(
+ "{\"keyHeaderName\": \"x-shard-key\", \"initialAssignmentTimeout\": \"-1s\"}"))
+ .contains("'initialAssignmentTimeout' must not be negative");
+ }
+
+ @Test
+ public void parse_unparseableInitialAssignmentTimeout_isReportedAsFailure() throws IOException {
+ // Durations are JSON strings ending in "s"; a bare number is a ClassCastException inside
+ // JsonUtil, which the provider turns into an error rather than letting it escape.
+ assertThat(
+ parseError("{\"keyHeaderName\": \"x-shard-key\", \"initialAssignmentTimeout\": 60}"))
+ .isEqualTo("Failed parsing configuration for autosharding_experimental");
+ }
+
+ @Test
+ public void parse_wronglyTypedField_isReportedAsFailure() throws IOException {
+ assertThat(parseError("{\"keyHeaderName\": 42}"))
+ .isEqualTo("Failed parsing configuration for autosharding_experimental");
+ }
+
+ private AutoShardingLoadBalancerConfig parseSuccessfully(String json) throws IOException {
+ ConfigOrError configOrError = provider.parseLoadBalancingPolicyConfig(parseJsonObject(json));
+ assertThat(configOrError.getError()).isNull();
+ return (AutoShardingLoadBalancerConfig) configOrError.getConfig();
+ }
+
+ /** Parses a config expected to be rejected, returning the error description. */
+ private String parseError(String json) throws IOException {
+ ConfigOrError configOrError = provider.parseLoadBalancingPolicyConfig(parseJsonObject(json));
+ assertThat(configOrError.getConfig()).isNull();
+ Status error = configOrError.getError();
+ assertThat(error.getCode()).isEqualTo(Status.Code.UNAVAILABLE);
+ return error.getDescription();
+ }
+
+ @SuppressWarnings("unchecked")
+ private static Map parseJsonObject(String json) throws IOException {
+ return (Map) JsonParser.parse(json);
+ }
+}
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..1ef8ddd3b5d
--- /dev/null
+++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingLoadBalancerTest.java
@@ -0,0 +1,1203 @@
+/*
+ * 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.assertThrows;
+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 String OTHER_KEY_HEADER = "x-other-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 emptyKeyHeaderName_isRejected() {
+ // Not a supported mode: with no key header every RPC would carry the empty key and the whole
+ // channel would end up on whichever slice covers it. C++ rejects this at config-parse time.
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ new AutoShardingLoadBalancerConfig(
+ CHANNEL_FACTORY_KEY, TARGET, "", true, ASSIGNMENT_TIMEOUT_NANOS));
+ }
+
+ @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");
+ }
+
+ @Test
+ public void channelFactoryKeyChangedWhileEndpointsEmpty_stillCreatesTheNewChannel() {
+ deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a");
+ assertThat(channelFactory.keys).containsExactly(CHANNEL_FACTORY_KEY);
+
+ // An empty endpoint set says nothing about the configuration, so the new key still takes
+ // effect. Recording the config without acting on it would make the next update look
+ // unchanged, stranding the policy on the old sharding service for good.
+ deliverAddresses(config(OTHER_CHANNEL_FACTORY_KEY, true));
+
+ assertThat(channelFactory.keys)
+ .containsExactly(CHANNEL_FACTORY_KEY, OTHER_CHANNEL_FACTORY_KEY)
+ .inOrder();
+ assertThat(channelFactory.isReleased(0)).isTrue();
+
+ deliverAddresses(config(OTHER_CHANNEL_FACTORY_KEY, true), "a");
+
+ // Already applied above; the restored endpoints must not cause a second channel.
+ assertThat(channelFactory.keys).hasSize(2);
+ }
+
+ @Test
+ public void keyHeaderNameChangedWhileEndpointsEmpty_stillTakesEffect() throws Exception {
+ deliverAddresses(configWithKeyHeader(KEY_HEADER), "a", "b");
+ deliverAssignment(1, slice("", "a"), slice("m", "b"));
+
+ deliverAddresses(configWithKeyHeader(OTHER_KEY_HEADER));
+ deliverAddresses(configWithKeyHeader(OTHER_KEY_HEADER), "a", "b");
+ reportReady("a");
+ reportReady("b");
+
+ // Read under the new header, "z" is past "m" and belongs to "b". Were the policy still
+ // reading the old header, it would find nothing, and the empty key would send this to "a".
+ assertThat(pickedHost(pick(OTHER_KEY_HEADER, "z"))).isEqualTo("b");
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // 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();
+
+ deliverAddresses(retargetedConfig("other-target"), "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 {
+ acceptAddresses(
+ ResolvedAddresses.newBuilder()
+ .setAddresses(endpoints("a"))
+ .setAttributes(attributesWithLocality("us-central1-a"))
+ .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(endpoints("a"))
+ .setAttributes(attributesWithLocality("us-central1-a"))
+ .setLoadBalancingPolicyConfig(localityConfig)
+ .build());
+ takeRequest();
+
+ acceptAddresses(
+ ResolvedAddresses.newBuilder()
+ .setAddresses(endpoints("a"))
+ .setAttributes(attributesWithLocality("us-central1-b"))
+ .setLoadBalancingPolicyConfig(localityConfig)
+ .build());
+
+ assertThat(channelFactory.keys).containsExactly(CHANNEL_FACTORY_KEY);
+ assertThat(service.streamCount.get()).isEqualTo(2);
+ assertThat(takeRequest().getInitialClientConfig().getTarget())
+ .isEqualTo("target/us-central1-b");
+ }
+
+ @Test
+ public void targetWithLocalityToken_noLocality_substitutesEmptyString() throws Exception {
+ deliverAddresses(retargetedConfig("target/%s"), "a");
+
+ assertThat(takeRequest().getInitialClientConfig().getTarget()).isEqualTo("target/");
+ }
+
+ @Test
+ public void targetWithLocalityToken_endpointsSpanLocalities_substitutesEmptyString()
+ throws Exception {
+ // The policy is doing its own locality picking, so it sees endpoints from every locality and
+ // no locality attribute. gRFC A119 says the token is not meant to be used here; the target
+ // must not end up depending on which endpoint the resolver happened to list first.
+ acceptAddresses(
+ ResolvedAddresses.newBuilder()
+ .setAddresses(
+ ImmutableList.of(
+ endpointInLocality("a", "us-central1-a"),
+ endpointInLocality("b", "us-central1-b")))
+ .setAttributes(attributesWithChannelFactory())
+ .setLoadBalancingPolicyConfig(retargetedConfig("target/%s"))
+ .build());
+
+ assertThat(takeRequest().getInitialClientConfig().getTarget()).isEqualTo("target/");
+ }
+
+ @Test
+ public void targetWithLocalityToken_endpointReordering_doesNotRecreateTheClient()
+ throws Exception {
+ AutoShardingLoadBalancerConfig localityConfig = retargetedConfig("target/%s");
+ acceptAddresses(
+ ResolvedAddresses.newBuilder()
+ .setAddresses(
+ ImmutableList.of(
+ endpointInLocality("a", "us-central1-a"),
+ endpointInLocality("b", "us-central1-b")))
+ .setAttributes(attributesWithChannelFactory())
+ .setLoadBalancingPolicyConfig(localityConfig)
+ .build());
+ takeRequest();
+
+ // Same endpoints, other order. Sourcing the locality from the first endpoint would change the
+ // target here and reconnect the client to a different sharding resource.
+ acceptAddresses(
+ ResolvedAddresses.newBuilder()
+ .setAddresses(
+ ImmutableList.of(
+ endpointInLocality("b", "us-central1-b"),
+ endpointInLocality("a", "us-central1-a")))
+ .setAttributes(attributesWithChannelFactory())
+ .setLoadBalancingPolicyConfig(localityConfig)
+ .build());
+
+ assertThat(service.streamCount.get()).isEqualTo(1);
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // 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");
+ }
+
+ @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
+ // ---------------------------------------------------------------------------------------------
+
+ @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 nameResolutionError_withEndpointsButNoneReady_reportsTheResolverError()
+ throws Exception {
+ deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a");
+ deliverAssignment(1, slice("", "a"));
+ reportTransientFailure("a");
+
+ syncContext.execute(
+ () -> loadBalancer.handleNameResolutionError(Status.UNAVAILABLE.withDescription("boom")));
+
+ // We are not serving, so the resolver failure is the more useful thing to report. Leaving it
+ // out would make RPCs blame the endpoints we can no longer refresh.
+ assertThat(currentState).isEqualTo(TRANSIENT_FAILURE);
+ assertThat(pick("k").getStatus().getDescription()).contains("boom");
+ }
+
+ @Test
+ public void nameResolutionError_whileAwaitingInitialAssignment_keepsQueueing() {
+ deliverAddresses(config(CHANNEL_FACTORY_KEY, true), "a");
+ // Not READY, so only the initial assignment wait can be holding these RPCs.
+ reportTransientFailure("a");
+ assertThat(currentState).isEqualTo(CONNECTING);
+
+ syncContext.execute(
+ () -> loadBalancer.handleNameResolutionError(Status.UNAVAILABLE.withDescription("boom")));
+
+ // gRFC A119 holds RPCs until the initial assignment timer fires; a failed refresh of
+ // endpoints we already have must not cut that short.
+ assertThat(currentState).isEqualTo(CONNECTING);
+ assertThat(pick("k").getStatus().isOk()).isTrue();
+ assertThat(pick("k").getSubchannel()).isNull();
+ }
+
+ @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);
+ }
+
+ /** 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);
+ }
+
+ /** The default config with a different {@code key_header_name}. */
+ private AutoShardingLoadBalancerConfig configWithKeyHeader(String keyHeaderName) {
+ return new AutoShardingLoadBalancerConfig(
+ CHANNEL_FACTORY_KEY, TARGET, keyHeaderName, true, ASSIGNMENT_TIMEOUT_NANOS);
+ }
+
+ private Attributes attributesWithChannelFactory() {
+ return Attributes.newBuilder()
+ .set(AutoShardingAttributes.ATTR_CHANNEL_FACTORY, channelFactory)
+ .build();
+ }
+
+ /** Resolver attributes as they look under a locality picker, which supplies the locality. */
+ private Attributes attributesWithLocality(String locality) {
+ return attributesWithChannelFactory().toBuilder()
+ .set(AutoShardingAttributes.ATTR_LOCALITY, locality)
+ .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);
+ }
+
+ 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<>();
+ 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) {
+ return pick(KEY_HEADER, key);
+ }
+
+ private PickResult pick(String headerName, String key) {
+ Metadata headers = new Metadata();
+ headers.put(Metadata.Key.of(headerName, 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..8c3084443b4
--- /dev/null
+++ b/autosharding/src/test/java/io/grpc/autosharding/AutoshardingClientTest.java
@@ -0,0 +1,585 @@
+/*
+ * 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.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 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 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 {
+ channel = newChannelToFakeService();
+ client = newClient(channel, TARGET);
+ }
+
+ @After
+ public void tearDown() {
+ // Must happen before GrpcCleanupRule shuts the channels down, otherwise a client keeps
+ // retrying against a terminating channel.
+ syncContext.execute(
+ () -> {
+ for (AutoshardingClient created : clients) {
+ created.shutdown();
+ }
+ });
+ }
+
+ @Test
+ public void start_opensStreamAndSendsInitialClientConfig() throws Exception {
+ start(client);
+
+ 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 chunksBufferedUntilMetadata_thenAssignmentDeliveredAndAcked() throws Exception {
+ start(client);
+ 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 {
+ start(client);
+ 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 noUsableSlices_nackedAndReportedAsAnError() throws Exception {
+ start(client);
+ takeRequest();
+ StreamObserver serverStream = takeServerStream();
+
+ // 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));
+
+ 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(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 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));
+
+ Assignment assignment = takeAssignment();
+ assertThat(assignment.getEndpointNames()).containsExactly("host-b");
+ }
+
+ @Test
+ public void staleGeneration_nackedAndNotDelivered() throws Exception {
+ start(client);
+ 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));
+
+ 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(errors).isEmpty();
+ assertThat(client.getLatestGeneration()).isEqualTo(5);
+ }
+
+ @Test
+ public void olderGeneration_nackedAndNotDelivered() throws Exception {
+ start(client);
+ 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));
+
+ WatchShardingAssignmentRequest nack = takeRequest();
+ assertThat(nack.getAssignmentAck().getGeneration()).isEqualTo(4);
+ assertThat(nack.getAssignmentAck().getAccepted()).isFalse();
+ assertThat(assignments).isEmpty();
+ assertThat(client.getLatestGeneration()).isEqualTo(5);
+ }
+
+ @Test
+ public void loadReportingConfig_ignored() throws Exception {
+ start(client);
+ 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 {
+ start(client);
+ 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 {
+ start(client);
+ 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 {
+ start(client);
+ takeRequest();
+ StreamObserver serverStream = takeServerStream();
+
+ serverStream.onCompleted();
+ fireRetryTimer();
+
+ takeRequest();
+ assertThat(service.streamCount.get()).isEqualTo(2);
+ }
+
+ @Test
+ public void backoffSequence_onlyResetAfterGoodAssignment() throws Exception {
+ start(client);
+ 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);
+ }
+
+ /**
+ * 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 newClient_startsFromGenerationZero() throws Exception {
+ start(client);
+ takeRequest();
+ StreamObserver serverStream = takeServerStream();
+ serverStream.onNext(chunkResponse(chunkWithEndpoint("host-a", "", null, 0)));
+ serverStream.onNext(metadataResponse(9));
+ takeAssignment();
+ takeRequest(); // ACK
+ assertThat(client.getLatestGeneration()).isEqualTo(9);
+ syncContext.execute(client::shutdown);
+
+ 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(replacement.getLatestGeneration()).isEqualTo(0);
+ assertThat(service.streamCount.get()).isEqualTo(2);
+ }
+
+ @Test
+ public void shutdown_cancelsStreamAndStopsReconnecting() throws Exception {
+ start(client);
+ takeRequest();
+
+ syncContext.execute(client::shutdown);
+
+ assertThat(service.streamCount.get()).isEqualTo(1);
+ assertThat(fakeClock.numPendingTasks()).isEqualTo(0);
+ }
+
+ @Test
+ public void shutdown_isIdempotent() throws Exception {
+ start(client);
+ takeRequest();
+
+ syncContext.execute(client::shutdown);
+ syncContext.execute(client::shutdown);
+
+ assertThat(service.streamCount.get()).isEqualTo(1);
+ }
+
+ @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. */
+ 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 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();
+ }
+
+ 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 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<>();
+ 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
new file mode 100644
index 00000000000..13fbe2e7899
--- /dev/null
+++ b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java
@@ -0,0 +1,589 @@
+/*
+ * 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.mockito.Mockito.mock;
+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 java.net.SocketAddress;
+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;
+
+/** Unit tests for {@link EndpointMap}. */
+@RunWith(JUnit4.class)
+public class EndpointMapTest {
+
+ private final SynchronizationContext syncContext =
+ 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;
+
+ @Before
+ public void setUp() {
+ when(helper.getSynchronizationContext()).thenReturn(syncContext);
+ endpointMap = new EndpointMap(helper, childProvider, () -> stateUpdates.add(1));
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // 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 indexOf_unknownHostname_returnsMinusOne() {
+ endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY);
+
+ assertThat(endpointMap.indexOf("nope")).isEqualTo(-1);
+ }
+
+ @Test
+ public void updateEndpoints_duplicateHostnames_keepsOneEntry() {
+ endpointMap.updateEndpoints(endpoints("a", "b", "a"), Attributes.EMPTY);
+
+ // Indices stay dense so that they remain valid offsets into toPickerEndpoints().
+ assertThat(endpointMap.size()).isEqualTo(2);
+ 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.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 updateEndpoints_doesNotNotifyListenerWhileRebuilding() {
+ // New children publish their initial IDLE state from inside updateEndpoints(). Forwarding
+ // those would make the LB policy publish one picker per endpoint for a single update.
+ endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY);
+
+ assertThat(stateUpdates).isEmpty();
+ }
+
+ @Test
+ public void updateEndpoints_childReenteringDuringUpdate_seesTheCompleteNewEndpointSet() {
+ endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY);
+ activate(0);
+ activate(1);
+
+ // Surviving children are handed their new addresses from inside updateEndpoints() and can
+ // call straight back in. Record how the map looks from in there.
+ List observedSizes = new ArrayList<>();
+ List observedIndicesOfC = new ArrayList<>();
+ childProvider.onAccept =
+ () -> {
+ observedSizes.add(endpointMap.size());
+ observedIndicesOfC.add(endpointMap.indexOf("c"));
+ };
+
+ endpointMap.updateEndpoints(endpoints("a", "b", "c"), Attributes.EMPTY);
+
+ // Both callbacks see all three endpoints and the final indices, never a partial rebuild.
+ assertThat(observedSizes).containsExactly(3, 3);
+ assertThat(observedIndicesOfC).containsExactly(2, 2);
+ }
+
+ @Test
+ public void childStateUpdate_notifiesListener() {
+ endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY);
+ activate(0);
+ stateUpdates.clear();
+
+ reportState(0, READY);
+
+ assertThat(stateUpdates).hasSize(1);
+ assertThat(stateAt(0)).isEqualTo(READY);
+ }
+
+ @Test
+ 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));
+
+ assertThat(stateUpdates).isEmpty();
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Connecting lazily
+ // ---------------------------------------------------------------------------------------------
+
+ @Test
+ public void endpointsStartIdleWithoutCreatingChildren() {
+ endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY);
+
+ assertThat(childProvider.children).isEmpty();
+ assertThat(stateAt(0)).isEqualTo(IDLE);
+ assertThat(stateAt(1)).isEqualTo(IDLE);
+ }
+
+ @Test
+ public void pickerEndpoint_requestConnection_createsChildAndConnects() {
+ endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY);
+
+ endpointMap.toPickerEndpoints().get(0).requestConnection();
+
+ assertThat(childProvider.children).hasSize(1);
+ assertThat(childProvider.children.get(0).requestConnectionCount).isEqualTo(1);
+ }
+
+ @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);
+ }
+
+ @Test
+ public void pickerEndpoint_requestConnectionAfterEndpointRemoved_isNoOp() {
+ endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY);
+ PickerEndpoint stale = endpointMap.toPickerEndpoints().get(0);
+
+ endpointMap.updateEndpoints(endpoints("b"), Attributes.EMPTY);
+ stale.requestConnection();
+
+ assertThat(childProvider.children).isEmpty();
+ }
+
+ @Test
+ public void pickerEndpoint_requestConnectionAfterShutdown_isNoOp() {
+ endpointMap.updateEndpoints(endpoints("a"), Attributes.EMPTY);
+ PickerEndpoint stale = endpointMap.toPickerEndpoints().get(0);
+
+ endpointMap.shutdown();
+ stale.requestConnection();
+
+ assertThat(childProvider.children).isEmpty();
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Aggregated connectivity state (gRFC A42 rules)
+ // ---------------------------------------------------------------------------------------------
+
+ @Test
+ public void aggregate_empty_isTransientFailure() {
+ assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(TRANSIENT_FAILURE);
+ }
+
+ @Test
+ public void aggregate_anyReady_isReady() {
+ setUpStates(TRANSIENT_FAILURE, TRANSIENT_FAILURE, READY);
+
+ assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(READY);
+ }
+
+ @Test
+ public void aggregate_twoTransientFailures_isTransientFailure() {
+ setUpStates(TRANSIENT_FAILURE, TRANSIENT_FAILURE, IDLE);
+
+ assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(TRANSIENT_FAILURE);
+ }
+
+ @Test
+ public void aggregate_anyConnecting_isConnecting() {
+ setUpStates(IDLE, CONNECTING);
+
+ assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(CONNECTING);
+ }
+
+ @Test
+ public void aggregate_oneTransientFailureAmongMany_isConnecting() {
+ setUpStates(TRANSIENT_FAILURE, IDLE);
+
+ assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(CONNECTING);
+ }
+
+ @Test
+ public void aggregate_soleEndpointInTransientFailure_isTransientFailure() {
+ setUpStates(TRANSIENT_FAILURE);
+
+ assertThat(endpointMap.aggregateConnectivityState()).isEqualTo(TRANSIENT_FAILURE);
+ }
+
+ @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 maybeWakeUpIdleEndpoint_connectsLowestIndexedIdleEndpoint() {
+ setUpStates(TRANSIENT_FAILURE, IDLE, IDLE);
+
+ endpointMap.maybeWakeUpIdleEndpoint();
+
+ assertThat(stateAt(1)).isEqualTo(CONNECTING);
+ assertThat(stateAt(2)).isEqualTo(IDLE);
+ }
+
+ @Test
+ public void maybeWakeUpIdleEndpoint_somethingAlreadyConnecting_doesNothing() {
+ setUpStates(CONNECTING, IDLE);
+ int childrenBefore = childProvider.children.size();
+
+ endpointMap.maybeWakeUpIdleEndpoint();
+
+ assertThat(childProvider.children).hasSize(childrenBefore);
+ assertThat(stateAt(1)).isEqualTo(IDLE);
+ }
+
+ @Test
+ public void maybeWakeUpIdleEndpoint_noIdleEndpoints_doesNothing() {
+ setUpStates(TRANSIENT_FAILURE, TRANSIENT_FAILURE);
+ int childrenBefore = childProvider.children.size();
+
+ endpointMap.maybeWakeUpIdleEndpoint();
+
+ assertThat(childProvider.children).hasSize(childrenBefore);
+ }
+
+ // ---------------------------------------------------------------------------------------------
+ // Shutdown
+ // ---------------------------------------------------------------------------------------------
+
+ @Test
+ public void shutdown_shutsDownChildrenAndEmptiesMap() {
+ endpointMap.updateEndpoints(endpoints("a", "b"), Attributes.EMPTY);
+ activate(0);
+ activate(1);
+
+ endpointMap.shutdown();
+
+ 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 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<>();
+
+ /** Run from every child's {@code acceptResolvedAddresses}, to exercise re-entrancy. */
+ Runnable onAccept;
+
+ @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, this);
+ children.add(child);
+ return child;
+ }
+ }
+
+ /** 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;
+ private final FakeChildProvider provider;
+ ResolvedAddresses lastAddresses;
+ int acceptCount;
+ int requestConnectionCount;
+ boolean shutdown;
+
+ FakeChild(Helper helper, FakeChildProvider provider) {
+ this.helper = helper;
+ this.provider = provider;
+ }
+
+ @Override
+ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+ lastAddresses = resolvedAddresses;
+ acceptCount++;
+ if (provider.onAccept != null) {
+ provider.onAccept.run();
+ }
+ 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);
+ }
+ }
+}