changed orca listener interface to use separate listener interfaces for per-request and out-of-band cases

Author: voidzcy

xds/src/main/java/io/grpc/xds/OrcaUtil.java

@@ -83,33 +83,35 @@ public ClientStreamTracer newClientStreamTracer(StreamInfo info, Metadata header
       };
 
   /**
-   * Creates a new {@link ClientStreamTracer.Factory} with provided {@link OrcaReportListener}
-   * installed to receive callback when a per-request ORCA report is received.
+   * Creates a new {@link ClientStreamTracer.Factory} with provided
+   * {@link OrcaPerRequestReportListener} installed to receive callback when a per-request
+   * ORCA report is received.
    *
    * @param listener contains the callback to be invoked when a per-request ORCA report is
    *     received.
    */
   public static ClientStreamTracer.Factory newOrcaClientStreamTracerFactory(
-      OrcaReportListener listener) {
+      OrcaPerRequestReportListener listener) {
     return newOrcaClientStreamTracerFactory(NOOP_CLIENT_STREAM_TRACER_FACTORY, listener);
   }
 
   /**
-   * Creates a new {@link ClientStreamTracer.Factory} with provided {@link OrcaReportListener}
-   * installed to receive callback when a per-request ORCA report is received.
+   * Creates a new {@link ClientStreamTracer.Factory} with provided
+   * {@link OrcaPerRequestReportListener} installed to receive callback when a per-request
+   * ORCA report is received.
    *
    * @param delegate the delegate factory to produce other client stream tracing.
    * @param listener contains the callback to be invoked when a per-request ORCA report is
    *     received.
    */
   public static ClientStreamTracer.Factory newOrcaClientStreamTracerFactory(
-      ClientStreamTracer.Factory delegate, OrcaReportListener listener) {
+      ClientStreamTracer.Factory delegate, OrcaPerRequestReportListener listener) {
     return new OrcaReportingTracerFactory(delegate, listener);
   }
 
   /**
-   * Creates a new {@link LoadBalancer.Helper} with provided {@link OrcaReportListener} installed to
-   * receive callback when an out-of-band ORCA report is received.
+   * Creates a new {@link LoadBalancer.Helper} with provided {@link OrcaOobReportListener}
+   * installed to receive callback when an out-of-band ORCA report is received.
    *
    * <p>Note the original {@code LoadBalancer} must call returned helper's
    * {@code Helper.createSubchannel()} from its SynchronizationContext, or it will throw.
@@ -124,7 +126,7 @@ public static ClientStreamTracer.Factory newOrcaClientStreamTracerFactory(
    */
   public static OrcaReportingHelperWrapper newOrcaReportingHelperWrapper(
       LoadBalancer.Helper delegate,
-      OrcaReportListener listener,
+      OrcaOobReportListener listener,
       BackoffPolicy.Provider backoffPolicyProvider,
       Supplier<Stopwatch> stopwatchSupplier) {
     final OrcaReportingHelper orcaHelper =
@@ -144,23 +146,38 @@ public Helper asHelper() {
   }
 
   /**
-   * The listener interface for receiving backend ORCA reports. The class that is interested in
-   * processing backend cost metrics implements this interface, and the object created with that
-   * class is registered with a component, using methods in {@link OrcaUtil}. When an ORCA report is
-   * received, that object's {@code onLoadReport} method is invoked.
+   * The listener interface for receiving per-request ORCA reports from backends. The class that
+   * is interested in processing backend cost metrics implements this interface, and the object
+   * created with that class is registered with a component, using methods in {@link OrcaUtil}.
+   * When an ORCA report is received, that object's {@code onLoadReport} method is invoked.
    */
-  public interface OrcaReportListener {
+  public interface OrcaPerRequestReportListener {
 
     /**
-     * Invoked when an ORCA report is received.
+     * Invoked when an per-request ORCA report is received.
      *
-     * <p>For out-of-band reporting, the actual reporting might be more frequently and the reports
-     * might contain more entries of named cost metrics than configured due to other load balancing
-     * polices requesting for more frequent and detailed reports.
+     * @param report load report in the format of ORCA format.
      */
     void onLoadReport(OrcaLoadReport report);
   }
 
+  /**
+   * The listener interface for receiving out-of-band ORCA reports from backends. The class that
+   * is interested in processing backend cost metrics implements this interface, and the object
+   * created with that class is registered with a component, using methods in {@link OrcaUtil}.
+   * When an ORCA report is received, that object's {@code onLoadReport} method is invoked.
+   */
+  public interface OrcaOobReportListener {
+
+    /**
+     * Invoked when an out-of-band ORCA report is received.
+     *
+     * @param subchannel the subchannel over which the connection to a backend is established.
+     * @param report load report in the format of ORCA protocol.
+     */
+    void onLoadReport(Subchannel subchannel, OrcaLoadReport report);
+  }
+
   /**
    * The blueprint for {@link LoadBalancer.Helper} with the capability of allowing {@link
    * LoadBalancer}s interested in receiving out-of-band ORCA reports to update the reporting
@@ -204,10 +221,10 @@ static final class OrcaReportingTracerFactory extends ClientStreamTracer.Factory
     private static final CallOptions.Key<OrcaReportBroker> ORCA_REPORT_BROKER_KEY =
         CallOptions.Key.create("internal-orca-report-broker");
     private final ClientStreamTracer.Factory delegate;
-    private final OrcaReportListener listener;
+    private final OrcaPerRequestReportListener listener;
 
     OrcaReportingTracerFactory(
-        ClientStreamTracer.Factory delegate, OrcaReportListener listener) {
+        ClientStreamTracer.Factory delegate, OrcaPerRequestReportListener listener) {
       this.delegate = checkNotNull(delegate, "delegate");
       this.listener = checkNotNull(listener, "listener");
     }
@@ -250,19 +267,19 @@ public void inboundTrailers(Metadata trailers) {
   }
 
   /**
-   * An {@link OrcaReportBroker} instance holds registered {@link OrcaReportListener}s and invoke
-   * all of them when an {@link OrcaLoadReport} is received.
+   * A container class to hold registered {@link OrcaPerRequestReportListener}s and invoke all
+   * of them when an {@link OrcaLoadReport} is received.
    */
   private static final class OrcaReportBroker {
 
-    private final List<OrcaReportListener> listeners = new ArrayList<>();
+    private final List<OrcaPerRequestReportListener> listeners = new ArrayList<>();
 
-    void addListener(OrcaReportListener listener) {
+    void addListener(OrcaPerRequestReportListener listener) {
       listeners.add(listener);
     }
 
     void onReport(OrcaLoadReport report) {
-      for (OrcaReportListener listener : listeners) {
+      for (OrcaPerRequestReportListener listener : listeners) {
         listener.onLoadReport(report);
       }
     }
@@ -275,12 +292,12 @@ void onReport(OrcaLoadReport report) {
    */
   private static final class OrcaReportingHelper
       extends ForwardingLoadBalancerHelper
-      implements OrcaReportListener {
+      implements OrcaOobReportListener {
 
     private static final CreateSubchannelArgs.Key<OrcaReportingState> ORCA_REPORTING_STATE_KEY =
         CreateSubchannelArgs.Key.create("internal-orca-reporting-state");
     private final LoadBalancer.Helper delegate;
-    private final OrcaReportListener listener;
+    private final OrcaOobReportListener listener;
     private final SynchronizationContext syncContext;
     private final BackoffPolicy.Provider backoffPolicyProvider;
     private final Supplier<Stopwatch> stopwatchSupplier;
@@ -289,7 +306,7 @@ private static final class OrcaReportingHelper
     private OrcaReportingConfig orcaConfig;
 
     OrcaReportingHelper(LoadBalancer.Helper delegate,
-        OrcaReportListener listener,
+        OrcaOobReportListener listener,
         BackoffPolicy.Provider backoffPolicyProvider,
         Supplier<Stopwatch> stopwatchSupplier) {
       this.delegate = checkNotNull(delegate, "delegate");
@@ -315,15 +332,14 @@ public Subchannel createSubchannel(CreateSubchannelArgs args) {
         // create subchannel to.
         orcaState = new OrcaReportingState(
             this,
-            new OrcaReportBroker(),
             args.getStateListener(),
             syncContext,
             delegate().getScheduledExecutorService());
         orcaStates.add(orcaState);
         args = args.toBuilder().addOption(ORCA_REPORTING_STATE_KEY, orcaState).build();
         augmented = true;
       }
-      orcaState.broker.addListener(this);
+      orcaState.listeners.add(this);
       Subchannel subchannel;
       if (augmented) {
         subchannel = super.createSubchannel(args.toBuilder().setStateListener(orcaState).build());
@@ -345,9 +361,9 @@ void setReportingConfig(OrcaReportingConfig config) {
     }
 
     @Override
-    public void onLoadReport(OrcaLoadReport report) {
+    public void onLoadReport(Subchannel subchannel, OrcaLoadReport report) {
       if (orcaConfig != null) {
-        listener.onLoadReport(report);
+        listener.onLoadReport(subchannel, report);
       }
     }
 
@@ -362,7 +378,7 @@ private final class OrcaReportingState implements SubchannelStateListener {
       private final SubchannelStateListener stateListener;
       private final SynchronizationContext syncContext;
       private final ScheduledExecutorService timeService;
-      private final OrcaReportBroker broker;
+      private final List<OrcaOobReportListener> listeners = new ArrayList<>();
       @Nullable
       private Subchannel subchannel;
       @Nullable
@@ -387,12 +403,10 @@ public void run() {
 
       OrcaReportingState(
           OrcaReportingHelper orcaHelper,
-          OrcaReportBroker broker,
           SubchannelStateListener stateListener,
           SynchronizationContext syncContext,
           ScheduledExecutorService timeService) {
         this.orcaHelper = checkNotNull(orcaHelper, "orcaHelper");
-        this.broker = checkNotNull(broker, "broker");
         this.stateListener = checkNotNull(stateListener, "stateListener");
         this.syncContext = checkNotNull(syncContext, "syncContext");
         this.timeService = checkNotNull(timeService, "timeService");
@@ -544,7 +558,9 @@ void handleResponse(OrcaLoadReport response) {
           callHasResponded = true;
           backoffPolicy = null;
           subchannelLogger.log(ChannelLogLevel.DEBUG, "Received an ORCA report: {0}", response);
-          broker.onReport(response);
+          for (OrcaOobReportListener listener : listeners) {
+            listener.onLoadReport(subchannel, response);
+          }
           call.request(1);
         }
 

xds/src/test/java/io/grpc/xds/OrcaUtilOobReportingTest.java

@@ -64,7 +64,7 @@
 import io.grpc.internal.FakeClock;
 import io.grpc.stub.StreamObserver;
 import io.grpc.testing.GrpcCleanupRule;
-import io.grpc.xds.OrcaUtil.OrcaReportListener;
+import io.grpc.xds.OrcaUtil.OrcaOobReportListener;
 import io.grpc.xds.OrcaUtil.OrcaReportingConfig;
 import io.grpc.xds.OrcaUtil.OrcaReportingHelperWrapper;
 import java.net.SocketAddress;
@@ -123,7 +123,7 @@ public void uncaughtException(Thread t, Throwable e) {
   private final FakeClock fakeClock = new FakeClock();
   private final Helper origHelper = mock(Helper.class, delegatesTo(new FakeHelper()));
   @Mock
-  private OrcaReportListener mockOrcaListener;
+  private OrcaOobReportListener mockOrcaListener;
   @Mock private BackoffPolicy.Provider backoffPolicyProvider;
   @Mock private BackoffPolicy backoffPolicy1;
   @Mock private BackoffPolicy backoffPolicy2;
@@ -302,7 +302,7 @@ public void typicalWorkflow() {
       OrcaLoadReport report = OrcaLoadReport.getDefaultInstance();
       serverCall.responseObserver.onNext(report);
       assertLog(subchannel.logs, "DEBUG: Received an ORCA report: " + report);
-      verify(mockOrcaListener, times(i + 1)).onLoadReport(eq(report));
+      verify(mockOrcaListener).onLoadReport(same(subchannel), eq(report));
     }
 
     for (int i = 0; i < NUM_SUBCHANNELS; i++) {
@@ -367,7 +367,7 @@ public void orcReportingDisabledWhenServiceNotImplemented() {
     OrcaLoadReport report = OrcaLoadReport.getDefaultInstance();
     serverCall.responseObserver.onNext(report);
     assertLog(subchannel.logs, "DEBUG: Received an ORCA report: " + report);
-    verify(mockOrcaListener).onLoadReport(eq(report));
+    verify(mockOrcaListener).onLoadReport(same(subchannel), eq(report));
 
     verifyZeroInteractions(backoffPolicyProvider);
   }
@@ -416,7 +416,7 @@ public void orcaReportingStreamClosedAndRetried() {
     OrcaLoadReport report = OrcaLoadReport.getDefaultInstance();
     orcaServiceImp.calls.peek().responseObserver.onNext(report);
     assertLog(subchannel.logs, "DEBUG: Received an ORCA report: " + report);
-    inOrder.verify(mockOrcaListener).onLoadReport(eq(report));
+    inOrder.verify(mockOrcaListener).onLoadReport(same(subchannel), eq(report));
 
     // Server closes the ORCA reporting RPC after a response, will restart immediately.
     orcaServiceImp.calls.poll().responseObserver.onCompleted();
@@ -447,10 +447,10 @@ public void orcaReportingStreamClosedAndRetried() {
    */
   @Test
   public void twoLevelPoliciesReceiveSameReport() {
-    OrcaReportListener parentListener = mockOrcaListener;
+    OrcaOobReportListener parentListener = mockOrcaListener;
     OrcaReportingHelperWrapper parentHelperWrapper = orcaHelperWrapper;
     parentHelperWrapper.setReportingConfig(ORCA_REPORTING_CONFIG);
-    OrcaReportListener childListener = mock(OrcaReportListener.class);
+    OrcaOobReportListener childListener = mock(OrcaOobReportListener.class);
     OrcaReportingHelperWrapper childHelperWrapper =
         OrcaUtil.newOrcaReportingHelperWrapper(
             parentHelperWrapper.asHelper(),
@@ -475,8 +475,8 @@ public void twoLevelPoliciesReceiveSameReport() {
     assertLog(subchannel.logs, "DEBUG: Received an ORCA report: " + report);
     ArgumentCaptor<OrcaLoadReport> parentReportCaptor = ArgumentCaptor.forClass(null);
     ArgumentCaptor<OrcaLoadReport> childReportCaptor = ArgumentCaptor.forClass(null);
-    verify(parentListener).onLoadReport(parentReportCaptor.capture());
-    verify(childListener).onLoadReport(childReportCaptor.capture());
+    verify(parentListener).onLoadReport(same(subchannel), parentReportCaptor.capture());
+    verify(childListener).onLoadReport(same(subchannel), childReportCaptor.capture());
     assertThat(parentReportCaptor.getValue()).isEqualTo(report);
     assertThat(childReportCaptor.getValue()).isSameInstanceAs(parentReportCaptor.getValue());
   }
@@ -491,7 +491,7 @@ public void twoLevelPoliciesReceiveSameReport() {
   public void reportMostEntriesAndMostFrequentIntervalRequested() {
     OrcaReportingHelperWrapper parentHelperWrapper = orcaHelperWrapper;
     parentHelperWrapper.setReportingConfig(ORCA_REPORTING_CONFIG);
-    OrcaReportListener childListener = mock(OrcaReportListener.class);
+    OrcaOobReportListener childListener = mock(OrcaOobReportListener.class);
     OrcaReportingConfig config =
         OrcaReportingConfig.newBuilder()
             .setReportInterval(12, TimeUnit.NANOSECONDS)

xds/src/test/java/io/grpc/xds/OrcaUtilPerRequestReportingTest.java

@@ -28,7 +28,7 @@
 import io.envoyproxy.udpa.data.orca.v1.OrcaLoadReport;
 import io.grpc.ClientStreamTracer;
 import io.grpc.Metadata;
-import io.grpc.xds.OrcaUtil.OrcaReportListener;
+import io.grpc.xds.OrcaUtil.OrcaPerRequestReportListener;
 import io.grpc.xds.OrcaUtil.OrcaReportingTracerFactory;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -50,7 +50,7 @@ public class OrcaUtilPerRequestReportingTest {
    */
   @Test
   public void singlePolicyPerRequestListener() {
-    OrcaReportListener mockListener = mock(OrcaReportListener.class);
+    OrcaPerRequestReportListener mockListener = mock(OrcaPerRequestReportListener.class);
     // Use a mocked noop stream tracer factory as the original stream tracer factory.
     ClientStreamTracer.Factory fakeDelegateFactory = mock(ClientStreamTracer.Factory.class);
     ClientStreamTracer fakeTracer = mock(ClientStreamTracer.class);
@@ -92,26 +92,25 @@ public void singlePolicyPerRequestListener() {
    */
   @Test
   public void twoLevelPoliciesPerRequestListeners() {
-    OrcaReportListener childListener = mock(OrcaReportListener.class);
-    ClientStreamTracer.Factory childFactory =
+    OrcaPerRequestReportListener parentListener = mock(OrcaPerRequestReportListener.class);
+    ClientStreamTracer.Factory parentFactory =
         mock(ClientStreamTracer.Factory.class,
-            delegatesTo(OrcaUtil.newOrcaClientStreamTracerFactory(childListener)));
+            delegatesTo(OrcaUtil.newOrcaClientStreamTracerFactory(parentListener)));
 
-    OrcaReportListener parentListener = mock(OrcaReportListener.class);
-    ClientStreamTracer.Factory parentFactory =
-        OrcaUtil.newOrcaClientStreamTracerFactory(childFactory, parentListener);
-    // Parent factory will augment the StreamInfo with a broker added and pass it to the child
-    // factory.
-    ClientStreamTracer parentTracer =
-        parentFactory.newClientStreamTracer(STREAM_INFO, new Metadata());
+    OrcaPerRequestReportListener childListener = mock(OrcaPerRequestReportListener.class);
+    ClientStreamTracer.Factory childFactory =
+        OrcaUtil.newOrcaClientStreamTracerFactory(parentFactory, childListener);
+    // Child factory will augment the StreamInfo and pass it to the parent factory.
+    ClientStreamTracer childTracer =
+        childFactory.newClientStreamTracer(STREAM_INFO, new Metadata());
     ArgumentCaptor<ClientStreamTracer.StreamInfo> streamInfoCaptor = ArgumentCaptor.forClass(null);
-    verify(childFactory).newClientStreamTracer(streamInfoCaptor.capture(), any(Metadata.class));
-    ClientStreamTracer.StreamInfo childStreamInfo = streamInfoCaptor.getValue();
-    assertThat(childStreamInfo).isNotEqualTo(STREAM_INFO);
+    verify(parentFactory).newClientStreamTracer(streamInfoCaptor.capture(), any(Metadata.class));
+    ClientStreamTracer.StreamInfo parentStreamInfo = streamInfoCaptor.getValue();
+    assertThat(parentStreamInfo).isNotEqualTo(STREAM_INFO);
 
     // When the trailer does not contain ORCA report, no listener callback will be invoked.
     Metadata trailer = new Metadata();
-    parentTracer.inboundTrailers(trailer);
+    childTracer.inboundTrailers(trailer);
     verify(parentListener, never()).onLoadReport(any(OrcaLoadReport.class));
     verify(childListener, never()).onLoadReport(any(OrcaLoadReport.class));
 
@@ -121,7 +120,7 @@ public void twoLevelPoliciesPerRequestListeners() {
     trailer.put(
         OrcaReportingTracerFactory.ORCA_ENDPOINT_LOAD_METRICS_KEY,
         OrcaLoadReport.getDefaultInstance());
-    parentTracer.inboundTrailers(trailer);
+    childTracer.inboundTrailers(trailer);
     ArgumentCaptor<OrcaLoadReport> parentReportCap = ArgumentCaptor.forClass(null);
     ArgumentCaptor<OrcaLoadReport> childReportCap = ArgumentCaptor.forClass(null);
     verify(parentListener).onLoadReport(parentReportCap.capture());