fix memory leaks due to not calling grpc_shutdown in bm_cq_multiple_threads

3 files changed, 18 insertions, 2 deletions

diff --git a/src/core/lib/iomgr/ev_posix.c b/src/core/lib/iomgr/ev_posix.c
index a5ae04cb5b..cff77e641c 100644
--- a/src/core/lib/iomgr/ev_posix.c
+++ b/src/core/lib/iomgr/ev_posix.c
@@ -121,6 +121,10 @@ void grpc_set_event_engine_test_only(
   g_event_engine = ev_engine;
 }
 
+const grpc_event_engine_vtable *grpc_get_event_engine_test_only() {
+  return g_event_engine;
+}
+
 /* Call this only after calling grpc_event_engine_init() */
 const char *grpc_get_poll_strategy_name() { return g_poll_strategy_name; }
 
diff --git a/src/core/lib/iomgr/ev_posix.h b/src/core/lib/iomgr/ev_posix.h
index 54c4f2ee11..0de7333843 100644
--- a/src/core/lib/iomgr/ev_posix.h
+++ b/src/core/lib/iomgr/ev_posix.h
@@ -153,7 +153,9 @@ void grpc_pollset_set_del_fd(grpc_exec_ctx *exec_ctx,
 typedef int (*grpc_poll_function_type)(struct pollfd *, nfds_t, int);
 extern grpc_poll_function_type grpc_poll_function;
 
-/* This should be used for testing purposes ONLY */
+/* WARNING: The following two functions should be used for testing purposes
+ * ONLY */
 void grpc_set_event_engine_test_only(const grpc_event_engine_vtable *);
+const grpc_event_engine_vtable *grpc_get_event_engine_test_only();
 
 #endif /* GRPC_CORE_LIB_IOMGR_EV_POSIX_H */
diff --git a/test/cpp/microbenchmarks/bm_cq_multiple_threads.cc b/test/cpp/microbenchmarks/bm_cq_multiple_threads.cc
index 1e3830a556..f79db15a47 100644
--- a/test/cpp/microbenchmarks/bm_cq_multiple_threads.cc
+++ b/test/cpp/microbenchmarks/bm_cq_multiple_threads.cc
@@ -41,6 +41,7 @@ namespace testing {
 static void* g_tag = (void*)(intptr_t)10;  // Some random number
 static grpc_completion_queue* g_cq;
 static grpc_event_engine_vtable g_vtable;
+static const grpc_event_engine_vtable* g_old_vtable;
 
 static void pollset_shutdown(grpc_exec_ctx* exec_ctx, grpc_pollset* ps,
                              grpc_closure* closure) {
@@ -72,7 +73,7 @@ static grpc_error* pollset_work(grpc_exec_ctx* exec_ctx, grpc_pollset* ps,
                                 grpc_pollset_worker** worker, gpr_timespec now,
                                 gpr_timespec deadline) {
   if (gpr_time_cmp(deadline, gpr_time_0(GPR_CLOCK_MONOTONIC)) == 0) {
-    gpr_log(GPR_ERROR, "no-op");
+    gpr_log(GPR_DEBUG, "no-op");
     return GRPC_ERROR_NONE;
   }
 
@@ -98,7 +99,12 @@ static void init_engine_vtable() {
 
 static void setup() {
   grpc_init();
+
+  /* Override the event engine with our test event engine (g_vtable); but before
+   * that, save the current event engine in g_old_vtable. We will have to set
+   * g_old_vtable back before calling grpc_shutdown() */
   init_engine_vtable();
+  g_old_vtable = grpc_get_event_engine_test_only();
   grpc_set_event_engine_test_only(&g_vtable);
 
   g_cq = grpc_completion_queue_create_for_next(NULL);
@@ -115,6 +121,10 @@ static void teardown() {
   }
 
   grpc_completion_queue_destroy(g_cq);
+
+  /* Restore the old event engine before calling grpc_shutdown */
+  grpc_set_event_engine_test_only(g_old_vtable);
+  grpc_shutdown();
 }
 
 /* A few notes about Multi-threaded benchmarks: