grpc-js: Add ChildLoadBalancerHandler and use it for refactoring

packages/grpc-js/src/load-balancer-child-handler.ts

@@ -0,0 +1,144 @@
+/*
+ * Copyright 2020 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.
+ *
+ */
+
+import {
+  LoadBalancer,
+  ChannelControlHelper,
+  createLoadBalancer,
+} from './load-balancer';
+import { SubchannelAddress, Subchannel } from './subchannel';
+import { LoadBalancingConfig } from './load-balancing-config';
+import { ChannelOptions } from './channel-options';
+import { ConnectivityState } from './channel';
+import { Picker } from './picker';
+
+const TYPE_NAME = 'child_load_balancer_helper';
+
+export class ChildLoadBalancerHandler implements LoadBalancer {
+  private currentChild: LoadBalancer | null = null;
+  private pendingChild: LoadBalancer | null = null;
+
+  private ChildPolicyHelper = class {
+    private child: LoadBalancer | null = null;
+    constructor(private parent: ChildLoadBalancerHandler) {}
+    createSubchannel(
+      subchannelAddress: SubchannelAddress,
+      subchannelArgs: ChannelOptions
+    ): Subchannel {
+      return this.parent.channelControlHelper.createSubchannel(
+        subchannelAddress,
+        subchannelArgs
+      );
+    }
+    updateState(connectivityState: ConnectivityState, picker: Picker): void {
+      if (this.calledByPendingChild()) {
+        if (connectivityState !== ConnectivityState.READY) {
+          return;
+        }
+        this.parent.currentChild?.destroy();
+        this.parent.currentChild = this.parent.pendingChild;
+        this.parent.pendingChild = null;
+      } else if (!this.calledByCurrentChild()) {
+        return;
+      }
+      this.parent.channelControlHelper.updateState(connectivityState, picker);
+    }
+    requestReresolution(): void {
+      const latestChild = this.parent.pendingChild ?? this.parent.currentChild;
+      if (this.child === latestChild) {
+        this.parent.channelControlHelper.requestReresolution();
+      }
+    }
+    setChild(newChild: LoadBalancer) {
+      this.child = newChild;
+    }
+    private calledByPendingChild(): boolean {
+      return this.child === this.parent.pendingChild;
+    }
+    private calledByCurrentChild(): boolean {
+      return this.child === this.parent.currentChild;
+    }
+  };
+
+  constructor(private readonly channelControlHelper: ChannelControlHelper) {}
+
+  /**
+   * Prerequisites: lbConfig !== null and lbConfig.name is registered
+   * @param addressList
+   * @param lbConfig
+   * @param attributes
+   */
+  updateAddressList(
+    addressList: SubchannelAddress[],
+    lbConfig: LoadBalancingConfig,
+    attributes: { [key: string]: unknown }
+  ): void {
+    let childToUpdate: LoadBalancer;
+    if (
+      this.currentChild === null ||
+      this.currentChild.getTypeName() !== lbConfig.name
+    ) {
+      const newHelper = new this.ChildPolicyHelper(this);
+      const newChild = createLoadBalancer(lbConfig.name, newHelper)!;
+      newHelper.setChild(newChild);
+      if (this.currentChild === null) {
+        this.currentChild = newChild;
+        childToUpdate = this.currentChild;
+      } else {
+        if (this.pendingChild) {
+          this.pendingChild.destroy();
+        }
+        this.pendingChild = newChild;
+        childToUpdate = this.pendingChild;
+      }
+    } else {
+      if (this.pendingChild === null) {
+        childToUpdate = this.currentChild;
+      } else {
+        childToUpdate = this.pendingChild;
+      }
+    }
+    childToUpdate.updateAddressList(addressList, lbConfig, attributes);
+  }
+  exitIdle(): void {
+    if (this.currentChild) {
+      this.currentChild.resetBackoff();
+      if (this.pendingChild) {
+        this.pendingChild.resetBackoff();
+      }
+    }
+  }
+  resetBackoff(): void {
+    if (this.currentChild) {
+      this.currentChild.resetBackoff();
+      if (this.pendingChild) {
+        this.pendingChild.resetBackoff();
+      }
+    }
+  }
+  destroy(): void {
+    if (this.currentChild) {
+      this.currentChild.destroy();
+    }
+    if (this.pendingChild) {
+      this.pendingChild.destroy();
+    }
+  }
+  getTypeName(): string {
+    return TYPE_NAME;
+  }
+}

