PL-300: Adding InfoClassSchema Class that represents the configuratio…

…n state for each table

Summary:
Adding InfoClassSchema implementation. This class contains all the
information regarding the tables that will be created and the elements within
the table that will be subscribed too.

Test Plan: Added unit tests

Reviewers: #engineering, zasgar, oazizi

Reviewed By: #engineering, zasgar, oazizi

Differential Revision: https://phab.corp.pixielabs.ai/D247

GitOrigin-RevId: 9ece889

.gitignore

@@ -31,6 +31,7 @@ bazel-*
 # VScode files
 **/.vscode/tags
 **/.vscode/launch.json
+**/.vscode/settings.json
 **/.vscode/tasks.json
 **/.vscode/.vs_code_bazel_build/**
 

src/data_collector/BUILD.bazel

@@ -1,6 +1,22 @@
 package(default_visibility = ["//visibility:public"])
 
-load("//bazel:pl_build_system.bzl", "pl_cc_binary")
+load("//bazel:pl_build_system.bzl", "pl_cc_binary", "pl_cc_library", "pl_cc_test")
+
+pl_cc_library(
+    name = "cc_library",
+    srcs = glob(
+        ["*.cc"],
+        exclude = [
+            "**/*_test.cc",
+            "**/*_benchmark.cc",
+        ],
+    ),
+    hdrs = glob(["*.h"]),
+    deps = [
+        "//src/common/types:types_cc_proto",
+        "//src/data_collector/proto:collector_config_cc_proto",
+    ],
+)
 
 pl_cc_binary(
     name = "data_collector_benchmark",
@@ -11,3 +27,9 @@ pl_cc_binary(
         "@com_google_benchmark//:benchmark_main",
     ],
 )
+
+pl_cc_test(
+    name = "info_class_schema_test",
+    srcs = ["info_class_schema_test.cc"],
+    deps = [":cc_library"],
+)

src/data_collector/info_class_schema.cc

@@ -0,0 +1,40 @@
+#include "src/data_collector/info_class_schema.h"
+
+namespace pl {
+namespace datacollector {
+
+std::atomic<uint64_t> InfoClassSchema::global_id_ = 0;
+
+datacollectorpb::Element InfoClassElement::ToProto() const {
+  datacollectorpb::Element element_proto;
+  element_proto.set_name(name_);
+  element_proto.set_state(state_);
+  element_proto.set_type(data_type_);
+  return element_proto;
+}
+
+datacollectorpb::InfoClass InfoClassSchema::ToProto() const {
+  datacollectorpb::InfoClass info_class_proto;
+  // Populate the proto with Elements.
+  for (auto element : elements_) {
+    datacollectorpb::Element* element_proto_ptr = info_class_proto.add_elements();
+    element_proto_ptr->MergeFrom(element.ToProto());
+  }
+
+  // Add metadata.
+  // TODO(kgandhi): For M2, only add the source name. Later on add other information
+  // from the SourceConnector.
+  auto metadata_map = info_class_proto.mutable_metadata();
+  std::string key = "source";
+  std::string value = source_.name();
+  (*metadata_map)[key] = value;
+
+  // Add all the other fields for the proto.
+  info_class_proto.set_name(name_);
+  info_class_proto.set_id(id_);
+
+  return info_class_proto;
+}
+
+}  // namespace datacollector
+}  // namespace pl

src/data_collector/info_class_schema.h

@@ -0,0 +1,121 @@
+#pragma once
+
+#include <gflags/gflags.h>
+#include <glog/logging.h>
+
+#include <atomic>
+#include <string>
+#include <vector>
+
+#include "src/common/status.h"
+#include "src/common/types/types.pb.h"
+#include "src/data_collector/proto/collector_config.pb.h"
+#include "src/data_collector/source_connector.h"
+
+namespace pl {
+namespace datacollector {
+
+using datacollectorpb::Element_State;
+using types::DataType;
+
+class InfoClassElement {
+ public:
+  InfoClassElement(const std::string& name, const DataType& type, const Element_State& state)
+      : name_(name), data_type_(type), state_(state) {}
+  virtual ~InfoClassElement() = default;
+
+  /**
+   * @brief Generate a proto message based on the InfoClassElement.
+   *
+   * @return datacollectorpb::Element
+   */
+  datacollectorpb::Element ToProto() const;
+  void SetState(const Element_State& state) { state_ = state; }
+  const std::string& name() const { return name_; }
+  const DataType& data_type() const { return data_type_; }
+  const Element_State& state() const { return state_; }
+
+ private:
+  std::string name_;
+  DataType data_type_;
+  Element_State state_;
+};
+
+class InfoClassSchema {
+ public:
+  /**
+   * @brief Construct a new Info Class Schema object
+   * SourceConnector constructs InfoClassSchema objects with and adds Elements to it
+   *
+   * @param name Name of the InfoClass
+   * @param source Pointer to the SourceConnector that created the InfoClassSchema object.
+   * This is required to identify an InfoClassSchema parent source and also to generate
+   * the publish proto.
+   */
+  InfoClassSchema(const std::string& name, const SourceConnector& source)
+      : name_(name), source_(source) {
+    id_ = global_id_++;
+  }
+  virtual ~InfoClassSchema() = default;
+
+  /**
+   * @brief Add an Element to the InfoClassSchema
+   *
+   * @param name Name of the Element
+   * @param type Data type of the element
+   * @param state Subscription state for the Element
+   */
+  void AddElement(const InfoClassElement& element) { elements_.push_back(element); }
+
+  /**
+   * @brief Get number of elements in the InfoClassSchema
+   *
+   * @return size_t Number of elements
+   */
+  size_t NumElements() const { return elements_.size(); }
+
+  /**
+   * @brief Set the subscription state of an Element
+   *
+   * @param index Element to update in the InfoClassSchema
+   * @param state Subscription state
+   * @return Status
+   */
+  void UpdateElementSubscription(size_t index, const Element_State& state) {
+    DCHECK(index < elements_.size());
+    elements_[index].SetState(state);
+  }
+
+  /**
+   * @brief Get an Element object
+   *
+   * @param index
+   * @return InfoClassElement
+   */
+  const InfoClassElement& GetElement(size_t index) const {
+    DCHECK(index < elements_.size());
+    return elements_[index];
+  }
+
+  /**
+   * @brief Generate a proto message based on the InfoClassSchema.
+   *
+   * @return datacollectorpb::InfoClass
+   */
+  datacollectorpb::InfoClass ToProto() const;
+
+  const std::string& name() const { return name_; }
+  const SourceConnector& source() const { return source_; }
+
+  uint64_t id() { return id_; }
+
+ private:
+  static std::atomic<uint64_t> global_id_;
+  std::vector<InfoClassElement> elements_;
+  std::string name_;
+  const SourceConnector& source_;
+  uint64_t id_;
+};
+
+}  // namespace datacollector
+}  // namespace pl

