Add sample labels flag to perf_to_profile

src/BUILD

@@ -143,6 +143,7 @@ cc_library(
         ":perf_data_converter",
         "//src/quipper:base",
         "//src/quipper:perf_data_cc_proto",
+        "@abseil-cpp//absl/strings",
     ],
 )
 

src/perf_to_profile.cc

@@ -14,8 +14,10 @@ int main(int argc, char** argv) {
   std::string input, output;
   bool overwriteOutput = false;
   bool allowUnalignedJitMappings = false;
+  uint32_t sampleLabels = perftools::kNoLabels;
   if (!ParseArguments(argc, const_cast<const char**>(argv), &input, &output,
-                      &overwriteOutput, &allowUnalignedJitMappings)) {
+                      &overwriteOutput, &allowUnalignedJitMappings,
+                      &sampleLabels)) {
     PrintUsage();
     return EXIT_FAILURE;
   }
@@ -25,7 +27,7 @@ int main(int argc, char** argv) {
     options |= perftools::ConversionOptions::kAllowUnalignedJitMappings;
   }
   std::string data = ReadFileToString(input);
-  const auto profiles = StringToProfiles(data, perftools::kNoLabels, options);
+  const auto profiles = StringToProfiles(data, sampleLabels, options);
 
   // With kNoOptions, all of the PID profiles should be merged into a
   // single one.

src/perf_to_profile_lib.cc

@@ -8,7 +8,52 @@
 #include "src/perf_to_profile_lib.h"
 
 #include <sys/stat.h>
