From 80dc06714b9511ed8a8a542ff1d3a6926d2aee0c Mon Sep 17 00:00:00 2001
From: yifeizhuang <zivy@google.com>
Date: Mon, 25 Jul 2022 16:28:58 -0700
Subject: [PATCH] example-orca: fix lost streamTracerFactor (#9411)

---
 ...tomBackendMetricsLoadBalancerProvider.java | 28 ++++++++++++-------
 services/BUILD.bazel                          |  1 +
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/examples/example-orca/src/main/java/io/grpc/examples/orca/CustomBackendMetricsLoadBalancerProvider.java b/examples/example-orca/src/main/java/io/grpc/examples/orca/CustomBackendMetricsLoadBalancerProvider.java
index 3a5e4f97a89..c42fb7cdc10 100644
--- a/examples/example-orca/src/main/java/io/grpc/examples/orca/CustomBackendMetricsLoadBalancerProvider.java
+++ b/examples/example-orca/src/main/java/io/grpc/examples/orca/CustomBackendMetricsLoadBalancerProvider.java
@@ -124,16 +124,24 @@ public LoadBalancer.PickResult pickSubchannel(LoadBalancer.PickSubchannelArgs ar
           return result;
         }
         // Installs ORCA per-query metrics reporting listener.
-        return LoadBalancer.PickResult.withSubchannel(
-            result.getSubchannel(),
-            OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(
-                new OrcaPerRequestUtil.OrcaPerRequestReportListener() {
-                  @Override
-                  public void onLoadReport(MetricReport orcaLoadReport) {
-                    System.out.println("Example load balancer received per-rpc metrics report:\n"
-                        + orcaLoadReport);
-                  }
-                }));
+        final OrcaPerRequestUtil.OrcaPerRequestReportListener orcaListener =
+            new OrcaPerRequestUtil.OrcaPerRequestReportListener() {
+          @Override
+          public void onLoadReport(MetricReport orcaLoadReport) {
+            System.out.println("Example load balancer received per-rpc metrics report:\n"
+                + orcaLoadReport);
+          }
+        };
+        if (result.getStreamTracerFactory() == null) {
+          return LoadBalancer.PickResult.withSubchannel(
+              result.getSubchannel(),
+              OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(orcaListener));
+        } else {
+            return LoadBalancer.PickResult.withSubchannel(
+                result.getSubchannel(),
+                OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(
+                    result.getStreamTracerFactory(), orcaListener));
+        }
       }
     }
   }
diff --git a/services/BUILD.bazel b/services/BUILD.bazel
index 78a953d0882..2854b666ba1 100644
--- a/services/BUILD.bazel
+++ b/services/BUILD.bazel
@@ -43,6 +43,7 @@ java_library(
         "//api",
         "//context",
         "@com_google_code_findbugs_jsr305//jar",
+        "@com_google_errorprone_error_prone_annotations//jar",
         "@com_google_guava_guava//jar",
     ],
 )