packages/grpc-js/src/load-balancer-pick-first.ts

@@ -129,7 +129,7 @@ export class PickFirstLoadBalancer implements LoadBalancer {
    * @param channelControlHelper `ChannelControlHelper` instance provided by
    *     this load balancer's owner.
    */
-  constructor(private channelControlHelper: ChannelControlHelper) {
+  constructor(private readonly channelControlHelper: ChannelControlHelper) {
     this.subchannelStateCounts = {
       [ConnectivityState.CONNECTING]: 0,
       [ConnectivityState.IDLE]: 0,
@@ -384,7 +384,7 @@ export class PickFirstLoadBalancer implements LoadBalancer {
 
   updateAddressList(
     addressList: SubchannelAddress[],
-    lbConfig: LoadBalancingConfig | null
+    lbConfig: LoadBalancingConfig
   ): void {
     // lbConfig has no useful information for pick first load balancing
     /* To avoid unnecessary churn, we only do something with this address list
@@ -436,10 +436,6 @@ export class PickFirstLoadBalancer implements LoadBalancer {
   getTypeName(): string {
     return TYPE_NAME;
   }
-
-  replaceChannelControlHelper(channelControlHelper: ChannelControlHelper) {
-    this.channelControlHelper = channelControlHelper;
-  }
 }
 
 export function setup(): void {

packages/grpc-js/src/load-balancer-round-robin.ts

@@ -94,7 +94,7 @@ export class RoundRobinLoadBalancer implements LoadBalancer {
 
   private currentReadyPicker: RoundRobinPicker | null = null;
 
-  constructor(private channelControlHelper: ChannelControlHelper) {
+  constructor(private readonly channelControlHelper: ChannelControlHelper) {
     this.subchannelStateCounts = {
       [ConnectivityState.CONNECTING]: 0,
       [ConnectivityState.IDLE]: 0,
@@ -188,7 +188,7 @@ export class RoundRobinLoadBalancer implements LoadBalancer {
 
   updateAddressList(
     addressList: SubchannelAddress[],
-    lbConfig: LoadBalancingConfig | null
+    lbConfig: LoadBalancingConfig
   ): void {
     this.resetSubchannelList();
     trace(
@@ -228,11 +228,6 @@ export class RoundRobinLoadBalancer implements LoadBalancer {
   getTypeName(): string {
     return TYPE_NAME;
   }
-  replaceChannelControlHelper(
-    channelControlHelper: ChannelControlHelper
-  ): void {
-    this.channelControlHelper = channelControlHelper;
-  }
 }
 
 export function setup() {

packages/grpc-js/src/load-balancer.ts

@@ -67,7 +67,7 @@ export interface LoadBalancer {
    */
   updateAddressList(
     addressList: SubchannelAddress[],
-    lbConfig: LoadBalancingConfig | null,
+    lbConfig: LoadBalancingConfig,
     attributes: { [key: string]: unknown }
   ): void;
   /**
@@ -91,11 +91,6 @@ export interface LoadBalancer {
    * balancer implementation class was registered with.
    */
   getTypeName(): string;
-  /**
-   * Replace the existing ChannelControlHelper with a new one
-   * @param channelControlHelper The new ChannelControlHelper to use from now on
-   */
-  replaceChannelControlHelper(channelControlHelper: ChannelControlHelper): void;
 }
 
 export interface LoadBalancerConstructor {
@@ -128,6 +123,17 @@ export function isLoadBalancerNameRegistered(typeName: string): boolean {
   return typeName in registeredLoadBalancerTypes;
 }
 
+export function getFirstUsableConfig(
+  configs: LoadBalancingConfig[]
+): LoadBalancingConfig | null {
+  for (const config of configs) {
+    if (config.name in registeredLoadBalancerTypes) {
+      return config;
+    }
+  }
+  return null;
+}
+
 export function registerAll() {
   load_balancer_pick_first.setup();
   load_balancer_round_robin.setup();

packages/grpc-js/src/load-balancing-config.ts

@@ -25,6 +25,8 @@
  * runtime */
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
+export type PickFirstConfig = {};
+
 export type RoundRobinConfig = {};
 
 export interface XdsConfig {
@@ -37,11 +39,69 @@ export interface GrpcLbConfig {
   childPolicy: LoadBalancingConfig[];
 }
 
-export interface LoadBalancingConfig {
-  /* Exactly one of these must be set for a config to be valid */
-  round_robin?: RoundRobinConfig;
-  xds?: XdsConfig;
-  grpclb?: GrpcLbConfig;
+export interface PriorityChild {
+  config: LoadBalancingConfig[];
+}
+
+export interface PriorityLbConfig {
+  children: Map<string, PriorityChild>;
+  priorities: string[];
+}
+
+export interface PickFirstLoadBalancingConfig {
+  name: 'pick_first';
+  pick_first: PickFirstConfig;
+}
+
+export interface RoundRobinLoadBalancingConfig {
+  name: 'round_robin';
+  round_robin: RoundRobinConfig;
+}
+
+export interface XdsLoadBalancingConfig {
+  name: 'xds';
+  xds: XdsConfig;
+}
+
+export interface GrpcLbLoadBalancingConfig {
+  name: 'grpclb';
+  grpclb: GrpcLbConfig;
+}
+
+export interface PriorityLoadBalancingConfig {
+  name: 'priority';
+  priority: PriorityLbConfig;
+}
+
+export type LoadBalancingConfig =
+  | PickFirstLoadBalancingConfig
+  | RoundRobinLoadBalancingConfig
+  | XdsLoadBalancingConfig
+  | GrpcLbLoadBalancingConfig
+  | PriorityLoadBalancingConfig;
+
+export function isRoundRobinLoadBalancingConfig(
+  lbconfig: LoadBalancingConfig
+): lbconfig is RoundRobinLoadBalancingConfig {
+  return lbconfig.name === 'round_robin';
+}
+
+export function isXdsLoadBalancingConfig(
+  lbconfig: LoadBalancingConfig
+): lbconfig is XdsLoadBalancingConfig {
+  return lbconfig.name === 'xds';
+}
+
+export function isGrpcLbLoadBalancingConfig(
+  lbconfig: LoadBalancingConfig
+): lbconfig is GrpcLbLoadBalancingConfig {
+  return lbconfig.name === 'grpclb';
+}
+
+export function isPriorityLoadBalancingConfig(
+  lbconfig: LoadBalancingConfig
+): lbconfig is PriorityLoadBalancingConfig {
+  return lbconfig.name === 'priority';
 }
 
 /* In these functions we assume the input came from a JSON object. Therefore we
@@ -97,17 +157,26 @@ export function validateConfig(obj: any): LoadBalancingConfig {
       throw new Error('Multiple load balancing policies configured');
     }
     if (obj['round_robin'] instanceof Object) {
-      return { round_robin: {} };
+      return {
+        name: 'round_robin',
+        round_robin: {},
+      };
     }
   }
   if ('xds' in obj) {
     if ('grpclb' in obj) {
       throw new Error('Multiple load balancing policies configured');
     }
-    return { xds: validateXdsConfig(obj.xds) };
+    return {
+      name: 'xds',
+      xds: validateXdsConfig(obj.xds),
+    };
   }
   if ('grpclb' in obj) {
-    return { grpclb: validateGrpcLbConfig(obj.grpclb) };
+    return {
+      name: 'grpclb',
+      grpclb: validateGrpcLbConfig(obj.grpclb),
+    };
   }
   throw new Error('No recognized load balancing policy configured');
 }