-#include <sstream>
+#include <optional>
+#include <string_view>
+
+#include "absl/strings/str_split.h"
+
+namespace {
+
+// Returns the SampleLabels value for the given label name, or std::nullopt if
+// unrecognized. Names match the label key constants in perf_data_converter.h.
+std::optional<perftools::SampleLabels> LabelNameToValue(
+    std::string_view name) {
+  if (name == "pid") return perftools::kPidLabel;
+  if (name == "tid") return perftools::kTidLabel;
+  if (name == "timestamp_ns") return perftools::kTimestampNsLabel;
+  if (name == "execution_mode") return perftools::kExecutionModeLabel;
+  if (name == "comm") return perftools::kCommLabel;
+  if (name == "thread_type") return perftools::kThreadTypeLabel;
+  if (name == "thread_comm") return perftools::kThreadCommLabel;
+  if (name == "cgroup") return perftools::kCgroupLabel;
+  if (name == "code_page_size") return perftools::kCodePageSizeLabel;
+  if (name == "data_page_size") return perftools::kDataPageSizeLabel;
+  if (name == "cpu") return perftools::kCpuLabel;
+  if (name == "cache_latency") return perftools::kCacheLatencyLabel;
+  if (name == "data_src") return perftools::kDataSrcLabel;
+  if (name == "total_latency") return perftools::kTotalLatencyLabel;
+  if (name == "issue_latency") return perftools::kIssueLatencyLabel;
+  if (name == "translation_latency") return perftools::kTranslationLatencyLabel;
+  return std::nullopt;
+}
+
+// Parses a comma-separated list of label names into a bitmask of
+// SampleLabels values. Returns false if any name is unrecognized.
+bool ParseSampleLabels(const char* arg, uint32_t* sample_labels) {
+  *sample_labels = perftools::kNoLabels;
+  for (std::string_view name : absl::StrSplit(arg, ',')) {
+    auto value = LabelNameToValue(name);
+    if (!value.has_value()) {
+      LOG(ERROR) << "Unknown sample label: " << name;
+      return false;
+    }
+    *sample_labels |= *value;
+  }
+  return true;
+}
+
+}  // namespace
 
 bool FileExists(const std::string& path) {
   struct stat file_stat;
@@ -52,22 +97,31 @@ void CreateFile(const std::string& path, std::ofstream* file,
 
 void PrintUsage() {
   LOG(INFO) << "Usage:";
-  LOG(INFO) << "perf_to_profile -i <input perf data> -o <output profile> [-f]";
+  LOG(INFO) << "perf_to_profile -i <input perf data> -o <output profile> [-f] "
+            << "[-l <labels>]";
   LOG(INFO) << "If the -f option is given, overwrite the existing output "
             << "profile.";
   LOG(INFO) << "If the -j option is given, allow unaligned MMAP events "
             << "required by perf data from VMs with JITs.";
+  LOG(INFO) << "If the -l option is given, add the specified sample labels to "
+            << "the output profile. Labels are a comma-separated list of: "
+            << "pid, tid, timestamp_ns, execution_mode, comm, thread_type, "
+            << "thread_comm, cgroup, code_page_size, data_page_size, cpu, "
+            << "cache_latency, data_src, total_latency, issue_latency, "
+            << "translation_latency.";
 }
 
 bool ParseArguments(int argc, const char* argv[], std::string* input,
                     std::string* output, bool* overwrite_output,
-                    bool* allow_unaligned_jit_mappings) {
+                    bool* allow_unaligned_jit_mappings,
+                    uint32_t* sample_labels) {
   *input = "";
   *output = "";
   *overwrite_output = false;
   *allow_unaligned_jit_mappings = false;
+  *sample_labels = perftools::kNoLabels;
   int opt;
-  while ((opt = getopt(argc, const_cast<char* const*>(argv), ":jfi:o:")) !=
+  while ((opt = getopt(argc, const_cast<char* const*>(argv), ":jfi:o:l:")) !=
          -1) {
     switch (opt) {
       case 'i':
@@ -82,8 +136,13 @@ bool ParseArguments(int argc, const char* argv[], std::string* input,
       case 'j':
         *allow_unaligned_jit_mappings = true;
         break;
+      case 'l':
+        if (!ParseSampleLabels(optarg, sample_labels)) {
+          return false;
+        }
+        break;
       case ':':
-        LOG(ERROR) << "Must provide arguments for flags -i and -o";
+        LOG(ERROR) << "Must provide arguments for flags -i, -o, and -l";
         return false;
       case '?':
         LOG(ERROR) << "Invalid option: " << static_cast<char>(optopt);

src/perf_to_profile_lib.h

@@ -33,12 +33,13 @@ perftools::ProcessProfiles StringToProfiles(
 void CreateFile(const std::string& path, std::ofstream* file,
                 bool overwrite_output);
 
-// Parses arguments, stores the results in |input|, |output|
-// |overwrite_output| and |allow_unaligned_jit_mappings|, and returns true if
-// arguments parsed successfully and false otherwise.
+// Parses arguments, stores the results in |input|, |output|,
+// |overwrite_output|, |allow_unaligned_jit_mappings|, and |sample_labels|,
+// and returns true if arguments parsed successfully and false otherwise.
 bool ParseArguments(int argc, const char* argv[], std::string* input,
                     std::string* output, bool* overwrite_output,
-                    bool* allow_unaligned_jit_mappings);
+                    bool* allow_unaligned_jit_mappings,
+                    uint32_t* sample_labels);
 
 // Prints the usage of the tool.
 void PrintUsage();

src/perf_to_profile_lib_test.cc

@@ -23,6 +23,7 @@ TEST(PerfToProfileTest, ParseArguments) {
     std::string expected_output;
     bool expected_overwrite_output;
     bool allow_unaligned_jit_mappings;
+    uint32_t expected_sample_labels;
     bool want_error;
   };
 
@@ -34,6 +35,7 @@ TEST(PerfToProfileTest, ParseArguments) {
       .expected_output = "output_profile",
       .expected_overwrite_output = true,
       .allow_unaligned_jit_mappings = false,
+      .expected_sample_labels = perftools::kNoLabels,
       .want_error = false});
   tests.push_back(
       Test{.desc = "With input and output flags",
@@ -42,6 +44,7 @@ TEST(PerfToProfileTest, ParseArguments) {
            .expected_output = "output_profile",
            .expected_overwrite_output = false,
            .allow_unaligned_jit_mappings = false,
+           .expected_sample_labels = perftools::kNoLabels,
            .want_error = false});
   tests.push_back(Test{
       .desc = "With input and output flags and jit-support",
@@ -50,13 +53,45 @@ TEST(PerfToProfileTest, ParseArguments) {
       .expected_output = "output_profile",
       .expected_overwrite_output = false,
       .allow_unaligned_jit_mappings = true,
+      .expected_sample_labels = perftools::kNoLabels,
       .want_error = false});
+  tests.push_back(Test{
+      .desc = "With pid label",
+      .argv = {"<exec>", "-i", "input_perf_file", "-o", "output_profile",
+               "-l", "pid"},
+      .expected_input = "input_perf_file",
+      .expected_output = "output_profile",
+      .expected_overwrite_output = false,
+      .allow_unaligned_jit_mappings = false,
+      .expected_sample_labels = perftools::kPidLabel,
+      .want_error = false});
+  tests.push_back(Test{
+      .desc = "With pid and tid labels",
+      .argv = {"<exec>", "-i", "input_perf_file", "-o", "output_profile",
+               "-l", "pid,tid"},
+      .expected_input = "input_perf_file",
+      .expected_output = "output_profile",
+      .expected_overwrite_output = false,
+      .allow_unaligned_jit_mappings = false,
+      .expected_sample_labels = perftools::kPidAndTidLabels,
+      .want_error = false});
+  tests.push_back(Test{
+      .desc = "With unknown label",
+      .argv = {"<exec>", "-i", "input_perf_file", "-o", "output_profile",
+               "-l", "bogus"},
+      .expected_input = "",
+      .expected_output = "",
+      .expected_overwrite_output = false,
+      .allow_unaligned_jit_mappings = false,
+      .expected_sample_labels = perftools::kNoLabels,
+      .want_error = true});
   tests.push_back(Test{.desc = "With only overwrite flag",
                        .argv = {"<exec>", "-f"},
                        .expected_input = "",
                        .expected_output = "",
                        .expected_overwrite_output = false,
                        .allow_unaligned_jit_mappings = false,
+                       .expected_sample_labels = perftools::kNoLabels,
                        .want_error = true});
   tests.push_back(Test{
       .desc = "With input, output, and invalid flags",
@@ -65,30 +100,35 @@ TEST(PerfToProfileTest, ParseArguments) {
       .expected_output = "",
       .expected_overwrite_output = false,
       .allow_unaligned_jit_mappings = false,
+      .expected_sample_labels = perftools::kNoLabels,
       .want_error = true});
   tests.push_back(Test{.desc = "With an invalid flag",
                        .argv = {"<exec>", "-F"},
                        .expected_input = "",
                        .expected_output = "",
                        .expected_overwrite_output = false,
                        .allow_unaligned_jit_mappings = false,
+                       .expected_sample_labels = perftools::kNoLabels,
                        .want_error = true});
   for (auto test : tests) {
     std::string input;
     std::string output;
     bool overwrite_output;
     bool allow_unaligned_jit_mappings;
+    uint32_t sample_labels;
     LOG(INFO) << "Testing: " << test.desc;
     EXPECT_THAT(
         ParseArguments(test.argv.size(), test.argv.data(), &input, &output,
-                       &overwrite_output, &allow_unaligned_jit_mappings),
+                       &overwrite_output, &allow_unaligned_jit_mappings,
+                       &sample_labels),
         Eq(!test.want_error));
     if (!test.want_error) {
       EXPECT_THAT(input, Eq(test.expected_input));
       EXPECT_THAT(output, Eq(test.expected_output));
       EXPECT_THAT(overwrite_output, Eq(test.expected_overwrite_output));
       EXPECT_THAT(allow_unaligned_jit_mappings,
                   Eq(test.allow_unaligned_jit_mappings));
+      EXPECT_THAT(sample_labels, Eq(test.expected_sample_labels));
     }
     optind = 1;
   }