Fix zookeeper resolver plugin bugs

4 files changed, 36 insertions, 16 deletions

diff --git a/include/grpc/grpc.h b/include/grpc/grpc.h
index fbe6be99e9..943982554a 100644
--- a/include/grpc/grpc.h
+++ b/include/grpc/grpc.h
@@ -362,6 +362,13 @@ typedef struct grpc_op {
     deinitialization order isn't guaranteed. */
 void grpc_register_plugin(void (*init)(void), void (*deinit)(void));
 
+/** Frees the memory used by all the plugin information.
+
+    While grpc_init and grpc_shutdown can be called multiple times, the plugins
+    won't be unregistered and their memory cleaned up unless you call that
+    function. Using atexit(grpc_unregister_all_plugins) is a valid method. */
+void grpc_unregister_all_plugins();
+
 /* Propagation bits: this can be bitwise or-ed to form propagation_mask for
  * grpc_call */
 /** Propagate deadline */
diff --git a/src/core/client_config/resolvers/zookeeper_resolver.c b/src/core/client_config/resolvers/zookeeper_resolver.c
index ea4f0f996d..96c59eacf9 100644
--- a/src/core/client_config/resolvers/zookeeper_resolver.c
+++ b/src/core/client_config/resolvers/zookeeper_resolver.c
@@ -267,7 +267,6 @@ static char *zookeeper_parse_address(char *buffer, int buffer_len) {
     }
     if (host != NULL && port != NULL) {
       gpr_asprintf(&address, "%s:%s", host, port);
-      gpr_log(GPR_DEBUG, address);
     }
     grpc_json_destroy(json);
   }
diff --git a/src/core/surface/init.c b/src/core/surface/init.c
index 10f3634f93..053878a87b 100644
--- a/src/core/surface/init.c
+++ b/src/core/surface/init.c
@@ -84,6 +84,17 @@ void grpc_register_plugin(void (*init)(void), void (*deinit)(void)) {
   g_plugins_head->next = old_head;
 }
 
+void grpc_unregister_all_plugins() {
+  grpc_plugin *plugin;
+  grpc_plugin *next;
+
+  for (plugin = g_plugins_head; plugin != NULL; plugin = next) {
+    next = plugin->next;
+    gpr_free(plugin);
+  }
+  g_plugins_head = NULL;
+}
+
 void grpc_init(void) {
   grpc_plugin *plugin;
   gpr_once_init(&g_basic_init, do_basic_init);
@@ -139,7 +150,6 @@ void grpc_shutdown(void) {
         plugin->deinit();
       }
       next = plugin->next;
-      gpr_free(plugin);
     }
   }
   gpr_mu_unlock(&g_init_mu);
diff --git a/test/cpp/end2end/zookeeper_test.cc b/test/cpp/end2end/zookeeper_test.cc
index a00828de5f..27adcf9a96 100644
--- a/test/cpp/end2end/zookeeper_test.cc
+++ b/test/cpp/end2end/zookeeper_test.cc
@@ -75,8 +75,9 @@ class ZookeeperTest : public ::testing::Test {
 
     // Setup two servers
     int port1 = grpc_pick_unused_port_or_die();
-    int port2 = grpc_pick_unused_port_or_die();
     server1_ = SetUpServer(port1);
+
+    int port2 = grpc_pick_unused_port_or_die();
     server2_ = SetUpServer(port2);
 
     // Register service /test in zookeeper
@@ -93,17 +94,7 @@ class ZookeeperTest : public ::testing::Test {
     RegisterService("/test/2", value);
   }
 
-  std::unique_ptr<Server> SetUpServer(int port) {
-    string server_address = "localhost:" + std::to_string(port);
-
-    ServerBuilder builder;
-    builder.AddListeningPort(server_address, InsecureServerCredentials());
-    builder.RegisterService(&service_);
-    std::unique_ptr<Server> server = builder.BuildAndStart();
-    return server;
-  }
-
-  // Require zookeeper server running beforehand
+  // Require zookeeper server running
   void SetUpZookeeper() {
     // Find zookeeper server address in environment
     // Default is localhost:2181
@@ -124,6 +115,19 @@ class ZookeeperTest : public ::testing::Test {
 
     // Register zookeeper name resolver in grpc
     grpc_zookeeper_register();
+
+    // Unregister all plugins when exit
+    atexit(grpc_unregister_all_plugins);
+  }
+
+  std::unique_ptr<Server> SetUpServer(int port) {
+    string server_address = "localhost:" + std::to_string(port);
+
+    ServerBuilder builder;
+    builder.AddListeningPort(server_address, InsecureServerCredentials());
+    builder.RegisterService(&service_);
+    std::unique_ptr<Server> server = builder.BuildAndStart();
+    return server;
   }
 
   void RegisterService(string name, string value) {
@@ -169,7 +173,7 @@ class ZookeeperTest : public ::testing::Test {
 };
 
 // Test zookeeper state change between two RPCs
-// TODO(ctiller): Handle leaked objects
+// TODO(ctiller): leaked objects
 TEST_F(ZookeeperTest, ZookeeperStateChangeTwoRpc) {
   ResetStub();
 
@@ -184,8 +188,8 @@ TEST_F(ZookeeperTest, ZookeeperStateChangeTwoRpc) {
   EXPECT_TRUE(s1.ok());
 
   // Zookeeper state change
-  gpr_log(GPR_DEBUG, "Zookeeper state change");
   DeleteService("/test/1");
+  gpr_log(GPR_DEBUG, "Zookeeper state change");
   sleep(1);
 
   // Second RPC