xds: Custom LB configs to support UDPA TypeStruct (grpc#9198)

* xds: Custom LB configs to support UDPA TypeStruct

The legacy com.github.udpa.udpa.type.v1.TypedStruct proto should be supported in addition to the current com.github.udpa.udpa.type.v1.TypedStruct one.

Co-authored-by: Sergii Tkachenko <hi@sergii.org>

xds/src/main/java/io/grpc/xds/LoadBalancerConfigFactory.java

@@ -16,12 +16,12 @@
 
 package io.grpc.xds;
 
-import com.github.xds.type.v3.TypedStruct;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.protobuf.Any;
 import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.Struct;
 import com.google.protobuf.util.JsonFormat;
 import io.envoyproxy.envoy.config.cluster.v3.Cluster;
 import io.envoyproxy.envoy.config.cluster.v3.Cluster.LeastRequestLbConfig;
@@ -167,9 +167,14 @@ static class LoadBalancingPolicyConverter {
                 recursionDepth);
           } else if (typedConfig.is(RoundRobin.class)) {
             serviceConfig = convertRoundRobinConfig();
-          } else if (typedConfig.is(TypedStruct.class)) {
-            serviceConfig = convertCustomConfig(typedConfig.unpack(TypedStruct.class));
+          } else if (typedConfig.is(com.github.xds.type.v3.TypedStruct.class)) {
+            serviceConfig = convertCustomConfig(
+                typedConfig.unpack(com.github.xds.type.v3.TypedStruct.class));
+          } else if (typedConfig.is(com.github.udpa.udpa.type.v1.TypedStruct.class)) {
+            serviceConfig = convertCustomConfig(
+                typedConfig.unpack(com.github.udpa.udpa.type.v1.TypedStruct.class));
           }
+
           // TODO: support least_request once it is added to the envoy protos.
         } catch (InvalidProtocolBufferException e) {
           throw new ResourceInvalidException(
@@ -227,29 +232,52 @@ static class LoadBalancingPolicyConverter {
     }
 
     /**
-     * Converts a custom LB config {@link Any} configuration to service config format.
+     * Converts a custom TypedStruct LB config to service config format.
+     */
+    @SuppressWarnings("unchecked")
+    private static ImmutableMap<String, ?> convertCustomConfig(
+        com.github.xds.type.v3.TypedStruct configTypedStruct)
+        throws ResourceInvalidException {
+      return ImmutableMap.of(parseCustomConfigTypeName(configTypedStruct.getTypeUrl()),
+          (Map<String, ?>) parseCustomConfigJson(configTypedStruct.getValue()));
+    }
+
+    /**
+     * Converts a custom UDPA (legacy) TypedStruct LB config to service config format.
      */
     @SuppressWarnings("unchecked")
-    private static ImmutableMap<String, ?> convertCustomConfig(TypedStruct configTypedStruct)
+    private static ImmutableMap<String, ?> convertCustomConfig(
+        com.github.udpa.udpa.type.v1.TypedStruct configTypedStruct)
+        throws ResourceInvalidException {
+      return ImmutableMap.of(parseCustomConfigTypeName(configTypedStruct.getTypeUrl()),
+          (Map<String, ?>) parseCustomConfigJson(configTypedStruct.getValue()));
+    }
+
+    /**
+     * Print the config Struct into JSON and then parse that into our internal representation.
+     */
+    private static Object parseCustomConfigJson(Struct configStruct)
         throws ResourceInvalidException {
       Object rawJsonConfig = null;
       try {
-        rawJsonConfig = JsonParser.parse(JsonFormat.printer().print(configTypedStruct.getValue()));
+        rawJsonConfig = JsonParser.parse(JsonFormat.printer().print(configStruct));
       } catch (IOException e) {
         throw new ResourceInvalidException("Unable to parse custom LB config JSON", e);
       }
 
       if (!(rawJsonConfig instanceof Map)) {
         throw new ResourceInvalidException("Custom LB config does not contain a JSON object");
       }
+      return rawJsonConfig;
+    }
+
 
-      String customConfigTypeName = configTypedStruct.getTypeUrl();
+    private static String parseCustomConfigTypeName(String customConfigTypeName) {
       if (customConfigTypeName.contains("/")) {
         customConfigTypeName = customConfigTypeName.substring(
             customConfigTypeName.lastIndexOf("/") + 1);
       }
-
-      return ImmutableMap.of(customConfigTypeName, (Map<String, ?>) rawJsonConfig);
+      return customConfigTypeName;
     }
 
     // Used to signal that the LB config goes too deep.

xds/src/test/java/io/grpc/xds/LoadBalancerConfigFactoryTest.java

@@ -81,6 +81,13 @@ public class LoadBalancerConfigFactoryTest {
                       Struct.newBuilder().putFields(CUSTOM_POLICY_FIELD_KEY,
                           Value.newBuilder().setNumberValue(CUSTOM_POLICY_FIELD_VALUE).build()))
                   .build()))).build();
+  private static final Policy CUSTOM_POLICY_UDPA = Policy.newBuilder().setTypedExtensionConfig(
+      TypedExtensionConfig.newBuilder().setTypedConfig(Any.pack(
+          com.github.udpa.udpa.type.v1.TypedStruct.newBuilder().setTypeUrl(
+                  "type.googleapis.com/" + CUSTOM_POLICY_NAME).setValue(
+                  Struct.newBuilder().putFields(CUSTOM_POLICY_FIELD_KEY,
+                      Value.newBuilder().setNumberValue(CUSTOM_POLICY_FIELD_VALUE).build()))
+              .build()))).build();
   private static final FakeCustomLoadBalancerProvider CUSTOM_POLICY_PROVIDER
       = new FakeCustomLoadBalancerProvider();
 
@@ -211,6 +218,18 @@ public void customLbInWrr_providerRegistered() throws ResourceInvalidException {
     assertThat(newLbConfig(cluster, false, true)).isEqualTo(VALID_CUSTOM_CONFIG_IN_WRR);
   }
 
+  // When a provider for the endpoint picking custom policy is available, the configuration should
+  // use it. This one uses the legacy UDPA TypedStruct that is also supported.
+  @Test
+  public void customLbInWrr_providerRegistered_udpa() throws ResourceInvalidException {
+    LoadBalancerRegistry.getDefaultRegistry().register(CUSTOM_POLICY_PROVIDER);
+
+    Cluster cluster = Cluster.newBuilder().setLoadBalancingPolicy(LoadBalancingPolicy.newBuilder()
+        .addPolicies(buildWrrPolicy(CUSTOM_POLICY_UDPA, ROUND_ROBIN_POLICY))).build();
+
+    assertThat(newLbConfig(cluster, false, true)).isEqualTo(VALID_CUSTOM_CONFIG_IN_WRR);
+  }
+
   // When a provider for the custom wrr_locality child policy is NOT available, we should fall back
   // to the round_robin that is also provided.
   @Test