Update stirling_profiler binary to use UnitConnector. (#1284)

Summary: Converting the `stirling_profiler` binary to use
`UnitConnector` provides some cleanup and also removes the dependency on
`//src/stirling:cc_library`.

Type of change: /kind cleanup

Test Plan: Ran `scripts/sudo_bazel_run.sh
//src/stirling/binaries:stirling_profiler` and saw the stack trace
output written to `stdout` as expected.

---------

Signed-off-by: Pete Stevenson <jps@pixielabs.ai>

Author: Pete Stevenson

src/stirling/binaries/BUILD.bazel

@@ -96,7 +96,7 @@ pl_cc_binary(
     name = "stirling_profiler",
     srcs = ["stirling_profiler.cc"],
     deps = [
-        "//src/stirling:cc_library",
+        "//src/stirling/source_connectors/perf_profiler:cc_library",
     ],
 )
 

src/stirling/binaries/stirling_profiler.cc

@@ -22,111 +22,112 @@
 
 #include "src/common/base/base.h"
 #include "src/shared/upid/upid.h"
-#include "src/stirling/core/pub_sub_manager.h"
-#include "src/stirling/core/source_registry.h"
+#include "src/stirling/core/unit_connector.h"
 #include "src/stirling/source_connectors/perf_profiler/perf_profile_connector.h"
 #include "src/stirling/source_connectors/perf_profiler/stack_traces_table.h"
-#include "src/stirling/stirling.h"
-
-using ::px::ProcessStatsMonitor;
 
 using ::px::Status;
-using ::px::StatusOr;
 
-using ::px::stirling::IndexPublication;
-using ::px::stirling::PerfProfileConnector;
-using ::px::stirling::SourceRegistry;
-using ::px::stirling::Stirling;
-using ::px::stirling::stirlingpb::InfoClass;
-using ::px::stirling::stirlingpb::Publish;
+DEFINE_uint32(time, 30, "Number of seconds to run the profiler.");
 
-using ::px::md::UPID;
-using ::px::types::ColumnWrapperRecordBatch;
-using ::px::types::TabletID;
+namespace px {
+namespace stirling {
+
+class Profiler : public UnitConnector<PerfProfileConnector> {
+ public:
+  Status PrintData() {
+    // Build the stack traces histogram.
+    PX_RETURN_IF_ERROR(BuildHistogram());
+
+    // Print the stack traces histogram.
+    // TODO(jps): replace this with a pprof proto file writer.
+    // 15x: libc.so;main;foo;bar
+    // 12x: libc.so;main;foo;qux
+    for (const auto& [str, count] : histo_) {
+      LOG(INFO) << count << "x: " << str;
+    }
+    return Status::OK();
+  }
 
-DEFINE_uint32(time, 30, "Number of seconds to run the profiler.");
-DEFINE_uint32(pid, 0, "PID to profile. Leave unspecified to profile everything.");
-
-// Put this in global space, so we can kill it in the signal handler.
-Stirling* g_stirling = nullptr;
-ProcessStatsMonitor* g_process_stats_monitor = nullptr;
-std::atomic<bool> g_data_received = false;
-
-Status StirlingWrapperCallback(uint64_t /* table_id */, TabletID /* tablet_id */,
-                               std::unique_ptr<ColumnWrapperRecordBatch> record_batch) {
-  auto& upid_col = (*record_batch)[px::stirling::kStackTraceUPIDIdx];
-  auto& stack_trace_str_col = (*record_batch)[px::stirling::kStackTraceStackTraceStrIdx];
-  auto& count_col = (*record_batch)[px::stirling::kStackTraceCountIdx];
-
-  std::string out;
-  for (size_t i = 0; i < stack_trace_str_col->Size(); ++i) {
-    UPID upid(upid_col->Get<px::types::UInt128Value>(i).val);
-
-    if (FLAGS_pid == upid.pid() || FLAGS_pid == 0) {
-      std::cout << stack_trace_str_col->Get<px::types::StringValue>(i);
-      std::cout << " ";
-      std::cout << count_col->Get<px::types::Int64Value>(i).val;
-      std::cout << "\n";
+ private:
+  Status BuildHistogram() {
+    PX_ASSIGN_OR_RETURN(const auto& records, ConsumeRecords(0));
+
+    const uint64_t num_rows = records[kStackTraceStackTraceStrIdx]->Size();
+    const auto traces_column = records[kStackTraceStackTraceStrIdx];
+    const auto counts_column = records[kStackTraceCountIdx];
+
+    // Build the stack traces histogram.
+    for (uint64_t row_idx = 0; row_idx < num_rows; ++row_idx) {
+      const std::string& stack_trace_str = traces_column->Get<types::StringValue>(row_idx);
+      const int64_t count = counts_column->Get<types::Int64Value>(row_idx).val;
+      histo_[stack_trace_str] += count;
     }
+    return Status::OK();
   }
 
-  g_data_received = true;
+  // A local stack trace histo (for convenience, to be populated after all samples are collected).
+  absl::flat_hash_map<std::string, uint64_t> histo_;
+};
 
-  return Status::OK();
-}
+}  // namespace stirling
+}  // namespace px
+
+std::unique_ptr<px::stirling::Profiler> g_profiler;
 
 void SignalHandler(int signum) {
   std::cerr << "\n\nStopping, might take a few seconds ..." << std::endl;
-  // Important to call Stop(), because it releases BPF resources,
+
+  // Important to call Stop(), because it releases eBPF resources,
   // which would otherwise leak.
-  if (g_stirling != nullptr) {
-    g_stirling->Stop();
-  }
-  if (g_process_stats_monitor != nullptr) {
-    g_process_stats_monitor->PrintCPUTime();
+  if (g_profiler != nullptr) {
+    PX_UNUSED(g_profiler->Stop());
+    g_profiler = nullptr;
   }
+
   exit(signum);
 }
 
+Status RunProfiler() {
+  // Bring up eBPF.
+  PX_RETURN_IF_ERROR(g_profiler->Init());
+
+  // Separate thread to periodically wake up and read the eBPF perf buffer & maps.
+  PX_RETURN_IF_ERROR(g_profiler->Start());
+
+  // Collect data for the user specified amount of time.
+  sleep(FLAGS_time);
+
+  // Stop collecting data and do a final read out of eBPF perf buffer & maps.
+  PX_RETURN_IF_ERROR(g_profiler->Stop());
+
+  // Print the info. We will replace this with a pprof proto file write out.
+  PX_RETURN_IF_ERROR(g_profiler->PrintData());
+
+  // Phew. We are outta here.
+  return Status::OK();
+}
+
 int main(int argc, char** argv) {
   // Register signal handlers to clean-up on exit.
+  signal(SIGHUP, SignalHandler);
   signal(SIGINT, SignalHandler);
   signal(SIGQUIT, SignalHandler);
   signal(SIGTERM, SignalHandler);
-  signal(SIGHUP, SignalHandler);
 
   px::EnvironmentGuard env_guard(&argc, argv);
 
-  // Make Stirling.
-  auto registry = std::make_unique<SourceRegistry>();
-  registry->RegisterOrDie<PerfProfileConnector>();
-  std::unique_ptr<Stirling> stirling = Stirling::Create(std::move(registry));
-  g_stirling = stirling.get();
-  stirling->RegisterDataPushCallback(StirlingWrapperCallback);
-
-  // Enable use of USR1/USR2 for controlling debug.
-  stirling->RegisterUserDebugSignalHandlers();
+  // Need to do this after env setup.
+  g_profiler = std::make_unique<px::stirling::Profiler>();
 
-  // Start measuring process stats after init.
-  ProcessStatsMonitor process_stats_monitor;
-  g_process_stats_monitor = &process_stats_monitor;
+  // Run the profiler (in more detail: setup, collect data, and tear down).
+  const auto status = RunProfiler();
 
-  // Run Stirling.
-  std::thread run_thread = std::thread(&Stirling::Run, stirling.get());
-
-  // Run for the specified amount of time.
-  std::this_thread::sleep_for(std::chrono::seconds(FLAGS_time));
-
-  // This is not likely because a table push is triggered immediately. But, just in case,
-  // provide some help if no data was received.
-  LOG_IF(WARNING, !g_data_received) << "No data received from profiler. Try increasing -time or "
-                                       "reducing -stirling_profiler_table_update_period_seconds.";
+  // Something happened, log that.
+  LOG_IF(WARNING, !status.ok()) << status.msg();
 
   // Cleanup.
-  stirling->Stop();
-
-  // Wait for the thread to return.
-  run_thread.join();
+  g_profiler = nullptr;
 
-  return 0;
+  return status.ok() ? 0 : -1;
 }