From f6a0028fe51cca8541db8a1acea30b2b7fb08b22 Mon Sep 17 00:00:00 2001
From: yifeizhuang <yifeizhuang@gmail.com>
Date: Wed, 8 Feb 2023 09:50:32 -0800
Subject: [PATCH] orca: support remove listener in OrcaOobUtil (#9881)

---
 .../java/io/grpc/xds/orca/OrcaOobUtil.java    | 19 +++++++++----
 .../io/grpc/xds/orca/OrcaOobUtilTest.java     | 28 +++++++++++++++++++
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/xds/src/main/java/io/grpc/xds/orca/OrcaOobUtil.java b/xds/src/main/java/io/grpc/xds/orca/OrcaOobUtil.java
index 016c4ba0eb..c90a9f58d3 100644
--- a/xds/src/main/java/io/grpc/xds/orca/OrcaOobUtil.java
+++ b/xds/src/main/java/io/grpc/xds/orca/OrcaOobUtil.java
@@ -180,7 +180,8 @@ public final class OrcaOobUtil {
   /**
    *  Update {@link OrcaOobReportListener} to receive Out-of-Band metrics report for the
    *  particular subchannel connection, and set the configuration of receiving ORCA reports,
-   *  such as the interval of receiving reports.
+   *  such as the interval of receiving reports. Set listener to null to remove listener, and the
+   *  config will have no effect.
    *
    * <p>This method needs to be called from the SynchronizationContext returned by the wrapped
    * helper's {@link Helper#getSynchronizationContext()}.
@@ -194,8 +195,9 @@ public final class OrcaOobUtil {
    * @param subchannel the server connected by this subchannel to receive the metrics.
    *
    * @param listener the callback upon receiving backend metrics from the Out-Of-Band stream.
+   *                 Setting to null to removes the listener from the subchannel.
    *
-   * @param config the configuration to be set.
+   * @param config the configuration to be set. It has no effect when listener is null.
    *
    */
   public static void setListener(Subchannel subchannel, OrcaOobReportListener listener,
@@ -305,18 +307,23 @@ public final class OrcaOobUtil {
             if (oldListener != null) {
               configs.remove(oldListener);
             }
+            if (listener != null) {
+              configs.put(listener, config);
+            }
             orcaSubchannel.reportListener = listener;
-            setReportingConfig(listener, config);
+            setReportingConfig(config);
           }
         });
       }
 
-      private void setReportingConfig(OrcaOobReportListener listener, OrcaReportingConfig config) {
+      private void setReportingConfig(OrcaReportingConfig config) {
         boolean reconfigured = false;
-        configs.put(listener, config);
         // Real reporting interval is the minimum of intervals requested by all participating
         // helpers.
-        if (overallConfig == null) {
+        if (configs.isEmpty()) {
+          overallConfig = null;
+          reconfigured = true;
+        } else if (overallConfig == null) {
           overallConfig = config.toBuilder().build();
           reconfigured = true;
         } else {
diff --git a/xds/src/test/java/io/grpc/xds/orca/OrcaOobUtilTest.java b/xds/src/test/java/io/grpc/xds/orca/OrcaOobUtilTest.java
index 5c8de10e2d..61092856cf 100644
--- a/xds/src/test/java/io/grpc/xds/orca/OrcaOobUtilTest.java
+++ b/xds/src/test/java/io/grpc/xds/orca/OrcaOobUtilTest.java
@@ -546,6 +546,34 @@ public class OrcaOobUtilTest {
     }
   }
 
+  @Test
+  public void removeListener() {
+    Subchannel created = createSubchannel(orcaHelper, 0, Attributes.EMPTY);
+    OrcaOobUtil.setListener(created, null, SHORT_INTERVAL_CONFIG);
+    deliverSubchannelState(0, ConnectivityStateInfo.forNonError(READY));
+    verify(mockStateListeners[0])
+            .onSubchannelState(eq(ConnectivityStateInfo.forNonError(READY)));
+
+    assertThat(orcaServiceImps[0].calls).isEmpty();
+    assertThat(subchannels[0].logs).isEmpty();
+    assertThat(unwrap(created)).isSameInstanceAs(subchannels[0]);
+
+    OrcaOobUtil.setListener(created, mockOrcaListener0, SHORT_INTERVAL_CONFIG);
+    assertThat(orcaServiceImps[0].calls).hasSize(1);
+    assertLog(subchannels[0].logs,
+            "DEBUG: Starting ORCA reporting for " + subchannels[0].getAllAddresses());
+    assertThat(orcaServiceImps[0].calls.peek().request)
+            .isEqualTo(buildOrcaRequestFromConfig(SHORT_INTERVAL_CONFIG));
+
+    OrcaOobUtil.setListener(created, null, null);
+    assertThat(orcaServiceImps[0].calls.poll().cancelled).isTrue();
+    assertThat(orcaServiceImps[0].calls).isEmpty();
+    assertThat(subchannels[0].logs).isEmpty();
+    assertThat(fakeClock.getPendingTasks()).isEmpty();
+    verifyNoMoreInteractions(mockOrcaListener0);
+    verifyNoInteractions(backoffPolicyProvider);
+  }
+
   @Test
   public void updateReportingIntervalBeforeCreatingSubchannel() {
     Subchannel created = createSubchannel(orcaHelper, 0, Attributes.EMPTY);