Fix child balancer's helper not wrapped by GracefulSwitchLoadBalancer.

grpc · Sep 3, 2020 · ae61247 · ae61247
1 parent 1356a41
commit ae61247
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/xds/src/main/java/io/grpc/xds/EdsLoadBalancer2.java b/xds/src/main/java/io/grpc/xds/EdsLoadBalancer2.java
@@ -144,7 +144,7 @@ private EdsLbState(@Nullable String edsServiceName, @Nullable String lrsServerNa
 
     @Override
     public LoadBalancer newLoadBalancer(Helper helper) {
-      return new ChildLbState();
+      return new ChildLbState(helper);
     }
 
     @Override
@@ -175,7 +175,7 @@ private final class ChildLbState extends LoadBalancer implements EndpointWatcher
       @Nullable
       private LoadBalancer lb;
 
-      private ChildLbState() {
+      private ChildLbState(Helper helper) {
         logger.log(
             XdsLogLevel.INFO,
             "Start endpoint watcher on {0} with xDS client {1}", resourceName, xdsClient);
@@ -191,7 +191,7 @@ private ChildLbState() {
           loadStatsStore = null;
           attributes = resolvedAddresses.getAttributes();
         }
-        lbHelper = new DropHandlingLbHelper();
+        lbHelper = new DropHandlingLbHelper(helper);
       }
 
       @Override
@@ -316,8 +316,13 @@ public void onError(Status error) {
       }
 
       private final class DropHandlingLbHelper extends ForwardingLoadBalancerHelper {
+        private final Helper helper;
         private List<DropOverload> dropPolicies = Collections.emptyList();
 
+        private DropHandlingLbHelper(Helper helper) {
+          this.helper = helper;
+        }
+
         @Override
         public void updateBalancingState(
             ConnectivityState newState, final SubchannelPicker newPicker) {