Skip to content

Commit 2c9b067

Browse files
kannanjgithubkannanjgithub
authored
core: Handle new PF enabled (#13094)
Make tests detect and handle when new PickFirst LB is enabled.
1 parent 38db48e commit 2c9b067

5 files changed

Lines changed: 158 additions & 63 deletions

File tree

‎grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java‎

Lines changed: 125 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@
8282
import io.grpc.inprocess.InProcessServerBuilder;
8383
import io.grpc.internal.BackoffPolicy;
8484
import io.grpc.internal.FakeClock;
85+
import io.grpc.internal.PickFirstLoadBalancerProvider;
8586
import io.grpc.lb.v1.ClientStats;
8687
import io.grpc.lb.v1.ClientStatsPerToken;
8788
import io.grpc.lb.v1.FallbackResponse;
@@ -1921,10 +1922,16 @@ public void grpclbWorking_pickFirstMode() throws Exception {
19211922
// With delegation, the child pick_first creates the subchannel
19221923
inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
19231924
CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
1924-
assertThat(createSubchannelArgs.getAddresses())
1925-
.containsExactly(
1926-
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
1927-
new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
1925+
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
1926+
assertThat(createSubchannelArgs.getAddresses())
1927+
.containsExactly(
1928+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
1929+
} else {
1930+
assertThat(createSubchannelArgs.getAddresses())
1931+
.containsExactly(
1932+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
1933+
new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
1934+
}
19281935

19291936
// Child pick_first eagerly connects, so we start in CONNECTING
19301937
inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
@@ -1939,18 +1946,20 @@ public void grpclbWorking_pickFirstMode() throws Exception {
19391946
// Child pick_first eagerly calls requestConnection()
19401947
verify(subchannel).requestConnection();
19411948

1942-
// TRANSIENT_FAILURE
1943-
Status error = Status.UNAVAILABLE.withDescription("Simulated connection error");
1944-
deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure(error));
1945-
// The child LB will notify our helper, which updates grpclb state
1946-
inOrder.verify(helper, atLeast(1))
1947-
.updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
1948-
RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue();
1949-
assertThat(picker1.dropList).containsExactly(null, null);
1950-
ChildLbPickerEntry failureEntry = (ChildLbPickerEntry) picker1.pickList.get(0);
1951-
PickResult failureResult =
1952-
failureEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class));
1953-
assertThat(failureResult.getStatus()).isEqualTo(error);
1949+
if (!PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
1950+
// TRANSIENT_FAILURE
1951+
Status error = Status.UNAVAILABLE.withDescription("Simulated connection error");
1952+
deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure(error));
1953+
// The child LB will notify our helper, which updates grpclb state
1954+
inOrder.verify(helper, atLeast(1))
1955+
.updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
1956+
RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue();
1957+
assertThat(picker1.dropList).containsExactly(null, null);
1958+
ChildLbPickerEntry failureEntry = (ChildLbPickerEntry) picker1.pickList.get(0);
1959+
PickResult failureResult =
1960+
failureEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class));
1961+
assertThat(failureResult.getStatus()).isEqualTo(error);
1962+
}
19541963

