Skip to content

Commit

Permalink
grpc-js: Enable LB policies to override subchannel credentials
Browse files Browse the repository at this point in the history
  • Loading branch information
@murgatroid99
murgatroid99 committed Sep 9, 2024
1 parent 21f4708 commit 57f0a71
Show file tree
Hide file tree
Showing 18 changed files with 88 additions and 60 deletions.
4 changes: 2 additions & 2 deletions packages/grpc-js-xds/interop/xds-interop-client.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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)) {
Expand Down
6 changes: 3 additions & 3 deletions packages/grpc-js-xds/src/load-balancer-cds.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
6 changes: 3 additions & 3 deletions packages/grpc-js-xds/src/load-balancer-priority.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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(
Expand Down
6 changes: 3 additions & 3 deletions packages/grpc-js-xds/src/load-balancer-ring-hash.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
{
Expand Down Expand Up @@ -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);
Expand Down
6 changes: 3 additions & 3 deletions packages/grpc-js-xds/src/load-balancer-weighted-target.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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) {
Expand Down
10 changes: 5 additions & 5 deletions packages/grpc-js-xds/src/load-balancer-xds-cluster-impl.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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)) {
Expand Down Expand Up @@ -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)) {
Expand Down
6 changes: 3 additions & 3 deletions packages/grpc-js-xds/src/load-balancer-xds-cluster-manager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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) {
Expand Down
6 changes: 3 additions & 3 deletions packages/grpc-js-xds/src/load-balancer-xds-cluster-resolver.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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() {
Expand Down
6 changes: 3 additions & 3 deletions packages/grpc-js-xds/src/load-balancer-xds-wrr-locality.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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)) {
Expand Down
6 changes: 3 additions & 3 deletions packages/grpc-js-xds/test/test-custom-lb-policies.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand All @@ -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)) {
Expand Down
6 changes: 4 additions & 2 deletions packages/grpc-js/src/internal-channel.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -349,6 +350,7 @@ export class InternalChannel {
this.resolvingLoadBalancer = new ResolvingLoadBalancer(
this.target,
channelControlHelper,
credentials,
options,
(serviceConfig, configSelector) => {
if (serviceConfig.retryThrottling) {
Expand Down
10 changes: 7 additions & 3 deletions packages/grpc-js/src/load-balancer-child-handler.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand All @@ -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 {
Expand Down Expand Up @@ -86,6 +89,7 @@ export class ChildLoadBalancerHandler implements LoadBalancer {

constructor(
private readonly channelControlHelper: ChannelControlHelper,
private readonly credentials: ChannelCredentials,
private readonly options: ChannelOptions
) {}

Expand Down Expand Up @@ -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;
Expand Down
9 changes: 7 additions & 2 deletions packages/grpc-js/src/load-balancer-outlier-detection.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand Down Expand Up @@ -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);
Expand All @@ -505,6 +509,7 @@ export class OutlierDetectionLoadBalancer implements LoadBalancer {
}
},
}),
credentials,
options
);
this.ejectionTimer = setInterval(() => {}, 0);
Expand Down
Loading

0 comments on commit 57f0a71

Please sign in to comment.