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..b7ab4e71cea --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/AutoShardingPicker.java @@ -0,0 +1,156 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import io.grpc.ConnectivityState; +import io.grpc.InternalMetadata; +import io.grpc.LoadBalancer.PickResult; +import io.grpc.LoadBalancer.PickSubchannelArgs; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.Metadata; +import io.grpc.Status; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ThreadLocalRandom; + +final class AutoShardingPicker extends SubchannelPicker { + private static final byte[] EMPTY_BYTES = new byte[0]; + + private static final InternalMetadata.TrustedAsciiMarshaller RAW_ASCII_MARSHALLER = + new InternalMetadata.TrustedAsciiMarshaller() { + @Override + public byte[] toAsciiString(byte[] value) { + return value; + } + + @Override + public byte[] parseAsciiString(byte[] serialized) { + return serialized; + } + }; + + private final SliceMap sliceMap; + private final List endpoints; + private final boolean[] sliceInFallback; + private final boolean fallbackEnabled; + private final Metadata.Key keyHeader; + + AutoShardingPicker( + SliceMap sliceMap, + List endpoints, + boolean fallbackEnabled, + String keyHeaderName) { + this.sliceMap = sliceMap; + this.endpoints = Collections.unmodifiableList(new ArrayList<>(endpoints)); + this.fallbackEnabled = fallbackEnabled; + + if (keyHeaderName == null || keyHeaderName.isEmpty()) { + this.keyHeader = null; + } else if (keyHeaderName.endsWith(Metadata.BINARY_HEADER_SUFFIX)) { + this.keyHeader = Metadata.Key.of(keyHeaderName, Metadata.BINARY_BYTE_MARSHALLER); + } else { + this.keyHeader = InternalMetadata.keyOf(keyHeaderName, RAW_ASCII_MARSHALLER); + } + + this.sliceInFallback = new boolean[sliceMap.getSlices().size()]; + for (int i = 0; i < sliceInFallback.length; i++) { + this.sliceInFallback[i] = isPoolInFallback(sliceMap.getSlices().get(i).endpoints); + } + } + + private boolean isPoolInFallback(List indices) { + if (indices.isEmpty()) { + return true; + } + for (int idx : indices) { + if (endpoints.get(idx).state != ConnectivityState.TRANSIENT_FAILURE) { + return false; + } + } + return true; + } + + @Override + public PickResult pickSubchannel(PickSubchannelArgs args) { + byte[] key = extractKeyBytes(args.getHeaders()); + Integer sliceIdx = sliceMap.lookup(key); + + if (sliceIdx == null) { + if (fallbackEnabled) { + return pickFromEndpointIndices(sliceMap.getFallbackPool(), args); + } else { + return PickResult.withError( + Status.UNAVAILABLE.withDescription( + "No sharding assignment available and fallback disabled")); + } + } + + if (sliceInFallback[sliceIdx] && fallbackEnabled) { + return pickFromEndpointIndices(sliceMap.getFallbackPool(), args); + } + + SliceMap.SliceEntry sliceEntry = sliceMap.getSlices().get(sliceIdx); + return pickFromEndpointIndices(sliceEntry.endpoints, args); + } + + private PickResult pickFromEndpointIndices( + List indices, PickSubchannelArgs args) { + if (indices.isEmpty()) { + return PickResult.withError( + Status.UNAVAILABLE.withDescription("No valid endpoints in slice and fallback disabled")); + } + + int size = indices.size(); + int firstIndex = ThreadLocalRandom.current().nextInt(size); + boolean requestedConnection = false; + boolean foundConnecting = false; + + for (int i = 0; i < size; i++) { + int epIdx = indices.get((firstIndex + i) % size); + PickerEndpoint endpoint = endpoints.get(epIdx); + + if (endpoint.state == ConnectivityState.READY) { + return endpoint.picker.pickSubchannel(args); + } + + if (endpoint.state == ConnectivityState.CONNECTING) { + foundConnecting = true; + } else if (!requestedConnection && endpoint.state == ConnectivityState.IDLE) { + if (endpoint.requestConnection != null) { + endpoint.requestConnection.run(); + } + requestedConnection = true; + } + } + + if (requestedConnection || foundConnecting) { + return PickResult.withNoResult("connecting", "Waiting for endpoint connection"); + } + + int firstEpIdx = indices.get(firstIndex); + return endpoints.get(firstEpIdx).picker.pickSubchannel(args); + } + + private byte[] extractKeyBytes(Metadata headers) { + if (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..54310a1bdb6 --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/PickerEndpoint.java @@ -0,0 +1,36 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import io.grpc.ConnectivityState; +import io.grpc.LoadBalancer.SubchannelPicker; + +/** + * Immutable snapshot of endpoint state used by the AutoShardingPicker. + */ +final class PickerEndpoint { + final ConnectivityState state; + final SubchannelPicker picker; + final Runnable requestConnection; + + PickerEndpoint( + ConnectivityState state, SubchannelPicker picker, Runnable requestConnection) { + this.state = state; + this.picker = picker; + this.requestConnection = requestConnection; + } +} diff --git a/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java new file mode 100644 index 00000000000..5444c252556 --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/SliceMap.java @@ -0,0 +1,105 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import javax.annotation.Nullable; + +final class SliceMap { + + static final class SliceEntry { + final byte[] startKey; + final List endpoints; + + SliceEntry(byte[] startKey, List endpoints) { + this.startKey = startKey; + this.endpoints = Collections.unmodifiableList(new ArrayList<>(endpoints)); + } + } + + private static final byte[] EMPTY_BYTES = new byte[0]; + + private final List slices; + private final List fallbackPool; + private final long generation; + + SliceMap(List slices, List fallbackPool, long generation) { + List sortedSlices = new ArrayList<>(slices); + sortedSlices.sort((e1, e2) -> compareUnsigned(e1.startKey, e2.startKey)); + this.slices = Collections.unmodifiableList(sortedSlices); + this.fallbackPool = Collections.unmodifiableList(new ArrayList<>(fallbackPool)); + this.generation = generation; + } + + /** + * Looks up the matching slice index for the given key. + * Returns null if slices is empty (e.g. startup/fallback case where there are no assignments). + */ + @Nullable + Integer lookup(@Nullable byte[] key) { + if (slices.isEmpty()) { + return null; + } + byte[] searchKey = key != null ? key : EMPTY_BYTES; + int low = 0; + int high = slices.size() - 1; + + while (low <= high) { + int mid = (low + high) >>> 1; + int cmp = compareUnsigned(slices.get(mid).startKey, searchKey); + + if (cmp < 0) { + low = mid + 1; + } else if (cmp > 0) { + high = mid - 1; + } else { + return mid; // Exact match on startKey + } + } + + if (low == 0) { + // Key is smaller than first slice's startKey + return null; + } + return low - 1; + } + + private static int compareUnsigned(byte[] a, byte[] b) { + int minLength = Math.min(a.length, b.length); + for (int i = 0; i < minLength; i++) { + int result = (a[i] & 0xFF) - (b[i] & 0xFF); + if (result != 0) { + return result; + } + } + return a.length - b.length; + } + + List getSlices() { + return slices; + } + + List getFallbackPool() { + return fallbackPool; + } + + long getGeneration() { + return generation; + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java new file mode 100644 index 00000000000..f1d29062da7 --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/AutoShardingPickerTest.java @@ -0,0 +1,283 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import static com.google.common.truth.Truth.assertThat; + +import io.grpc.CallOptions; +import io.grpc.ConnectivityState; +import io.grpc.LoadBalancer.PickDetailsConsumer; +import io.grpc.LoadBalancer.PickResult; +import io.grpc.LoadBalancer.PickSubchannelArgs; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.Status; +import io.grpc.autosharding.SliceMap.SliceEntry; +import io.grpc.internal.PickSubchannelArgsImpl; +import io.grpc.testing.TestMethodDescriptors; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AutoShardingPickerTest { + + private static final MethodDescriptor METHOD = TestMethodDescriptors.voidMethod(); + private static final Runnable NOOP = new Runnable() { + @Override + public void run() {} + }; + private static final PickDetailsConsumer NOOP_CONSUMER = new PickDetailsConsumer() {}; + + private PickSubchannelArgs createArgs(Metadata headers) { + return new PickSubchannelArgsImpl(METHOD, headers, CallOptions.DEFAULT, NOOP_CONSUMER); + } + + private static class FakePicker extends SubchannelPicker { + private final PickResult result; + + FakePicker(PickResult result) { + this.result = result; + } + + @Override + public PickResult pickSubchannel(PickSubchannelArgs args) { + return result; + } + } + + @Test + public void pick_noSliceMap_fallbackEnabled_picksFromFallbackPool() { + PickResult readyResult = PickResult.withNoResult(); // using as token + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(readyResult), NOOP); + + SliceMap emptySliceMap = new SliceMap( + Collections.emptyList(), Collections.singletonList(0), 1L); + AutoShardingPicker picker = new AutoShardingPicker( + emptySliceMap, Collections.singletonList(ep0), true, "x-slice-key"); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("x-slice-key", Metadata.ASCII_STRING_MARSHALLER), "user123"); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(readyResult); + } + + @Test + public void pick_noSliceMap_fallbackDisabled_returnsUnavailableError() { + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP); + + SliceMap emptySliceMap = new SliceMap( + Collections.emptyList(), Collections.singletonList(0), 1L); + AutoShardingPicker picker = new AutoShardingPicker( + emptySliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + Metadata headers = new Metadata(); + PickResult result = picker.pickSubchannel(createArgs(headers)); + + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()) + .contains("No sharding assignment available and fallback disabled"); + } + + @Test + public void pick_sliceFound_readyEndpoint_returnsPickResult() { + PickResult expectedResult = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(expectedResult), NOOP); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("x-slice-key", Metadata.ASCII_STRING_MARSHALLER), "anyKey"); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(expectedResult); + } + + @Test + public void pick_sliceFound_idleEndpoint_triggersConnectionAndQueues() { + AtomicInteger connectCalls = new AtomicInteger(0); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.IDLE, + new FakePicker(PickResult.withNoResult()), + connectCalls::incrementAndGet); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + + assertThat(connectCalls.get()).isEqualTo(1); + assertThat(result.hasResult()).isFalse(); + } + + @Test + public void pick_sliceFound_connectingEndpoint_queuesPick() { + AtomicInteger connectCalls = new AtomicInteger(0); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.CONNECTING, + new FakePicker(PickResult.withNoResult()), + connectCalls::incrementAndGet); + + SliceEntry slice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + + assertThat(connectCalls.get()).isEqualTo(0); + assertThat(result.hasResult()).isFalse(); + } + + @Test + public void pick_sliceFound_allTransientFailure_fallbackEnabled_picksFromFallbackPool() { + PickResult fallbackReadyResult = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.TRANSIENT_FAILURE, + new FakePicker(PickResult.withError(Status.UNAVAILABLE.withDescription("ep0 down"))), + NOOP); + PickerEndpoint ep1 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP); + + // Slice 0 only has ep0 (which is down) + SliceEntry slice0 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + // Fallback pool has ep1 (which is ready) + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice0), Collections.singletonList(1), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Arrays.asList(ep0, ep1), true, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result).isSameInstanceAs(fallbackReadyResult); + } + + @Test + public void pick_sliceFound_allTransientFailure_fallbackDisabled_delegatesToEndpointPicker() { + Status epError = Status.UNAVAILABLE.withDescription("connection refused to ep0"); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.TRANSIENT_FAILURE, + new FakePicker(PickResult.withError(epError)), + NOOP); + + SliceEntry slice0 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice0), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result.getStatus()).isEqualTo(epError); + } + + @Test + public void pick_binaryHeader_extractedProperly() { + PickResult ready0 = PickResult.withNoResult(); + PickResult ready1 = PickResult.withNoResult(); + + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready0), NOOP); + PickerEndpoint ep1 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(ready1), NOOP); + + SliceEntry s0 = new SliceEntry(new byte[] {0x00}, Collections.singletonList(0)); + SliceEntry s1 = new SliceEntry(new byte[] {0x50}, Collections.singletonList(1)); + SliceMap sliceMap = new SliceMap(Arrays.asList(s0, s1), Arrays.asList(0, 1), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Arrays.asList(ep0, ep1), false, "slice-key-bin"); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("slice-key-bin", Metadata.BINARY_BYTE_MARSHALLER), + new byte[] {0x60}); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(ready1); + } + + @Test + public void pick_emptySliceEndpoints_fallbackDisabled_returnsUnavailable() { + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(PickResult.withNoResult()), NOOP); + + SliceEntry emptySlice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.emptyList()); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(emptySlice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), false, "x-slice-key"); + + PickResult result = picker.pickSubchannel(createArgs(new Metadata())); + assertThat(result.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()) + .contains("No valid endpoints in slice and fallback disabled"); + } + + @Test + public void pick_emptySliceEndpoints_fallbackEnabled_routesToFallbackPool() { + PickResult fallbackReadyResult = PickResult.withNoResult(); + PickerEndpoint ep0 = new PickerEndpoint( + ConnectivityState.READY, new FakePicker(fallbackReadyResult), NOOP); + + // Gap slice with empty endpoints list + SliceEntry gapSlice = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.emptyList()); + // Fallback pool has ep0 + SliceMap sliceMap = new SliceMap( + Collections.singletonList(gapSlice), Collections.singletonList(0), 1L); + + AutoShardingPicker picker = new AutoShardingPicker( + sliceMap, Collections.singletonList(ep0), true, "x-key"); + + Metadata headers = new Metadata(); + headers.put( + Metadata.Key.of("x-key", Metadata.ASCII_STRING_MARSHALLER), "anyKey"); + + PickResult result = picker.pickSubchannel(createArgs(headers)); + assertThat(result).isSameInstanceAs(fallbackReadyResult); + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java new file mode 100644 index 00000000000..2ec753cf4e9 --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/SliceMapTest.java @@ -0,0 +1,139 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import io.grpc.autosharding.SliceMap.SliceEntry; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class SliceMapTest { + + @Test + public void lookup_emptySlices_returnsNull() { + SliceMap sliceMap = new SliceMap(Collections.emptyList(), Arrays.asList(0, 1), 1L); + assertThat(sliceMap.lookup(new byte[] {1, 2, 3})).isNull(); + assertThat(sliceMap.lookup(null)).isNull(); + assertThat(sliceMap.lookup(new byte[0])).isNull(); + } + + @Test + public void lookup_singleSlice() { + byte[] startKey = new byte[0]; // Covers ["" .. inf) + SliceEntry slice = new SliceEntry(startKey, Arrays.asList(0, 1)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(slice), Arrays.asList(0, 1), 10L); + + assertThat(sliceMap.lookup(new byte[0])).isEqualTo(0); + assertThat(sliceMap.lookup("foo".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup(null)).isEqualTo(0); + } + + @Test + public void lookup_multipleSlices() { + // Slices: ["" .. "m"), ["m" .. "t"), ["t" .. inf) + SliceEntry s1 = new SliceEntry( + "".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceEntry s2 = new SliceEntry( + "m".getBytes(StandardCharsets.UTF_8), Collections.singletonList(1)); + SliceEntry s3 = new SliceEntry( + "t".getBytes(StandardCharsets.UTF_8), Collections.singletonList(2)); + + SliceMap sliceMap = new SliceMap(Arrays.asList(s3, s1, s2), Arrays.asList(0, 1, 2), 5L); + + // Exact matches + assertThat(sliceMap.lookup("".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("m".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("t".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + + // In-between matches + assertThat(sliceMap.lookup("a".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("l".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("n".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("s".getBytes(StandardCharsets.UTF_8))).isEqualTo(1); + assertThat(sliceMap.lookup("u".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + assertThat(sliceMap.lookup("zzz".getBytes(StandardCharsets.UTF_8))).isEqualTo(2); + } + + @Test + public void lookup_keySmallerThanFirstSlice_returnsNull() { + // Slice starts at "m" + SliceEntry s1 = new SliceEntry( + "m".getBytes(StandardCharsets.UTF_8), Collections.singletonList(0)); + SliceMap sliceMap = new SliceMap( + Collections.singletonList(s1), Collections.singletonList(0), 1L); + + assertThat(sliceMap.lookup("a".getBytes(StandardCharsets.UTF_8))).isNull(); + assertThat(sliceMap.lookup("".getBytes(StandardCharsets.UTF_8))).isNull(); + assertThat(sliceMap.lookup(null)).isNull(); + assertThat(sliceMap.lookup("m".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + assertThat(sliceMap.lookup("z".getBytes(StandardCharsets.UTF_8))).isEqualTo(0); + } + + @Test + public void lookup_unsignedByteComparison() { + // Test that 0x80 is treated as greater than 0x7F (unsigned) + byte[] key1 = new byte[] {0x7F}; + byte[] key2 = new byte[] {(byte) 0x80}; + byte[] key3 = new byte[] {(byte) 0xFF}; + + SliceEntry s1 = new SliceEntry(new byte[0], Collections.singletonList(0)); + SliceEntry s2 = new SliceEntry(key1, Collections.singletonList(1)); + SliceEntry s3 = new SliceEntry(key2, Collections.singletonList(2)); + SliceEntry s4 = new SliceEntry(key3, Collections.singletonList(3)); + + SliceMap sliceMap = new SliceMap( + Arrays.asList(s4, s2, s1, s3), Arrays.asList(0, 1, 2, 3), 1L); + + assertThat(sliceMap.lookup(new byte[] {0x10})).isEqualTo(0); + assertThat(sliceMap.lookup(new byte[] {0x7F})).isEqualTo(1); + assertThat(sliceMap.lookup(new byte[] {(byte) 0x80})).isEqualTo(2); + assertThat(sliceMap.lookup(new byte[] {(byte) 0x90})).isEqualTo(2); + assertThat(sliceMap.lookup(new byte[] {(byte) 0xFF})).isEqualTo(3); + assertThat(sliceMap.lookup(new byte[] {(byte) 0xFF, 0x01})).isEqualTo(3); + } + + @Test + public void gettersAndImmutability() { + List slices = new ArrayList<>(); + slices.add(new SliceEntry(new byte[] {1}, Arrays.asList(0, 1))); + List fallback = new ArrayList<>(Arrays.asList(0, 1)); + + SliceMap sliceMap = new SliceMap(slices, fallback, 42L); + + assertThat(sliceMap.getGeneration()).isEqualTo(42L); + assertThat(sliceMap.getFallbackPool()).containsExactly(0, 1).inOrder(); + assertThat(sliceMap.getSlices()).hasSize(1); + assertThat(sliceMap.getSlices().get(0).endpoints).containsExactly(0, 1).inOrder(); + + // Verify immutability + assertThrows(UnsupportedOperationException.class, () -> sliceMap.getSlices().clear()); + assertThrows(UnsupportedOperationException.class, () -> sliceMap.getFallbackPool().clear()); + assertThrows( + UnsupportedOperationException.class, + () -> sliceMap.getSlices().get(0).endpoints.clear()); + } +}