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.

Author: temawi

util/src/main/java/io/grpc/util/OutlierDetectionLoadBalancer.java

@@ -257,6 +257,14 @@ public void start(SubchannelStateListener listener) {
       super.start(new OutlierDetectionSubchannelStateListener(listener));
     }
 
+    @Override
+    public void shutdown() {
+      if (addressTracker != null) {
+        addressTracker.removeSubchannel(this);
+      }
+      super.shutdown();
+    }
+
     @Override
     public Attributes getAttributes() {
       if (addressTracker != null) {

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 LoadBalancer newLoadBalancer(Helper helper) {
       "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 void acceptResolvedAddresses() {
     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 void handleNameResolutionError(Status error) {
     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