grpc-js: Enable LB policies to override subchannel credentials

murgatroid99 · murgatroid99 · commit 57f0a713992b · 2024-09-09T11:19:50.000-07:00
diff --git a/packages/grpc-js-xds/interop/xds-interop-client.ts b/packages/grpc-js-xds/interop/xds-interop-client.ts
@@ -88,7 +88,7 @@ const RPC_BEHAVIOR_CHILD_CONFIG = parseLoadBalancingConfig({round_robin: {}});
 class RpcBehaviorLoadBalancer implements LoadBalancer {
   private child: ChildLoadBalancerHandler;
   private latestConfig: RpcBehaviorLoadBalancingConfig | null = null;
-  constructor(channelControlHelper: ChannelControlHelper, options: grpc.ChannelOptions) {
+  constructor(channelControlHelper: ChannelControlHelper, credentials: grpc.ChannelCredentials, options: grpc.ChannelOptions) {
     const childChannelControlHelper = createChildChannelControlHelper(channelControlHelper, {
       updateState: (connectivityState, picker) => {
         if (connectivityState === grpc.connectivityState.READY && this.latestConfig) {
@@ -97,7 +97,7 @@ class RpcBehaviorLoadBalancer implements LoadBalancer {
         channelControlHelper.updateState(connectivityState, picker);
       }
     });
-    this.child = new ChildLoadBalancerHandler(childChannelControlHelper, options);
+    this.child = new ChildLoadBalancerHandler(childChannelControlHelper, credentials, options);
   }
   updateAddressList(endpointList: Endpoint[], lbConfig: TypedLoadBalancingConfig, attributes: { [key: string]: unknown; }): void {
     if (!(lbConfig instanceof RpcBehaviorLoadBalancingConfig)) {
diff --git a/packages/grpc-js-xds/src/load-balancer-cds.ts b/packages/grpc-js-xds/src/load-balancer-cds.ts
@@ -15,7 +15,7 @@
  *
  */
 
-import { connectivityState, status, Metadata, logVerbosity, experimental, LoadBalancingConfig, ChannelOptions } from '@grpc/grpc-js';
+import { connectivityState, status, Metadata, logVerbosity, experimental, LoadBalancingConfig, ChannelOptions, ChannelCredentials } from '@grpc/grpc-js';
 import { getSingletonXdsClient, Watcher, XdsClient } from './xds-client';
 import { Cluster__Output } from './generated/envoy/config/cluster/v3/Cluster';
 import Endpoint = experimental.Endpoint;
@@ -155,8 +155,8 @@ export class CdsLoadBalancer implements LoadBalancer {
 
   private updatedChild = false;
 
-  constructor(private readonly channelControlHelper: ChannelControlHelper, options: ChannelOptions) {
-    this.childBalancer = new XdsClusterResolverChildPolicyHandler(channelControlHelper, options);
+  constructor(private readonly channelControlHelper: ChannelControlHelper, credentials: ChannelCredentials, options: ChannelOptions) {
+    this.childBalancer = new XdsClusterResolverChildPolicyHandler(channelControlHelper, credentials, options);
   }
 
   private reportError(errorMessage: string) {
diff --git a/packages/grpc-js-xds/src/load-balancer-priority.ts b/packages/grpc-js-xds/src/load-balancer-priority.ts
@@ -15,7 +15,7 @@
  *
  */
 
-import { connectivityState as ConnectivityState, status as Status, Metadata, logVerbosity as LogVerbosity, experimental, LoadBalancingConfig, ChannelOptions } from '@grpc/grpc-js';
+import { connectivityState as ConnectivityState, status as Status, Metadata, logVerbosity as LogVerbosity, experimental, LoadBalancingConfig, ChannelOptions, ChannelCredentials } from '@grpc/grpc-js';
 import LoadBalancer = experimental.LoadBalancer;
 import ChannelControlHelper = experimental.ChannelControlHelper;
 import registerLoadBalancerType = experimental.registerLoadBalancerType;
@@ -197,7 +197,7 @@ export class PriorityLoadBalancer implements LoadBalancer {
             this.parent.channelControlHelper.requestReresolution();
           }
         }
-      }), parent.options);
+      }), parent.credentials, parent.options);
       this.picker = new QueuePicker(this.childBalancer);
       this.startFailoverTimer();
     }
@@ -323,7 +323,7 @@ export class PriorityLoadBalancer implements LoadBalancer {
 
   private updatesPaused = false;
 
-  constructor(private channelControlHelper: ChannelControlHelper, private options: ChannelOptions) {}
+  constructor(private channelControlHelper: ChannelControlHelper, private credentials: ChannelCredentials, private options: ChannelOptions) {}
 
   private updateState(state: ConnectivityState, picker: Picker) {
     trace(
diff --git a/packages/grpc-js-xds/src/load-balancer-ring-hash.ts b/packages/grpc-js-xds/src/load-balancer-ring-hash.ts
@@ -15,7 +15,7 @@
  *
  */
 
-import { experimental, logVerbosity, connectivityState, status, Metadata, ChannelOptions, LoadBalancingConfig } from '@grpc/grpc-js';
+import { experimental, logVerbosity, connectivityState, status, Metadata, ChannelOptions, LoadBalancingConfig, ChannelCredentials } from '@grpc/grpc-js';
 import { isLocalityEndpoint } from './load-balancer-priority';
 import TypedLoadBalancingConfig = experimental.TypedLoadBalancingConfig;
 import LeafLoadBalancer = experimental.LeafLoadBalancer;
@@ -226,7 +226,7 @@ class RingHashLoadBalancer implements LoadBalancer {
   private currentState: connectivityState = connectivityState.IDLE;
   private ring: RingEntry[] = [];
   private ringHashSizeCap = DEFAULT_RING_SIZE_CAP;
-  constructor(private channelControlHelper: ChannelControlHelper, private options: ChannelOptions) {
+  constructor(private channelControlHelper: ChannelControlHelper, private credentials: ChannelCredentials, private options: ChannelOptions) {
     this.childChannelControlHelper = createChildChannelControlHelper(
       channelControlHelper,
       {
@@ -407,7 +407,7 @@ class RingHashLoadBalancer implements LoadBalancer {
       } else {
         this.leafMap.set(
           endpoint,
-          new LeafLoadBalancer(endpoint, this.childChannelControlHelper, this.options)
+          new LeafLoadBalancer(endpoint, this.childChannelControlHelper, this.credentials, this.options)
         );
       }
       const weight = this.leafWeightMap.get(endpoint);
diff --git a/packages/grpc-js-xds/src/load-balancer-weighted-target.ts b/packages/grpc-js-xds/src/load-balancer-weighted-target.ts
@@ -15,7 +15,7 @@
  *
  */
 
-import { connectivityState as ConnectivityState, status as Status, Metadata, logVerbosity, experimental, LoadBalancingConfig, ChannelOptions } from "@grpc/grpc-js";
+import { connectivityState as ConnectivityState, status as Status, Metadata, logVerbosity, experimental, LoadBalancingConfig, ChannelOptions, ChannelCredentials } from "@grpc/grpc-js";
 import { isLocalityEndpoint, LocalityEndpoint } from "./load-balancer-priority";
 import TypedLoadBalancingConfig = experimental.TypedLoadBalancingConfig;
 import LoadBalancer = experimental.LoadBalancer;
@@ -178,7 +178,7 @@ export class WeightedTargetLoadBalancer implements LoadBalancer {
         updateState: (connectivityState: ConnectivityState, picker: Picker) => {
           this.updateState(connectivityState, picker);
         },
-      }), parent.options);
+      }), parent.credentials, parent.options);
 
       this.picker = new QueuePicker(this.childBalancer);
     }
@@ -243,7 +243,7 @@ export class WeightedTargetLoadBalancer implements LoadBalancer {
   private targetList: string[] = [];
   private updatesPaused = false;
 
-  constructor(private channelControlHelper: ChannelControlHelper, private options: ChannelOptions) {}
+  constructor(private channelControlHelper: ChannelControlHelper, private credentials: ChannelCredentials, private options: ChannelOptions) {}
 
   private maybeUpdateState() {
     if (!this.updatesPaused) {
diff --git a/packages/grpc-js-xds/src/load-balancer-xds-cluster-impl.ts b/packages/grpc-js-xds/src/load-balancer-xds-cluster-impl.ts
@@ -15,7 +15,7 @@
  *
  */
 
-import { experimental, logVerbosity, status as Status, Metadata, connectivityState, ChannelOptions } from "@grpc/grpc-js";
+import { experimental, logVerbosity, status as Status, Metadata, connectivityState, ChannelOptions, ChannelCredentials } from "@grpc/grpc-js";
 import { validateXdsServerConfig, XdsServerConfig } from "./xds-bootstrap";
 import { getSingletonXdsClient, XdsClient, XdsClusterDropStats, XdsClusterLocalityStats } from "./xds-client";
 import { LocalityEndpoint } from "./load-balancer-priority";
@@ -253,13 +253,13 @@ class XdsClusterImplBalancer implements LoadBalancer {
   private clusterDropStats: XdsClusterDropStats | null = null;
   private xdsClient: XdsClient | null = null;
 
-  constructor(private readonly channelControlHelper: ChannelControlHelper, options: ChannelOptions) {
+  constructor(private readonly channelControlHelper: ChannelControlHelper, credentials: ChannelCredentials, options: ChannelOptions) {
       this.childBalancer = new ChildLoadBalancerHandler(createChildChannelControlHelper(channelControlHelper, {
-        createSubchannel: (subchannelAddress, subchannelArgs) => {
+        createSubchannel: (subchannelAddress, subchannelArgs, credentialsOverride) => {
           if (!this.xdsClient || !this.latestConfig || !this.lastestEndpointList) {
             throw new Error('xds_cluster_impl: invalid state: createSubchannel called with xdsClient or latestConfig not populated');
           }
-          const wrapperChild = channelControlHelper.createSubchannel(subchannelAddress, subchannelArgs);
+          const wrapperChild = channelControlHelper.createSubchannel(subchannelAddress, subchannelArgs, credentialsOverride);
           let locality: Locality__Output | null = null;
           for (const endpoint of this.lastestEndpointList) {
             if (endpointHasAddress(endpoint, subchannelAddress)) {
@@ -290,7 +290,7 @@ class XdsClusterImplBalancer implements LoadBalancer {
             channelControlHelper.updateState(connectivityState, picker);
           }
         }
-      }), options);
+      }), credentials, options);
     }
   updateAddressList(endpointList: Endpoint[], lbConfig: TypedLoadBalancingConfig, attributes: { [key: string]: unknown; }): void {
     if (!(lbConfig instanceof XdsClusterImplLoadBalancingConfig)) {
diff --git a/packages/grpc-js-xds/src/load-balancer-xds-cluster-manager.ts b/packages/grpc-js-xds/src/load-balancer-xds-cluster-manager.ts
@@ -15,7 +15,7 @@
  *
  */
 
-import { connectivityState as ConnectivityState, status as Status, experimental, logVerbosity, Metadata, status, ChannelOptions } from "@grpc/grpc-js/";
+import { connectivityState as ConnectivityState, status as Status, experimental, logVerbosity, Metadata, status, ChannelOptions, ChannelCredentials } from "@grpc/grpc-js/";
 
 import TypedLoadBalancingConfig = experimental.TypedLoadBalancingConfig;
 import LoadBalancer = experimental.LoadBalancer;
@@ -131,7 +131,7 @@ class XdsClusterManager implements LoadBalancer {
         updateState: (connectivityState: ConnectivityState, picker: Picker) => {
           this.updateState(connectivityState, picker);
         },
-      }), parent.options);
+      }), parent.credentials, parent.options);
 
       this.picker = new QueuePicker(this.childBalancer);
     }
@@ -167,7 +167,7 @@ class XdsClusterManager implements LoadBalancer {
   // Shutdown is a placeholder value that will never appear in normal operation.
   private currentState: ConnectivityState = ConnectivityState.SHUTDOWN;
   private updatesPaused = false;
-  constructor(private channelControlHelper: ChannelControlHelper, private options: ChannelOptions) {}
+  constructor(private channelControlHelper: ChannelControlHelper, private credentials: ChannelCredentials, private options: ChannelOptions) {}
 
   private maybeUpdateState() {
     if (!this.updatesPaused) {
diff --git a/packages/grpc-js-xds/src/load-balancer-xds-cluster-resolver.ts b/packages/grpc-js-xds/src/load-balancer-xds-cluster-resolver.ts
@@ -15,7 +15,7 @@
  *
  */
 
-import { ChannelOptions, LoadBalancingConfig, Metadata, connectivityState, experimental, logVerbosity, status } from "@grpc/grpc-js";
+import { ChannelCredentials, ChannelOptions, LoadBalancingConfig, Metadata, connectivityState, experimental, logVerbosity, status } from "@grpc/grpc-js";
 import { registerLoadBalancerType } from "@grpc/grpc-js/build/src/load-balancer";
 import { EXPERIMENTAL_DUALSTACK_ENDPOINTS, EXPERIMENTAL_OUTLIER_DETECTION } from "./environment";
 import { Locality__Output } from "./generated/envoy/config/core/v3/Locality";
@@ -249,14 +249,14 @@ export class XdsClusterResolver implements LoadBalancer {
   private xdsClient: XdsClient | null = null;
   private childBalancer: ChildLoadBalancerHandler;
 
-  constructor(private readonly channelControlHelper: ChannelControlHelper, options: ChannelOptions) {
+  constructor(private readonly channelControlHelper: ChannelControlHelper, credentials: ChannelCredentials, options: ChannelOptions) {
     this.childBalancer = new ChildLoadBalancerHandler(experimental.createChildChannelControlHelper(channelControlHelper, {
       requestReresolution: () => {
         for (const entry of this.discoveryMechanismList) {
           entry.resolver?.updateResolution();
         }
       }
-    }), options);
+    }), credentials, options);
   }
 
   private maybeUpdateChild() {
diff --git a/packages/grpc-js-xds/src/load-balancer-xds-wrr-locality.ts b/packages/grpc-js-xds/src/load-balancer-xds-wrr-locality.ts
@@ -17,7 +17,7 @@
 
 // https://github.com/grpc/proposal/blob/master/A52-xds-custom-lb-policies.md
 
-import { ChannelOptions, LoadBalancingConfig, experimental, logVerbosity } from "@grpc/grpc-js";
+import { ChannelCredentials, ChannelOptions, LoadBalancingConfig, experimental, logVerbosity } from "@grpc/grpc-js";
 import { loadProtosWithOptionsSync } from "@grpc/proto-loader/build/src/util";
 import { WeightedTargetRaw } from "./load-balancer-weighted-target";
 import { isLocalityEndpoint } from "./load-balancer-priority";
@@ -73,8 +73,8 @@ class XdsWrrLocalityLoadBalancingConfig implements TypedLoadBalancingConfig {
 
 class XdsWrrLocalityLoadBalancer implements LoadBalancer {
   private childBalancer: ChildLoadBalancerHandler;
-  constructor(private readonly channelControlHelper: ChannelControlHelper, options: ChannelOptions) {
-    this.childBalancer = new ChildLoadBalancerHandler(channelControlHelper, options);
+  constructor(private readonly channelControlHelper: ChannelControlHelper, credentials: ChannelCredentials, options: ChannelOptions) {
+    this.childBalancer = new ChildLoadBalancerHandler(channelControlHelper, credentials, options);
   }
   updateAddressList(endpointList: Endpoint[], lbConfig: TypedLoadBalancingConfig, attributes: { [key: string]: unknown; }): void {
     if (!(lbConfig instanceof XdsWrrLocalityLoadBalancingConfig)) {
diff --git a/packages/grpc-js-xds/test/test-custom-lb-policies.ts b/packages/grpc-js-xds/test/test-custom-lb-policies.ts
@@ -24,7 +24,7 @@ import { ControlPlaneServer } from "./xds-server";
 import * as assert from 'assert';
 import { WrrLocality } from "../src/generated/envoy/extensions/load_balancing_policies/wrr_locality/v3/WrrLocality";
 import { TypedStruct } from "../src/generated/xds/type/v3/TypedStruct";
-import { ChannelOptions, connectivityState, experimental, logVerbosity } from "@grpc/grpc-js";
+import { ChannelCredentials, ChannelOptions, connectivityState, experimental, logVerbosity } from "@grpc/grpc-js";
 
 import TypedLoadBalancingConfig = experimental.TypedLoadBalancingConfig;
 import LoadBalancer = experimental.LoadBalancer;
@@ -84,7 +84,7 @@ const RPC_BEHAVIOR_CHILD_CONFIG = parseLoadBalancingConfig({round_robin: {}});
 class RpcBehaviorLoadBalancer implements LoadBalancer {
   private child: ChildLoadBalancerHandler;
   private latestConfig: RpcBehaviorLoadBalancingConfig | null = null;
-  constructor(channelControlHelper: ChannelControlHelper, options: ChannelOptions) {
+  constructor(channelControlHelper: ChannelControlHelper, credentials: ChannelCredentials, options: ChannelOptions) {
     const childChannelControlHelper = createChildChannelControlHelper(channelControlHelper, {
       updateState: (state, picker) => {
         if (state === connectivityState.READY && this.latestConfig) {
@@ -93,7 +93,7 @@ class RpcBehaviorLoadBalancer implements LoadBalancer {
         channelControlHelper.updateState(state, picker);
       }
     });
-    this.child = new ChildLoadBalancerHandler(childChannelControlHelper, options);
+    this.child = new ChildLoadBalancerHandler(childChannelControlHelper, credentials, options);
   }
   updateAddressList(endpointList: Endpoint[], lbConfig: TypedLoadBalancingConfig, attributes: { [key: string]: unknown; }): void {
     if (!(lbConfig instanceof RpcBehaviorLoadBalancingConfig)) {
diff --git a/packages/grpc-js/src/internal-channel.ts b/packages/grpc-js/src/internal-channel.ts
@@ -294,13 +294,14 @@ export class InternalChannel {
     const channelControlHelper: ChannelControlHelper = {
       createSubchannel: (
         subchannelAddress: SubchannelAddress,
-        subchannelArgs: ChannelOptions
+        subchannelArgs: ChannelOptions,
+        credentialsOverride: ChannelCredentials | null
       ) => {
         const subchannel = this.subchannelPool.getOrCreateSubchannel(
           this.target,
           subchannelAddress,
           Object.assign({}, this.options, subchannelArgs),
-          this.credentials
+          credentialsOverride ?? this.credentials
         );
         subchannel.throttleKeepalive(this.keepaliveTime);
         if (this.channelzEnabled) {
@@ -349,6 +350,7 @@ export class InternalChannel {
     this.resolvingLoadBalancer = new ResolvingLoadBalancer(
       this.target,
       channelControlHelper,
+      credentials,
       options,
       (serviceConfig, configSelector) => {
         if (serviceConfig.retryThrottling) {
diff --git a/packages/grpc-js/src/load-balancer-child-handler.ts b/packages/grpc-js/src/load-balancer-child-handler.ts
@@ -27,6 +27,7 @@ import { ConnectivityState } from './connectivity-state';
 import { Picker } from './picker';
 import type { ChannelRef, SubchannelRef } from './channelz';
 import { SubchannelInterface } from './subchannel-interface';
+import { ChannelCredentials } from './channel-credentials';
 
 const TYPE_NAME = 'child_load_balancer_helper';
 
@@ -40,11 +41,13 @@ export class ChildLoadBalancerHandler implements LoadBalancer {
     constructor(private parent: ChildLoadBalancerHandler) {}
     createSubchannel(
       subchannelAddress: SubchannelAddress,
-      subchannelArgs: ChannelOptions
+      subchannelArgs: ChannelOptions,
+      credentialsOverride: ChannelCredentials | null
     ): SubchannelInterface {
       return this.parent.channelControlHelper.createSubchannel(
         subchannelAddress,
-        subchannelArgs
+        subchannelArgs,
+        credentialsOverride
       );
     }
     updateState(connectivityState: ConnectivityState, picker: Picker): void {
@@ -86,6 +89,7 @@ export class ChildLoadBalancerHandler implements LoadBalancer {
 
   constructor(
     private readonly channelControlHelper: ChannelControlHelper,
+    private readonly credentials: ChannelCredentials,
     private readonly options: ChannelOptions
   ) {}
 
@@ -114,7 +118,7 @@ export class ChildLoadBalancerHandler implements LoadBalancer {
       this.configUpdateRequiresNewPolicyInstance(this.latestConfig, lbConfig)
     ) {
       const newHelper = new this.ChildPolicyHelper(this);
-      const newChild = createLoadBalancer(lbConfig, newHelper, this.options)!;
+      const newChild = createLoadBalancer(lbConfig, newHelper, this.credentials, this.options)!;
       newHelper.setChild(newChild);
       if (this.currentChild === null) {
         this.currentChild = newChild;
diff --git a/packages/grpc-js/src/load-balancer-outlier-detection.ts b/packages/grpc-js/src/load-balancer-outlier-detection.ts
@@ -43,6 +43,7 @@ import {
 } from './subchannel-interface';
 import * as logging from './logging';
 import { LoadBalancingConfig } from './service-config';
+import { ChannelCredentials } from './channel-credentials';
 
 const TRACER_NAME = 'outlier_detection';
 
@@ -469,17 +470,20 @@ export class OutlierDetectionLoadBalancer implements LoadBalancer {
 
   constructor(
     channelControlHelper: ChannelControlHelper,
+    credentials: ChannelCredentials,
     options: ChannelOptions
   ) {
     this.childBalancer = new ChildLoadBalancerHandler(
       createChildChannelControlHelper(channelControlHelper, {
         createSubchannel: (
           subchannelAddress: SubchannelAddress,
-          subchannelArgs: ChannelOptions
+          subchannelArgs: ChannelOptions,
+          credentialsOverride: ChannelCredentials | null
         ) => {
           const originalSubchannel = channelControlHelper.createSubchannel(
             subchannelAddress,
-            subchannelArgs
+            subchannelArgs,
+            credentialsOverride
           );
           const mapEntry =
             this.entryMap.getForSubchannelAddress(subchannelAddress);
@@ -505,6 +509,7 @@ export class OutlierDetectionLoadBalancer implements LoadBalancer {
           }
         },
       }),
+      credentials,
       options
     );
     this.ejectionTimer = setInterval(() => {}, 0);
diff --git a/packages/grpc-js/src/load-balancer-pick-first.ts b/packages/grpc-js/src/load-balancer-pick-first.ts
diff --git a/packages/grpc-js/src/load-balancer-round-robin.ts b/packages/grpc-js/src/load-balancer-round-robin.ts
diff --git a/packages/grpc-js/src/load-balancer.ts b/packages/grpc-js/src/load-balancer.ts
diff --git a/packages/grpc-js/src/resolving-load-balancer.ts b/packages/grpc-js/src/resolving-load-balancer.ts
diff --git a/packages/grpc-js/test/test-pick-first.ts b/packages/grpc-js/test/test-pick-first.ts

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@`
`15`	`15`	`*`
`16`	`16`	`*/`
`17`	`17`
`18`		`-import { connectivityState as ConnectivityState, status as Status, Metadata, logVerbosity as LogVerbosity, experimental, LoadBalancingConfig, ChannelOptions } from '@grpc/grpc-js';`
	`18`	`+import { connectivityState as ConnectivityState, status as Status, Metadata, logVerbosity as LogVerbosity, experimental, LoadBalancingConfig, ChannelOptions, ChannelCredentials } from '@grpc/grpc-js';`
`19`	`19`	`import LoadBalancer = experimental.LoadBalancer;`
`20`	`20`	`import ChannelControlHelper = experimental.ChannelControlHelper;`
`21`	`21`	`import registerLoadBalancerType = experimental.registerLoadBalancerType;`
`@@ -197,7 +197,7 @@ export class PriorityLoadBalancer implements LoadBalancer {`
`197`	`197`	`this.parent.channelControlHelper.requestReresolution();`
`198`	`198`	`}`
`199`	`199`	`}`
`200`		`- }), parent.options);`
	`200`	`+ }), parent.credentials, parent.options);`
`201`	`201`	`this.picker = new QueuePicker(this.childBalancer);`
`202`	`202`	`this.startFailoverTimer();`
`203`	`203`	`}`
`@@ -323,7 +323,7 @@ export class PriorityLoadBalancer implements LoadBalancer {`
`323`	`323`
`324`	`324`	`private updatesPaused = false;`
`325`	`325`
`326`		`- constructor(private channelControlHelper: ChannelControlHelper, private options: ChannelOptions) {}`
	`326`	`+ constructor(private channelControlHelper: ChannelControlHelper, private credentials: ChannelCredentials, private options: ChannelOptions) {}`
`327`	`327`
`328`	`328`	`private updateState(state: ConnectivityState, picker: Picker) {`
`329`	`329`	`trace(`