From 91ca04cf6f29e9a8d46adff61e9ffc4ccc48350d Mon Sep 17 00:00:00 2001
From: Esun Kim <veblush@google.com>
Date: Thu, 17 Sep 2020 12:41:14 -0700
Subject: [PATCH] Add TestGrpcScope

---
 test/core/util/test_config.cc                 | 33 ++++++++++++++-----
 test/core/util/test_config.h                  | 13 ++++++++
 .../microbenchmarks/bm_opencensus_plugin.cc   | 13 +++++---
 3 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/test/core/util/test_config.cc b/test/core/util/test_config.cc
index 3b5e77c86d7..a98bf7e2b1a 100644
--- a/test/core/util/test_config.cc
+++ b/test/core/util/test_config.cc
@@ -380,6 +380,19 @@ void grpc_test_init(int /*argc*/, char** /*argv*/) {
   srand(seed());
 }
 
+bool grpc_wait_until_shutdown(int64_t time_s) {
+  gpr_timespec deadline = grpc_timeout_seconds_to_deadline(time_s);
+  while (grpc_is_initialized()) {
+    grpc_maybe_wait_for_async_shutdown();
+    gpr_sleep_until(gpr_time_add(gpr_now(GPR_CLOCK_REALTIME),
+                                 gpr_time_from_millis(1, GPR_TIMESPAN)));
+    if (gpr_time_cmp(gpr_now(GPR_CLOCK_MONOTONIC), deadline) > 0) {
+      return false;
+    }
+  }
+  return true;
+}
+
 namespace grpc {
 namespace testing {
 
@@ -390,15 +403,8 @@ TestEnvironment::TestEnvironment(int argc, char** argv) {
 TestEnvironment::~TestEnvironment() {
   // This will wait until gRPC shutdown has actually happened to make sure
   // no gRPC resources (such as thread) are active. (timeout = 10s)
-  gpr_timespec deadline = grpc_timeout_seconds_to_deadline(10);
-  while (grpc_is_initialized()) {
-    grpc_maybe_wait_for_async_shutdown();
-    gpr_sleep_until(gpr_time_add(gpr_now(GPR_CLOCK_REALTIME),
-                                 gpr_time_from_millis(1, GPR_TIMESPAN)));
-    if (gpr_time_cmp(gpr_now(GPR_CLOCK_MONOTONIC), deadline) > 0) {
-      gpr_log(GPR_ERROR, "Timeout in waiting for gRPC shutdown");
-      break;
-    }
+  if (!grpc_wait_until_shutdown(10)) {
+    gpr_log(GPR_ERROR, "Timeout in waiting for gRPC shutdown");
   }
   if (BuiltUnderMsan()) {
     // This is a workaround for MSAN. MSAN doesn't like having shutdown thread
@@ -412,5 +418,14 @@ TestEnvironment::~TestEnvironment() {
   gpr_log(GPR_INFO, "TestEnvironment ends");
 }
 
+TestGrpcScope::TestGrpcScope() { grpc_init(); }
+
+TestGrpcScope::~TestGrpcScope() {
+  grpc_shutdown();
+  if (!grpc_wait_until_shutdown(10)) {
+    gpr_log(GPR_ERROR, "Timeout in waiting for gRPC shutdown");
+  }
+}
+
 }  // namespace testing
 }  // namespace grpc
diff --git a/test/core/util/test_config.h b/test/core/util/test_config.h
index 112af3176f9..49960c67dde 100644
--- a/test/core/util/test_config.h
+++ b/test/core/util/test_config.h
@@ -40,6 +40,10 @@ gpr_timespec grpc_timeout_milliseconds_to_deadline(int64_t time_ms);
 // Prefer TestEnvironment below.
 void grpc_test_init(int argc, char** argv);
 
+// Wait until gRPC is fully shut down.
+// Returns if grpc is shutdown
+bool grpc_wait_until_shutdown(int64_t time_s);
+
 namespace grpc {
 namespace testing {
 
@@ -51,6 +55,15 @@ class TestEnvironment {
   ~TestEnvironment();
 };
 
+// A TestGrpcScope makes sure that
+// - when it's created, gRPC will be initialized
+// - when it's destroyed, gRPC will shutdown and it waits until shutdown
+class TestGrpcScope {
+ public:
+  TestGrpcScope();
+  ~TestGrpcScope();
+};
+
 }  // namespace testing
 }  // namespace grpc
 
diff --git a/test/cpp/microbenchmarks/bm_opencensus_plugin.cc b/test/cpp/microbenchmarks/bm_opencensus_plugin.cc
index 33553073c02..75f5a50bc2d 100644
--- a/test/cpp/microbenchmarks/bm_opencensus_plugin.cc
+++ b/test/cpp/microbenchmarks/bm_opencensus_plugin.cc
@@ -83,8 +83,7 @@ class EchoServerThread final {
 };
 
 static void BM_E2eLatencyCensusDisabled(benchmark::State& state) {
-  grpc::testing::TestEnvironment env(0, {});
-
+  grpc::testing::TestGrpcScope grpc_scope;
   EchoServerThread server;
   std::unique_ptr<grpc::testing::EchoTestService::Stub> stub =
       grpc::testing::EchoTestService::NewStub(grpc::CreateChannel(
@@ -100,14 +99,13 @@ static void BM_E2eLatencyCensusDisabled(benchmark::State& state) {
 BENCHMARK(BM_E2eLatencyCensusDisabled);
 
 static void BM_E2eLatencyCensusEnabled(benchmark::State& state) {
-  grpc::testing::TestEnvironment env(0, {});
-
   // Now start the test by registering the plugin (once in the execution)
   RegisterOnce();
   // This we can safely repeat, and doing so clears accumulated data to avoid
   // initialization costs varying between runs.
   grpc::RegisterOpenCensusViewsForExport();
 
+  grpc::testing::TestGrpcScope grpc_scope;
   EchoServerThread server;
   std::unique_ptr<grpc::testing::EchoTestService::Stub> stub =
       grpc::testing::EchoTestService::NewStub(grpc::CreateChannel(
@@ -122,4 +120,9 @@ static void BM_E2eLatencyCensusEnabled(benchmark::State& state) {
 }
 BENCHMARK(BM_E2eLatencyCensusEnabled);
 
-BENCHMARK_MAIN();
+int main(int argc, char** argv) {
+  grpc::testing::TestEnvironment env(argc, argv);
+  ::benchmark::Initialize(&argc, argv);
+  if (::benchmark::ReportUnrecognizedArguments(argc, argv)) return 1;
+  ::benchmark::RunSpecifiedBenchmarks();
+}