reviewer feedback

8 files changed, 27 insertions, 26 deletions

diff --git a/src/core/ext/filters/client_channel/connector.h b/src/core/ext/filters/client_channel/connector.h
index 484cc1b3c3..a5ef9bc255 100644
--- a/src/core/ext/filters/client_channel/connector.h
+++ b/src/core/ext/filters/client_channel/connector.h
@@ -49,8 +49,8 @@ typedef struct {
   /** channel arguments (to be passed to the filters) */
   grpc_channel_args* channel_args;
 
-  /** socket node of the connected transport */
-  grpc_core::channelz::SocketNode* socket_node;
+  /** socket node of the connected transport. 0 if not availible */
+  intptr_t socket_uuid;
 } grpc_connect_out_args;
 
 struct grpc_connector_vtable {
diff --git a/src/core/ext/filters/client_channel/subchannel.cc b/src/core/ext/filters/client_channel/subchannel.cc
index e66b0711cf..0817b1dd39 100644
--- a/src/core/ext/filters/client_channel/subchannel.cc
+++ b/src/core/ext/filters/client_channel/subchannel.cc
@@ -826,9 +826,7 @@ static bool publish_transport_locked(grpc_subchannel* c) {
     GRPC_ERROR_UNREF(error);
     return false;
   }
-  intptr_t socket_uuid = c->connecting_result.socket_node == nullptr
-                             ? 0
-                             : c->connecting_result.socket_node->uuid();
+  intptr_t socket_uuid = c->connecting_result.socket_uuid;
   memset(&c->connecting_result, 0, sizeof(c->connecting_result));
 
   if (c->disconnected) {
diff --git a/src/core/ext/transport/chttp2/client/chttp2_connector.cc b/src/core/ext/transport/chttp2/client/chttp2_connector.cc
index 62a07d2ba6..42a2e2e896 100644
--- a/src/core/ext/transport/chttp2/client/chttp2_connector.cc
+++ b/src/core/ext/transport/chttp2/client/chttp2_connector.cc
@@ -117,8 +117,9 @@ static void on_handshake_done(void* arg, grpc_error* error) {
                                           c->args.interested_parties);
     c->result->transport =
         grpc_create_chttp2_transport(args->args, args->endpoint, true);
-    c->result->socket_node =
+    grpc_core::RefCountedPtr<grpc_core::channelz::SocketNode> socket_node =
         grpc_chttp2_transport_get_socket_node(c->result->transport);
+    c->result->socket_uuid = socket_node == nullptr ? 0 : socket_node->uuid();
     GPR_ASSERT(c->result->transport);
     // TODO(roth): We ideally want to wait until we receive HTTP/2
     // settings from the server before we consider the connection
diff --git a/src/core/ext/transport/chttp2/transport/chttp2_transport.cc b/src/core/ext/transport/chttp2/transport/chttp2_transport.cc
index 73e43131a0..9b6574b612 100644
--- a/src/core/ext/transport/chttp2/transport/chttp2_transport.cc
+++ b/src/core/ext/transport/chttp2/transport/chttp2_transport.cc
@@ -3145,11 +3145,11 @@ static const grpc_transport_vtable vtable = {sizeof(grpc_chttp2_stream),
 
 static const grpc_transport_vtable* get_vtable(void) { return &vtable; }
 
-grpc_core::channelz::SocketNode* grpc_chttp2_transport_get_socket_node(
-    grpc_transport* transport) {
+grpc_core::RefCountedPtr<grpc_core::channelz::SocketNode>
+grpc_chttp2_transport_get_socket_node(grpc_transport* transport) {
   grpc_chttp2_transport* t =
       reinterpret_cast<grpc_chttp2_transport*>(transport);
-  return t->channelz_socket.get();
+  return t->channelz_socket;
 }
 
 grpc_transport* grpc_create_chttp2_transport(
diff --git a/src/core/ext/transport/chttp2/transport/chttp2_transport.h b/src/core/ext/transport/chttp2/transport/chttp2_transport.h
index b9929b1662..c22cfb0ad7 100644
--- a/src/core/ext/transport/chttp2/transport/chttp2_transport.h
+++ b/src/core/ext/transport/chttp2/transport/chttp2_transport.h
@@ -36,8 +36,8 @@ grpc_transport* grpc_create_chttp2_transport(
     const grpc_channel_args* channel_args, grpc_endpoint* ep, bool is_client,
     grpc_resource_user* resource_user = nullptr);
 
-grpc_core::channelz::SocketNode* grpc_chttp2_transport_get_socket_node(
-    grpc_transport* transport);
+grpc_core::RefCountedPtr<grpc_core::channelz::SocketNode>
+grpc_chttp2_transport_get_socket_node(grpc_transport* transport);
 
 /// Takes ownership of \a read_buffer, which (if non-NULL) contains
 /// leftover bytes previously read from the endpoint (e.g., by handshakers).
diff --git a/src/core/lib/surface/server.cc b/src/core/lib/surface/server.cc
index 4c63b6bc39..e257caa8f9 100644
--- a/src/core/lib/surface/server.cc
+++ b/src/core/lib/surface/server.cc
@@ -109,7 +109,7 @@ struct channel_data {
   uint32_t registered_method_max_probes;
   grpc_closure finish_destroy_channel_closure;
   grpc_closure channel_connectivity_changed;
-  grpc_core::channelz::SocketNode* socket_node;
+  grpc_core::RefCountedPtr<grpc_core::channelz::SocketNode> socket_node;
 };
 
 typedef struct shutdown_tag {
@@ -462,6 +462,9 @@ static void finish_destroy_channel(void* cd, grpc_error* error) {
   channel_data* chand = static_cast<channel_data*>(cd);
   grpc_server* server = chand->server;
   GRPC_CHANNEL_INTERNAL_UNREF(chand->channel, "server");
+  if (chand->socket_node != nullptr) {
+    chand->socket_node->Unref();
+  }
   server_unref(server);
 }
 
@@ -1155,11 +1158,11 @@ void grpc_server_get_pollsets(grpc_server* server, grpc_pollset*** pollsets,
   *pollsets = server->pollsets;
 }
 
-void grpc_server_setup_transport(grpc_server* s, grpc_transport* transport,
-                                 grpc_pollset* accepting_pollset,
-                                 const grpc_channel_args* args,
-                                 grpc_core::channelz::SocketNode* socket_node,
-                                 grpc_resource_user* resource_user) {
+void grpc_server_setup_transport(
+    grpc_server* s, grpc_transport* transport, grpc_pollset* accepting_pollset,
+    const grpc_channel_args* args,
+    grpc_core::RefCountedPtr<grpc_core::channelz::SocketNode> socket_node,
+    grpc_resource_user* resource_user) {
   size_t num_registered_methods;
   size_t alloc;
   registered_method* rm;
@@ -1261,9 +1264,8 @@ void grpc_server_populate_server_sockets(
   gpr_mu_lock(&s->mu_global);
   channel_data* c = nullptr;
   for (c = s->root_channel_data.next; c != &s->root_channel_data; c = c->next) {
-    grpc_core::channelz::SocketNode* socket_node = c->socket_node;
-    if (socket_node && socket_node->uuid() >= start_idx) {
-      server_sockets->push_back(socket_node);
+    if (c->socket_node != nullptr && c->socket_node->uuid() >= start_idx) {
+      server_sockets->push_back(c->socket_node.get());
     }
   }
   gpr_mu_unlock(&s->mu_global);
diff --git a/src/core/lib/surface/server.h b/src/core/lib/surface/server.h
index 8e8903d76b..393bb24214 100644
--- a/src/core/lib/surface/server.h
+++ b/src/core/lib/surface/server.h
@@ -44,11 +44,11 @@ void grpc_server_add_listener(grpc_server* server, void* listener,
 
 /* Setup a transport - creates a channel stack, binds the transport to the
    server */
-void grpc_server_setup_transport(grpc_server* server, grpc_transport* transport,
-                                 grpc_pollset* accepting_pollset,
-                                 const grpc_channel_args* args,
-                                 grpc_core::channelz::SocketNode* socket_node,
-                                 grpc_resource_user* resource_user = nullptr);
+void grpc_server_setup_transport(
+    grpc_server* server, grpc_transport* transport,
+    grpc_pollset* accepting_pollset, const grpc_channel_args* args,
+    grpc_core::RefCountedPtr<grpc_core::channelz::SocketNode> socket_node,
+    grpc_resource_user* resource_user = nullptr);
 
 /* fills in the uuids of all sockets used for connections on this server */
 void grpc_server_populate_server_sockets(
diff --git a/test/core/end2end/tests/channelz.cc b/test/core/end2end/tests/channelz.cc
index 922783aa0d..49a0bc8011 100644
--- a/test/core/end2end/tests/channelz.cc
+++ b/test/core/end2end/tests/channelz.cc
@@ -260,7 +260,7 @@ static void test_channelz(grpc_end2end_test_config config) {
   gpr_free(json);
 
   json = channelz_server->RenderServerSockets(0);
-  GPR_ASSERT(nullptr != strstr(json, "\"socketRef\":"));
+  GPR_ASSERT(nullptr != strstr(json, "\"end\":true"));
   gpr_free(json);
 
   end_test(&f);