simpleperf: support --cpu option in record/stat command.

--cpu option is used to record on selected cpus.

Change-Id: If5bb9b42a064d2ff69fbeec77906fc79943dddc1

10 files changed, 72 insertions, 30 deletions

diff --git a/simpleperf/cmd_record.cpp b/simpleperf/cmd_record.cpp
index e5342dde..8cfbfdff 100644
--- a/simpleperf/cmd_record.cpp
+++ b/simpleperf/cmd_record.cpp
@@ -68,6 +68,9 @@ class RecordCommand : public Command {
             "    --call-graph fp | dwarf[,<dump_stack_size>]\n"
             "                 Enable call graph recording. Use frame pointer or dwarf as the\n"
             "                 method to parse call graph in stack. Default is dwarf,8192.\n"
+            "    --cpu cpu_item1,cpu_item2,...\n"
+            "                 Collect samples only on the selected cpus. cpu_item can be cpu\n"
+            "                 number like 1, or cpu range like 0-3.\n"
             "    -e event1[:modifier1],event2[:modifier2],...\n"
             "                 Select the event list to sample. Use `simpleperf list` to find\n"
             "                 all possible event names. Modifiers can be added to define\n"
@@ -154,6 +157,7 @@ class RecordCommand : public Command {
   bool post_unwind_;
   bool child_inherit_;
   std::vector<pid_t> monitored_threads_;
+  std::vector<int> cpus_;
   std::vector<EventTypeAndModifier> measured_event_types_;
   EventSelectionSet event_selection_set_;
 
@@ -207,11 +211,11 @@ bool RecordCommand::Run(const std::vector<std::string>& args) {
   // 3. Open perf_event_files, create memory mapped buffers for perf_event_files, add prepare poll
   //    for perf_event_files.
   if (system_wide_collection_) {
-    if (!event_selection_set_.OpenEventFilesForAllCpus()) {
+    if (!event_selection_set_.OpenEventFilesForCpus(cpus_)) {
       return false;
     }
   } else {
-    if (!event_selection_set_.OpenEventFilesForThreadsOnAllCpus(monitored_threads_)) {
+    if (!event_selection_set_.OpenEventFilesForThreadsOnCpus(monitored_threads_, cpus_)) {
       return false;
     }
   }
@@ -321,6 +325,11 @@ bool RecordCommand::ParseOptions(const std::vector<std::string>& args,
         LOG(ERROR) << "unexpected argument for --call-graph option: " << args[i];
         return false;
       }
+    } else if (args[i] == "--cpu") {
+      if (!NextArgumentOrError(args, &i)) {
+        return false;
+      }
+      cpus_ = GetCpusFromString(args[i]);
     } else if (args[i] == "-e") {
       if (!NextArgumentOrError(args, &i)) {
         return false;
diff --git a/simpleperf/cmd_record_test.cpp b/simpleperf/cmd_record_test.cpp
index 78dae8c7..72c871b9 100644
--- a/simpleperf/cmd_record_test.cpp
+++ b/simpleperf/cmd_record_test.cpp
@@ -167,3 +167,8 @@ TEST(record_cmd, more_than_one_event_types) {
   ASSERT_TRUE(RecordCmd()->Run({"-e", "cpu-cycles,cpu-clock", "sleep", "1"}));
   ASSERT_TRUE(RecordCmd()->Run({"-e", "cpu-cycles", "-e", "cpu-clock", "sleep", "1"}));
 }
+
+TEST(record_cmd, cpu_option) {
+  ASSERT_TRUE(RecordCmd()->Run({"--cpu", "0", "sleep", "1"}));
+  ASSERT_TRUE(RecordCmd()->Run({"--cpu", "0", "-a", "sleep", "1"}));
+}
diff --git a/simpleperf/cmd_stat.cpp b/simpleperf/cmd_stat.cpp
index 245a9967..5e82ba8e 100644
--- a/simpleperf/cmd_stat.cpp
+++ b/simpleperf/cmd_stat.cpp
@@ -55,6 +55,9 @@ class StatCommand : public Command {
                 "Usage: simpleperf stat [options] [command [command-args]]\n"
                 "    Gather performance counter information of running [command].\n"
                 "    -a           Collect system-wide information.\n"
+                "    --cpu cpu_item1,cpu_item2,...\n"
+                "                 Collect information only on the selected cpus. cpu_item can\n"
+                "                 be a cpu number like 1, or a cpu range like 0-3.\n"
                 "    -e event1[:modifier1],event2[:modifier2],...\n"
                 "                 Select the event list to count. Use `simpleperf list` to find\n"
                 "                 all possible event names. Modifiers can be added to define\n"
@@ -89,6 +92,7 @@ class StatCommand : public Command {
   bool system_wide_collection_;
   bool child_inherit_;
   std::vector<pid_t> monitored_threads_;
+  std::vector<int> cpus_;
   std::vector<EventTypeAndModifier> measured_event_types_;
   EventSelectionSet event_selection_set_;
 
@@ -130,11 +134,11 @@ bool StatCommand::Run(const std::vector<std::string>& args) {
 
   // 3. Open perf_event_files.
   if (system_wide_collection_) {
-    if (!event_selection_set_.OpenEventFilesForAllCpus()) {
+    if (!event_selection_set_.OpenEventFilesForCpus(cpus_)) {
       return false;
     }
   } else {
-    if (!event_selection_set_.OpenEventFilesForThreads(monitored_threads_)) {
+    if (!event_selection_set_.OpenEventFilesForThreadsOnCpus(monitored_threads_, cpus_)) {
       return false;
     }
   }
@@ -174,6 +178,11 @@ bool StatCommand::ParseOptions(const std::vector<std::string>& args,
   for (i = 0; i < args.size() && args[i].size() > 0 && args[i][0] == '-'; ++i) {
     if (args[i] == "-a") {
       system_wide_collection_ = true;
+    } else if (args[i] == "--cpu") {
+      if (!NextArgumentOrError(args, &i)) {
+        return false;
+      }
+      cpus_ = GetCpusFromString(args[i]);
     } else if (args[i] == "-e") {
       if (!NextArgumentOrError(args, &i)) {
         return false;
diff --git a/simpleperf/cmd_stat_test.cpp b/simpleperf/cmd_stat_test.cpp
index c6c4ef74..8de50fc0 100644
--- a/simpleperf/cmd_stat_test.cpp
+++ b/simpleperf/cmd_stat_test.cpp
@@ -69,3 +69,8 @@ TEST(stat_cmd, existing_threads) {
 TEST(stat_cmd, no_monitored_threads) {
   ASSERT_FALSE(StatCmd()->Run({""}));
 }
+
+TEST(stat_cmd, cpu_option) {
+  ASSERT_TRUE(StatCmd()->Run({"--cpu", "0", "sleep", "1"}));
+  ASSERT_TRUE(StatCmd()->Run({"--cpu", "0", "-a", "sleep", "1"}));
+}
diff --git a/simpleperf/environment.cpp b/simpleperf/environment.cpp
index 4e3c212e..560fa89b 100644
--- a/simpleperf/environment.cpp
+++ b/simpleperf/environment.cpp
@@ -21,6 +21,7 @@
 #include <stdlib.h>
 
 #include <limits>
+#include <set>
 #include <unordered_map>
 #include <vector>
 
@@ -43,27 +44,29 @@ std::vector<int> GetOnlineCpus() {
   LineReader reader(fp);
   char* line;
   if ((line = reader.ReadLine()) != nullptr) {
-    result = GetOnlineCpusFromString(line);
+    result = GetCpusFromString(line);
   }
   CHECK(!result.empty()) << "can't get online cpu information";
   return result;
 }
 
-std::vector<int> GetOnlineCpusFromString(const std::string& s) {
-  std::vector<int> result;
+std::vector<int> GetCpusFromString(const std::string& s) {
+  std::set<int> cpu_set;
   bool have_dash = false;
   const char* p = s.c_str();
   char* endp;
+  int last_cpu;
   long cpu;
   // Parse line like: 0,1-3, 5, 7-8
   while ((cpu = strtol(p, &endp, 10)) != 0 || endp != p) {
-    if (have_dash && result.size() > 0) {
-      for (int t = result.back() + 1; t < cpu; ++t) {
-        result.push_back(t);
+    if (have_dash && !cpu_set.empty()) {
+      for (int t = last_cpu + 1; t < cpu; ++t) {
+        cpu_set.insert(t);
       }
     }
     have_dash = false;
-    result.push_back(cpu);
+    cpu_set.insert(cpu);
+    last_cpu = cpu;
     p = endp;
     while (!isdigit(*p) && *p != '\0') {
       if (*p == '-') {
@@ -72,7 +75,7 @@ std::vector<int> GetOnlineCpusFromString(const std::string& s) {
       ++p;
     }
   }
-  return result;
+  return std::vector<int>(cpu_set.begin(), cpu_set.end());
 }
 
 bool ProcessKernelSymbols(const std::string& symbol_file,
diff --git a/simpleperf/environment.h b/simpleperf/environment.h
index a0c085da..853ac440 100644
--- a/simpleperf/environment.h
+++ b/simpleperf/environment.h
@@ -27,6 +27,7 @@
 #include "build_id.h"
 
 std::vector<int> GetOnlineCpus();
+std::vector<int> GetCpusFromString(const std::string& s);
 
 constexpr char DEFAULT_KERNEL_MMAP_NAME[] = "[kernel.kallsyms]_text";
 
@@ -76,8 +77,6 @@ bool GetValidThreadsFromThreadString(const std::string& tid_str, std::set<pid_t>
 bool GetExecPath(std::string* exec_path);
 
 // Expose the following functions for unit tests.
-std::vector<int> GetOnlineCpusFromString(const std::string& s);
-
 struct KernelSymbol {
   uint64_t addr;
   char type;
diff --git a/simpleperf/environment_test.cpp b/simpleperf/environment_test.cpp
index 3cf81fa6..1257cd5d 100644
--- a/simpleperf/environment_test.cpp
+++ b/simpleperf/environment_test.cpp
@@ -21,10 +21,11 @@
 
 #include "environment.h"
 
-TEST(environment, GetOnlineCpusFromString) {
-  ASSERT_EQ(GetOnlineCpusFromString(""), std::vector<int>());
-  ASSERT_EQ(GetOnlineCpusFromString("0-2"), std::vector<int>({0, 1, 2}));
-  ASSERT_EQ(GetOnlineCpusFromString("0,2-3"), std::vector<int>({0, 2, 3}));
+TEST(environment, GetCpusFromString) {
+  ASSERT_EQ(GetCpusFromString(""), std::vector<int>());
+  ASSERT_EQ(GetCpusFromString("0-2"), std::vector<int>({0, 1, 2}));
+  ASSERT_EQ(GetCpusFromString("0,2-3"), std::vector<int>({0, 2, 3}));
+  ASSERT_EQ(GetCpusFromString("1,0-3,3,4"), std::vector<int>({0, 1, 2, 3, 4}));
 }
 
 static bool FindKernelSymbol(const KernelSymbol& sym1, const KernelSymbol& sym2) {
diff --git a/simpleperf/event_attr.cpp b/simpleperf/event_attr.cpp
index b0d75ce6..d0b66139 100644
--- a/simpleperf/event_attr.cpp
+++ b/simpleperf/event_attr.cpp
@@ -87,7 +87,8 @@ perf_event_attr CreateDefaultPerfEventAttr(const EventType& event_type) {
   // PerfCounter in event_fd.h.
   attr.read_format =
       PERF_FORMAT_TOTAL_TIME_ENABLED | PERF_FORMAT_TOTAL_TIME_RUNNING | PERF_FORMAT_ID;
-  attr.sample_type |= PERF_SAMPLE_IP | PERF_SAMPLE_TID | PERF_SAMPLE_TIME | PERF_SAMPLE_PERIOD;
+  attr.sample_type |=
+      PERF_SAMPLE_IP | PERF_SAMPLE_TID | PERF_SAMPLE_TIME | PERF_SAMPLE_PERIOD | PERF_SAMPLE_CPU;
 
   if (attr.type == PERF_TYPE_TRACEPOINT) {
     attr.sample_freq = 0;
diff --git a/simpleperf/event_selection_set.cpp b/simpleperf/event_selection_set.cpp
index e42b89a4..1a9de630 100644
--- a/simpleperf/event_selection_set.cpp
+++ b/simpleperf/event_selection_set.cpp
@@ -166,18 +166,29 @@ void EventSelectionSet::SetInherit(bool enable) {
   }
 }
 
-bool EventSelectionSet::OpenEventFilesForAllCpus() {
-  return OpenEventFilesForThreadsOnAllCpus({-1});
+static bool CheckIfCpusOnline(const std::vector<int>& cpus) {
+  std::vector<int> online_cpus = GetOnlineCpus();
+  for (const auto& cpu : cpus) {
+    if (std::find(online_cpus.begin(), online_cpus.end(), cpu) == online_cpus.end()) {
+      LOG(ERROR) << "cpu " << cpu << " is not online.";
+      return false;
+    }
+  }
+  return true;
 }
 
-bool EventSelectionSet::OpenEventFilesForThreads(const std::vector<pid_t>& threads) {
-  return OpenEventFiles(threads, {-1});
+bool EventSelectionSet::OpenEventFilesForCpus(const std::vector<int>& cpus) {
+  return OpenEventFilesForThreadsOnCpus({-1}, cpus);
 }
 
-bool EventSelectionSet::OpenEventFilesForThreadsOnAllCpus(const std::vector<pid_t>& threads) {
-  std::vector<int> cpus = GetOnlineCpus();
-  if (cpus.empty()) {
-    return false;
+bool EventSelectionSet::OpenEventFilesForThreadsOnCpus(const std::vector<pid_t>& threads,
+                                                       std::vector<int> cpus) {
+  if (!cpus.empty()) {
+    if (!CheckIfCpusOnline(cpus)) {
+      return false;
+    }
+  } else {
+    cpus = GetOnlineCpus();
   }
   return OpenEventFiles(threads, cpus);
 }
diff --git a/simpleperf/event_selection_set.h b/simpleperf/event_selection_set.h
index 54cf3cd3..fed3c0bc 100644
--- a/simpleperf/event_selection_set.h
+++ b/simpleperf/event_selection_set.h
@@ -68,9 +68,8 @@ class EventSelectionSet {
   bool EnableDwarfCallChainSampling(uint32_t dump_stack_size);
   void SetInherit(bool enable);
 
-  bool OpenEventFilesForAllCpus();
-  bool OpenEventFilesForThreads(const std::vector<pid_t>& threads);
-  bool OpenEventFilesForThreadsOnAllCpus(const std::vector<pid_t>& threads);
+  bool OpenEventFilesForCpus(const std::vector<int>& cpus);
+  bool OpenEventFilesForThreadsOnCpus(const std::vector<pid_t>& threads, std::vector<int> cpus);
   bool EnableEvents();
   bool ReadCounters(std::vector<CountersInfo>* counters);
   void PreparePollForEventFiles(std::vector<pollfd>* pollfds);