rls: delegating helper for rls child policies (#6904)

rls/src/main/java/io/grpc/rls/internal/ChildLoadBalancerHelper.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2020 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.rls.internal;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import io.grpc.ConnectivityState;
+import io.grpc.LoadBalancer.Helper;
+import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.util.ForwardingLoadBalancerHelper;
+import javax.annotation.Nonnull;
+
+/**
+ * A delegating {@link Helper} for the child load blanacer. The child load-balancer notifies the
+ * higher level load-blancer with aggregated status instead of each individual child load-blanacer's
+ * state.
+ */
+final class ChildLoadBalancerHelper extends ForwardingLoadBalancerHelper {
+
+  private final String target;
+  private final Helper rlsHelper;
+  private final SubchannelStateManager subchannelStateManager;
+  private final SubchannelPicker picker;
+
+  private ChildLoadBalancerHelper(
+      String target,
+      Helper rlsHelper,
+      SubchannelStateManager subchannelStateManager,
+      SubchannelPicker picker) {
+    this.target = checkNotNull(target, "target");
+    this.rlsHelper = checkNotNull(rlsHelper, "rlsHelper");
+    this.subchannelStateManager = checkNotNull(subchannelStateManager, "subchannelStateManager");
+    this.picker = checkNotNull(picker, "picker");
+  }
+
+  @Override
+  protected Helper delegate() {
+    return rlsHelper;
+  }
+
+  /**
+   * Updates balancing state from one or more subchannels tracked in the {@link
+   * SubchannelStateManager}. The passed picker will be ignored, instead the picker which governs
+   * many subchannels/pickers will be reported to the parent load-balancer.
+   */
+  @Override
+  public void updateBalancingState(
+      @Nonnull ConnectivityState newState,
+      @Nonnull SubchannelPicker unused) {
+    subchannelStateManager.updateState(target, newState);
+    super.updateBalancingState(subchannelStateManager.getAggregatedState(), picker);
+  }
+
+  static final class ChildLoadBalancerHelperProvider {
+    private final Helper helper;
+    private final SubchannelStateManager subchannelStateManager;
+    private final SubchannelPicker picker;
+
+    ChildLoadBalancerHelperProvider(
+        Helper helper, SubchannelStateManager subchannelStateManager, SubchannelPicker picker) {
+      this.helper = checkNotNull(helper, "helper");
+      this.subchannelStateManager = checkNotNull(subchannelStateManager, "subchannelStateManager");
+      this.picker = checkNotNull(picker, "picker");
+    }
+
+    ChildLoadBalancerHelper forTarget(String target) {
+      return new ChildLoadBalancerHelper(target, helper, subchannelStateManager, picker);
+    }
+  }
+}

rls/src/test/java/io/grpc/rls/internal/ChildLoadBalancerHelperTest.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2020 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.rls.internal;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+
+import io.grpc.ConnectivityState;
+import io.grpc.LoadBalancer.Helper;
+import io.grpc.LoadBalancer.SubchannelPicker;
+import io.grpc.rls.internal.ChildLoadBalancerHelper.ChildLoadBalancerHelperProvider;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+
+@RunWith(JUnit4.class)
+public class ChildLoadBalancerHelperTest {
+
+  private final Helper helper = mock(Helper.class);
+  private final SubchannelStateManager subchannelStateManager = new SubchannelStateManagerImpl();
+  private final SubchannelPicker picker = mock(SubchannelPicker.class);
+  private final ChildLoadBalancerHelperProvider provider =
+      new ChildLoadBalancerHelperProvider(helper, subchannelStateManager, picker);
+
+  @Test
+  public void childLoadBalancerHelper_shouldReportsSubchannelState() {
+    InOrder inOrder = Mockito.inOrder(helper);
+    String target1 = "foo.com";
+    ChildLoadBalancerHelper childLbHelper1 = provider.forTarget(target1);
+    SubchannelPicker picker1 = mock(SubchannelPicker.class);
+    String target2 = "bar.com";
+    ChildLoadBalancerHelper childLbHelper2 = provider.forTarget(target2);
+    SubchannelPicker picker2 = mock(SubchannelPicker.class);
+
+    assertThat(subchannelStateManager.getState(target1)).isNull();
+    assertThat(subchannelStateManager.getState(target2)).isNull();
+
+    childLbHelper1.updateBalancingState(ConnectivityState.TRANSIENT_FAILURE, picker1);
+    inOrder.verify(helper).updateBalancingState(ConnectivityState.TRANSIENT_FAILURE, picker);
+    assertThat(subchannelStateManager.getState(target1))
+        .isEqualTo(ConnectivityState.TRANSIENT_FAILURE);
+
+    childLbHelper2.updateBalancingState(ConnectivityState.CONNECTING, picker2);
+    inOrder.verify(helper).updateBalancingState(ConnectivityState.CONNECTING, picker);
+    assertThat(subchannelStateManager.getState(target2)).isEqualTo(ConnectivityState.CONNECTING);
+
+    childLbHelper1.updateBalancingState(ConnectivityState.READY, picker1);
+    inOrder.verify(helper).updateBalancingState(ConnectivityState.READY, picker);
+    assertThat(subchannelStateManager.getState(target1)).isEqualTo(ConnectivityState.READY);
+
+    childLbHelper1.updateBalancingState(ConnectivityState.SHUTDOWN, picker1);
+    inOrder.verify(helper).updateBalancingState(ConnectivityState.CONNECTING, picker);
+    assertThat(subchannelStateManager.getState(target1)).isNull();
+  }
+}