19551964
// READY
19561965
deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
@@ -1977,11 +1986,15 @@ public void grpclbWorking_pickFirstMode() throws Exception {
19771986
assertThat(mockSubchannels).isEmpty();
19781987

19791988
// The child LB policy internally calls updateAddresses on the subchannel
1980-
verify(subchannel).updateAddresses(
1981-
eq(Arrays.asList(
1982-
new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")),
1983-
new EquivalentAddressGroup(backends2.get(2).addr,
1984-
eagAttrsWithToken("token0004")))));
1989+
List<EquivalentAddressGroup> expectedAddresses =
1990+
PickFirstLoadBalancerProvider.isEnabledNewPickFirst()
1991+
? Collections.singletonList(
1992+
new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")))
1993+
: Arrays.asList(
1994+
new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")),
1995+
new EquivalentAddressGroup(backends2.get(2).addr,
1996+
eagAttrsWithToken("token0004")));
1997+
verify(subchannel).updateAddresses(eq(expectedAddresses));
19851998
inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture());
19861999
RoundRobinPicker picker3 = (RoundRobinPicker) pickerCaptor.getValue();
19872000
assertThat(picker3.dropList).containsExactly(
@@ -2004,7 +2017,8 @@ public void grpclbWorking_pickFirstMode() throws Exception {
20042017
assertThat(pick.getSubchannel()).isNull();
20052018
verify(subchannel, times(2)).requestConnection();
20062019
balancer.requestConnection();
2007-
verify(subchannel, times(3)).requestConnection();
2020+
int expectedRequests = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 2 : 3;
2021+
verify(subchannel, times(expectedRequests)).requestConnection();
20082022

20092023
// PICK_FIRST doesn't use subchannelPool
20102024
verify(subchannelPool, never())
@@ -2045,10 +2059,16 @@ public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception {
20452059
// The child pick_first creates the first subchannel
20462060
inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
20472061
CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
2048-
assertThat(createSubchannelArgs.getAddresses())
2049-
.containsExactly(
2050-
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
2051-
new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
2062+
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
2063+
assertThat(createSubchannelArgs.getAddresses())
2064+
.containsExactly(
2065+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
2066+
} else {
2067+
assertThat(createSubchannelArgs.getAddresses())
2068+
.containsExactly(
2069+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
2070+
new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
2071+
}
20522072

20532073
// Child pick_first eagerly connects, so initial state is CONNECTING
20542074
inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
@@ -2104,7 +2124,18 @@ public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception {
21042124
lbResponseObserver.onNext(buildLbResponse(backends2));
21052125

21062126
// A NEW child LB and NEW subchannel are created upon recovery
2107-
inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class));
2127+
inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
2128+
CreateSubchannelArgs createSubchannelArgs2 = createSubchannelArgsCaptor.getValue();
2129+
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
2130+
assertThat(createSubchannelArgs2.getAddresses())
2131+
.containsExactly(
2132+
new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")));
2133+
} else {
2134+
assertThat(createSubchannelArgs2.getAddresses())
2135+
.containsExactly(
2136+
new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")),
2137+
new EquivalentAddressGroup(backends2.get(2).addr, eagAttrsWithToken("token0004")));
2138+
}
21082139
assertThat(mockSubchannels).hasSize(1);
21092140
Subchannel subchannel2 = mockSubchannels.poll();
21102141
inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
@@ -2180,8 +2211,13 @@ private void pickFirstModeFallback(long timeout) throws Exception {
21802211
// Entering fallback mode - child LB is created for fallback backends
21812212
inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
21822213
CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
2183-
assertThat(createSubchannelArgs.getAddresses())
2184-
.containsExactly(backendList.get(0), backendList.get(1));
2214+
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
2215+
assertThat(createSubchannelArgs.getAddresses())
2216+
.containsExactly(backendList.get(0));
2217+
} else {
2218+
assertThat(createSubchannelArgs.getAddresses())
2219+
.containsExactly(backendList.get(0), backendList.get(1));
2220+
}
21852221

21862222
assertThat(mockSubchannels).hasSize(1);
21872223
Subchannel subchannel = mockSubchannels.poll();
@@ -2212,16 +2248,44 @@ private void pickFirstModeFallback(long timeout) throws Exception {
22122248
lbResponseObserver.onNext(buildInitialResponse());
22132249
lbResponseObserver.onNext(buildLbResponse(backends1));
22142250

2215-
// subchannel should be updated, NOT recreated
2216-
inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class));
2217-
assertThat(mockSubchannels).isEmpty();
2218-
// The child LB internally calls updateAddresses on the existing subchannel
2219-
verify(subchannel).updateAddresses(
2220-
eq(Arrays.asList(
2221-
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
2222-
new EquivalentAddressGroup(backends1.get(1).addr,
2223-
eagAttrsWithToken("token0002")))));
2224-
inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture());
2251+
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
2252+
// In new PF, active fallback subchannel is shut down because its address is not in backends1
2253+
verify(subchannel).shutdown();
2254+
// And LB state transitions to IDLE with RequestConnectionPicker
2255+
inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(IDLE), pickerCaptor.capture());
2256+
RoundRobinPicker pickerIdle = (RoundRobinPicker) pickerCaptor.getValue();
2257+
ChildLbPickerEntry idleEntry = (ChildLbPickerEntry) pickerIdle.pickList.get(0);
2258+
2259+
// Picking while IDLE triggers requestConnection(), creating a subchannel for backends1.get(0)
2260+
PickSubchannelArgs args = mock(PickSubchannelArgs.class);
2261+
PickResult pick = idleEntry.getChildPicker().pickSubchannel(args);
2262+
assertThat(pick.getSubchannel()).isNull(); // BUFFERing while IDLE
2263+
2264+
inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
2265+
CreateSubchannelArgs createSubchannelArgs2 = createSubchannelArgsCaptor.getValue();
2266+
assertThat(createSubchannelArgs2.getAddresses())
2267+
.containsExactly(
2268+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
2269+
assertThat(mockSubchannels).hasSize(1);
2270+
subchannel = mockSubchannels.poll();
2271+
verify(subchannel).requestConnection();
2272+
2273+
// subchannel becomes READY
2274+
deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
2275+
inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture());
2276+
} else {
2277+
// subchannel should be updated, NOT recreated
2278+
inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class));
2279+
assertThat(mockSubchannels).isEmpty();
2280+
// The child LB internally calls updateAddresses on the existing subchannel
2281+
verify(subchannel).updateAddresses(
2282+
eq(Arrays.asList(
2283+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
2284+
new EquivalentAddressGroup(backends1.get(1).addr,
2285+
eagAttrsWithToken("token0002")))));
2286+
inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture());
2287+
}
2288+
22252289
RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue();
22262290
assertThat(picker2.dropList).containsExactly(null, null);
22272291

