Rewrite registry to know type

src/core/lib/channel/channelz.cc

@@ -92,14 +92,14 @@ ChannelNode::ChannelNode(grpc_channel* channel, size_t channel_tracer_max_nodes)
     : channel_(channel), target_(nullptr), channel_uuid_(-1) {
   trace_.Init(channel_tracer_max_nodes);
   target_ = UniquePtr<char>(grpc_channel_get_target(channel_));
-  channel_uuid_ = ChannelzRegistry::Register(this);
+  channel_uuid_ = ChannelzRegistry::RegisterChannelNode(this);
   gpr_atm_no_barrier_store(&last_call_started_millis_,
                            (gpr_atm)ExecCtx::Get()->Now());
 }
 
 ChannelNode::~ChannelNode() {
   trace_.Destroy();
-  ChannelzRegistry::Unregister(channel_uuid_);
+  ChannelzRegistry::UnregisterChannelNode(channel_uuid_);
 }
 
 void ChannelNode::RecordCallStarted() {

src/core/lib/channel/channelz_registry.cc

@@ -25,10 +25,12 @@
 
 #include <grpc/support/alloc.h>
 #include <grpc/support/log.h>
+#include <grpc/support/sync.h>
 
 #include <cstring>
 
 namespace grpc_core {
+namespace channelz {
 namespace {
 
 // singleton instance of the registry.
@@ -49,12 +51,32 @@ ChannelzRegistry::ChannelzRegistry() { gpr_mu_init(&mu_); }
 
 ChannelzRegistry::~ChannelzRegistry() { gpr_mu_destroy(&mu_); }
 
-void ChannelzRegistry::InternalUnregister(intptr_t uuid) {
+intptr_t ChannelzRegistry::InternalRegisterEntry(const RegistryEntry& entry) {
+  mu_guard guard(&mu_);
+  entities_.push_back(entry);
+  intptr_t uuid = entities_.size();
+  return uuid;
+}
+
+void ChannelzRegistry::InternalUnregisterEntry(intptr_t uuid, EntityType type) {
   GPR_ASSERT(uuid >= 1);
-  gpr_mu_lock(&mu_);
+  mu_guard guard(&mu_);
   GPR_ASSERT(static_cast<size_t>(uuid) <= entities_.size());
-  entities_[uuid - 1] = nullptr;
-  gpr_mu_unlock(&mu_);
+  GPR_ASSERT(entities_[uuid - 1].type == type);
+  entities_[uuid - 1].object = nullptr;
+}
+
+void* ChannelzRegistry::InternalGetEntry(intptr_t uuid, EntityType type) {
+  mu_guard guard(&mu_);
+  if (uuid < 1 || uuid > static_cast<intptr_t>(entities_.size())) {
+    return nullptr;
+  }
+  if (entities_[uuid - 1].type == type) {
+    return entities_[uuid - 1].object;
+  } else {
+    return nullptr;
+  }
 }
 
+}  // namespace channelz
 }  // namespace grpc_core

src/core/lib/channel/channelz_registry.h

@@ -27,6 +27,7 @@
 #include <stdint.h>
 
 namespace grpc_core {
+namespace channelz {
 
 // singleton registry object to track all objects that are needed to support
 // channelz bookkeeping. All objects share globally distributed uuids.
@@ -38,23 +39,31 @@ class ChannelzRegistry {
   // To be callen in grpc_shutdown();
   static void Shutdown();
 
-  // globally registers a channelz Object. Returns its unique uuid
-  template <typename Object>
-  static intptr_t Register(Object* object) {
-    return Default()->InternalRegister(object);
+  static intptr_t RegisterChannelNode(ChannelNode* channel_node) {
+    RegistryEntry entry(channel_node, EntityType::kChannelNode);
+    return Default()->InternalRegisterEntry(entry);
   }
-
-  // globally unregisters the object that is associated to uuid.
-  static void Unregister(intptr_t uuid) { Default()->InternalUnregister(uuid); }
-
-  // if object with uuid has previously been registered, returns the
-  // Object associated with that uuid. Else returns nullptr.
-  template <typename Object>
-  static Object* Get(intptr_t uuid) {
-    return Default()->InternalGet<Object>(uuid);
+  static void UnregisterChannelNode(intptr_t uuid) {
+    Default()->InternalUnregisterEntry(uuid, EntityType::kChannelNode);
+  }
+  static ChannelNode* GetChannelNode(intptr_t uuid) {
+    void* gotten = Default()->InternalGetEntry(uuid, EntityType::kChannelNode);
+    return gotten == nullptr ? nullptr : static_cast<ChannelNode*>(gotten);
   }
 
  private:
+  enum class EntityType {
+    kChannelNode,
+    kUnset,
+  };
+
+  struct RegistryEntry {
+    RegistryEntry(void* object_in, EntityType type_in)
+        : object(object_in), type(type_in) {}
+    void* object;
+    EntityType type;
+  };
+
   GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_NEW
   GPRC_ALLOW_CLASS_TO_USE_NON_PUBLIC_DELETE
 
@@ -64,40 +73,23 @@ class ChannelzRegistry {
   // Returned the singleton instance of ChannelzRegistry;
   static ChannelzRegistry* Default();
 
-  // globally registers a channelz Object. Returns its unique uuid
-  template <typename Object>
-  intptr_t InternalRegister(Object* object) {
-    gpr_mu_lock(&mu_);
-    entities_.push_back(static_cast<void*>(object));
-    intptr_t uuid = entities_.size();
-    gpr_mu_unlock(&mu_);
-    return uuid;
-  }
+  // globally registers an Entry. Returns its unique uuid
+  intptr_t InternalRegisterEntry(const RegistryEntry& entry);
 
-  // globally unregisters the object that is associated to uuid.
-  void InternalUnregister(intptr_t uuid);
-
-  // if object with uuid has previously been registered, returns the
-  // Object associated with that uuid. Else returns nullptr.
-  template <typename Object>
-  Object* InternalGet(intptr_t uuid) {
-    gpr_mu_lock(&mu_);
-    if (uuid < 1 || uuid > static_cast<intptr_t>(entities_.size())) {
-      gpr_mu_unlock(&mu_);
-      return nullptr;
-    }
-    Object* ret = static_cast<Object*>(entities_[uuid - 1]);
-    gpr_mu_unlock(&mu_);
-    return ret;
-  }
+  // globally unregisters the object that is associated to uuid. Also does
+  // sanity check that an object doesn't try to unregister the wrong type.
+  void InternalUnregisterEntry(intptr_t uuid, EntityType type);
 
-  // private members
+  // if object with uuid has previously been registered as the correct type,
+  // returns the void* associated with that uuid. Else returns nullptr.
+  void* InternalGetEntry(intptr_t uuid, EntityType type);
 
   // protects entities_ and uuid_
   gpr_mu mu_;
-  InlinedVector<void*, 20> entities_;
+  InlinedVector<RegistryEntry, 20> entities_;
 };
 
+}  // namespace channelz
 }  // namespace grpc_core
 
 #endif /* GRPC_CORE_LIB_CHANNEL_CHANNELZ_REGISTRY_H */

src/core/lib/surface/init.cc

@@ -127,7 +127,7 @@ void grpc_init(void) {
     grpc_slice_intern_init();
     grpc_mdctx_global_init();
     grpc_channel_init_init();
-    grpc_core::ChannelzRegistry::Init();
+    grpc_core::channelz::ChannelzRegistry::Init();
     grpc_security_pre_init();
     grpc_core::ExecCtx::GlobalInit();
     grpc_iomgr_init();
@@ -176,7 +176,7 @@ void grpc_shutdown(void) {
       grpc_mdctx_global_shutdown();
       grpc_handshaker_factory_registry_shutdown();
       grpc_slice_intern_shutdown();
-      grpc_core::ChannelzRegistry::Shutdown();
+      grpc_core::channelz::ChannelzRegistry::Shutdown();
       grpc_stats_shutdown();
       grpc_core::Fork::GlobalShutdown();
     }

test/core/channel/channelz_registry_test.cc

@@ -19,17 +19,20 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include <grpc/grpc.h>
 #include <gtest/gtest.h>
 
 #include <grpc/support/alloc.h>
 #include <grpc/support/log.h>
 
 #include "src/core/lib/channel/channel_trace.h"
+#include "src/core/lib/channel/channelz.h"
 #include "src/core/lib/channel/channelz_registry.h"
 #include "src/core/lib/gpr/useful.h"
 #include "src/core/lib/gprpp/memory.h"
 #include "src/core/lib/iomgr/exec_ctx.h"
 #include "src/core/lib/json/json.h"
+#include "src/core/lib/surface/channel.h"
 
 #include "test/core/util/test_config.h"
 
@@ -37,27 +40,55 @@
 #include <string.h>
 
 namespace grpc_core {
+namespace channelz {
 namespace testing {
+namespace {
+
+class ChannelFixture {
+ public:
+  ChannelFixture() {
+    grpc_arg client_a[1];
+    client_a[0].type = GRPC_ARG_INTEGER;
+    client_a[0].key = const_cast<char*>(GRPC_ARG_ENABLE_CHANNELZ);
+    client_a[0].value.integer = true;
+    grpc_channel_args client_args = {GPR_ARRAY_SIZE(client_a), client_a};
+    channel_ =
+        grpc_insecure_channel_create("fake_target", &client_args, nullptr);
+  }
+
+  ~ChannelFixture() { grpc_channel_destroy(channel_); }
+
+  grpc_channel* channel() { return channel_; }
+
+ private:
+  grpc_channel* channel_;
+};
+
+}  // namespace
 
 // Tests basic ChannelTrace functionality like construction, adding trace, and
 // lookups by uuid.
 TEST(ChannelzRegistryTest, UuidStartsAboveZeroTest) {
-  int object_to_register;
-  intptr_t uuid = ChannelzRegistry::Register(&object_to_register);
+  ChannelFixture channel;
+  ChannelNode* channelz_channel =
+      grpc_channel_get_channelz_node(channel.channel());
+  intptr_t uuid = ChannelzRegistry::RegisterChannelNode(channelz_channel);
   EXPECT_GT(uuid, 0) << "First uuid chose must be greater than zero. Zero if "
                         "reserved according to "
                         "https://github.com/grpc/proposal/blob/master/"
                         "A14-channelz.md";
-  ChannelzRegistry::Unregister(uuid);
+  ChannelzRegistry::UnregisterChannelNode(uuid);
 }
 
 TEST(ChannelzRegistryTest, UuidsAreIncreasing) {
-  int object_to_register;
+  ChannelFixture channel;
+  ChannelNode* channelz_channel =
+      grpc_channel_get_channelz_node(channel.channel());
   std::vector<intptr_t> uuids;
   uuids.reserve(10);
   for (int i = 0; i < 10; ++i) {
     // reregister the same object. It's ok since we are just testing uuids
-    uuids.push_back(ChannelzRegistry::Register(&object_to_register));
+    uuids.push_back(ChannelzRegistry::RegisterChannelNode(channelz_channel));
   }
   for (size_t i = 1; i < uuids.size(); ++i) {
     EXPECT_LT(uuids[i - 1], uuids[i]) << "Uuids must always be increasing";
@@ -65,60 +96,39 @@ TEST(ChannelzRegistryTest, UuidsAreIncreasing) {
 }
 
 TEST(ChannelzRegistryTest, RegisterGetTest) {
-  int object_to_register = 42;
-  intptr_t uuid = ChannelzRegistry::Register(&object_to_register);
-  int* retrieved = ChannelzRegistry::Get<int>(uuid);
-  EXPECT_EQ(&object_to_register, retrieved);
-}
-
-TEST(ChannelzRegistryTest, MultipleTypeTest) {
-  int int_to_register = 42;
-  intptr_t int_uuid = ChannelzRegistry::Register(&int_to_register);
-  std::string str_to_register = "hello world";
-  intptr_t str_uuid = ChannelzRegistry::Register(&str_to_register);
-  int* retrieved_int = ChannelzRegistry::Get<int>(int_uuid);
-  std::string* retrieved_str = ChannelzRegistry::Get<std::string>(str_uuid);
-  EXPECT_EQ(&int_to_register, retrieved_int);
-  EXPECT_EQ(&str_to_register, retrieved_str);
+  ChannelFixture channel;
+  ChannelNode* channelz_channel =
+      grpc_channel_get_channelz_node(channel.channel());
+  intptr_t uuid = ChannelzRegistry::RegisterChannelNode(channelz_channel);
+  ChannelNode* retrieved = ChannelzRegistry::GetChannelNode(uuid);
+  EXPECT_EQ(channelz_channel, retrieved);
 }
 
 TEST(ChannelzRegistryTest, RegisterManyItems) {
-  int object_to_register = 42;
+  ChannelFixture channel;
+  ChannelNode* channelz_channel =
+      grpc_channel_get_channelz_node(channel.channel());
   for (int i = 0; i < 100; i++) {
-    intptr_t uuid = ChannelzRegistry::Register(&object_to_register);
-    int* retrieved = ChannelzRegistry::Get<int>(uuid);
-    EXPECT_EQ(&object_to_register, retrieved);
-  }
+    intptr_t uuid = ChannelzRegistry::RegisterChannelNode(channelz_channel);
+    ChannelNode* retrieved = ChannelzRegistry::GetChannelNode(uuid);
+    EXPECT_EQ(channelz_channel, retrieved);
   }
-
-namespace {
-class Foo {
- public:
-  int bar;
-};
-}  // namespace
-
-TEST(ChannelzRegistryTest, CustomObjectTest) {
-  Foo* foo = New<Foo>();
-  foo->bar = 1024;
-  intptr_t uuid = ChannelzRegistry::Register(foo);
-  Foo* retrieved = ChannelzRegistry::Get<Foo>(uuid);
-  EXPECT_EQ(foo, retrieved);
-  Delete(foo);
 }
 
 TEST(ChannelzRegistryTest, NullIfNotPresentTest) {
-  int object_to_register = 42;
-  intptr_t uuid = ChannelzRegistry::Register(&object_to_register);
+  ChannelFixture channel;
+  ChannelNode* channelz_channel =
+      grpc_channel_get_channelz_node(channel.channel());
+  intptr_t uuid = ChannelzRegistry::RegisterChannelNode(channelz_channel);
   // try to pull out a uuid that does not exist.
-  int* nonexistant = ChannelzRegistry::Get<int>(uuid + 1);
+  ChannelNode* nonexistant = ChannelzRegistry::GetChannelNode(uuid + 1);
   EXPECT_EQ(nonexistant, nullptr);
-  int* retrieved = ChannelzRegistry::Get<int>(uuid);
-  EXPECT_EQ(object_to_register, *retrieved);
-  EXPECT_EQ(&object_to_register, retrieved);
+  ChannelNode* retrieved = ChannelzRegistry::GetChannelNode(uuid);
+  EXPECT_EQ(channelz_channel, retrieved);
 }
 
 }  // namespace testing
+}  // namespace channelz
 }  // namespace grpc_core
 
 int main(int argc, char** argv) {