service: HealthCheckingLoadBalancerFactory using its own attribute (#…

…6704)

api/src/main/java/io/grpc/LoadBalancer.java

@@ -121,6 +121,10 @@ public abstract class LoadBalancer {
   public static final Attributes.Key<Map<String, ?>> ATTR_LOAD_BALANCING_CONFIG =
       Attributes.Key.create("io.grpc.LoadBalancer.loadBalancingConfig");
 
+  @Internal
+  @NameResolver.ResolutionResultAttr
+  public static final Attributes.Key<Map<String, ?>> ATTR_HEALTH_CHECKING_CONFIG =
+      Attributes.Key.create("health-checking-config");
   private int recursionCount;
 
   /**

core/src/main/java/io/grpc/internal/ManagedChannelImpl.java

@@ -1409,11 +1409,19 @@ public void run() {
           if (NameResolverListener.this.helper == ManagedChannelImpl.this.lbHelper) {
             Attributes effectiveAttrs = attrs;
             if (effectiveServiceConfig != validServiceConfig) {
-              effectiveAttrs = attrs.toBuilder()
-                  .set(
-                      GrpcAttributes.NAME_RESOLVER_SERVICE_CONFIG,
-                      effectiveServiceConfig.rawServiceConfig)
-                  .build();
+              Attributes.Builder attrsBuilder = attrs.toBuilder();
+              attrsBuilder.set(
+                  GrpcAttributes.NAME_RESOLVER_SERVICE_CONFIG,
+                  effectiveServiceConfig.rawServiceConfig);
+              Map<String, ?> healthCheckingConfig =
+                  effectiveServiceConfig
+                      .managedChannelServiceConfig
+                      .getHealthCheckingConfig();
+              if (healthCheckingConfig != null) {
+                attrsBuilder
+                    .set(LoadBalancer.ATTR_HEALTH_CHECKING_CONFIG, healthCheckingConfig);
+              }
+              effectiveAttrs = attrsBuilder.build();
             }
 
             Status handleResult = helper.lb.tryHandleResolvedAddresses(

core/src/main/java/io/grpc/internal/ManagedChannelServiceConfig.java

@@ -43,16 +43,23 @@ final class ManagedChannelServiceConfig {
   private final Throttle retryThrottling;
   @Nullable
   private final Object loadBalancingConfig;
+  @Nullable
+  private final Map<String, ?> healthCheckingConfig;
 
   ManagedChannelServiceConfig(
       Map<String, MethodInfo> serviceMethodMap,
       Map<String, MethodInfo> serviceMap,
       @Nullable Throttle retryThrottling,
-      @Nullable Object loadBalancingConfig) {
+      @Nullable Object loadBalancingConfig,
+      @Nullable Map<String, ?> healthCheckingConfig) {
     this.serviceMethodMap = Collections.unmodifiableMap(new HashMap<>(serviceMethodMap));
     this.serviceMap = Collections.unmodifiableMap(new HashMap<>(serviceMap));
     this.retryThrottling = retryThrottling;
     this.loadBalancingConfig = loadBalancingConfig;
+    this.healthCheckingConfig =
+        healthCheckingConfig != null
+            ? Collections.unmodifiableMap(new HashMap<>(healthCheckingConfig))
+            : null;
   }
 
   /** Returns an empty {@link ManagedChannelServiceConfig}. */
@@ -62,7 +69,8 @@ static ManagedChannelServiceConfig empty() {
             new HashMap<String, MethodInfo>(),
             new HashMap<String, MethodInfo>(),
             /* retryThrottling= */ null,
-            /* loadBalancingConfig= */ null);
+            /* loadBalancingConfig= */ null,
+            /* healthCheckingConfig= */ null);
   }
 
   /**
@@ -80,6 +88,8 @@ static ManagedChannelServiceConfig fromServiceConfig(
     }
     Map<String, MethodInfo> serviceMethodMap = new HashMap<>();
     Map<String, MethodInfo> serviceMap = new HashMap<>();
+    Map<String, ?> healthCheckingConfig =
+        ServiceConfigUtil.getHealthCheckedService(serviceConfig);
 
     // Try and do as much validation here before we swap out the existing configuration.  In case
     // the input is invalid, we don't want to lose the existing configuration.
@@ -88,8 +98,13 @@ static ManagedChannelServiceConfig fromServiceConfig(
 
     if (methodConfigs == null) {
       // this is surprising, but possible.
-      return new ManagedChannelServiceConfig(
-          serviceMethodMap, serviceMap, retryThrottling, loadBalancingConfig);
+      return
+          new ManagedChannelServiceConfig(
+              serviceMethodMap,
+              serviceMap,
+              retryThrottling,
+              loadBalancingConfig,
+              healthCheckingConfig);
     }
 
     for (Map<String, ?> methodConfig : methodConfigs) {
@@ -122,8 +137,13 @@ static ManagedChannelServiceConfig fromServiceConfig(
       }
     }
 
-    return new ManagedChannelServiceConfig(
-        serviceMethodMap, serviceMap, retryThrottling, loadBalancingConfig);
+    return
+        new ManagedChannelServiceConfig(
+            serviceMethodMap,
+            serviceMap,
+            retryThrottling,
+            loadBalancingConfig,
+            healthCheckingConfig);
   }
 
   /**
@@ -133,6 +153,11 @@ Map<String, MethodInfo> getServiceMap() {
     return serviceMap;
   }
 
+  @Nullable
+  Map<String, ?> getHealthCheckingConfig() {
+    return healthCheckingConfig;
+  }
+
   /**
    * Returns the per-method configuration for the channel.
    */

