grpc-js-xds: Add support for pick_first in xDS config

Author: murgatroid99

packages/grpc-js-xds/gulpfile.ts

@@ -63,6 +63,7 @@ const compile = checkTask(() => execNpmCommand('compile'));
 const runTests = checkTask(() => {
   process.env.GRPC_EXPERIMENTAL_XDS_FEDERATION = 'true';
   process.env.GRPC_EXPERIMENTAL_XDS_CUSTOM_LB_CONFIG = 'true';
+  process.env.GRPC_EXPERIMENTAL_PICKFIRST_LB_CONFIG = 'true';
   if (Number(process.versions.node.split('.')[0]) > 14) {
     process.env.GRPC_XDS_EXPERIMENTAL_ENABLE_RING_HASH = 'true';
   }

packages/grpc-js-xds/src/environment.ts

@@ -21,3 +21,4 @@ export const EXPERIMENTAL_RETRY = (process.env.GRPC_XDS_EXPERIMENTAL_ENABLE_RETR
 export const EXPERIMENTAL_FEDERATION = (process.env.GRPC_EXPERIMENTAL_XDS_FEDERATION ?? 'false') === 'true';
 export const EXPERIMENTAL_CUSTOM_LB_CONFIG = (process.env.GRPC_EXPERIMENTAL_XDS_CUSTOM_LB_CONFIG ?? 'false') === 'true';
 export const EXPERIMENTAL_RING_HASH = (process.env.GRPC_XDS_EXPERIMENTAL_ENABLE_RING_HASH ?? 'false') === 'true';
+export const EXPERIMENTAL_PICK_FIRST = (process.env.GRPC_EXPERIMENTAL_PICKFIRST_LB_CONFIG ?? 'false') === 'true';

packages/grpc-js-xds/src/index.ts

@@ -29,6 +29,7 @@ import * as fault_injection_filter from './http-filter/fault-injection-filter';
 import * as csds from './csds';
 import * as round_robin_lb from './lb-policy-registry/round-robin';
 import * as typed_struct_lb from './lb-policy-registry/typed-struct';
+import * as pick_first_lb from './lb-policy-registry/pick-first';
 
 /**
  * Register the "xds:" name scheme with the @grpc/grpc-js library.
@@ -48,4 +49,5 @@ export function register() {
   csds.setup();
   round_robin_lb.setup();
   typed_struct_lb.setup();
+  pick_first_lb.setup();
 }

packages/grpc-js-xds/src/lb-policy-registry/pick-first.ts

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2023 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.
+ *
+ */
+
+// https://github.com/grpc/proposal/blob/master/A62-pick-first.md#pick_first-via-xds-1
+
+import { LoadBalancingConfig } from "@grpc/grpc-js";
+import { LoadBalancingPolicy__Output } from "../generated/envoy/config/cluster/v3/LoadBalancingPolicy";
+import { TypedExtensionConfig__Output } from "../generated/envoy/config/core/v3/TypedExtensionConfig";
+import { loadProtosWithOptionsSync } from "@grpc/proto-loader/build/src/util";
+import { Any__Output } from "../generated/google/protobuf/Any";
+import { PickFirst__Output } from "../generated/envoy/extensions/load_balancing_policies/pick_first/v3/PickFirst";
+import { EXPERIMENTAL_PICK_FIRST } from "../environment";
+import { registerLbPolicy } from "../lb-policy-registry";
+
+const PICK_FIRST_TYPE_URL = 'type.googleapis.com/envoy.extensions.load_balancing_policies.pick_first.v3.PickFirst';
+
+const resourceRoot = loadProtosWithOptionsSync([
+  'envoy/extensions/load_balancing_policies/pick_first/v3/pick_first.proto'], {
+    keepCase: true,
+    includeDirs: [
+      // Paths are relative to src/build/lb-policy-registry
+      __dirname + '/../../../deps/envoy-api/',
+      __dirname + '/../../../deps/xds/',
+      __dirname + '/../../../deps/protoc-gen-validate'
+    ],
+  }
+);
+
+const toObjectOptions = {
+  longs: String,
+  enums: String,
+  defaults: true,
+  oneofs: true
+}
+
+function decodePickFirstConfig(message: Any__Output): PickFirst__Output {
+  const name = message.type_url.substring(message.type_url.lastIndexOf('/') + 1);
+  const type = resourceRoot.lookup(name);
+  if (type) {
+    const decodedMessage = (type as any).decode(message.value);
+    return decodedMessage.$type.toObject(decodedMessage, toObjectOptions) as PickFirst__Output;
+  } else {
+    throw new Error(`TypedStruct parsing error: unexpected type URL ${message.type_url}`);
+  }
+}
+
+function convertToLoadBalancingPolicy(protoPolicy: TypedExtensionConfig__Output, selectChildPolicy: (childPolicy: LoadBalancingPolicy__Output) => LoadBalancingConfig): LoadBalancingConfig | null {
+  if (protoPolicy.typed_config?.type_url !== PICK_FIRST_TYPE_URL) {
+    throw new Error(`Pick first LB policy parsing error: unexpected type URL ${protoPolicy.typed_config?.type_url}`);
+  }
+  const pickFirstMessage = decodePickFirstConfig(protoPolicy.typed_config);
+  return {
+    pick_first: {
+      shuffleAddressList: pickFirstMessage.shuffle_address_list
+    }
+  };
+}
+
+export function setup() {
+  if (EXPERIMENTAL_PICK_FIRST) {
+    registerLbPolicy(PICK_FIRST_TYPE_URL, convertToLoadBalancingPolicy);
+  }
+}

packages/grpc-js-xds/test/test-custom-lb-policies.ts

@@ -38,6 +38,7 @@ import PickResultType = experimental.PickResultType;
 import createChildChannelControlHelper = experimental.createChildChannelControlHelper;
 import parseLoadBalancingConfig = experimental.parseLoadBalancingConfig;
 import registerLoadBalancerType = experimental.registerLoadBalancerType;
+import { PickFirst } from "../src/generated/envoy/extensions/load_balancing_policies/pick_first/v3/PickFirst";
 
 const LB_POLICY_NAME = 'test.RpcBehaviorLoadBalancer';
 
@@ -297,5 +298,28 @@ describe('Custom LB policies', () => {
         done();
       });
     }, reason => done(reason));
-  })
+  });
+  it('Should handle pick_first', done => {
+    const lbPolicy: PickFirst & AnyExtension = {
+      '@type': 'type.googleapis.com/envoy.extensions.load_balancing_policies.pick_first.v3.PickFirst',
+      shuffle_address_list: true
+    };
+    const cluster = new FakeEdsCluster('cluster1', 'endpoint1', [{backends: [new Backend()], locality:{region: 'region1'}}], lbPolicy);
+    const routeGroup = new FakeRouteGroup('listener1', 'route1', [{cluster: cluster}]);
+    routeGroup.startAllBackends().then(() => {
+      xdsServer.setEdsResource(cluster.getEndpointConfig());
+      xdsServer.setCdsResource(cluster.getClusterConfig());
+      xdsServer.setRdsResource(routeGroup.getRouteConfiguration());
+      xdsServer.setLdsResource(routeGroup.getListener());
+      xdsServer.addResponseListener((typeUrl, responseState) => {
+        if (responseState.state === 'NACKED') {
+          client.stopCalls();
+          assert.fail(`Client NACKED ${typeUrl} resource with message ${responseState.errorMessage}`);
+        }
+      })
+      client = XdsTestClient.createFromServer('listener1', xdsServer);
+      client.sendOneCall(done);
+    }, reason => done(reason));
+  });
+
 });

packages/grpc-js-xds/test/xds-server.ts

@@ -45,6 +45,7 @@ const loadedProtos = loadPackageDefinition(loadSync(
     'envoy/extensions/load_balancing_policies/round_robin/v3/round_robin.proto',
     'envoy/extensions/load_balancing_policies/wrr_locality/v3/wrr_locality.proto',
     'envoy/extensions/load_balancing_policies/ring_hash/v3/ring_hash.proto',
+    'envoy/extensions/load_balancing_policies/pick_first/v3/pick_first.proto',
     'xds/type/v3/typed_struct.proto'
   ],
   {