From 2ade1dbe89f8d2b029445eeebcf20728b75d4a16 Mon Sep 17 00:00:00 2001
From: Muxi Yan <muxi@users.noreply.github.com>
Date: Mon, 27 Apr 2020 12:23:57 -0700
Subject: [PATCH] Fix client_lb_end2end_test flake

---
 test/cpp/end2end/client_lb_end2end_test.cc | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/test/cpp/end2end/client_lb_end2end_test.cc b/test/cpp/end2end/client_lb_end2end_test.cc
index 295afb312b6..0e526a0ec80 100644
--- a/test/cpp/end2end/client_lb_end2end_test.cc
+++ b/test/cpp/end2end/client_lb_end2end_test.cc
@@ -234,9 +234,6 @@ class ClientLbEnd2endTest : public ::testing::Test {
     for (size_t i = 0; i < servers_.size(); ++i) {
       servers_[i]->Shutdown();
     }
-    // Explicitly destroy all the members so that we can make sure grpc_shutdown
-    // has finished by the end of this function, and thus all the registered
-    // LB policy factories are removed.
     servers_.clear();
     creds_.reset();
     grpc_shutdown_blocking();
@@ -1640,12 +1637,19 @@ class ClientLbInterceptTrailingMetadataTest : public ClientLbEnd2endTest {
  protected:
   void SetUp() override {
     ClientLbEnd2endTest::SetUp();
-    grpc_core::RegisterInterceptRecvTrailingMetadataLoadBalancingPolicy(
-        ReportTrailerIntercepted, this);
+    current_test_instance_ = this;
   }
 
   void TearDown() override { ClientLbEnd2endTest::TearDown(); }
 
+  static void SetUpTestCase() {
+    grpc_init();
+    grpc_core::RegisterInterceptRecvTrailingMetadataLoadBalancingPolicy(
+        ReportTrailerIntercepted, nullptr);
+  }
+
+  static void TearDownTestCase() { grpc_shutdown_blocking(); }
+
   int trailers_intercepted() {
     grpc::internal::MutexLock lock(&mu_);
     return trailers_intercepted_;
@@ -1660,8 +1664,7 @@ class ClientLbInterceptTrailingMetadataTest : public ClientLbEnd2endTest {
   static void ReportTrailerIntercepted(
       void* arg, const grpc_core::LoadBalancingPolicy::BackendMetricData*
                      backend_metric_data) {
-    ClientLbInterceptTrailingMetadataTest* self =
-        static_cast<ClientLbInterceptTrailingMetadataTest*>(arg);
+    ClientLbInterceptTrailingMetadataTest* self = current_test_instance_;
     grpc::internal::MutexLock lock(&self->mu_);
     self->trailers_intercepted_++;
     if (backend_metric_data != nullptr) {
@@ -1684,11 +1687,15 @@ class ClientLbInterceptTrailingMetadataTest : public ClientLbEnd2endTest {
     }
   }
 
+  static ClientLbInterceptTrailingMetadataTest* current_test_instance_;
   grpc::internal::Mutex mu_;
   int trailers_intercepted_ = 0;
   std::unique_ptr<udpa::data::orca::v1::OrcaLoadReport> load_report_;
 };
 
+ClientLbInterceptTrailingMetadataTest*
+    ClientLbInterceptTrailingMetadataTest::current_test_instance_ = nullptr;
+
 TEST_F(ClientLbInterceptTrailingMetadataTest, InterceptsRetriesDisabled) {
   const int kNumServers = 1;
   const int kNumRpcs = 10;