From 517805f20bffd3ef274faf86e2995c1ed823e758 Mon Sep 17 00:00:00 2001
From: Terry Wilson <terrymwilson@gmail.com>
Date: Thu, 16 Jun 2022 13:11:53 -0700
Subject: [PATCH] xds: weighted target to delay picker updates while updating
 children

---
 .../grpc/xds/WeightedTargetLoadBalancer.java  | 33 +++++----
 .../xds/WeightedTargetLoadBalancerTest.java   | 69 +++++++++++++++++++
 2 files changed, 88 insertions(+), 14 deletions(-)

diff --git a/xds/src/main/java/io/grpc/xds/WeightedTargetLoadBalancer.java b/xds/src/main/java/io/grpc/xds/WeightedTargetLoadBalancer.java
index ee8c0308fce..5edc7bbf20f 100644
--- a/xds/src/main/java/io/grpc/xds/WeightedTargetLoadBalancer.java
+++ b/xds/src/main/java/io/grpc/xds/WeightedTargetLoadBalancer.java
@@ -28,7 +28,6 @@
 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;
@@ -49,13 +48,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();
+  // Set to true if currently in the process of handling resolved addresses.
+  private boolean resolvingAddresses;
 
   WeightedTargetLoadBalancer(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,6 +62,15 @@ final class WeightedTargetLoadBalancer extends LoadBalancer {
 
   @Override
   public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+    try {
+      resolvingAddresses = true;
+      handleResolvedAddressesInternal(resolvedAddresses);
+    } finally {
+      resolvingAddresses = false;
+    }
+  }
+
+  public void handleResolvedAddressesInternal(ResolvedAddresses resolvedAddresses) {
     logger.log(XdsLogLevel.DEBUG, "Received resolution result: {0}", resolvedAddresses);
     Object lbConfig = resolvedAddresses.getLoadBalancingPolicyConfig();
     checkNotNull(lbConfig, "missing weighted_target lb config");
@@ -191,17 +199,14 @@ private ChildHelper(String name) {
     @Override
     public void updateBalancingState(final ConnectivityState newState,
         final SubchannelPicker newPicker) {
-      syncContext.execute(new Runnable() {
-        @Override
-        public void run() {
-          if (!childBalancers.containsKey(name)) {
-            return;
-          }
-          currentState = newState;
-          currentPicker = newPicker;
-          updateOverallBalancingState();
-        }
-      });
+      currentState = newState;
+      currentPicker = newPicker;
+
+      // If we are already in the process of resolving addresses, the overall balancing state
+      // will be updated at the end of it, and we don't need to trigger that update here.
+      if (!resolvingAddresses && childBalancers.containsKey(name)) {
+        updateOverallBalancingState();
+      }
     }
 
     @Override
diff --git a/xds/src/test/java/io/grpc/xds/WeightedTargetLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/WeightedTargetLoadBalancerTest.java
index 7d9d30385e4..9bd33f0f182 100644
--- a/xds/src/test/java/io/grpc/xds/WeightedTargetLoadBalancerTest.java
+++ b/xds/src/test/java/io/grpc/xds/WeightedTargetLoadBalancerTest.java
@@ -402,4 +402,73 @@ public void raceBetweenShutdownAndChildLbBalancingStateUpdate() {
     weightedChildHelper0.updateBalancingState(READY, mock(SubchannelPicker.class));
     verifyNoMoreInteractions(helper);
   }
+
+  // When the ChildHelper is asked to update the overall balancing state, it should not do that if
+  // the update was triggered by the parent LB that will handle triggering the overall state update.
+  @Test
+  public void noDuplicateOverallBalancingStateUpdate() {
+    FakeLoadBalancerProvider fakeLbProvider = new FakeLoadBalancerProvider();
+
+    Map<String, WeightedPolicySelection> targets = ImmutableMap.of(
+        "target0", new WeightedPolicySelection(
+            weights[0], new PolicySelection(fakeLbProvider, configs[0])),
+        "target3", new WeightedPolicySelection(
+            weights[3], new PolicySelection(fakeLbProvider, configs[3])));
+    weightedTargetLb.handleResolvedAddresses(
+        ResolvedAddresses.newBuilder()
+            .setAddresses(ImmutableList.<EquivalentAddressGroup>of())
+            .setLoadBalancingPolicyConfig(new WeightedTargetConfig(targets))
+            .build());
+
+    // Both of the two child LB policies will call the helper to update the balancing state.
+    // But since those calls happen during the handling of teh resolved addresses of the parent
+    // WeightedTargetLLoadBalancer, the overall balancing state should only be updated once.
+    verify(helper, times(1)).updateBalancingState(any(), any());
+
+  }
+
+  private static class FakeLoadBalancerProvider extends LoadBalancerProvider {
+
+    @Override
+    public boolean isAvailable() {
+      return true;
+    }
+
+    @Override
+    public int getPriority() {
+      return 5;
+    }
+
+    @Override
+    public String getPolicyName() {
+      return "foo";
+    }
+
+    @Override
+    public LoadBalancer newLoadBalancer(Helper helper) {
+      return new FakeLoadBalancer(helper);
+    }
+  }
+
+  static class FakeLoadBalancer extends LoadBalancer {
+
+    private Helper helper;
+
+    FakeLoadBalancer(Helper helper) {
+      this.helper = helper;
+    }
+
+    @Override
+    public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
+      helper.updateBalancingState(TRANSIENT_FAILURE, new ErrorPicker(Status.INTERNAL));
+    }
+
+    @Override
+    public void handleNameResolutionError(Status error) {
+    }
+
+    @Override
+    public void shutdown() {
+    }
+  }
 }