core: Avoid implicit requestConnection in PickFirst

This makes the behavior more clear.

2 files changed, 35 insertions, 20 deletions

diff --git a/core/src/main/java/io/grpc/PickFirstBalancerFactory.java b/core/src/main/java/io/grpc/PickFirstBalancerFactory.java
index 86677d3e2..3c75556e4 100644
--- a/core/src/main/java/io/grpc/PickFirstBalancerFactory.java
+++ b/core/src/main/java/io/grpc/PickFirstBalancerFactory.java
@@ -95,23 +95,27 @@ public final class PickFirstBalancerFactory extends LoadBalancer.Factory {
         return;
       }
 
-      PickResult pickResult;
+      SubchannelPicker picker;
       switch (currentState) {
+        case IDLE:
+          picker = new RequestConnectionPicker(subchannel);
+          break;
         case CONNECTING:
-          pickResult = PickResult.withNoResult();
+          // It's safe to use RequestConnectionPicker here, so when coming from IDLE we could leave
+          // the current picker in-place. But ignoring the potential optimization is simpler.
+          picker = new Picker(PickResult.withNoResult());
           break;
         case READY:
-        case IDLE:
-          pickResult = PickResult.withSubchannel(subchannel);
+          picker = new Picker(PickResult.withSubchannel(subchannel));
           break;
         case TRANSIENT_FAILURE:
-          pickResult = PickResult.withError(stateInfo.getStatus());
+          picker = new Picker(PickResult.withError(stateInfo.getStatus()));
           break;
         default:
           throw new IllegalArgumentException("Unsupported state:" + currentState);
       }
 
-      helper.updateBalancingState(currentState, new Picker(pickResult));
+      helper.updateBalancingState(currentState, picker);
     }
 
     @Override
@@ -126,8 +130,7 @@ public final class PickFirstBalancerFactory extends LoadBalancer.Factory {
    * No-op picker which doesn't add any custom picking logic. It just passes already known result
    * received in constructor.
    */
-  @VisibleForTesting
-  static final class Picker extends SubchannelPicker {
+  private static final class Picker extends SubchannelPicker {
     private final PickResult result;
 
     Picker(PickResult result) {
@@ -138,13 +141,25 @@ public final class PickFirstBalancerFactory extends LoadBalancer.Factory {
     public PickResult pickSubchannel(PickSubchannelArgs args) {
       return result;
     }
+  }
+
+  /** Picker that requests connection during pick, and returns noResult. */
+  private static final class RequestConnectionPicker extends SubchannelPicker {
+    private final Subchannel subchannel;
+
+    RequestConnectionPicker(Subchannel subchannel) {
+      this.subchannel = checkNotNull(subchannel, "subchannel");
+    }
+
+    @Override
+    public PickResult pickSubchannel(PickSubchannelArgs args) {
+      subchannel.requestConnection();
+      return PickResult.withNoResult();
+    }
 
     @Override
     public void requestConnection() {
-      Subchannel subchannel = result.getSubchannel();
-      if (subchannel != null) {
-        subchannel.requestConnection();
-      }
+      subchannel.requestConnection();
     }
   }
 }
diff --git a/core/src/test/java/io/grpc/PickFirstLoadBalancerTest.java b/core/src/test/java/io/grpc/PickFirstLoadBalancerTest.java
index 5b7dbbf20..30c31e254 100644
--- a/core/src/test/java/io/grpc/PickFirstLoadBalancerTest.java
+++ b/core/src/test/java/io/grpc/PickFirstLoadBalancerTest.java
@@ -38,8 +38,8 @@ import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancer.PickResult;
 import io.grpc.LoadBalancer.PickSubchannelArgs;
 import io.grpc.LoadBalancer.Subchannel;
+import io.grpc.LoadBalancer.SubchannelPicker;
 import io.grpc.PickFirstBalancerFactory.PickFirstBalancer;
-import io.grpc.PickFirstBalancerFactory.Picker;
 import java.net.SocketAddress;
 import java.util.List;
 import org.junit.After;
@@ -65,7 +65,7 @@ public class PickFirstLoadBalancerTest {
   private Attributes affinity = Attributes.newBuilder().set(FOO, "bar").build();
 
   @Captor
-  private ArgumentCaptor<Picker> pickerCaptor;
+  private ArgumentCaptor<SubchannelPicker> pickerCaptor;
   @Captor
   private ArgumentCaptor<Attributes> attrsCaptor;
   @Mock
@@ -121,7 +121,8 @@ public class PickFirstLoadBalancerTest {
 
     verify(mockHelper).createSubchannel(anyListOf(EquivalentAddressGroup.class),
         any(Attributes.class));
-    verify(mockHelper).updateBalancingState(isA(ConnectivityState.class), isA(Picker.class));
+    verify(mockHelper)
+        .updateBalancingState(isA(ConnectivityState.class), isA(SubchannelPicker.class));
     // Updating the subchannel addresses is unnecessary, but doesn't hurt anything
     verify(mockHelper).updateSubchannelAddresses(
         eq(mockSubchannel), anyListOf(EquivalentAddressGroup.class));
@@ -201,7 +202,7 @@ public class PickFirstLoadBalancerTest {
 
     loadBalancer.handleNameResolutionError(Status.NOT_FOUND.withDescription("nameResolutionError"));
     inOrder.verify(mockHelper)
-        .updateBalancingState(any(ConnectivityState.class), any(Picker.class));
+        .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class));
     verify(mockSubchannel, never()).requestConnection();
 
     loadBalancer.handleResolvedAddressGroups(servers, affinity);
@@ -247,13 +248,12 @@ public class PickFirstLoadBalancerTest {
   @Test
   public void requestConnection() {
     loadBalancer.handleResolvedAddressGroups(servers, affinity);
-    verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
-    Picker picker = pickerCaptor.getValue();
+    loadBalancer.handleSubchannelState(mockSubchannel, ConnectivityStateInfo.forNonError(IDLE));
+    verify(mockHelper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
+    SubchannelPicker picker = pickerCaptor.getValue();
 
     verify(mockSubchannel).requestConnection();
-
     picker.requestConnection();
-
     verify(mockSubchannel, times(2)).requestConnection();
   }