1 files changed, 56 insertions, 39 deletions

diff --git a/src/core/lib/channel/channelz_registry.h b/src/core/lib/channel/channelz_registry.h
index a5a187a054..5d7c936726 100644
--- a/src/core/lib/channel/channelz_registry.h
+++ b/src/core/lib/channel/channelz_registry.h
@@ -22,11 +22,13 @@
 #include <grpc/impl/codegen/port_platform.h>
 
 #include "src/core/lib/channel/channel_trace.h"
+#include "src/core/lib/channel/channelz.h"
 #include "src/core/lib/gprpp/inlined_vector.h"
 
 #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.
@@ -35,26 +37,56 @@ class ChannelzRegistry {
   // To be called in grpc_init()
   static void Init();
 
-  // To be callen in grpc_shutdown();
+  // To be called 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);
+  // Register/Unregister/Get for ChannelNode
+  static intptr_t RegisterChannelNode(ChannelNode* channel_node) {
+    RegistryEntry entry(channel_node, EntityType::kChannelNode);
+    return Default()->InternalRegisterEntry(entry);
+  }
+  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);
   }
 
-  // globally unregisters the object that is associated to uuid.
-  static void Unregister(intptr_t uuid) { Default()->InternalUnregister(uuid); }
+  // Register/Unregister/Get for SubchannelNode
+  static intptr_t RegisterSubchannelNode(SubchannelNode* channel_node) {
+    RegistryEntry entry(channel_node, EntityType::kSubchannelNode);
+    return Default()->InternalRegisterEntry(entry);
+  }
+  static void UnregisterSubchannelNode(intptr_t uuid) {
+    Default()->InternalUnregisterEntry(uuid, EntityType::kSubchannelNode);
+  }
+  static SubchannelNode* GetSubchannelNode(intptr_t uuid) {
+    void* gotten =
+        Default()->InternalGetEntry(uuid, EntityType::kSubchannelNode);
+    return gotten == nullptr ? nullptr : static_cast<SubchannelNode*>(gotten);
+  }
 
-  // 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);
+  // Returns the allocated JSON string that represents the proto
+  // GetTopChannelsResponse as per channelz.proto.
+  static char* GetTopChannels(intptr_t start_channel_id) {
+    return Default()->InternalGetTopChannels(start_channel_id);
   }
 
  private:
+  enum class EntityType {
+    kChannelNode,
+    kSubchannelNode,
+    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 +96,25 @@ 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);
+
+  // 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);
 
-  // private members
+  char* InternalGetTopChannels(intptr_t start_channel_id);
 
   // 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 */