core/src/main/java/io/grpc/internal/ServiceConfigUtil.java

@@ -45,10 +45,10 @@ public final class ServiceConfigUtil {
   private ServiceConfigUtil() {}
 
   /**
-   * Fetch the health-checked service name from service config. {@code null} if can't find one.
+   * Fetches the health-checked service config from service config. {@code null} if can't find one.
    */
   @Nullable
-  public static String getHealthCheckedServiceName(@Nullable Map<String, ?> serviceConfig) {
+  public static Map<String, ?> getHealthCheckedService(@Nullable  Map<String, ?> serviceConfig) {
     if (serviceConfig == null) {
       return null;
     }
@@ -61,11 +61,20 @@ public static String getHealthCheckedServiceName(@Nullable Map<String, ?> servic
       }
     }
     */
-    Map<String, ?> healthCheck = JsonUtil.getObject(serviceConfig, "healthCheckConfig");
-    if (healthCheck == null) {
+    return JsonUtil.getObject(serviceConfig, "healthCheckConfig");
+  }
+
+  /**
+   * Fetches the health-checked service name from health-checked service config. {@code null} if
+   * can't find one.
+   */
+  @Nullable
+  public static String getHealthCheckedServiceName(
+      @Nullable Map<String, ?> healthCheckedServiceConfig) {
+    if (healthCheckedServiceConfig == null) {
       return null;
     }
-    return JsonUtil.getString(healthCheck, "serviceName");
+    return JsonUtil.getString(healthCheckedServiceConfig, "serviceName");
   }
 
   @Nullable

core/src/test/java/io/grpc/internal/ManagedChannelServiceConfigTest.java

@@ -0,0 +1,59 @@
+/*
+ * 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.internal;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import java.util.Collections;
+import java.util.Map;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ManagedChannelServiceConfigTest {
+
+  @Test
+  public void managedChannelServiceConfig_shouldParseHealthCheckingConfig() throws Exception {
+    Map<String, ?> rawServiceConfig =
+        parseConfig(
+            "{\"healthCheckConfig\": "
+                + "{\"serviceName\": \"COVID-19\", "
+                + "\"description\": \"I can't visit korea, because of you\"}}");
+    ManagedChannelServiceConfig mcsc =
+        ManagedChannelServiceConfig.fromServiceConfig(rawServiceConfig, true, 3, 4, null);
+    Map<String, ?> healthCheckingConfig = mcsc.getHealthCheckingConfig();
+    assertThat(healthCheckingConfig).isNotNull();
+    assertThat(healthCheckingConfig)
+        .containsExactly(
+            "serviceName", "COVID-19", "description", "I can't visit korea, because of you");
+  }
+
+  @Test
+  public void managedChannelServiceConfig_shouldHandleNoHealthCheckingConfig() throws Exception {
+    ManagedChannelServiceConfig mcsc =
+        ManagedChannelServiceConfig
+            .fromServiceConfig(Collections.<String, Object>emptyMap(), true, 3, 4, null);
+
+    assertThat(mcsc.getHealthCheckingConfig()).isNull();
+  }
+
+  @SuppressWarnings("unchecked")
+  private static Map<String, Object> parseConfig(String json) throws Exception {
+    return (Map<String, Object>) JsonParser.parse(json);
+  }
+}

core/src/test/java/io/grpc/internal/ServiceConfigStateTest.java

@@ -42,11 +42,13 @@ public class ServiceConfigStateTest {
       Collections.<String, MethodInfo>emptyMap(),
       Collections.<String, MethodInfo>emptyMap(),
       null,
+      null,
       null);
   private final ManagedChannelServiceConfig serviceConfig2 = new ManagedChannelServiceConfig(
       Collections.<String, MethodInfo>emptyMap(),
       Collections.<String, MethodInfo>emptyMap(),
       null,
+      null,
       null);
   private final ConfigOrError config1 = ConfigOrError.fromConfig(serviceConfig1);
   private final ConfigOrError config2 = ConfigOrError.fromConfig(serviceConfig2);
@@ -429,6 +431,7 @@ public void lookup_default_onPresent_onPresent() {
         Collections.<String, MethodInfo>emptyMap(),
         Collections.<String, MethodInfo>emptyMap(),
         null,
+        null,
         null);
     ConfigOrError config3 = ConfigOrError.fromConfig(serviceConfig3);
 

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

@@ -49,7 +49,6 @@
 import io.grpc.health.v1.HealthCheckResponse.ServingStatus;
 import io.grpc.health.v1.HealthGrpc;
 import io.grpc.internal.BackoffPolicy;
-import io.grpc.internal.GrpcAttributes;
 import io.grpc.internal.ServiceConfigUtil;
 import io.grpc.util.ForwardingLoadBalancer;
 import io.grpc.util.ForwardingLoadBalancerHelper;
@@ -183,7 +182,9 @@ protected LoadBalancer delegate() {
     @Override
     public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
       Map<String, ?> serviceConfig =
-          resolvedAddresses.getAttributes().get(GrpcAttributes.NAME_RESOLVER_SERVICE_CONFIG);
+          resolvedAddresses
+              .getAttributes()
+              .get(LoadBalancer.ATTR_HEALTH_CHECKING_CONFIG);
       String serviceName = ServiceConfigUtil.getHealthCheckedServiceName(serviceConfig);
       helper.setHealthCheckedService(serviceName);
       super.handleResolvedAddresses(resolvedAddresses);

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

@@ -69,7 +69,6 @@
 import io.grpc.inprocess.InProcessServerBuilder;
 import io.grpc.internal.BackoffPolicy;
 import io.grpc.internal.FakeClock;
-import io.grpc.internal.GrpcAttributes;
 import io.grpc.internal.ServiceConfigUtil;
 import io.grpc.services.HealthCheckingLoadBalancerFactory.SubchannelImpl;
 import io.grpc.stub.StreamObserver;
@@ -978,30 +977,21 @@ public void serviceConfigChangesServiceNameWhenRpcInactive() {
   }
 
   @Test
-  public void getHealthCheckedServiceName_nullServiceConfig() {
+  public void getHealthCheckedServiceName_nullHealthCheckConfig() {
     assertThat(ServiceConfigUtil.getHealthCheckedServiceName(null)).isNull();
   }
 
   @Test
-  public void getHealthCheckedServiceName_noHealthCheckConfig() {
-    assertThat(ServiceConfigUtil.getHealthCheckedServiceName(new HashMap<String, Void>())).isNull();
-  }
-
-  @Test
-  public void getHealthCheckedServiceName_healthCheckConfigMissingServiceName() {
-    HashMap<String, Object> serviceConfig = new HashMap<>();
+  public void getHealthCheckedServiceName_missingServiceName() {
     HashMap<String, Object> hcConfig = new HashMap<>();
-    serviceConfig.put("healthCheckConfig", hcConfig);
-    assertThat(ServiceConfigUtil.getHealthCheckedServiceName(serviceConfig)).isNull();
+    assertThat(ServiceConfigUtil.getHealthCheckedServiceName(hcConfig)).isNull();
   }
 
   @Test
   public void getHealthCheckedServiceName_healthCheckConfigHasServiceName() {
-    HashMap<String, Object> serviceConfig = new HashMap<>();
     HashMap<String, Object> hcConfig = new HashMap<>();
     hcConfig.put("serviceName", "FooService");
-    serviceConfig.put("healthCheckConfig", hcConfig);
-    assertThat(ServiceConfigUtil.getHealthCheckedServiceName(serviceConfig))
+    assertThat(ServiceConfigUtil.getHealthCheckedServiceName(hcConfig))
         .isEqualTo("FooService");
   }
 
@@ -1094,14 +1084,12 @@ public LoadBalancer newLoadBalancer(Helper helper) {
     assertThat(healthImpls[0].calls).hasSize(1);
   }
 
-  @SuppressWarnings("deprecation")  // TODO(creamsoup) migrate to parsed object
   private Attributes attrsWithHealthCheckService(@Nullable String serviceName) {
-    HashMap<String, Object> serviceConfig = new HashMap<>();
     HashMap<String, Object> hcConfig = new HashMap<>();
     hcConfig.put("serviceName", serviceName);
-    serviceConfig.put("healthCheckConfig", hcConfig);
     return Attributes.newBuilder()
-        .set(GrpcAttributes.NAME_RESOLVER_SERVICE_CONFIG, serviceConfig).build();
+        .set(LoadBalancer.ATTR_HEALTH_CHECKING_CONFIG, hcConfig)
+        .build();
   }
 
   private HealthCheckRequest makeRequest(String service) {