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..2edf86e9599 --- /dev/null +++ b/autosharding/src/main/java/io/grpc/autosharding/EndpointMap.java @@ -0,0 +1,257 @@ +/* + * 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 static io.grpc.ConnectivityState.IDLE; + +import com.google.common.base.MoreObjects; +import com.google.common.collect.ImmutableList; +import io.grpc.Attributes; +import io.grpc.ConnectivityState; +import io.grpc.EquivalentAddressGroup; +import io.grpc.LoadBalancer; +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.util.ForwardingLoadBalancerHelper; +import io.grpc.util.LazyLoadBalancer; +import java.util.Collection; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; +import javax.annotation.Nullable; +import javax.annotation.concurrent.NotThreadSafe; + +/** + * Manages the mapping from endpoint hostname to {@link EndpointHolder} and coordinates + * child load balancer lifecycle and connectivity state updates. + * + *

Threading model: This class is not thread-safe. All methods must be invoked from the + * {@link io.grpc.SynchronizationContext} by the parent load balancer. + */ +@NotThreadSafe +final class EndpointMap { + private final Map map = new LinkedHashMap<>(); + + @Nullable + EndpointHolder get(String hostname) { + return map.get(checkNotNull(hostname, "hostname")); + } + + void put(String hostname, EndpointHolder holder) { + map.put(checkNotNull(hostname, "hostname"), checkNotNull(holder, "holder")); + } + + @Nullable + EndpointHolder remove(String hostname) { + return map.remove(checkNotNull(hostname, "hostname")); + } + + Collection values() { + return map.values(); + } + + Set keySet() { + return map.keySet(); + } + + int size() { + return map.size(); + } + + boolean isEmpty() { + return map.isEmpty(); + } + + void clear() { + map.clear(); + } + + /** + * Re-assigns contiguous 0-based index values across all current endpoint holders. + */ + void reindex() { + int nextIdx = 0; + for (EndpointHolder holder : map.values()) { + holder.setIndex(nextIdx++); + } + } + + /** + * Shuts down all child load balancers and clears the map. + */ + void shutdownAll() { + for (EndpointHolder holder : map.values()) { + holder.shutdown(); + } + map.clear(); + } + + /** + * Builds an immutable snapshot list of {@link PickerEndpoint}s placed strictly at their + * corresponding {@link EndpointHolder#getIndex()} positions. + * + * @throws IllegalStateException if endpoint indices are not contiguous from 0 to N-1 + */ + ImmutableList toPickerEndpoints() { + int size = map.size(); + if (size == 0) { + return ImmutableList.of(); + } + PickerEndpoint[] array = new PickerEndpoint[size]; + for (EndpointHolder holder : map.values()) { + int idx = holder.getIndex(); + checkState( + idx >= 0 && idx < size, + "Endpoint holder index %s is out of bounds for size %s", + idx, + size); + checkState( + array[idx] == null, + "Duplicate endpoint holder index %s detected", + idx); + array[idx] = holder.toPickerEndpoint(); + } + return ImmutableList.copyOf(array); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("map", map) + .toString(); + } + + /** + * Holds the connectivity state, picker, and lazy child load balancer for a single endpoint. + */ + static final class EndpointHolder { + private int index; + private final LazyLoadBalancer childLb; + private final AtomicBoolean connectingScheduled = new AtomicBoolean(false); + private final Helper helper; + private ConnectivityState state = IDLE; + private SubchannelPicker picker = new FixedResultPicker(PickResult.withNoResult()); + + EndpointHolder( + int index, + Helper helper, + LoadBalancer.Factory pickFirstFactory, + @Nullable Runnable stateUpdateCallback) { + this.index = index; + this.helper = checkNotNull(helper, "helper"); + this.childLb = new LazyLoadBalancer( + new ChildHelper(helper, stateUpdateCallback), + checkNotNull(pickFirstFactory, "pickFirstFactory")); + } + + int getIndex() { + return index; + } + + void setIndex(int index) { + this.index = index; + } + + ConnectivityState getState() { + return state; + } + + SubchannelPicker getPicker() { + return picker; + } + + LazyLoadBalancer getChildLb() { + return childLb; + } + + PickerEndpoint toPickerEndpoint() { + return new PickerEndpoint(state, picker, this::exitIdle); + } + + private void exitIdle() { + if (connectingScheduled.compareAndSet(false, true)) { + helper.getSynchronizationContext().execute(() -> { + connectingScheduled.set(false); + childLb.requestConnection(); + }); + } + } + + void updateAddresses(List eags, Attributes attributes) { + ResolvedAddresses childAddresses = ResolvedAddresses.newBuilder() + .setAddresses(ImmutableList.copyOf(checkNotNull(eags, "eags"))) + .setAttributes(checkNotNull(attributes, "attributes")) + .build(); + childLb.acceptResolvedAddresses(childAddresses); + } + + void requestConnection() { + childLb.requestConnection(); + } + + void shutdown() { + childLb.shutdown(); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("index", index) + .add("state", state) + .add("childLb", childLb) + .toString(); + } + + private final class ChildHelper extends ForwardingLoadBalancerHelper { + private final Helper delegateHelper; + @Nullable private final Runnable stateUpdateCallback; + + ChildHelper(Helper delegateHelper, @Nullable Runnable stateUpdateCallback) { + this.delegateHelper = checkNotNull(delegateHelper, "delegateHelper"); + this.stateUpdateCallback = stateUpdateCallback; + } + + @Override + protected Helper delegate() { + return delegateHelper; + } + + @Override + public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) { + state = checkNotNull(newState, "newState"); + picker = checkNotNull(newPicker, "newPicker"); + if (stateUpdateCallback != null) { + stateUpdateCallback.run(); + } + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("delegateHelper", delegateHelper) + .toString(); + } + } + } +} diff --git a/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java new file mode 100644 index 00000000000..5dd6f4c5332 --- /dev/null +++ b/autosharding/src/test/java/io/grpc/autosharding/EndpointMapTest.java @@ -0,0 +1,280 @@ +/* + * Copyright 2026 The gRPC Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.grpc.autosharding; + +import static com.google.common.truth.Truth.assertThat; +import static io.grpc.ConnectivityState.IDLE; +import static io.grpc.ConnectivityState.READY; +import static io.grpc.ConnectivityState.TRANSIENT_FAILURE; +import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.common.collect.ImmutableList; +import io.grpc.Attributes; +import io.grpc.EquivalentAddressGroup; +import io.grpc.LoadBalancer; +import io.grpc.LoadBalancer.Helper; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.LoadBalancerProvider; +import io.grpc.SynchronizationContext; +import io.grpc.autosharding.EndpointMap.EndpointHolder; +import java.net.SocketAddress; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; + +@RunWith(JUnit4.class) +public class EndpointMapTest { + + private final Helper mockHelper = mock(Helper.class); + private final LoadBalancerProvider mockProvider = mock(LoadBalancerProvider.class); + private final LoadBalancer mockDelegate = mock(LoadBalancer.class); + private final SynchronizationContext syncContext = + new SynchronizationContext((t, e) -> { + throw new AssertionError("Unhandled exception in syncContext", e); + }); + + private EndpointMap endpointMap; + private final AtomicInteger stateChangeCount = new AtomicInteger(0); + + @Before + public void setUp() { + when(mockHelper.getSynchronizationContext()).thenReturn(syncContext); + when(mockProvider.newLoadBalancer(any())).thenReturn(mockDelegate); + endpointMap = new EndpointMap(); + } + + private EndpointHolder createHolder(int index) { + return new EndpointHolder(index, mockHelper, mockProvider, stateChangeCount::incrementAndGet); + } + + @Test + public void basicMapOperations() { + assertThat(endpointMap.isEmpty()).isTrue(); + assertThat(endpointMap.size()).isEqualTo(0); + + EndpointHolder h1 = createHolder(0); + EndpointHolder h2 = createHolder(1); + + endpointMap.put("host1", h1); + endpointMap.put("host2", h2); + + assertThat(endpointMap.isEmpty()).isFalse(); + assertThat(endpointMap.size()).isEqualTo(2); + assertThat(endpointMap.get("host1")).isSameInstanceAs(h1); + assertThat(endpointMap.get("host2")).isSameInstanceAs(h2); + assertThat(endpointMap.get("unknown")).isNull(); + assertThat(endpointMap.keySet()).containsExactly("host1", "host2").inOrder(); + assertThat(endpointMap.values()).containsExactly(h1, h2).inOrder(); + + EndpointHolder removed = endpointMap.remove("host1"); + assertThat(removed).isSameInstanceAs(h1); + assertThat(endpointMap.size()).isEqualTo(1); + assertThat(endpointMap.get("host1")).isNull(); + } + + @Test + public void nullChecks() { + EndpointHolder h = createHolder(0); + + assertThrows(NullPointerException.class, () -> endpointMap.get(null)); + assertThrows(NullPointerException.class, () -> endpointMap.put(null, h)); + assertThrows(NullPointerException.class, () -> endpointMap.put("host", null)); + assertThrows(NullPointerException.class, () -> endpointMap.remove(null)); + + assertThrows( + NullPointerException.class, + () -> new EndpointHolder(0, null, mockProvider, null)); + assertThrows( + NullPointerException.class, + () -> new EndpointHolder(0, mockHelper, null, null)); + + assertThrows( + NullPointerException.class, + () -> h.updateAddresses(null, Attributes.EMPTY)); + assertThrows( + NullPointerException.class, + () -> h.updateAddresses(Collections.emptyList(), null)); + } + + @Test + public void reindex_updatesIndicesContiguously() { + EndpointHolder h0 = createHolder(0); + EndpointHolder h1 = createHolder(1); + EndpointHolder h2 = createHolder(2); + + endpointMap.put("host0", h0); + endpointMap.put("host1", h1); + endpointMap.put("host2", h2); + + // Remove middle element + endpointMap.remove("host1"); + assertThat(h0.getIndex()).isEqualTo(0); + assertThat(h2.getIndex()).isEqualTo(2); + + endpointMap.reindex(); + assertThat(h0.getIndex()).isEqualTo(0); + assertThat(h2.getIndex()).isEqualTo(1); + } + + @Test + public void endpointHolder_childHelperUpdatesStateAndTriggersCallback() { + EndpointHolder holder = createHolder(0); + assertThat(holder.getState()).isEqualTo(IDLE); + + // Capture child helper passed to LazyChildLoadBalancer + ArgumentCaptor helperCaptor = ArgumentCaptor.forClass(Helper.class); + verify(mockProvider, org.mockito.Mockito.never()).newLoadBalancer(any()); + + // Trigger connection to create child helper and delegate + holder.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + holder.requestConnection(); + + verify(mockProvider).newLoadBalancer(helperCaptor.capture()); + Helper childHelper = helperCaptor.getValue(); + + // Reset counter before state update to verify callback fires on update + stateChangeCount.set(0); + + // Simulate child balancer updating state + SubchannelPicker testPicker = mock(SubchannelPicker.class); + childHelper.updateBalancingState(READY, testPicker); + + assertThat(holder.getState()).isEqualTo(READY); + assertThat(holder.getPicker()).isSameInstanceAs(testPicker); + assertThat(stateChangeCount.get()).isEqualTo(1); + } + + @Test + public void toPickerEndpoints_buildsImmutableListMatchingHoldersByIndex() { + EndpointHolder h0 = createHolder(0); + EndpointHolder h1 = createHolder(1); + + ArgumentCaptor helperCaptor = ArgumentCaptor.forClass(Helper.class); + + // Trigger connections so child helpers are passed to provider + h0.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + h0.requestConnection(); + + h1.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + h1.requestConnection(); + + verify(mockProvider, times(2)).newLoadBalancer(helperCaptor.capture()); + Helper childHelper0 = helperCaptor.getAllValues().get(0); + Helper childHelper1 = helperCaptor.getAllValues().get(1); + + SubchannelPicker picker0 = mock(SubchannelPicker.class); + SubchannelPicker picker1 = mock(SubchannelPicker.class); + + childHelper0.updateBalancingState(READY, picker0); + childHelper1.updateBalancingState(TRANSIENT_FAILURE, picker1); + + // Insert in reverse index order to verify explicit index placement + endpointMap.put("host1", h1); + endpointMap.put("host0", h0); + + ImmutableList pickerEndpoints = endpointMap.toPickerEndpoints(); + assertThat(pickerEndpoints).hasSize(2); + assertThat(pickerEndpoints.get(0).getState()).isEqualTo(READY); + assertThat(pickerEndpoints.get(0).getPicker()).isSameInstanceAs(picker0); + assertThat(pickerEndpoints.get(1).getState()).isEqualTo(TRANSIENT_FAILURE); + assertThat(pickerEndpoints.get(1).getPicker()).isSameInstanceAs(picker1); + } + + @Test + public void toPickerEndpoints_emptyMap_returnsEmptyList() { + assertThat(endpointMap.toPickerEndpoints()).isEmpty(); + } + + @Test + public void toPickerEndpoints_duplicateOrOutOfBoundsIndex_throwsIllegalStateException() { + EndpointHolder h0 = createHolder(0); + EndpointHolder h0Duplicate = createHolder(0); + + endpointMap.put("host0", h0); + endpointMap.put("host1", h0Duplicate); + + assertThrows(IllegalStateException.class, () -> endpointMap.toPickerEndpoints()); + + endpointMap.clear(); + EndpointHolder hOutOfBounds = createHolder(5); + endpointMap.put("host0", hOutOfBounds); + + assertThrows(IllegalStateException.class, () -> endpointMap.toPickerEndpoints()); + } + + @Test + public void shutdownAll_cleansUpAllHoldersAndClearsMap() { + EndpointHolder h0 = createHolder(0); + EndpointHolder h1 = createHolder(1); + + endpointMap.put("host0", h0); + endpointMap.put("host1", h1); + + // Trigger connections so delegates exist + h0.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + h0.requestConnection(); + + endpointMap.shutdownAll(); + assertThat(endpointMap.isEmpty()).isTrue(); + verify(mockDelegate).shutdown(); + } + + @Test + public void toPickerEndpoint_requestConnection_wakesUpChildBalancerOnSyncContext() { + EndpointHolder holder = createHolder(0); + holder.updateAddresses( + Collections.singletonList(new EquivalentAddressGroup(new SocketAddress() {})), + Attributes.EMPTY); + + PickerEndpoint pickerEndpoint = holder.toPickerEndpoint(); + verify(mockProvider, org.mockito.Mockito.never()).newLoadBalancer(any()); + + // Trigger connection through PickerEndpoint (simulate AutoShardingPicker encountering IDLE) + pickerEndpoint.requestConnection(); + + verify(mockProvider).newLoadBalancer(any()); + verify(mockDelegate).acceptResolvedAddresses(any()); + verify(mockDelegate).requestConnection(); + } + + @Test + public void toString_containsDebugFields() { + EndpointHolder h = createHolder(3); + endpointMap.put("host3", h); + + assertThat(endpointMap.toString()).contains("host3"); + assertThat(h.toString()).contains("index=3"); + assertThat(h.toString()).contains("state=IDLE"); + } +}