From ba78e4ff475550e1fa3fd09fbcee326c3c6dfd5a Mon Sep 17 00:00:00 2001
From: Yijie Ma <yijiem@google.com>
Date: Thu, 18 Apr 2024 18:04:10 -0700
Subject: [PATCH] [test] Wait for `RecordEnd` in `FakeClientCallAttemptTracer`
 destruction (#36396)

<!--

If you know who should review your pull request, please assign it to that
person, otherwise the pull request would get assigned randomly.

If your pull request is for a specific language, please add the appropriate
lang label.

-->

Closes #36396

COPYBARA_INTEGRATE_REVIEW=https://github.com/grpc/grpc/pull/36396 from yijiem:xds-cluster-end2end-test-twin 26062ea1b4584d4408cc274d96b90a37fab5c70b
PiperOrigin-RevId: 626207011
---
 src/cpp/ext/otel/otel_plugin.cc                 |  2 +-
 test/core/channel/call_tracer_test.cc           |  1 +
 test/core/util/fake_stats_plugin.h              | 17 +++++++++--------
 .../cpp/end2end/xds/xds_cluster_end2end_test.cc |  7 -------
 4 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/src/cpp/ext/otel/otel_plugin.cc b/src/cpp/ext/otel/otel_plugin.cc
index de527591a0d..5546b21c8a3 100644
--- a/src/cpp/ext/otel/otel_plugin.cc
+++ b/src/cpp/ext/otel/otel_plugin.cc
@@ -104,7 +104,7 @@ class OpenTelemetryPlugin::NPCMetricsKeyValueIterable
                        bool(opentelemetry::nostd::string_view,
                             opentelemetry::common::AttributeValue)>
                            callback) const noexcept override {
-    for (size_t i = 0; i < label_keys_.size(); i++) {
+    for (size_t i = 0; i < label_keys_.size(); ++i) {
       if (!callback(AbslStrViewToOpenTelemetryStrView(label_keys_[i]),
                     AbslStrViewToOpenTelemetryStrView(label_values_[i]))) {
         return false;
diff --git a/test/core/channel/call_tracer_test.cc b/test/core/channel/call_tracer_test.cc
index d83093cb128..4950e8d0276 100644
--- a/test/core/channel/call_tracer_test.cc
+++ b/test/core/channel/call_tracer_test.cc
@@ -94,6 +94,7 @@ TEST_F(CallTracerTest, MultipleClientCallAttemptTracers) {
   attempt_tracer->RecordAnnotation("Test");
   EXPECT_EQ(annotation_logger_,
             std::vector<std::string>({"Test", "Test", "Test"}));
+  attempt_tracer->RecordEnd(gpr_timespec());
 }
 
 TEST_F(CallTracerTest, BasicServerCallTracerTest) {
diff --git a/test/core/util/fake_stats_plugin.h b/test/core/util/fake_stats_plugin.h
index 4b3a4018fa3..8a32e4992c0 100644
--- a/test/core/util/fake_stats_plugin.h
+++ b/test/core/util/fake_stats_plugin.h
@@ -31,6 +31,7 @@
 #include "src/core/lib/channel/metrics.h"
 #include "src/core/lib/channel/promise_based_filter.h"
 #include "src/core/lib/channel/tcp_tracer.h"
+#include "src/core/lib/gprpp/ref_counted.h"
 
 namespace grpc_core {
 
@@ -59,12 +60,12 @@ void RegisterFakeStatsPlugin();
 class FakeClientCallTracer : public ClientCallTracer {
  public:
   class FakeClientCallAttemptTracer
-      : public ClientCallTracer::CallAttemptTracer {
+      : public ClientCallTracer::CallAttemptTracer,
+        public RefCounted<FakeClientCallAttemptTracer> {
    public:
     explicit FakeClientCallAttemptTracer(
         std::vector<std::string>* annotation_logger)
         : annotation_logger_(annotation_logger) {}
-    ~FakeClientCallAttemptTracer() override {}
     void RecordSendInitialMetadata(
         grpc_metadata_batch* /*send_initial_metadata*/) override {}
     void RecordSendTrailingMetadata(
@@ -83,7 +84,7 @@ class FakeClientCallTracer : public ClientCallTracer {
         grpc_metadata_batch* /*recv_trailing_metadata*/,
         const grpc_transport_stream_stats* /*transport_stream_stats*/)
         override {}
-    void RecordEnd(const gpr_timespec& /*latency*/) override {}
+    void RecordEnd(const gpr_timespec& /*latency*/) override { Unref(); }
     void RecordAnnotation(absl::string_view annotation) override {
       annotation_logger_->push_back(std::string(annotation));
     }
@@ -113,9 +114,10 @@ class FakeClientCallTracer : public ClientCallTracer {
       : annotation_logger_(annotation_logger) {}
   ~FakeClientCallTracer() override {}
   CallAttemptTracer* StartNewAttempt(bool /*is_transparent_retry*/) override {
-    call_attempt_tracers_.emplace_back(
-        new FakeClientCallAttemptTracer(annotation_logger_));
-    return call_attempt_tracers_.back().get();
+    auto call_attempt_tracer =
+        MakeRefCounted<FakeClientCallAttemptTracer>(annotation_logger_);
+    call_attempt_tracers_.emplace_back(call_attempt_tracer);
+    return call_attempt_tracer.release();  // Released in RecordEnd().
   }
 
   void RecordAnnotation(absl::string_view annotation) override {
@@ -132,8 +134,7 @@ class FakeClientCallTracer : public ClientCallTracer {
 
  private:
   std::vector<std::string>* annotation_logger_;
-  std::vector<std::unique_ptr<FakeClientCallAttemptTracer>>
-      call_attempt_tracers_;
+  std::vector<RefCountedPtr<FakeClientCallAttemptTracer>> call_attempt_tracers_;
 };
 
 #define GRPC_ARG_INJECT_FAKE_CLIENT_CALL_TRACER_FACTORY \
diff --git a/test/cpp/end2end/xds/xds_cluster_end2end_test.cc b/test/cpp/end2end/xds/xds_cluster_end2end_test.cc
index 41780a96313..d52e767053b 100644
--- a/test/cpp/end2end/xds/xds_cluster_end2end_test.cc
+++ b/test/cpp/end2end/xds/xds_cluster_end2end_test.cc
@@ -355,13 +355,6 @@ TEST_P(CdsTest, MetricLabels) {
                           "mynamespace"),
           ::testing::Pair(OptionalLabelKey::kLocality,
                           LocalityNameString("locality1"))));
-  // TODO(yashkt, yijiem): This shutdown shouldn't actually be necessary. The
-  // only reason it's here is to add a delay before
-  // fake_client_call_tracer_factory goes out of scope, since there may be
-  // lingering callbacks in the call stack that are using the CallAttemptTracer
-  // even after we get here, which would then cause a crash.  Find a cleaner way
-  // to fix this.
-  balancer_->Shutdown();
 }
 
 //