Add flag protection

Author: temawi

core/src/main/java/io/grpc/internal/PickFirstLoadBalancer.java

@@ -22,7 +22,9 @@
 import static io.grpc.ConnectivityState.SHUTDOWN;
 import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
+import com.google.common.base.Strings;
 import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.EquivalentAddressGroup;
@@ -44,6 +46,10 @@ final class PickFirstLoadBalancer extends LoadBalancer {
   private final Helper helper;
   private Subchannel subchannel;
   private ConnectivityState currentState = IDLE;
+  @VisibleForTesting
+  static boolean enablePickFirstConfig =
+      !Strings.isNullOrEmpty(System.getenv("GRPC_EXPERIMENTAL_PICKFIRST_LB_CONFIG"))
+          && Boolean.parseBoolean(System.getenv("GRPC_EXPERIMENTAL_PICKFIRST_LB_CONFIG"));
 
   PickFirstLoadBalancer(Helper helper) {
     this.helper = checkNotNull(helper, "helper");
@@ -61,7 +67,7 @@ public boolean acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
 
     // We can optionally be configured to shuffle the address list. This can help better distribute
     // the load.
-    if (resolvedAddresses.getLoadBalancingPolicyConfig() != null) {
+    if (enablePickFirstConfig && resolvedAddresses.getLoadBalancingPolicyConfig() != null) {
       PickFirstLoadBalancerConfig config
           = (PickFirstLoadBalancerConfig) resolvedAddresses.getLoadBalancingPolicyConfig();
       if (config.shuffleAddressList != null && config.shuffleAddressList) {

core/src/test/java/io/grpc/internal/PickFirstLoadBalancerTest.java

@@ -121,6 +121,7 @@ public void setUp() {
   @After
   public void tearDown() throws Exception {
     verifyNoMoreInteractions(mockArgs);
+    PickFirstLoadBalancer.enablePickFirstConfig = false;
   }
 
   @Test
@@ -143,6 +144,7 @@ public void pickAfterResolved() throws Exception {
 
   @Test
   public void pickAfterResolved_shuffle() throws Exception {
+    PickFirstLoadBalancer.enablePickFirstConfig = true;
     loadBalancer.acceptResolvedAddresses(
         ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(affinity)
             .setLoadBalancingPolicyConfig(new PickFirstLoadBalancerConfig(true, 123L)).build());