xds: fix LB policy address and balancing state update propagations (#7772)

Delaying handleResolvedAddresses() for propagating configs to the child LB policy can be problematic. For example, if channel shutdown has been enqueued when calling child policy's handleResolvedAddresses() is being enqueued (e.g., receiving updates from XdsClient), it should not be executed. Otherwise, subchannels may be created by LBs that have already been shut down.

This change fixes LB config propagations in LB policies that manage a group of child LBs and delay the propagation for avoiding reentrancy. LB policies will always directly propagate child LB config/addresses updates directly. On the other hand, upcalls from child LB policies for balancing state updates will be queued and executed later.

xds/src/main/java/io/grpc/xds/ClusterManagerLoadBalancer.java

@@ -84,21 +84,18 @@ class ClusterManagerLoadBalancer extends LoadBalancer {
       } else {
         childLbStates.get(name).reactivate(childPolicyProvider);
       }
-      final LoadBalancer childLb = childLbStates.get(name).lb;
+      LoadBalancer childLb = childLbStates.get(name).lb;
-      final ResolvedAddresses childAddresses =
+      ResolvedAddresses childAddresses =
           resolvedAddresses.toBuilder().setLoadBalancingPolicyConfig(childConfig).build();
-      syncContext.execute(new Runnable() {
+      childLb.handleResolvedAddresses(childAddresses);
-        @Override
-        public void run() {
-          childLb.handleResolvedAddresses(childAddresses);
-        }
-      });
     }
     for (String name : childLbStates.keySet()) {
       if (!newChildPolicies.containsKey(name)) {
         childLbStates.get(name).deactivate();
       }
     }
+    // Must update channel picker before return so that new RPCs will not be routed to deleted
+    // clusters and resolver can remove them in service config.
     updateOverallBalancingState();
   }
 
