From ee3ffef3ee0eca9d63424b15e6c31d6e729ef4ed Mon Sep 17 00:00:00 2001
From: Eric Anderson <ejona@google.com>
Date: Sat, 17 Aug 2024 11:23:01 -0700
Subject: [PATCH] core: In PF, disjoint update while READY should transition to
 IDLE

This is the same as if we received a GOAWAY. We wait for the next RPC to
begin connecting again. This is InternalSubchannel's behavior.
---
 .../internal/PickFirstLeafLoadBalancer.java   | 18 ++++++++----------
 .../PickFirstLeafLoadBalancerTest.java        | 19 ++++++++++++++-----
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/core/src/main/java/io/grpc/internal/PickFirstLeafLoadBalancer.java b/core/src/main/java/io/grpc/internal/PickFirstLeafLoadBalancer.java
index 8d9c5aa8a1..344012fef2 100644
--- a/core/src/main/java/io/grpc/internal/PickFirstLeafLoadBalancer.java
+++ b/core/src/main/java/io/grpc/internal/PickFirstLeafLoadBalancer.java
@@ -152,20 +152,18 @@ final class PickFirstLeafLoadBalancer extends LoadBalancer {
       }
     }
 
-    if (oldAddrs.size() == 0 || rawConnectivityState == CONNECTING
-        || rawConnectivityState == READY) {
-      // start connection attempt at first address
+    if (oldAddrs.size() == 0) {
+      // Make tests happy; they don't properly assume starting in CONNECTING
       rawConnectivityState = CONNECTING;
       updateBalancingState(CONNECTING, new Picker(PickResult.withNoResult()));
-      cancelScheduleTask();
-      requestConnection();
+    }
 
-    } else if (rawConnectivityState == IDLE) {
-      // start connection attempt at first address when requested
-      SubchannelPicker picker = new RequestConnectionPicker(this);
-      updateBalancingState(IDLE, picker);
+    if (rawConnectivityState == READY) {
+      // connect from beginning when prompted
+      rawConnectivityState = IDLE;
+      updateBalancingState(IDLE, new RequestConnectionPicker(this));
 
-    } else if (rawConnectivityState == TRANSIENT_FAILURE) {
+    } else if (rawConnectivityState == CONNECTING || rawConnectivityState == TRANSIENT_FAILURE) {
       // start connection attempt at first address
       cancelScheduleTask();
       requestConnection();
diff --git a/core/src/test/java/io/grpc/internal/PickFirstLeafLoadBalancerTest.java b/core/src/test/java/io/grpc/internal/PickFirstLeafLoadBalancerTest.java
index a6e3f99ab3..9c22c4a708 100644
--- a/core/src/test/java/io/grpc/internal/PickFirstLeafLoadBalancerTest.java
+++ b/core/src/test/java/io/grpc/internal/PickFirstLeafLoadBalancerTest.java
@@ -1121,10 +1121,17 @@ public class PickFirstLeafLoadBalancerTest {
     loadBalancer.acceptResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(newServers).setAttributes(affinity).build());
     inOrder.verify(mockSubchannel1).shutdown();
-    inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
+    inOrder.verify(mockHelper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
+    inOrder.verify(mockSubchannel3, never()).start(stateListenerCaptor.capture());
+
+    // Trigger connection creation
+    picker = pickerCaptor.getValue();
+    assertEquals(PickResult.withNoResult(), picker.pickSubchannel(mockArgs));
     inOrder.verify(mockSubchannel3).start(stateListenerCaptor.capture());
     SubchannelStateListener stateListener3 = stateListenerCaptor.getValue();
     inOrder.verify(mockSubchannel3).requestConnection();
+    stateListener3.onSubchannelState(ConnectivityStateInfo.forNonError(CONNECTING));
+    inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
 
     if (enableHappyEyeballs) {
       forwardTimeByConnectionDelay();
@@ -1171,17 +1178,19 @@ public class PickFirstLeafLoadBalancerTest {
     loadBalancer.acceptResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(newestServers).setAttributes(affinity).build());
     inOrder.verify(mockSubchannel3).shutdown();
-    inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
-    inOrder.verify(mockSubchannel1n2).start(stateListenerCaptor.capture());
-    stateListener = stateListenerCaptor.getValue();
-    assertEquals(CONNECTING, loadBalancer.getConcludedConnectivityState());
+    inOrder.verify(mockHelper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
+    assertEquals(IDLE, loadBalancer.getConcludedConnectivityState());
     picker = pickerCaptor.getValue();
 
     // Calling pickSubchannel() twice gave the same result
     assertEquals(picker.pickSubchannel(mockArgs), picker.pickSubchannel(mockArgs));
 
     // But the picker calls requestConnection() only once
+    inOrder.verify(mockSubchannel1n2).start(stateListenerCaptor.capture());
+    stateListener = stateListenerCaptor.getValue();
     inOrder.verify(mockSubchannel1n2).requestConnection();
+    stateListener.onSubchannelState(ConnectivityStateInfo.forNonError(CONNECTING));
+    inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
     assertEquals(PickResult.withNoResult(), pickerCaptor.getValue().pickSubchannel(mockArgs));
     assertEquals(CONNECTING, loadBalancer.getConcludedConnectivityState());