services: HealthCheckingLoadBalancerUtil and HealthCheckingLoadBalancerProvider (#5026)

HealthCheckingLoadBalancerUtil is the public wrapper utility that helps
turn a LoadBalancerFactory into a health-checking capable one.

HealthCheckingRoundRobinLoadBalancerProvider overrides the
RoundRobinLoadBalancerProvider from grpc-core.

services/src/main/java/io/grpc/services/HealthCheckingLoadBalancerUtil.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2018 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.services;
+
+import io.grpc.ExperimentalApi;
+import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.Factory;
+import io.grpc.LoadBalancer.Helper;
+import io.grpc.internal.ExponentialBackoffPolicy;
+import io.grpc.internal.TimeProvider;
+
+/**
+ * Utility for enabling
+ * <a href="https://github.com/grpc/proposal/blob/master/A17-client-side-health-checking.md">
+ * client-side health checking</a> for {@link LoadBalancer}s.
+ */
+@ExperimentalApi("https://github.com/grpc/grpc-java/issues/5025")
+public final class HealthCheckingLoadBalancerUtil {
+  private HealthCheckingLoadBalancerUtil() {
+  }
+
+  /**
+   * Creates a health-checking-capable LoadBalancer.  This method is used to implement
+   * health-checking-capable {@link Factory}s, which will typically written this way:
+   *
+   * <pre>
+   * public class HealthCheckingFooLbFactory extends LoadBalancer.Factory {
+   *   // This is the original balancer implementation that doesn't have health checking
+   *   private final LoadBalancer.Factory fooLbFactory;
+   *
+   *   ...
+   *
+   *   // Returns the health-checking-capable version of FooLb   
+   *   public LoadBalancer newLoadBalancer(Helper helper) {
+   *     return HealthCheckingLoadBalancerUtil.newHealthCheckingLoadBalancer(fooLbFactory, helper);
+   *   }
+   * }
+   * </pre>
+   *
+   * <p>As a requirement for the original LoadBalancer, it must call
+   * {@code Helper.createSubchannel()} from the {@link
+   * io.grpc.LoadBalancer.Helper#getSynchronizationContext() Synchronization Context}, or
+   * {@code createSubchannel()} will throw.
+   *
+   * @param factory the original factory that implements load-balancing logic without health
+   *        checking
+   * @param helper the helper passed to the resulting health-checking LoadBalancer.
+   */
+  public static LoadBalancer newHealthCheckingLoadBalancer(Factory factory, Helper helper) {
+    HealthCheckingLoadBalancerFactory hcFactory =
+        new HealthCheckingLoadBalancerFactory(
+            factory, new ExponentialBackoffPolicy.Provider(), TimeProvider.SYSTEM_TIME_PROVIDER);
+    return hcFactory.newLoadBalancer(helper);
+  }
+}

services/src/main/java/io/grpc/services/internal/HealthCheckingRoundRobinLoadBalancerProvider.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2018 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.services.internal;
+
+import io.grpc.Internal;
+import io.grpc.LoadBalancer;
+import io.grpc.LoadBalancer.Helper;
+import io.grpc.LoadBalancerProvider;
+import io.grpc.internal.RoundRobinLoadBalancerProvider;
+import io.grpc.services.HealthCheckingLoadBalancerUtil;
+
+/**
+ * The health-check-capable provider for the "round_robin" balancing policy.  This overrides
+ * the "round_robin" provided by grpc-core.
+ */
+@Internal
+public final class HealthCheckingRoundRobinLoadBalancerProvider extends LoadBalancerProvider {
+  private final RoundRobinLoadBalancerProvider rrProvider = new RoundRobinLoadBalancerProvider();
+
+  @Override
+  public boolean isAvailable() {
+    return rrProvider.isAvailable();
+  }
+
+  @Override
+  public int getPriority() {
+    return rrProvider.getPriority() + 1;
+  }
+
+  @Override
+  public String getPolicyName() {
+    return rrProvider.getPolicyName();
+  }
+
+  @Override
+  public LoadBalancer newLoadBalancer(Helper helper) {
+    return HealthCheckingLoadBalancerUtil.newHealthCheckingLoadBalancer(rrProvider, helper);
+  }
+}

services/src/main/resources/META-INF/services/io.grpc.LoadBalancerProvider

@@ -0,0 +1 @@
+io.grpc.services.internal.HealthCheckingRoundRobinLoadBalancerProvider

services/src/test/java/io/grpc/services/HealthCheckingLoadBalancerFactoryTest.java

@@ -133,6 +133,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
 
   @Mock
   private LoadBalancer origLb;
+  private LoadBalancer hcLb;
   @Captor
   ArgumentCaptor<Attributes> attrsCaptor;
   @Mock
@@ -175,7 +176,7 @@ public class HealthCheckingLoadBalancerFactoryTest {
 
     hcLbFactory = new HealthCheckingLoadBalancerFactory(
         origLbFactory, backoffPolicyProvider, clock.getTimeProvider());
-    final LoadBalancer hcLb = hcLbFactory.newLoadBalancer(origHelper);
+    hcLb = hcLbFactory.newLoadBalancer(origHelper);
     // Make sure all calls into the hcLb is from the syncContext
     hcLbEventDelivery = new LoadBalancer() {
         @Override
@@ -914,6 +915,33 @@ public class HealthCheckingLoadBalancerFactoryTest {
         .isEqualTo("FooService");
   }
 
+  @Test
+  public void util_newHealthCheckingLoadBalancer() {
+    Factory hcFactory =
+        new Factory() {
+          @Override
+          public LoadBalancer newLoadBalancer(Helper helper) {
+            return HealthCheckingLoadBalancerUtil.newHealthCheckingLoadBalancer(
+                origLbFactory, helper);
+          }
+        };
+
+    // hcLb and wrappedHelper are already set in setUp().  For this special test case, we
+    // clear wrappedHelper so that we can create hcLb again with the util.
+    wrappedHelper = null;
+    hcLb = hcFactory.newLoadBalancer(origHelper);
+
+    // Verify that HC works
+    Attributes resolutionAttrs = attrsWithHealthCheckService("BarService");
+    hcLbEventDelivery.handleResolvedAddressGroups(resolvedAddressList, resolutionAttrs);
+    verify(origLb).handleResolvedAddressGroups(same(resolvedAddressList), same(resolutionAttrs));
+    createSubchannel(0, Attributes.EMPTY);
+    assertThat(healthImpls[0].calls).isEmpty();
+    hcLbEventDelivery.handleSubchannelState(
+        subchannels[0], ConnectivityStateInfo.forNonError(READY));
+    assertThat(healthImpls[0].calls).hasSize(1);
+  }
+
   private Attributes attrsWithHealthCheckService(@Nullable String serviceName) {
     HashMap<String, Object> serviceConfig = new HashMap<String, Object>();
     HashMap<String, Object> hcConfig = new HashMap<String, Object>();

services/src/test/java/io/grpc/services/internal/HealthCheckingRoundRobinLoadBalancerProviderTest.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2018 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.services.internal;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import io.grpc.LoadBalancerProvider;
+import io.grpc.LoadBalancerRegistry;
+import io.grpc.internal.RoundRobinLoadBalancerProvider;
+import org.junit.Test;
+
+/**
+ * Tests for {@link InternalHealthCheckingRoundRobinLoadBalancerProvider}.
+ */
+public class HealthCheckingRoundRobinLoadBalancerProviderTest {
+  @Test
+  public void registry() {
+    LoadBalancerProvider hcRoundRobin =
+        LoadBalancerRegistry.getDefaultRegistry().getProvider("round_robin");
+    assertThat(hcRoundRobin).isInstanceOf(
+        HealthCheckingRoundRobinLoadBalancerProvider.class);
+  }
+
+  @Test
+  public void policyName() {
+    LoadBalancerProvider hcRoundRobin = new HealthCheckingRoundRobinLoadBalancerProvider();
+    assertThat(hcRoundRobin.getPolicyName())
+        .isEqualTo(new RoundRobinLoadBalancerProvider().getPolicyName());
+  }
+
+  @Test
+  public void priority() {
+    LoadBalancerProvider hcRoundRobin = new HealthCheckingRoundRobinLoadBalancerProvider();
+    assertThat(hcRoundRobin.getPriority())
+        .isEqualTo(new RoundRobinLoadBalancerProvider().getPriority() + 1);
+  }
+}