@@ -245,12 +242,20 @@ class ClusterManagerLoadBalancer extends LoadBalancer {
     private final class ChildLbStateHelper extends ForwardingLoadBalancerHelper {
 
       @Override
-      public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
+      public void updateBalancingState(final ConnectivityState newState,
-        currentState = newState;
+          final SubchannelPicker newPicker) {
-        currentPicker = newPicker;
+        syncContext.execute(new Runnable() {
-        if (!deactivated) {
+          @Override
-          updateOverallBalancingState();
+          public void run() {
-        }
+            currentState = newState;
+            currentPicker = newPicker;
+            // Subchannel picker and state are saved, but will only be propagated to the channel
+            // when the child instance exits deactivated state.
+            if (!deactivated) {
+              updateOverallBalancingState();
+            }
+          }
+        });
       }
 
       @Override

xds/src/main/java/io/grpc/xds/PriorityLoadBalancer.java

@@ -251,43 +251,42 @@ final class PriorityLoadBalancer extends LoadBalancer {
      * already exists.
      */
     void updateResolvedAddresses() {
-      final ResolvedAddresses addresses = resolvedAddresses;
+      PriorityLbConfig config =
-      syncContext.execute(
+          (PriorityLbConfig) resolvedAddresses.getLoadBalancingPolicyConfig();
-          new Runnable() {
+      PolicySelection childPolicySelection = config.childConfigs.get(priority);
-            @Override
+      LoadBalancerProvider lbProvider = childPolicySelection.getProvider();
-            public void run() {
+      String newPolicy = lbProvider.getPolicyName();
-              PriorityLbConfig config = (PriorityLbConfig) addresses.getLoadBalancingPolicyConfig();
+      if (!newPolicy.equals(policy)) {
-              PolicySelection childPolicySelection = config.childConfigs.get(priority);
+        policy = newPolicy;
-              LoadBalancerProvider lbProvider = childPolicySelection.getProvider();
+        lb.switchTo(lbProvider);
-              String newPolicy = lbProvider.getPolicyName();
+      }
-              if (!newPolicy.equals(policy)) {
+      lb.handleResolvedAddresses(
-                policy = newPolicy;
+          resolvedAddresses.toBuilder()
-                lb.switchTo(lbProvider);
+              .setAddresses(AddressFilter.filter(resolvedAddresses.getAddresses(), priority))
-              }
+              .setLoadBalancingPolicyConfig(childPolicySelection.getConfig())
-              lb.handleResolvedAddresses(
+              .build());
-                  addresses
-                      .toBuilder()
-                      .setAddresses(AddressFilter.filter(addresses.getAddresses(), priority))
-                      .setLoadBalancingPolicyConfig(childPolicySelection.getConfig())
-                      .build());
-            }
-          });
     }
 
     final class ChildHelper extends ForwardingLoadBalancerHelper {
       @Override
-      public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
+      public void updateBalancingState(final ConnectivityState newState,
-        connectivityState = newState;
+          final SubchannelPicker newPicker) {
-        picker = newPicker;
+        syncContext.execute(new Runnable() {
-        if (deletionTimer != null && deletionTimer.isPending()) {
+          @Override
-          return;
+          public void run() {
-        }
+            connectivityState = newState;
-        if (failOverTimer.isPending()) {
+            picker = newPicker;
-          if (newState.equals(READY) || newState.equals(TRANSIENT_FAILURE)) {
+            if (deletionTimer != null && deletionTimer.isPending()) {
-            failOverTimer.cancel();
+              return;
+            }
+            if (failOverTimer.isPending()) {
+              if (newState.equals(READY) || newState.equals(TRANSIENT_FAILURE)) {
+                failOverTimer.cancel();
+              }
+            }
+            tryNextPriority(true);
           }
-        }
+        });
-        tryNextPriority(true);
       }
 
       @Override

xds/src/main/java/io/grpc/xds/WeightedTargetLoadBalancer.java

@@ -28,6 +28,7 @@ import io.grpc.ConnectivityState;
 import io.grpc.InternalLogId;
 import io.grpc.LoadBalancer;
 import io.grpc.Status;
+import io.grpc.SynchronizationContext;
 import io.grpc.util.ForwardingLoadBalancerHelper;
 import io.grpc.util.GracefulSwitchLoadBalancer;
 import io.grpc.xds.WeightedRandomPicker.WeightedChildPicker;
@@ -48,11 +49,13 @@ final class WeightedTargetLoadBalancer extends LoadBalancer {
   private final Map<String, GracefulSwitchLoadBalancer> childBalancers = new HashMap<>();
   private final Map<String, ChildHelper> childHelpers = new HashMap<>();
   private final Helper helper;
+  private final SynchronizationContext syncContext;
 
   private Map<String, WeightedPolicySelection> targets = ImmutableMap.of();
 
   WeightedTargetLoadBalancer(Helper helper) {
-    this.helper = helper;
+    this.helper = checkNotNull(helper, "helper");
+    this.syncContext = checkNotNull(helper.getSynchronizationContext(), "syncContext");
     logger = XdsLogger.withLogId(
         InternalLogId.allocate("weighted-target-lb", helper.getAuthority()));
     logger.log(XdsLogLevel.INFO, "Created");
@@ -63,10 +66,8 @@ final class WeightedTargetLoadBalancer extends LoadBalancer {
     logger.log(XdsLogLevel.DEBUG, "Received resolution result: {0}", resolvedAddresses);
     Object lbConfig = resolvedAddresses.getLoadBalancingPolicyConfig();
     checkNotNull(lbConfig, "missing weighted_target lb config");
-
     WeightedTargetConfig weightedTargetConfig = (WeightedTargetConfig) lbConfig;
     Map<String, WeightedPolicySelection> newTargets = weightedTargetConfig.targets;
-
     for (String targetName : newTargets.keySet()) {
       WeightedPolicySelection weightedChildLbConfig = newTargets.get(targetName);
       if (!targets.containsKey(targetName)) {
@@ -81,9 +82,7 @@ final class WeightedTargetLoadBalancer extends LoadBalancer {
             .switchTo(weightedChildLbConfig.policySelection.getProvider());
       }
     }
-
     targets = newTargets;
-
     for (String targetName : targets.keySet()) {
       childBalancers.get(targetName).handleResolvedAddresses(
           resolvedAddresses.toBuilder()
@@ -101,6 +100,7 @@ final class WeightedTargetLoadBalancer extends LoadBalancer {
     }
     childBalancers.keySet().retainAll(targets.keySet());
     childHelpers.keySet().retainAll(targets.keySet());
+    updateOverallBalancingState();
   }
 
   @Override
@@ -180,10 +180,16 @@ final class WeightedTargetLoadBalancer extends LoadBalancer {
     SubchannelPicker currentPicker = BUFFER_PICKER;
 
     @Override
-    public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
+    public void updateBalancingState(final ConnectivityState newState,
-      currentState = newState;
+        final SubchannelPicker newPicker) {
-      currentPicker = newPicker;
+      syncContext.execute(new Runnable() {
-      updateOverallBalancingState();
+        @Override
+        public void run() {
+          currentState = newState;
+          currentPicker = newPicker;
+          updateOverallBalancingState();
+        }
+      });
     }
 
     @Override

xds/src/test/java/io/grpc/xds/ClusterImplLoadBalancerTest.java

@@ -41,6 +41,7 @@ import io.grpc.Metadata;
 import io.grpc.NameResolver;
 import io.grpc.Status;
 import io.grpc.Status.Code;
+import io.grpc.SynchronizationContext;
 import io.grpc.internal.ObjectPool;
 import io.grpc.internal.ServiceConfigUtil.PolicySelection;
 import io.grpc.xds.ClusterImplLoadBalancerProvider.ClusterImplConfig;
@@ -86,6 +87,13 @@ public class ClusterImplLoadBalancerTest {
   private static final String CLUSTER = "cluster-foo.googleapis.com";
   private static final String EDS_SERVICE_NAME = "service.googleapis.com";
   private static final String LRS_SERVER_NAME = "";
+  private final SynchronizationContext syncContext = new SynchronizationContext(
+      new Thread.UncaughtExceptionHandler() {
+        @Override
+        public void uncaughtException(Thread t, Throwable e) {
+          throw new AssertionError(e);
+        }
+      });
   private final Locality locality =
       new Locality("test-region", "test-zone", "test-subzone");
   private final PolicySelection roundRobin =
@@ -583,6 +591,12 @@ public class ClusterImplLoadBalancerTest {
   }
 
   private final class FakeLbHelper extends LoadBalancer.Helper {
+
+    @Override
+    public SynchronizationContext getSynchronizationContext() {
+      return syncContext;
+    }
+
     @Override
     public void updateBalancingState(
         @Nonnull ConnectivityState newState, @Nonnull SubchannelPicker newPicker) {

xds/src/test/java/io/grpc/xds/WeightedTargetLoadBalancerTest.java

@@ -23,18 +23,18 @@ import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
 import static io.grpc.xds.XdsSubchannelPickers.BUFFER_PICKER;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.atLeastOnce;
-import static org.mockito.Mockito.doReturn;
 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.when;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import io.grpc.Attributes;
-import io.grpc.ChannelLogger;
 import io.grpc.EquivalentAddressGroup;
 import io.grpc.LoadBalancer;
 import io.grpc.LoadBalancer.Helper;
@@ -45,6 +45,7 @@ import io.grpc.LoadBalancer.SubchannelPicker;
 import io.grpc.LoadBalancerProvider;
 import io.grpc.LoadBalancerRegistry;
 import io.grpc.Status;
+import io.grpc.SynchronizationContext;
 import io.grpc.internal.ServiceConfigUtil.PolicySelection;
 import io.grpc.xds.WeightedRandomPicker.WeightedChildPicker;
 import io.grpc.xds.WeightedTargetLoadBalancerProvider.WeightedPolicySelection;
@@ -68,6 +69,13 @@ import org.mockito.MockitoAnnotations;
 @RunWith(JUnit4.class)
 public class WeightedTargetLoadBalancerTest {
 
+  private final SynchronizationContext syncContext = new SynchronizationContext(
+      new Thread.UncaughtExceptionHandler() {
+        @Override
+        public void uncaughtException(Thread t, Throwable e) {
+          throw new AssertionError(e);
+        }
+      });
   private final LoadBalancerRegistry lbRegistry = new LoadBalancerRegistry();
   private final List<LoadBalancer> childBalancers = new ArrayList<>();
   private final List<Helper> childHelpers = new ArrayList<>();
@@ -143,8 +151,6 @@ public class WeightedTargetLoadBalancerTest {
 
   @Mock
   private Helper helper;
-  @Mock
-  private ChannelLogger channelLogger;
 
   private LoadBalancer weightedTargetLb;
   private int fooLbCreated;
@@ -153,8 +159,7 @@ public class WeightedTargetLoadBalancerTest {
   @Before
   public void setUp() {
     MockitoAnnotations.initMocks(this);
-
+    when(helper.getSynchronizationContext()).thenReturn(syncContext);
-    doReturn(channelLogger).when(helper).getChannelLogger();
     lbRegistry.register(fooLbProvider);
     lbRegistry.register(barLbProvider);
 
@@ -198,7 +203,7 @@ public class WeightedTargetLoadBalancerTest {
             .setAttributes(Attributes.newBuilder().set(fakeKey, fakeValue).build())
             .setLoadBalancingPolicyConfig(new WeightedTargetConfig(targets))
             .build());
-
+    verify(helper).updateBalancingState(eq(CONNECTING), eq(BUFFER_PICKER));
     assertThat(childBalancers).hasSize(4);
     assertThat(childHelpers).hasSize(4);
     assertThat(fooLbCreated).isEqualTo(2);
@@ -235,7 +240,7 @@ public class WeightedTargetLoadBalancerTest {
             .setAddresses(ImmutableList.<EquivalentAddressGroup>of())
             .setLoadBalancingPolicyConfig(new WeightedTargetConfig(newTargets))
             .build());
-
+    verify(helper, atLeast(2)).updateBalancingState(eq(CONNECTING), eq(BUFFER_PICKER));
     assertThat(childBalancers).hasSize(5);
     assertThat(childHelpers).hasSize(5);
     assertThat(fooLbCreated).isEqualTo(3); // One more foo LB created for target4
@@ -277,6 +282,7 @@ public class WeightedTargetLoadBalancerTest {
             .setAddresses(ImmutableList.<EquivalentAddressGroup>of())
             .setLoadBalancingPolicyConfig(new WeightedTargetConfig(targets))
             .build());
+    verify(helper).updateBalancingState(eq(CONNECTING), eq(BUFFER_PICKER));
 
     // Error after child balancers created.
     weightedTargetLb.handleNameResolutionError(Status.ABORTED);
@@ -303,6 +309,7 @@ public class WeightedTargetLoadBalancerTest {
             .setAddresses(ImmutableList.<EquivalentAddressGroup>of())
             .setLoadBalancingPolicyConfig(new WeightedTargetConfig(targets))
             .build());
+    verify(helper).updateBalancingState(eq(CONNECTING), eq(BUFFER_PICKER));
 
     // Subchannels to be created for each child balancer.
     final SubchannelPicker[] subchannelPickers = new SubchannelPicker[]{
@@ -316,7 +323,7 @@ public class WeightedTargetLoadBalancerTest {
     childHelpers.get(1).updateBalancingState(TRANSIENT_FAILURE, new ErrorPicker(Status.ABORTED));
     verify(helper, never()).updateBalancingState(
         eq(TRANSIENT_FAILURE), any(SubchannelPicker.class));
-    verify(helper).updateBalancingState(eq(CONNECTING), eq(BUFFER_PICKER));
+    verify(helper, times(2)).updateBalancingState(eq(CONNECTING), eq(BUFFER_PICKER));
 
     // Another child balancer goes to READY.
     childHelpers.get(2).updateBalancingState(READY, subchannelPickers[2]);