diff --git a/core/src/main/java/io/grpc/util/GracefulSwitchLoadBalancer.java b/core/src/main/java/io/grpc/util/GracefulSwitchLoadBalancer.java new file mode 100644 index 00000000000..5f0b9d9c70c --- /dev/null +++ b/core/src/main/java/io/grpc/util/GracefulSwitchLoadBalancer.java @@ -0,0 +1,162 @@ +/* + * Copyright 2019 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.util; + +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; + +import com.google.common.annotations.VisibleForTesting; +import io.grpc.ConnectivityState; +import io.grpc.ConnectivityStateInfo; +import io.grpc.ExperimentalApi; +import io.grpc.LoadBalancer; +import io.grpc.LoadBalancerProvider; +import io.grpc.Status; +import javax.annotation.Nullable; +import javax.annotation.concurrent.NotThreadSafe; + +/** + * A load balancer that gracefully swaps to a new lb policy. If the channel is currently in a state + * other than READY, the new policy will be swapped into place immediately. Otherwise, the channel + * will keep using the old policy until the new policy reports READY or the old policy exits READY. + */ +@ExperimentalApi("https://github.com/grpc/grpc-java/issues/5999") +@NotThreadSafe // Must be accessed in SynchronizationContext +public final class GracefulSwitchLoadBalancer extends ForwardingLoadBalancer { + private static final LoadBalancer NOOP_BALANCER = new LoadBalancer() { + @Override + public void handleNameResolutionError(Status error) {} + + @Override + public void shutdown() {} + }; + + @VisibleForTesting + static final SubchannelPicker BUFFER_PICKER = new SubchannelPicker() { + @Override + public PickResult pickSubchannel(PickSubchannelArgs args) { + return PickResult.withNoResult(); + } + + @Override + public String toString() { + return "BUFFER_PICKER"; + } + }; + + private final Helper helper; + + // While the new policy is not fully switched on, the pendingLb is handling new updates from name + // resolver, and the currentLb is updating channel state and picker for the given helper. + // The current fields are guaranteed to be set after the initial swapTo(). + // The pending fields are cleared when it becomes current. + @Nullable private String currentPolicyName; + private LoadBalancer currentLb = NOOP_BALANCER; + @Nullable private String pendingPolicyName; + private LoadBalancer pendingLb = NOOP_BALANCER; + private ConnectivityState pendingState; + private SubchannelPicker pendingPicker; + + private boolean currentLbIsReady; + + public GracefulSwitchLoadBalancer(Helper helper) { + this.helper = checkNotNull(helper, "helper"); + } + + /** Gracefully switch to a new load balancing policy. */ + public void switchTo(LoadBalancerProvider newLbProvider) { + checkNotNull(newLbProvider, "newLbProvider"); + + String newPolicyName = newLbProvider.getPolicyName(); + if (newPolicyName.equals(pendingPolicyName)) { + return; + } + pendingLb.shutdown(); + pendingLb = NOOP_BALANCER; + pendingPolicyName = null; + pendingState = ConnectivityState.CONNECTING; + pendingPicker = BUFFER_PICKER; + + if (newPolicyName.equals(currentPolicyName)) { + return; + } + + class PendingHelper extends ForwardingLoadBalancerHelper { + LoadBalancer lb; + + @Override + protected Helper delegate() { + return helper; + } + + @Override + public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) { + if (lb == pendingLb) { + checkState(currentLbIsReady, "there's pending lb while current lb has been out of READY"); + pendingState = newState; + pendingPicker = newPicker; + if (newState == ConnectivityState.READY) { + swap(); + } + } else if (lb == currentLb) { + currentLbIsReady = newState == ConnectivityState.READY; + if (!currentLbIsReady && pendingLb != NOOP_BALANCER) { + swap(); // current policy exits READY, so swap + } else { + helper.updateBalancingState(newState, newPicker); + } + } + } + } + + PendingHelper pendingHelper = new PendingHelper(); + pendingHelper.lb = newLbProvider.newLoadBalancer(pendingHelper); + pendingLb = pendingHelper.lb; + pendingPolicyName = newPolicyName; + if (!currentLbIsReady) { + swap(); // the old policy is not READY at the moment, so swap to the new one right now + } + } + + private void swap() { + helper.updateBalancingState(pendingState, pendingPicker); + currentLb.shutdown(); + currentLb = pendingLb; + currentPolicyName = pendingPolicyName; + pendingLb = NOOP_BALANCER; + pendingPolicyName = null; + } + + @Override + protected LoadBalancer delegate() { + return pendingLb == NOOP_BALANCER ? currentLb : pendingLb; + } + + @Override + @Deprecated + public void handleSubchannelState( + Subchannel subchannel, ConnectivityStateInfo stateInfo) { + throw new UnsupportedOperationException( + "handleSubchannelState() is not supported by " + this.getClass().getName()); + } + + @Override + public void shutdown() { + pendingLb.shutdown(); + currentLb.shutdown(); + } +} diff --git a/core/src/test/java/io/grpc/util/GracefulSwitchLoadBalancerTest.java b/core/src/test/java/io/grpc/util/GracefulSwitchLoadBalancerTest.java new file mode 100644 index 00000000000..b861d726885 --- /dev/null +++ b/core/src/test/java/io/grpc/util/GracefulSwitchLoadBalancerTest.java @@ -0,0 +1,480 @@ +/* + * Copyright 2019 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.util; + +import static com.google.common.truth.Truth.assertThat; +import static io.grpc.ConnectivityState.CONNECTING; +import static io.grpc.ConnectivityState.READY; +import static io.grpc.util.GracefulSwitchLoadBalancer.BUFFER_PICKER; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; + +import io.grpc.ConnectivityStateInfo; +import io.grpc.EquivalentAddressGroup; +import io.grpc.LoadBalancer; +import io.grpc.LoadBalancer.CreateSubchannelArgs; +import io.grpc.LoadBalancer.Helper; +import io.grpc.LoadBalancer.ResolvedAddresses; +import io.grpc.LoadBalancer.Subchannel; +import io.grpc.LoadBalancer.SubchannelPicker; +import io.grpc.LoadBalancerProvider; +import io.grpc.LoadBalancerRegistry; +import io.grpc.Status; +import java.net.SocketAddress; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.InOrder; + +/** + * Tests for {@link GracefulSwitchLoadBalancer}. + */ +@RunWith(JUnit4.class) +public class GracefulSwitchLoadBalancerTest { + @Rule + public final ExpectedException thrown = ExpectedException.none(); + + private final LoadBalancerRegistry lbRegistry = new LoadBalancerRegistry(); + // maps policy name to lb provide + private final Map lbProviders = new HashMap<>(); + // maps policy name to lb + private final Map balancers = new HashMap<>(); + private final Map helpers = new HashMap<>(); + private final Helper mockHelper = mock(Helper.class); + private final GracefulSwitchLoadBalancer gracefulSwitchLb = + new GracefulSwitchLoadBalancer(mockHelper); + private final String[] lbPolicies = {"lb_policy_0", "lb_policy_1", "lb_policy_2", "lb_policy_3"}; + + @Before + public void setUp() { + for (int i = 0; i < lbPolicies.length; i++) { + String lbPolicy = lbPolicies[i]; + LoadBalancerProvider lbProvider = new FakeLoadBalancerProvider(lbPolicy); + lbProviders.put(lbPolicy, lbProvider); + lbRegistry.register(lbProvider); + } + } + + @Test + public void canHandleEmptyAddressListFromNameResolutionForwardedToLatestPolicy() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + + assertThat(gracefulSwitchLb.canHandleEmptyAddressListFromNameResolution()).isFalse(); + doReturn(true).when(lb0).canHandleEmptyAddressListFromNameResolution(); + assertThat(gracefulSwitchLb.canHandleEmptyAddressListFromNameResolution()).isTrue(); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + + assertThat(gracefulSwitchLb.canHandleEmptyAddressListFromNameResolution()).isFalse(); + + doReturn(true).when(lb1).canHandleEmptyAddressListFromNameResolution(); + assertThat(gracefulSwitchLb.canHandleEmptyAddressListFromNameResolution()).isTrue(); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[2])); + LoadBalancer lb2 = balancers.get(lbPolicies[2]); + + assertThat(gracefulSwitchLb.canHandleEmptyAddressListFromNameResolution()).isFalse(); + + doReturn(true).when(lb2).canHandleEmptyAddressListFromNameResolution(); + assertThat(gracefulSwitchLb.canHandleEmptyAddressListFromNameResolution()).isTrue(); + } + + @Test + public void handleResolvedAddressesAndNameResolutionErrorForwardedToLatestPolicy() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + + ResolvedAddresses addresses = newFakeAddresses(); + gracefulSwitchLb.handleResolvedAddresses(addresses); + verify(lb0).handleResolvedAddresses(addresses); + gracefulSwitchLb.handleNameResolutionError(Status.DATA_LOSS); + verify(lb0).handleNameResolutionError(Status.DATA_LOSS); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + addresses = newFakeAddresses(); + gracefulSwitchLb.handleResolvedAddresses(addresses); + verify(lb0, never()).handleResolvedAddresses(addresses); + verify(lb1).handleResolvedAddresses(addresses); + gracefulSwitchLb.handleNameResolutionError(Status.ALREADY_EXISTS); + verify(lb0, never()).handleNameResolutionError(Status.ALREADY_EXISTS); + verify(lb1).handleNameResolutionError(Status.ALREADY_EXISTS); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[2])); + verify(lb1).shutdown(); + LoadBalancer lb2 = balancers.get(lbPolicies[2]); + addresses = newFakeAddresses(); + gracefulSwitchLb.handleResolvedAddresses(addresses); + verify(lb0, never()).handleResolvedAddresses(addresses); + verify(lb1, never()).handleResolvedAddresses(addresses); + verify(lb2).handleResolvedAddresses(addresses); + gracefulSwitchLb.handleNameResolutionError(Status.CANCELLED); + verify(lb0, never()).handleNameResolutionError(Status.CANCELLED); + verify(lb1, never()).handleNameResolutionError(Status.CANCELLED); + verify(lb2).handleNameResolutionError(Status.CANCELLED); + + verifyNoMoreInteractions(lb0, lb1, lb2); + } + + @Test + public void shutdownTriggeredWhenSwitchAndForwardedWhenSwitchLbShutdown() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + verify(lb1, never()).shutdown(); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[2])); + verify(lb1).shutdown(); + LoadBalancer lb2 = balancers.get(lbPolicies[2]); + verify(lb0, never()).shutdown(); + helpers.get(lb2).updateBalancingState(READY, mock(SubchannelPicker.class)); + verify(lb0).shutdown(); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[3])); + LoadBalancer lb3 = balancers.get(lbPolicies[3]); + verify(lb2, never()).shutdown(); + verify(lb3, never()).shutdown(); + + gracefulSwitchLb.shutdown(); + verify(lb2).shutdown(); + verify(lb3).shutdown(); + + verifyNoMoreInteractions(lb0, lb1, lb2, lb3); + } + + @Test + public void requestConnectionForwardedToLatestPolicies() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + + gracefulSwitchLb.requestConnection(); + verify(lb0).requestConnection(); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + gracefulSwitchLb.requestConnection(); + verify(lb1).requestConnection(); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[2])); + verify(lb1).shutdown(); + LoadBalancer lb2 = balancers.get(lbPolicies[2]); + gracefulSwitchLb.requestConnection(); + verify(lb2).requestConnection(); + + // lb2 reports READY + helpers.get(lb2).updateBalancingState(READY, mock(SubchannelPicker.class)); + verify(lb0).shutdown(); + + gracefulSwitchLb.requestConnection(); + verify(lb2, times(2)).requestConnection(); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[3])); + LoadBalancer lb3 = balancers.get(lbPolicies[3]); + gracefulSwitchLb.requestConnection(); + verify(lb3).requestConnection(); + + verifyNoMoreInteractions(lb0, lb1, lb2, lb3); + } + + @Test + public void createSubchannelForwarded() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + + CreateSubchannelArgs createSubchannelArgs = newFakeCreateSubchannelArgs(); + helper0.createSubchannel(createSubchannelArgs); + verify(mockHelper).createSubchannel(createSubchannelArgs); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + Helper helper1 = helpers.get(lb1); + createSubchannelArgs = newFakeCreateSubchannelArgs(); + helper1.createSubchannel(createSubchannelArgs); + verify(mockHelper).createSubchannel(createSubchannelArgs); + + createSubchannelArgs = newFakeCreateSubchannelArgs(); + helper0.createSubchannel(createSubchannelArgs); + verify(mockHelper).createSubchannel(createSubchannelArgs); + + verifyNoMoreInteractions(lb0, lb1); + } + + @Test + public void updateBalancingStateIsGraceful() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + verify(mockHelper).updateBalancingState(READY, picker); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + Helper helper1 = helpers.get(lb1); + picker = mock(SubchannelPicker.class); + helper1.updateBalancingState(CONNECTING, picker); + verify(mockHelper, never()).updateBalancingState(CONNECTING, picker); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[2])); + verify(lb1).shutdown(); + LoadBalancer lb2 = balancers.get(lbPolicies[2]); + Helper helper2 = helpers.get(lb2); + picker = mock(SubchannelPicker.class); + helper2.updateBalancingState(CONNECTING, picker); + verify(mockHelper, never()).updateBalancingState(CONNECTING, picker); + + // lb2 reports READY + SubchannelPicker picker2 = mock(SubchannelPicker.class); + helper2.updateBalancingState(READY, picker2); + verify(lb0).shutdown(); + verify(mockHelper).updateBalancingState(READY, picker2); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[3])); + LoadBalancer lb3 = balancers.get(lbPolicies[3]); + Helper helper3 = helpers.get(lb3); + SubchannelPicker picker3 = mock(SubchannelPicker.class); + helper3.updateBalancingState(CONNECTING, picker3); + verify(mockHelper, never()).updateBalancingState(CONNECTING, picker3); + + // lb2 out of READY + picker2 = mock(SubchannelPicker.class); + helper2.updateBalancingState(CONNECTING, picker2); + verify(mockHelper, never()).updateBalancingState(CONNECTING, picker2); + verify(mockHelper).updateBalancingState(CONNECTING, picker3); + verify(lb2).shutdown(); + + picker3 = mock(SubchannelPicker.class); + helper3.updateBalancingState(CONNECTING, picker3); + verify(mockHelper).updateBalancingState(CONNECTING, picker3); + + verifyNoMoreInteractions(lb0, lb1, lb2, lb3); + } + + @Test + public void switchWhileOldPolicyIsNotReady() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(CONNECTING, picker); + + verify(lb0, never()).shutdown(); + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + verify(lb0).shutdown(); + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + + Helper helper1 = helpers.get(lb1); + picker = mock(SubchannelPicker.class); + helper1.updateBalancingState(CONNECTING, picker); + verify(mockHelper).updateBalancingState(CONNECTING, picker); + + verify(lb1, never()).shutdown(); + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[2])); + verify(lb1).shutdown(); + LoadBalancer lb2 = balancers.get(lbPolicies[2]); + + verifyNoMoreInteractions(lb0, lb1, lb2); + } + + @Test + public void switchWhileOldPolicyGoesFromReadyToNotReady() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + verify(lb0, never()).shutdown(); + + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + Helper helper1 = helpers.get(lb1); + SubchannelPicker picker1 = mock(SubchannelPicker.class); + helper1.updateBalancingState(CONNECTING, picker1); + verify(mockHelper, never()).updateBalancingState(CONNECTING, picker1); + + picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(CONNECTING, picker); + verify(lb0).shutdown(); + verify(mockHelper, never()).updateBalancingState(CONNECTING, picker); + verify(mockHelper).updateBalancingState(CONNECTING, picker1); + + picker1 = mock(SubchannelPicker.class); + helper1.updateBalancingState(READY, picker1); + verify(mockHelper).updateBalancingState(READY, picker1); + + verifyNoMoreInteractions(lb0, lb1); + } + + @Test + public void switchWhileOldPolicyGoesFromReadyToNotReadyWhileNewPolicyStillIdle() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + InOrder inOrder = inOrder(lb0, mockHelper); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + verify(lb0, never()).shutdown(); + + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + Helper helper1 = helpers.get(lb1); + + picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(CONNECTING, picker); + + verify(mockHelper, never()).updateBalancingState(CONNECTING, picker); + inOrder.verify(mockHelper).updateBalancingState(CONNECTING, BUFFER_PICKER); + inOrder.verify(lb0).shutdown(); // shutdown after update + + picker = mock(SubchannelPicker.class); + helper1.updateBalancingState(CONNECTING, picker); + inOrder.verify(mockHelper).updateBalancingState(CONNECTING, picker); + + inOrder.verifyNoMoreInteractions(); + verifyNoMoreInteractions(lb1); + } + + @Test + public void newPolicyNameTheSameAsPendingPolicy_shouldHaveNoEffect() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + assertThat(balancers.get(lbPolicies[1])).isSameInstanceAs(lb1); + + verifyNoMoreInteractions(lb0, lb1); + } + + @Test + public void newPolicyNameTheSameAsCurrentPolicy_shouldShutdownPendingLb() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + LoadBalancer lb0 = balancers.get(lbPolicies[0]); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + assertThat(balancers.get(lbPolicies[0])).isSameInstanceAs(lb0); + + Helper helper0 = helpers.get(lb0); + SubchannelPicker picker = mock(SubchannelPicker.class); + helper0.updateBalancingState(READY, picker); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[1])); + LoadBalancer lb1 = balancers.get(lbPolicies[1]); + + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + verify(lb1).shutdown(); + assertThat(balancers.get(lbPolicies[0])).isSameInstanceAs(lb0); + + verifyNoMoreInteractions(lb0, lb1); + } + + @Deprecated + @Test + public void handleSubchannelState_shouldThrow() { + gracefulSwitchLb.switchTo(lbProviders.get(lbPolicies[0])); + Subchannel subchannel = mock(Subchannel.class); + ConnectivityStateInfo connectivityStateInfo = ConnectivityStateInfo.forNonError(READY); + thrown.expect(UnsupportedOperationException.class); + gracefulSwitchLb.handleSubchannelState(subchannel, connectivityStateInfo); + } + + private final class FakeLoadBalancerProvider extends LoadBalancerProvider { + + final String policyName; + + FakeLoadBalancerProvider(String policyName) { + this.policyName = policyName; + } + + @Override + public boolean isAvailable() { + return true; + } + + @Override + public int getPriority() { + return 5; + } + + @Override + public String getPolicyName() { + return policyName; + } + + @Override + public LoadBalancer newLoadBalancer(Helper helper) { + LoadBalancer balancer = mock(LoadBalancer.class); + balancers.put(policyName, balancer); + helpers.put(balancer, helper); + return balancer; + } + } + + private static ResolvedAddresses newFakeAddresses() { + return ResolvedAddresses + .newBuilder() + .setAddresses( + Collections.singletonList(new EquivalentAddressGroup(mock(SocketAddress.class)))) + .build(); + } + + private static CreateSubchannelArgs newFakeCreateSubchannelArgs() { + return CreateSubchannelArgs + .newBuilder() + .setAddresses(new EquivalentAddressGroup(mock(SocketAddress.class))) + .build(); + } +}