From a6575c3e734876a0eae68cabee40fde8741c7f37 Mon Sep 17 00:00:00 2001
From: Michael Lumish <mlumish@google.com>
Date: Tue, 30 Jul 2024 17:59:08 -0700
Subject: [PATCH] grpc-js: Simplify pick_first behavior

---
 .../grpc-js/src/load-balancer-pick-first.ts   | 66 ++++++++-----------
 1 file changed, 28 insertions(+), 38 deletions(-)

diff --git a/packages/grpc-js/src/load-balancer-pick-first.ts b/packages/grpc-js/src/load-balancer-pick-first.ts
index e042e116..933483fe 100644
--- a/packages/grpc-js/src/load-balancer-pick-first.ts
+++ b/packages/grpc-js/src/load-balancer-pick-first.ts
@@ -320,21 +320,16 @@ export class PickFirstLoadBalancer implements LoadBalancer {
 
   private removeCurrentPick() {
     if (this.currentPick !== null) {
-      /* Unref can cause a state change, which can cause a change in the value
-       * of this.currentPick, so we hold a local reference to make sure that
-       * does not impact this function. */
-      const currentPick = this.currentPick;
-      this.currentPick = null;
-      currentPick.unref();
-      currentPick.removeConnectivityStateListener(this.subchannelStateListener);
+      this.currentPick.removeConnectivityStateListener(this.subchannelStateListener);
       this.channelControlHelper.removeChannelzChild(
-        currentPick.getChannelzRef()
+        this.currentPick.getChannelzRef()
       );
-      if (this.reportHealthStatus) {
-        currentPick.removeHealthStateWatcher(
-          this.pickedSubchannelHealthListener
-        );
-      }
+      this.currentPick.removeHealthStateWatcher(
+        this.pickedSubchannelHealthListener
+      );
+      // Unref last, to avoid triggering listeners
+      this.currentPick.unref();
+      this.currentPick = null;
     }
   }
 
@@ -418,20 +413,25 @@ export class PickFirstLoadBalancer implements LoadBalancer {
     this.connectionDelayTimeout.unref?.();
   }
 
+  /**
+   * Declare that the specified subchannel should be used to make requests.
+   * This functions the same independent of whether subchannel is a member of
+   * this.children and whether it is equal to this.currentPick.
+   * Prerequisite: subchannel.getConnectivityState() === READY.
+   * @param subchannel
+   */
   private pickSubchannel(subchannel: SubchannelInterface) {
-    if (this.currentPick && subchannel.realSubchannelEquals(this.currentPick)) {
-      return;
-    }
     trace('Pick subchannel with address ' + subchannel.getAddress());
     this.stickyTransientFailureMode = false;
-    this.removeCurrentPick();
-    this.currentPick = subchannel;
+    /* Ref before removeCurrentPick and resetSubchannelList to avoid the
+     * refcount dropping to 0 during this process. */
     subchannel.ref();
-    if (this.reportHealthStatus) {
-      subchannel.addHealthStateWatcher(this.pickedSubchannelHealthListener);
-    }
     this.channelControlHelper.addChannelzChild(subchannel.getChannelzRef());
+    this.removeCurrentPick();
     this.resetSubchannelList();
+    subchannel.addConnectivityStateListener(this.subchannelStateListener);
+    subchannel.addHealthStateWatcher(this.pickedSubchannelHealthListener);
+    this.currentPick = subchannel;
     clearTimeout(this.connectionDelayTimeout);
     this.calculateAndReportNewState();
   }
@@ -448,20 +448,11 @@ export class PickFirstLoadBalancer implements LoadBalancer {
 
   private resetSubchannelList() {
     for (const child of this.children) {
-      if (
-        !(
-          this.currentPick &&
-          child.subchannel.realSubchannelEquals(this.currentPick)
-        )
-      ) {
-        /* The connectivity state listener is the same whether the subchannel
-         * is in the list of children or it is the currentPick, so if it is in
-         * both, removing it here would cause problems. In particular, that
-         * always happens immediately after the subchannel is picked. */
-        child.subchannel.removeConnectivityStateListener(
-          this.subchannelStateListener
-        );
-      }
+      /* Always remoev the connectivity state listener. If the subchannel is
+         getting picked, it will be re-added then. */
+      child.subchannel.removeConnectivityStateListener(
+        this.subchannelStateListener
+      );
       /* Refs are counted independently for the children list and the
        * currentPick, so we call unref whether or not the child is the
        * currentPick. Channelz child references are also refcounted, so
@@ -478,15 +469,13 @@ export class PickFirstLoadBalancer implements LoadBalancer {
   }
 
   private connectToAddressList(addressList: SubchannelAddress[]) {
+    trace('connectToAddressList([' + addressList.map(address => subchannelAddressToString(address)) + '])');
     const newChildrenList = addressList.map(address => ({
       subchannel: this.channelControlHelper.createSubchannel(address, {}),
       hasReportedTransientFailure: false,
     }));
-    trace('connectToAddressList([' + addressList.map(address => subchannelAddressToString(address)) + '])');
     for (const { subchannel } of newChildrenList) {
       if (subchannel.getConnectivityState() === ConnectivityState.READY) {
-        this.channelControlHelper.addChannelzChild(subchannel.getChannelzRef());
-        subchannel.addConnectivityStateListener(this.subchannelStateListener);
         this.pickSubchannel(subchannel);
         return;
       }
@@ -522,6 +511,7 @@ export class PickFirstLoadBalancer implements LoadBalancer {
     if (!(lbConfig instanceof PickFirstLoadBalancingConfig)) {
       return;
     }
+    this.requestedResolutionSinceLastUpdate = false;
     /* Previously, an update would be discarded if it was identical to the
      * previous update, to minimize churn. Now the DNS resolver is
      * rate-limited, so that is less of a concern. */