@@ -2247,7 +2311,8 @@ private void pickFirstModeFallback(long timeout) throws Exception {
22472311
verify(subchannel, times(2)).requestConnection();
22482312

22492313
balancer.requestConnection();
2250-
verify(subchannel, times(3)).requestConnection();
2314+
int expectedRequests = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 2 : 3;
2315+
verify(subchannel, times(expectedRequests)).requestConnection();
22512316

22522317
// PICK_FIRST doesn't use subchannelPool
22532318
verify(subchannelPool, never())
@@ -2334,10 +2399,16 @@ public void switchMode() throws Exception {
23342399
// PICK_FIRST Subchannel: child LB creates it
23352400
inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
23362401
CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
2337-
assertThat(createSubchannelArgs.getAddresses())
2338-
.containsExactly(
2339-
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
2340-
new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
2402+
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
2403+
assertThat(createSubchannelArgs.getAddresses())
2404+
.containsExactly(
2405+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
2406+
} else {
2407+
assertThat(createSubchannelArgs.getAddresses())
2408+
.containsExactly(
2409+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
2410+
new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
2411+
}
23412412

23422413
// Child pick_first eagerly connects, so initial state is CONNECTING (not IDLE)
23432414
inOrder.verify(helper, atLeast(1))
@@ -2424,10 +2495,16 @@ public void switchMode_nullLbPolicy() throws Exception {
24242495
// PICK_FIRST Subchannel: with delegation, child LB creates the subchannel
24252496
inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
24262497
CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
2427-
assertThat(createSubchannelArgs.getAddresses())
2428-
.containsExactly(
2429-
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
2430-
new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
2498+
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
2499+
assertThat(createSubchannelArgs.getAddresses())
2500+
.containsExactly(
2501+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
2502+
} else {
2503+
assertThat(createSubchannelArgs.getAddresses())
2504+
.containsExactly(
2505+
new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
2506+
new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
2507+
}
24312508

24322509
// Child pick_first eagerly connects, so state is CONNECTING (not IDLE)
24332510
inOrder.verify(helper, atLeast(1))

‎rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@
7373
import io.grpc.inprocess.InProcessServerBuilder;
7474
import io.grpc.internal.FakeClock;
7575
import io.grpc.internal.JsonParser;
76+
import io.grpc.internal.PickFirstLoadBalancerProvider;
7677
import io.grpc.internal.PickSubchannelArgsImpl;
7778
import io.grpc.internal.testing.StreamRecorder;
7879
import io.grpc.lookup.v1.RouteLookupServiceGrpc;
@@ -257,6 +258,9 @@ public void lb_working_withDefaultTarget_rlsResponding() throws Exception {
257258
InOrder inOrder = inOrder(helper);
258259
inOrder.verify(helper)
259260
.updateBalancingState(eq(ConnectivityState.CONNECTING), pickerCaptor.capture());
261+
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
262+
inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class));
263+
}
260264
SubchannelPicker picker = pickerCaptor.getValue();
261265
// Warm-up pick; will be queued
262266
PickResult res = picker.pickSubchannel(searchSubchannelArgs);

‎util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -255,7 +255,11 @@ public void pickAfterStateChange() throws Exception {
255255
assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus()).isEqualTo(error);
256256

257257
deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE));
258-
inOrder.verify(mockHelper).refreshNameResolution();
258+
if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
259+
inOrder.verify(mockHelper, never()).refreshNameResolution();
260+
} else {
261+
inOrder.verify(mockHelper).refreshNameResolution();
262+
}
259263
inOrder.verify(mockHelper, never())
260264
.updateBalancingState(eq(TRANSIENT_FAILURE), any(SubchannelPicker.class));
261265

‎xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,9 @@ public void pickAfterStateChange() throws Exception {
257257
.isEqualTo(PickResult.withNoResult());
258258

259259
deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE));
260-
inOrder.verify(helper).refreshNameResolution();
260+
if (!PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
261+
inOrder.verify(helper).refreshNameResolution();
262+
}
261263
assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE);
262264
assertThat(childLbState.getCurrentPicker().toString()).contains(error.toString());
263265

0 commit comments

Comments
 (0)