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.

2 files changed, 53 insertions, 1 deletions

diff --git a/util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java b/util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java
index 855e68524..e94811ac0 100644
--- a/util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java
+++ b/util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java
@@ -258,6 +258,14 @@ public final class OutlierDetectionLoadBalancer extends LoadBalancer {
     }
 
     @Override
+    public void shutdown() {
+      if (addressTracker != null) {
+        addressTracker.removeSubchannel(this);
+      }
+      super.shutdown();
+    }
+
+    @Override
     public Attributes getAttributes() {
       if (addressTracker != null) {
         return delegate.getAttributes().toBuilder().set(ADDRESS_TRACKER_ATTR_KEY, addressTracker)
diff --git a/util/src/test/java/io/grpc/util/OutlierDetectionLoadBalancerTest.java b/util/src/test/java/io/grpc/util/OutlierDetectionLoadBalancerTest.java
index e40a5a15f..c2a9c15d9 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(
@@ -267,6 +272,29 @@ public class OutlierDetectionLoadBalancerTest {
   }
 
   /**
+   * 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.
    */
   @Test
@@ -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