From da4618ace191f08b2b974dbfe585c1362a05fa8b Mon Sep 17 00:00:00 2001
From: Terry Wilson <tmwilson@google.com>
Date: Fri, 17 Nov 2023 15:04:45 -0800
Subject: [PATCH] util: Remove shutdown subchannels from OD tracking (#10683)

An OutlierDetectionLoadBalancer child load balancer might decided to
shut down any subchannel it is tracking. We need to make sure that those
subchannels are removed from the outlier detection tracker map to avoid
a memory leak.
---
 .../util/OutlierDetectionLoadBalancer.java    |  8 ++++
 .../OutlierDetectionLoadBalancerTest.java     | 46 ++++++++++++++++++-
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java b/util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java
index 855e685249..e94811ac0e 100644
--- a/util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java
+++ b/util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java
@@ -257,6 +257,14 @@ public final class OutlierDetectionLoadBalancer extends LoadBalancer {
       super.start(new OutlierDetectionSubchannelStateListener(listener));
     }
 
+    @Override
+    public void shutdown() {
+      if (addressTracker != null) {
+        addressTracker.removeSubchannel(this);
+      }
+      super.shutdown();
+    }
+
     @Override
     public Attributes getAttributes() {
       if (addressTracker != null) {
diff --git a/util/src/test/java/io/grpc/util/OutlierDetectionLoadBalancerTest.java b/util/src/test/java/io/grpc/util/OutlierDetectionLoadBalancerTest.java
index e40a5a15fc..c2a9c15d96 100644
--- a/util/src/test/java/io/grpc/util/OutlierDetectionLoadBalancerTest.java
+++ b/util/src/test/java/io/grpc/util/OutlierDetectionLoadBalancerTest.java
@@ -112,6 +112,8 @@ public class OutlierDetectionLoadBalancerTest {
   @Captor
   private ArgumentCaptor<ConnectivityState> stateCaptor;
 
+  private FakeLoadBalancer fakeChildLb;
+
   private final LoadBalancerProvider mockChildLbProvider = new StandardLoadBalancerProvider(
       "foo_policy") {
     @Override
@@ -123,7 +125,10 @@ public class OutlierDetectionLoadBalancerTest {
       "fake_policy") {
     @Override
     public LoadBalancer newLoadBalancer(Helper helper) {
-      return new FakeLoadBalancer(helper);
+      if (fakeChildLb == null) {
+        fakeChildLb = new FakeLoadBalancer(helper);
+      }
+      return fakeChildLb;
     }
   };
   private final LoadBalancerProvider roundRobinLbProvider = new StandardLoadBalancerProvider(
@@ -266,6 +271,29 @@ public class OutlierDetectionLoadBalancerTest {
     assertThat(task.getDelay(TimeUnit.NANOSECONDS)).isEqualTo(config.intervalNanos);
   }
 
+  /**
+   * The child LB might recreate subchannels leaving the ones we are tracking
+   * orphaned in the address tracker. Make sure subchannels that are shut down get
+   * removed from the tracker.
+   */
+  @Test
+  public void childLbRecreatesSubchannels() {
+    OutlierDetectionLoadBalancerConfig config = new OutlierDetectionLoadBalancerConfig.Builder()
+        .setSuccessRateEjection(new SuccessRateEjection.Builder().build())
+        .setChildPolicy(new PolicySelection(fakeLbProvider, null)).build();
+
+    loadBalancer.acceptResolvedAddresses(buildResolvedAddress(config, servers.get(0)));
+
+    assertThat(loadBalancer.trackerMap).hasSize(1);
+    AddressTracker addressTracker = (AddressTracker) loadBalancer.trackerMap.values().toArray()[0];
+    assertThat(addressTracker).isNotNull();
+    OutlierDetectionSubchannel trackedSubchannel
+        = (OutlierDetectionSubchannel) addressTracker.getSubchannels().toArray()[0];
+
+    fakeChildLb.recreateSubchannels();
+    assertThat(addressTracker.getSubchannels()).doesNotContain(trackedSubchannel);
+  }
+
   /**
    * Outlier detection first enabled, then removed.
    */
@@ -1227,6 +1255,22 @@ public class OutlierDetectionLoadBalancerTest {
     public void shutdown() {
     }
 
+    // Simulates a situation where a load balancer might recreate some of the subchannels it is
+    // tracking even if acceptResolvedAddresses() has not been called.
+    void recreateSubchannels() {
+      List<Subchannel> newSubchannelList = new ArrayList<>(subchannelList.size());
+      for (Subchannel subchannel : subchannelList) {
+        Subchannel newSubchannel = helper
+            .createSubchannel(
+                CreateSubchannelArgs.newBuilder().setAddresses(subchannel.getAddresses()).build());
+        newSubchannel.start(mock(SubchannelStateListener.class));
+        subchannel.shutdown();
+        newSubchannelList.add(newSubchannel);
+      }
+      subchannelList = newSubchannelList;
+      deliverSubchannelState(READY);
+    }
+
     void deliverSubchannelState(ConnectivityState state) {
       SubchannelPicker picker = new SubchannelPicker() {
         @Override