src/data_collector/info_class_schema_test.cc

@@ -0,0 +1,67 @@
+#include <gtest/gtest.h>
+
+#include "src/data_collector/info_class_schema.h"
+#include "src/data_collector/source_connector.h"
+
+namespace pl {
+namespace datacollector {
+
+using datacollectorpb::Element_State;
+using types::DataType;
+
+TEST(InfoClassElementTest, basic_test) {
+  InfoClassElement element("user_percentage", DataType::FLOAT64,
+                           Element_State::Element_State_COLLECTED_NOT_SUBSCRIBED);
+
+  EXPECT_EQ("user_percentage", element.name());
+  EXPECT_EQ(DataType::FLOAT64, element.data_type());
+  EXPECT_EQ(Element_State::Element_State_COLLECTED_NOT_SUBSCRIBED, element.state());
+
+  datacollectorpb::Element element_pb;
+  element_pb = element.ToProto();
+  EXPECT_EQ("user_percentage", element_pb.name());
+  EXPECT_EQ(DataType::FLOAT64, element_pb.type());
+  EXPECT_EQ(Element_State::Element_State_COLLECTED_NOT_SUBSCRIBED, element_pb.state());
+}
+
+TEST(InfoClassInfoSchemaTest, basic_test) {
+  SourceConnector source("eBPF_CPU_USAGE");
+  InfoClassSchema info_class_schema("cpu_usage", source);
+  InfoClassElement element_1("user_percentage", DataType::FLOAT64,
+                             Element_State::Element_State_COLLECTED_NOT_SUBSCRIBED);
+  InfoClassElement element_2("system_percentage", DataType::FLOAT64,
+                             Element_State::Element_State_COLLECTED_NOT_SUBSCRIBED);
+  InfoClassElement element_3("io_percentage", DataType::FLOAT64,
+                             Element_State::Element_State_COLLECTED_NOT_SUBSCRIBED);
+
+  info_class_schema.AddElement(element_1);
+  info_class_schema.AddElement(element_2);
+  info_class_schema.AddElement(element_3);
+
+  EXPECT_EQ(3, info_class_schema.NumElements());
+  EXPECT_EQ("cpu_usage", info_class_schema.name());
+  EXPECT_EQ("eBPF_CPU_USAGE", source.name());
+
+  datacollectorpb::InfoClass info_class_pb;
+  info_class_pb = info_class_schema.ToProto();
+  EXPECT_EQ(3, info_class_pb.elements_size());
+  EXPECT_EQ("cpu_usage", info_class_pb.name());
+  EXPECT_EQ("user_percentage", info_class_pb.elements(0).name());
+  EXPECT_EQ(DataType::FLOAT64, info_class_pb.elements(0).type());
+  EXPECT_EQ(0, info_class_pb.id());
+  auto metadata = info_class_pb.metadata();
+  EXPECT_EQ("eBPF_CPU_USAGE", metadata["source"]);
+
+  info_class_schema.UpdateElementSubscription(
+      0, Element_State::Element_State_COLLECTED_AND_SUBSCRIBED);
+  EXPECT_EQ(Element_State::Element_State_COLLECTED_AND_SUBSCRIBED,
+            info_class_schema.GetElement(0).state());
+
+  datacollectorpb::InfoClass subscribe_pb;
+  subscribe_pb = info_class_schema.ToProto();
+  EXPECT_EQ(Element_State::Element_State_COLLECTED_AND_SUBSCRIBED,
+            subscribe_pb.elements(0).state());
+}
+
+}  // namespace datacollector
+}  // namespace pl

src/data_collector/source_connector.cc

@@ -0,0 +1,9 @@
+#include "src/data_collector/source_connector.h"
+
+namespace pl {
+namespace datacollector {
+
+// Placeholder for SourceConnector code.
+
+}  // namespace datacollector
+}  // namespace pl

src/data_collector/source_connector.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#include <string>
+
+namespace pl {
+namespace datacollector {
+
+class SourceConnector {
+ public:
+  explicit SourceConnector(const std::string& name) : name_(name) {}
+  virtual ~SourceConnector() = default;
+
+  const std::string& name() const { return name_; }
+
+ private:
+  const std::string name_;
+};
+
+}  // namespace datacollector
+}  // namespace pl