diff --git a/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java new file mode 100644 index 00000000000..81eeb9b370c --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -0,0 +1,217 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; +import io.grpc.ConnectivityState; +import io.grpc.InternalMetadata; +import io.grpc.LoadBalancer.PickResult; +import io.grpc.LoadBalancer.PickSubchannelArgs; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.Metadata; +import io.grpc.Status; +import java.util.List; +import java.util.concurrent.ThreadLocalRandom; +import javax.annotation.Nullable; +import javax.annotation.concurrent.ThreadSafe; + +/** + * Subchannel picker for the auto-sharding load balancing policy. + * + *

Routes RPCs to backend endpoints based on a request metadata header key, matching against + * an immutable {@link SliceMap}. + */ +final class AutoShardingPicker extends SubchannelPicker { + private static final byte[] EMPTY_BYTES = new byte[0]; + + @ThreadSafe + @FunctionalInterface + interface ThreadSafeRandom { + int nextInt(int bound); + } + + private static final ThreadSafeRandom DEFAULT_RANDOM = + bound -> ThreadLocalRandom.current().nextInt(bound); + + private static final InternalMetadata.TrustedAsciiMarshaller RAW_ASCII_MARSHALLER = + new InternalMetadata.TrustedAsciiMarshaller() { + @Override + public byte[] toAsciiString(byte[] value) { + return value; + } + + @Override + public byte[] parseAsciiString(byte[] serialized) { + return serialized; + } + }; + + private final SliceMap sliceMap; + private final ImmutableList endpoints; + private final boolean[] sliceInFallback; + private final boolean fallbackEnabled; + @Nullable private final Metadata.Key keyHeader; + private final ThreadSafeRandom random; + + /** + * Pre-creates a {@link Metadata.Key} for the given key header name. + * + * @param keyHeaderName the metadata header name, or {@code null}/empty if no header routing + * @return the pre-computed {@link Metadata.Key}, or {@code null} if keyHeaderName is null/empty + */ + @Nullable + static Metadata.Key createKeyHeader(@Nullable String keyHeaderName) { + if (keyHeaderName == null || keyHeaderName.isEmpty()) { + return null; + } else if (keyHeaderName.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { + return Metadata.Key.of(keyHeaderName, Metadata.BINARY_BYTE_MARSHALLER); + } else { + return InternalMetadata.keyOf(keyHeaderName, RAW_ASCII_MARSHALLER); + } + } + + /** + * Constructs an {@link AutoShardingPicker}. + * + * @param sliceMap the pre-built, immutable mapping from key ranges to endpoint indices + * @param endpoints the list of endpoint snapshots corresponding 1:1 to endpoint indices + * @param fallbackEnabled whether fallback routing to all resolved endpoints is enabled + * @param keyHeader the pre-parsed metadata header key used to extract the routing key + */ + AutoShardingPicker( + SliceMap sliceMap, + List endpoints, + boolean fallbackEnabled, + @Nullable Metadata.Key keyHeader) { + this(sliceMap, endpoints, fallbackEnabled, keyHeader, DEFAULT_RANDOM); + } + + @VisibleForTesting + AutoShardingPicker( + SliceMap sliceMap, + List endpoints, + boolean fallbackEnabled, + @Nullable Metadata.Key keyHeader, + ThreadSafeRandom random) { + this.sliceMap = checkNotNull(sliceMap, "sliceMap"); + this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints")); + this.fallbackEnabled = fallbackEnabled; + this.keyHeader = keyHeader; + this.random = checkNotNull(random, "random"); + + boolean hasTransientFailure = false; + for (int i = 0; i < this.endpoints.size(); i++) { + if (this.endpoints.get(i).getState() == ConnectivityState.TRANSIENT_FAILURE) { + hasTransientFailure = true; + break; + } + } + + this.sliceInFallback = new boolean[sliceMap.getSlices().size()]; + if (!hasTransientFailure) { + for (int i = 0; i < sliceInFallback.length; i++) { + this.sliceInFallback[i] = sliceMap.getSlices().get(i).getEndpoints().isEmpty(); + } + } else { + for (int i = 0; i < sliceInFallback.length; i++) { + this.sliceInFallback[i] = isPoolInFallback(sliceMap.getSlices().get(i).getEndpoints()); + } + } + } + + private boolean isPoolInFallback(List indices) { + if (indices.isEmpty()) { + return true; + } + for (int idx : indices) { + if (endpoints.get(idx).getState() != ConnectivityState.TRANSIENT_FAILURE) { + return false; + } + } + return true; + } + + @Override + public PickResult pickSubchannel(PickSubchannelArgs args) { + byte[] key = extractKeyBytes(args.getHeaders()); + int sliceIdx = sliceMap.lookup(key); + + if (sliceIdx == -1) { + if (fallbackEnabled) { + return pickFromEndpointIndices(sliceMap.getFallbackPool(), args); + } else { + return PickResult.withError( + Status.UNAVAILABLE.withDescription( + "No sharding assignment available and fallback disabled")); + } + } + + if (sliceInFallback[sliceIdx] && fallbackEnabled) { + return pickFromEndpointIndices(sliceMap.getFallbackPool(), args); + } + + SliceMap.SliceEntry sliceEntry = sliceMap.getSlices().get(sliceIdx); + return pickFromEndpointIndices(sliceEntry.getEndpoints(), args); + } + + private PickResult pickFromEndpointIndices( + List indices, PickSubchannelArgs args) { + if (indices.isEmpty()) { + return PickResult.withError( + Status.UNAVAILABLE.withDescription("No valid endpoints in slice and fallback disabled")); + } + + int size = indices.size(); + int firstIndex = random.nextInt(size); + boolean requestedConnection = false; + boolean foundConnecting = false; + + for (int i = 0; i < size; i++) { + int epIdx = indices.get((firstIndex + i) % size); + PickerEndpoint endpoint = endpoints.get(epIdx); + + if (endpoint.getState() == ConnectivityState.READY) { + return endpoint.getPicker().pickSubchannel(args); + } + + if (endpoint.getState() == ConnectivityState.CONNECTING) { + foundConnecting = true; + } else if (!requestedConnection && endpoint.getState() == ConnectivityState.IDLE) { + endpoint.requestConnection(); + requestedConnection = true; + } + } + + if (requestedConnection || foundConnecting) { + return PickResult.withNoResult("connecting", "Waiting for endpoint connection"); + } + + int firstEpIdx = indices.get(firstIndex); + return endpoints.get(firstEpIdx).getPicker().pickSubchannel(args); + } + + private byte[] extractKeyBytes(Metadata headers) { + if (keyHeader != null) { + byte[] val = headers.get(keyHeader); + return val != null ? val : EMPTY_BYTES; + } + return EMPTY_BYTES; + } +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java new file mode 100644 index 00000000000..d85189d222e --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java @@ -0,0 +1,86 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.common.base.MoreObjects; +import io.grpc.ConnectivityState; +import io.grpc.LoadBalancer.SubchannelPicker; +import javax.annotation.Nullable; + +/** + * Immutable snapshot of endpoint state used by {@link AutoShardingPicker}. + */ +final class PickerEndpoint { + + /** + * Callback interface to trigger connection attempts on an IDLE endpoint's child balancer. + */ + @FunctionalInterface + interface ExitIdler { + /** + * Requests the child load balancer to exit IDLE and initiate a connection. + * + *

Implementations MUST be thread-safe, non-blocking, idempotent, and dispatch + * execution to the {@link io.grpc.SynchronizationContext}. + */ + void exitIdle(); + } + + private final ConnectivityState state; + private final SubchannelPicker picker; + @Nullable private final ExitIdler exitIdler; + + /** + * Constructs a {@link PickerEndpoint}. + * + * @param state the current connectivity state of the endpoint + * @param picker the latest subchannel picker for the endpoint + * @param exitIdler a callback to trigger an IDLE child balancer to start connecting + */ + PickerEndpoint( + ConnectivityState state, + SubchannelPicker picker, + @Nullable ExitIdler exitIdler) { + this.state = checkNotNull(state, "state"); + this.picker = checkNotNull(picker, "picker"); + this.exitIdler = exitIdler; + } + + ConnectivityState getState() { + return state; + } + + SubchannelPicker getPicker() { + return picker; + } + + void requestConnection() { + if (exitIdler != null) { + exitIdler.exitIdle(); + } + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("state", state) + .add("picker", picker) + .toString(); + } +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java new file mode 100644 index 00000000000..269c1b0ca7e --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -0,0 +1,135 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.common.collect.ImmutableList; +import com.google.common.primitives.UnsignedBytes; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import javax.annotation.Nullable; + +/** + * An immutable lookup structure mapping application routing keys to slice indices. + * + *

As defined in gRFC A119, the assignment provider guarantees that the assignment + * is pre-validated, gap-free, non-overlapping, and covers the entire keyspace {@code ["" .. inf)}. + * Any gaps returned by the autosharding server are filled as slice entries with an empty + * endpoints list. Therefore, each {@link SliceEntry} only needs to store {@code startKey} + * because the exclusive end key of slice {@code i} is implicitly the inclusive start key of + * slice {@code i + 1}. + */ +final class SliceMap { + + /** + * Represents a single key-range slice mapping to endpoint indices in the picker. + */ + static final class SliceEntry { + private final byte[] startKey; + private final ImmutableList endpoints; + + /** + * Constructs a {@link SliceEntry}. + * + * @param startKey the inclusive start key of the slice + * @param endpoints the list of endpoint indices assigned to this slice + */ + SliceEntry(byte[] startKey, List endpoints) { + this.startKey = checkNotNull(startKey, "startKey"); + this.endpoints = ImmutableList.copyOf(checkNotNull(endpoints, "endpoints")); + } + + byte[] getStartKey() { + return startKey; + } + + ImmutableList getEndpoints() { + return endpoints; + } + } + + private static final Comparator UNSIGNED_BYTES_COMPARATOR = + UnsignedBytes.lexicographicalComparator(); + private static final byte[] EMPTY_BYTES = new byte[0]; + + private final ImmutableList slices; + private final ImmutableList fallbackPool; + private final long generation; + + /** + * Constructs an immutable {@link SliceMap}. + * + * @param slices the pre-validated list of key-range slice entries + * @param fallbackPool the list of all available endpoint indices for fallback routing + * @param generation the snapshot generation number from the assignment + */ + SliceMap(List slices, List fallbackPool, long generation) { + List sortedSlices = new ArrayList<>(checkNotNull(slices, "slices")); + sortedSlices.sort( + (e1, e2) -> UNSIGNED_BYTES_COMPARATOR.compare(e1.getStartKey(), e2.getStartKey())); + this.slices = ImmutableList.copyOf(sortedSlices); + this.fallbackPool = ImmutableList.copyOf(checkNotNull(fallbackPool, "fallbackPool")); + this.generation = generation; + } + + /** + * Looks up the matching slice index for the given key. + * Returns -1 if slices is empty (e.g. startup/fallback case where there are no assignments) + * or if the key is smaller than the first slice's startKey. + */ + int lookup(@Nullable byte[] key) { + if (slices.isEmpty()) { + return -1; + } + byte[] searchKey = key != null ? key : EMPTY_BYTES; + int low = 0; + int high = slices.size() - 1; + + while (low <= high) { + int mid = (low + high) >>> 1; + int cmp = UNSIGNED_BYTES_COMPARATOR.compare(slices.get(mid).getStartKey(), searchKey); + + if (cmp < 0) { + low = mid + 1; + } else if (cmp > 0) { + high = mid - 1; + } else { + return mid; // Exact match on startKey + } + } + + if (low == 0) { + // Key is smaller than first slice's startKey + return -1; + } + return low - 1; + } + + ImmutableList getSlices() { + return slices; + } + + ImmutableList getFallbackPool() { + return fallbackPool; + } + + long getGeneration() { + return generation; + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java new file mode 100644 index 00000000000..7bced71cd94 --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -0,0 +1,384 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import static com.google.common.truth.Truth.assertThat; + +import io.grpc.CallOptions; +import io.grpc.ConnectivityState; +import io.grpc.LoadBalancer.PickDetailsConsumer; +import io.grpc.LoadBalancer.PickResult; +import io.grpc.LoadBalancer.PickSubchannelArgs; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Status; +import io.grpc.autosharding.PickerEndpoint.ExitIdler; +import io.grpc.autosharding.SliceMap.SliceEntry; +import io.grpc.internal.PickSubchannelArgsImpl; +import io.grpc.testing.TestMethodDescriptors; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AutoShardingPickerTest { + + private static final MethodDescriptor METHOD = TestMethodDescriptors.voidMethod(); + private static final ExitIdler NOOP_EXIT_IDLER = new ExitIdler() { + @Override + public void exitIdle() {} + }; + private static final PickDetailsConsumer NOOP_CONSUMER = new PickDetailsConsumer() {}; + + private PickSubchannelArgs createArgs(Metadata headers) { + return new PickSubchannelArgsImpl(METHOD, headers, CallOptions.DEFAULT, NOOP_CONSUMER); + } + + private static class FakePicker extends SubchannelPicker { + private final PickResult result; + + FakePicker(PickResult result) { + this.result = result; + } + + @Override + public PickResult pickSubchannel(PickSubchannelArgs args) { + return result; + } + } + + @Test + public void pick_noSliceMap_fallbackEnabled_picksFromFallbackPool() { + PickResult readyResult = PickResult.withNoResult(); // using as token + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(readyResult), NOOP_EXIT_IDLER); + + SliceMap emptySliceMap = new SliceMap( + Collections.emptyList(), Collections.singletonList(0), 1L); + AutoShardingPicker picker = new AutoShardingPicker( + emptySliceMap, + Collections.singletonList(ep0), + true, + AutoShardingPicker.createKeyHeader("x-slice-key")); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("x-slice-key", Metadata.ASCII_STRING_MARSHALLER), "user123"); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(readyResult); + } + + @Test + public void pick_noSliceMap_fallbackDisabled_returnsUnavailableError() { + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP_EXIT_IDLER); + + SliceMap emptySliceMap = new SliceMap( + Collections.emptyList(), Collections.singletonList(0), 1L); + AutoShardingPicker picker = new AutoShardingPicker( + emptySliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); + + Metadata headers = new Metadata(); + PickResult result = picker.pickSubchannel(createArgs(headers)); + + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()) + .contains("No sharding assignment available and fallback disabled"); + } + + @Test + public void pick_sliceFound_readyEndpoint_returnsPickResult() { + PickResult expectedResult = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(expectedResult), NOOP_EXIT_IDLER); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("x-slice-key", Metadata.ASCII_STRING_MARSHALLER), "anyKey"); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(expectedResult); + } + + @Test + public void pick_sliceFound_idleEndpoint_triggersConnectionAndQueues() { + AtomicInteger connectCalls = new AtomicInteger(0); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.IDLE, + new FakePicker(PickResult.withNoResult()), + connectCalls::incrementAndGet); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + + assertThat(connectCalls.get()).isEqualTo(1); + assertThat(result.hasResult()).isFalse(); + } + + @Test + public void pick_sliceFound_connectingEndpoint_queuesPick() { + AtomicInteger connectCalls = new AtomicInteger(0); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.CONNECTING, + new FakePicker(PickResult.withNoResult()), + connectCalls::incrementAndGet); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + + assertThat(connectCalls.get()).isEqualTo(0); + assertThat(result.hasResult()).isFalse(); + } + + @Test + public void pick_sliceFound_allTransientFailure_fallbackEnabled_picksFromFallbackPool() { + PickResult fallbackReadyResult = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.TRANSIENT_FAILURE, + new FakePicker(PickResult.withError(Status.UNAVAILABLE.withDescription("ep0 down"))), + NOOP_EXIT_IDLER); + PickerEndpoint ep1 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP_EXIT_IDLER); + + // Slice 0 only has ep0 (which is down) + SliceEntry slice0 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + // Fallback pool has ep1 (which is ready) + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice0), Collections.singletonList(1), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Arrays.asList(ep0, ep1), + true, + AutoShardingPicker.createKeyHeader("x-slice-key")); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result).isSameInstanceAs(fallbackReadyResult); + } + + @Test + public void pick_sliceFound_allTransientFailure_fallbackDisabled_delegatesToEndpointPicker() { + Status epError = Status.UNAVAILABLE.withDescription("connection refused to ep0"); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.TRANSIENT_FAILURE, + new FakePicker(PickResult.withError(epError)), + NOOP_EXIT_IDLER); + + SliceEntry slice0 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice0), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result.getStatus()).isEqualTo(epError); + } + + @Test + public void pick_binaryHeader_extractedProperly() { + PickResult ready0 = PickResult.withNoResult(); + PickResult ready1 = PickResult.withNoResult(); + + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready0), NOOP_EXIT_IDLER); + PickerEndpoint ep1 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready1), NOOP_EXIT_IDLER); + + SliceEntry s0 = new SliceEntry(new byte[] {0x00}, Collections.singletonList(0)); + SliceEntry s1 = new SliceEntry(new byte[] {0x50}, Collections.singletonList(1)); + SliceMap sliceMap = new SliceMap(Arrays.asList(s0, s1), Arrays.asList(0, 1), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Arrays.asList(ep0, ep1), + false, + AutoShardingPicker.createKeyHeader("slice-key-bin")); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("slice-key-bin", Metadata.BINARY_BYTE_MARSHALLER), + new byte[] {0x60}); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(ready1); + } + + @Test + public void pick_emptySliceEndpoints_fallbackDisabled_returnsUnavailable() { + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP_EXIT_IDLER); + + SliceEntry emptySlice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.emptyList()); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(emptySlice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Collections.singletonList(ep0), + false, + AutoShardingPicker.createKeyHeader("x-slice-key")); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()) + .contains("No valid endpoints in slice and fallback disabled"); + } + + @Test + public void pick_emptySliceEndpoints_fallbackEnabled_routesToFallbackPool() { + PickResult fallbackReadyResult = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP_EXIT_IDLER); + + // Gap slice with empty endpoints list + SliceEntry gapSlice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.emptyList()); + // Fallback pool has ep0 + SliceMap sliceMap = new SliceMap( + Collections.singletonList(gapSlice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, + Collections.singletonList(ep0), + true, + AutoShardingPicker.createKeyHeader("x-key")); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("x-key", Metadata.ASCII_STRING_MARSHALLER), "anyKey"); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(fallbackReadyResult); + } + + @Test + public void pickerEndpoint_gettersAndToString() { + FakePicker fakePicker = new FakePicker(PickResult.withNoResult()); + AtomicInteger count = new AtomicInteger(); + PickerEndpoint ep = new PickerEndpoint( + ConnectivityState.IDLE, fakePicker, count::incrementAndGet); + + assertThat(ep.getState()).isEqualTo(ConnectivityState.IDLE); + assertThat(ep.getPicker()).isSameInstanceAs(fakePicker); + assertThat(ep.toString()).contains("state=IDLE"); + + ep.requestConnection(); + assertThat(count.get()).isEqualTo(1); + } + + @Test + public void createKeyHeader_nullOrEmpty_returnsNull() { + assertThat(AutoShardingPicker.createKeyHeader(null)).isNull(); + assertThat(AutoShardingPicker.createKeyHeader("")).isNull(); + } + + @Test + public void createKeyHeader_asciiHeader() { + Metadata.Key key = AutoShardingPicker.createKeyHeader("x-slice-key"); + assertThat(key).isNotNull(); + assertThat(key.name()).isEqualTo("x-slice-key"); + } + + @Test + public void createKeyHeader_binaryHeader() { + Metadata.Key key = AutoShardingPicker.createKeyHeader("x-slice-key-bin"); + assertThat(key).isNotNull(); + assertThat(key.name()).isEqualTo("x-slice-key-bin"); + } + + @Test + public void pick_deterministicRandom_selectsExpectedEndpoint() { + PickResult ready0 = PickResult.withNoResult(); + PickResult ready1 = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready0), NOOP_EXIT_IDLER); + PickerEndpoint ep1 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready1), NOOP_EXIT_IDLER); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Arrays.asList(0, 1)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Arrays.asList(0, 1), 1L); + + // Test picking index 0 + AutoShardingPicker picker0 = new AutoShardingPicker( + sliceMap, + Arrays.asList(ep0, ep1), + false, + AutoShardingPicker.createKeyHeader("x-key"), + bound -> 0); + PickResult result0 = picker0.pickSubchannel(createArgs(new Metadata())); + assertThat(result0).isSameInstanceAs(ready0); + + // Test picking index 1 + AutoShardingPicker picker1 = new AutoShardingPicker( + sliceMap, + Arrays.asList(ep0, ep1), + false, + AutoShardingPicker.createKeyHeader("x-key"), + bound -> 1); + PickResult result1 = picker1.pickSubchannel(createArgs(new Metadata())); + assertThat(result1).isSameInstanceAs(ready1); + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java new file mode 100644 index 00000000000..2ad3d13998c --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java @@ -0,0 +1,138 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import static com.google.common.truth.Truth.assertThat; + +import io.grpc.autosharding.SliceMap.SliceEntry; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class SliceMapTest { + + @Test + public void lookup_emptySlices_returnsInvalidIndex() { + SliceMap sliceMap = new SliceMap(Collections.emptyList(), Arrays.asList(0, 1), 1L); + assertThat(sliceMap.lookup(new byte[] {1, 2, 3})).isEqualTo(-1); + assertThat(sliceMap.lookup(null)).isEqualTo(-1); + assertThat(sliceMap.lookup(new byte[0])).isEqualTo(-1); + } + + @Test + public void lookup_singleSlice() { + byte[] startKey = new byte[0]; // Covers ["" .. inf) + SliceEntry slice = new SliceEntry(startKey, Arrays.asList(0, 1)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Arrays.asList(0, 1), 10L); + + assertThat(sliceMap.lookup(new byte[0])).isEqualTo(0); + assertThat(sliceMap.lookup("foo".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup(null)).isEqualTo(0); + } + + @Test + public void lookup_multipleSlices() { + // Slices: ["" .. "m"), ["m" .. "t"), ["t" .. inf) + SliceEntry s1 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceEntry s2 = new SliceEntry( + "m".getBytes(StandardCharsets.UTF_8), Collections.singletonList(1)); + SliceEntry s3 = new SliceEntry( + "t".getBytes(StandardCharsets.UTF_8), Collections.singletonList(2)); + + SliceMap sliceMap = new SliceMap(Arrays.asList(s3, s1, s2), Arrays.asList(0, 1, 2), 5L); + + // Exact matches + assertThat(sliceMap.lookup("".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("m".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("t".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + + // In-between matches + assertThat(sliceMap.lookup("a".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("l".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("n".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("s".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("u".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + assertThat(sliceMap.lookup("zzz".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + } + + @Test + public void lookup_keySmallerThanFirstSlice_returnsInvalidIndex() { + // Slice starts at "m" + SliceEntry s1 = new SliceEntry( + "m".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(s1), Collections.singletonList(0), 1L); + + assertThat(sliceMap.lookup("a".getBytes(StandardCharsets.UTF_8))).isEqualTo(-1); + assertThat(sliceMap.lookup("".getBytes(StandardCharsets.UTF_8))).isEqualTo(-1); + assertThat(sliceMap.lookup(null)).isEqualTo(-1); + assertThat(sliceMap.lookup("m".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("z".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + } + + @Test + public void lookup_unsignedByteComparison() { + // Test that 0x80 is treated as greater than 0x7F (unsigned) + byte[] key1 = new byte[] {0x7F}; + byte[] key2 = new byte[] {(byte) 0x80}; + byte[] key3 = new byte[] {(byte) 0xFF}; + + SliceEntry s1 = new SliceEntry(new byte[0], Collections.singletonList(0)); + SliceEntry s2 = new SliceEntry(key1, Collections.singletonList(1)); + SliceEntry s3 = new SliceEntry(key2, Collections.singletonList(2)); + SliceEntry s4 = new SliceEntry(key3, Collections.singletonList(3)); + + SliceMap sliceMap = new SliceMap( + Arrays.asList(s4, s2, s1, s3), Arrays.asList(0, 1, 2, 3), 1L); + + assertThat(sliceMap.lookup(new byte[] {0x10})).isEqualTo(0); + assertThat(sliceMap.lookup(new byte[] {0x7F})).isEqualTo(1); + assertThat(sliceMap.lookup(new byte[] {(byte) 0x80})).isEqualTo(2); + assertThat(sliceMap.lookup(new byte[] {(byte) 0x90})).isEqualTo(2); + assertThat(sliceMap.lookup(new byte[] {(byte) 0xFF})).isEqualTo(3); + assertThat(sliceMap.lookup(new byte[] {(byte) 0xFF, 0x01})).isEqualTo(3); + } + + @Test + public void gettersAndImmutability() { + List slices = new ArrayList<>(); + slices.add(new SliceEntry(new byte[] {1}, Arrays.asList(0, 1))); + List fallback = new ArrayList<>(Arrays.asList(0, 1)); + + SliceMap sliceMap = new SliceMap(slices, fallback, 42L); + + assertThat(sliceMap.getGeneration()).isEqualTo(42L); + assertThat(sliceMap.getFallbackPool()).containsExactly(0, 1).inOrder(); + assertThat(sliceMap.getSlices()).hasSize(1); + assertThat(sliceMap.getSlices().get(0).getStartKey()).isEqualTo(new byte[] {1}); + assertThat(sliceMap.getSlices().get(0).getEndpoints()).containsExactly(0, 1).inOrder(); + + // Verify defensive copying: mutating input collections does not affect sliceMap + slices.clear(); + fallback.clear(); + assertThat(sliceMap.getSlices()).hasSize(1); + assertThat(sliceMap.getFallbackPool()).containsExactly(0, 1).inOrder(); + } +}