simpleperf: add --size-limit option in record cmd.

--size-limit option stops recording when the recorded data
reaches the size limit. It is used by
run_simpleperf_without_usb_connection.py to avoid taking too much disk
space.

Bug: http://b/74198167
Test: run simpleperf_unit_test.
Test: run test.py.
Change-Id: I11f0023c342c50e1cf8035430e6af1b3caa329e7
diff --git a/simpleperf/IOEventLoop.cpp b/simpleperf/IOEventLoop.cpp
index 662aba2..01f2acd 100644
--- a/simpleperf/IOEventLoop.cpp
+++ b/simpleperf/IOEventLoop.cpp
@@ -37,7 +37,8 @@
   }
 };
 
-IOEventLoop::IOEventLoop() : ebase_(nullptr), has_error_(false), use_precise_timer_(false) {}
+IOEventLoop::IOEventLoop()
+    : ebase_(nullptr), has_error_(false), use_precise_timer_(false), in_loop_(false) {}
 
 IOEventLoop::~IOEventLoop() {
   events_.clear();
@@ -163,8 +164,10 @@
 }
 
 bool IOEventLoop::RunLoop() {
+  in_loop_ = true;
   if (event_base_dispatch(ebase_) == -1) {
     LOG(ERROR) << "event_base_dispatch() failed";
+    in_loop_ = false;
     return false;
   }
   if (has_error_) {
@@ -174,9 +177,12 @@
 }
 
 bool IOEventLoop::ExitLoop() {
-  if (event_base_loopbreak(ebase_) == -1) {
-    LOG(ERROR) << "event_base_loopbreak() failed";
-    return false;
+  if (in_loop_) {
+    if (event_base_loopbreak(ebase_) == -1) {
+      LOG(ERROR) << "event_base_loopbreak() failed";
+      return false;
+    }
+    in_loop_ = false;
   }
   return true;
 }
diff --git a/simpleperf/IOEventLoop.h b/simpleperf/IOEventLoop.h
index 9dc73c3..4a84197 100644
--- a/simpleperf/IOEventLoop.h
+++ b/simpleperf/IOEventLoop.h
@@ -85,6 +85,7 @@
   std::vector<std::unique_ptr<IOEvent>> events_;
   bool has_error_;
   bool use_precise_timer_;
+  bool in_loop_;
 };
 
 #endif  // SIMPLE_PERF_IOEVENT_LOOP_H_
diff --git a/simpleperf/IOEventLoop_test.cpp b/simpleperf/IOEventLoop_test.cpp
index 3253502..cf231ba 100644
--- a/simpleperf/IOEventLoop_test.cpp
+++ b/simpleperf/IOEventLoop_test.cpp
@@ -215,3 +215,8 @@
   close(fd[0]);
   close(fd[1]);
 }
+
+TEST(IOEventLoop, exit_before_loop) {
+  IOEventLoop loop;
+  ASSERT_TRUE(loop.ExitLoop());
+}
diff --git a/simpleperf/cmd_debug_unwind.cpp b/simpleperf/cmd_debug_unwind.cpp
index 10c5f37..edfad9f 100644
--- a/simpleperf/cmd_debug_unwind.cpp
+++ b/simpleperf/cmd_debug_unwind.cpp
@@ -169,11 +169,7 @@
         return false;
       }
     } else if (args[i] == "--time") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      if (!android::base::ParseUint(args[i].c_str(), &selected_time_)) {
-        LOG(ERROR) << "Invalid option for " << args[i-1] << ": " << args[i];
+      if (!GetUintOption(args, &i, &selected_time_)) {
         return false;
       }
     } else {
diff --git a/simpleperf/cmd_record.cpp b/simpleperf/cmd_record.cpp
index c142feb..9e86965 100644
--- a/simpleperf/cmd_record.cpp
+++ b/simpleperf/cmd_record.cpp
@@ -28,7 +28,6 @@
 
 #include <android-base/logging.h>
 #include <android-base/file.h>
-#include <android-base/parsedouble.h>
 #include <android-base/parseint.h>
 #include <android-base/strings.h>
 #include <android-base/test_utils.h>
@@ -189,6 +188,8 @@
 "-o record_file_name    Set record file name, default is perf.data.\n"
 "--exit-with-parent            Stop recording when the process starting\n"
 "                              simpleperf dies.\n"
+"--size-limit SIZE[K|M|G]      Stop recording after SIZE bytes of records.\n"
+"                              Default is unlimited.\n"
 "--start_profiling_fd fd_no    After starting profiling, write \"STARTED\" to\n"
 "                              <fd_no>, then close <fd_no>.\n"
 "--symfs <dir>    Look for files with symbols relative to this directory.\n"
@@ -300,6 +301,7 @@
   bool in_app_context_;
   bool trace_offcpu_;
   bool exclude_kernel_callchain_;
+  uint64_t size_limit_in_bytes_ = 0;
 
   // For CallChainJoiner
   bool allow_callchain_joiner_;
@@ -565,13 +567,8 @@
     } else if (args[i] == "-b") {
       branch_sampling_ = branch_sampling_type_map["any"];
     } else if (args[i] == "-c" || args[i] == "-f") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      char* endptr;
-      uint64_t value = strtoull(args[i].c_str(), &endptr, 0);
-      if (*endptr != '\0' || value == 0) {
-        LOG(ERROR) << "Invalid option for " << args[i-1] << ": '" << args[i] << "'";
+      uint64_t value;
+      if (!GetUintOption(args, &i, &value, 1)) {
         return false;
       }
       if (args[i-1] == "-c") {
@@ -596,11 +593,9 @@
         fp_callchain_sampling_ = false;
         dwarf_callchain_sampling_ = true;
         if (strs.size() > 1) {
-          char* endptr;
-          uint64_t size = strtoull(strs[1].c_str(), &endptr, 0);
-          if (*endptr != '\0' || size > UINT_MAX) {
-            LOG(ERROR) << "invalid dump stack size in --call-graph option: "
-                       << strs[1];
+          uint64_t size;
+          if (!android::base::ParseUint(strs[1], &size)) {
+            LOG(ERROR) << "invalid dump stack size in --call-graph option: " << strs[1];
             return false;
           }
           if ((size & 7) != 0) {
@@ -642,12 +637,7 @@
       }
       cpus_ = GetCpusFromString(args[i]);
     } else if (args[i] == "--duration") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      if (!android::base::ParseDouble(args[i].c_str(), &duration_in_sec_,
-                                      1e-9)) {
-        LOG(ERROR) << "Invalid duration: " << args[i].c_str();
+      if (!GetDoubleOption(args, &i, &duration_in_sec_, 1e-9)) {
         return false;
       }
     } else if (args[i] == "-e") {
@@ -702,12 +692,11 @@
         branch_sampling_ |= it->second;
       }
     } else if (args[i] == "-m") {
-      if (!NextArgumentOrError(args, &i)) {
+      uint64_t pages;
+      if (!GetUintOption(args, &i, &pages)) {
         return false;
       }
-      char* endptr;
-      uint64_t pages = strtoull(args[i].c_str(), &endptr, 0);
-      if (*endptr != '\0' || !IsPowerOfTwo(pages)) {
+      if (!IsPowerOfTwo(pages)) {
         LOG(ERROR) << "Invalid mmap_pages: '" << args[i] << "'";
         return false;
       }
@@ -723,12 +712,7 @@
     } else if (args[i] == "--no-callchain-joiner") {
       allow_callchain_joiner_ = false;
     } else if (args[i] == "--callchain-joiner-min-matching-nodes") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      if (!android::base::ParseUint(args[i].c_str(), &callchain_joiner_min_matching_nodes_) ||
-          callchain_joiner_min_matching_nodes_ < 1u) {
-        LOG(ERROR) << "unexpected argument for " << args[i - 1] << " option";
+      if (!GetUintOption(args, &i, &callchain_joiner_min_matching_nodes_, 1)) {
         return false;
       }
     } else if (args[i] == "-o") {
@@ -754,12 +738,13 @@
         LOG(ERROR) << "unexpected option " << args[i];
         return false;
       }
-    } else if (args[i] == "--start_profiling_fd") {
-      if (!NextArgumentOrError(args, &i)) {
+    } else if (args[i] == "--size-limit") {
+      if (!GetUintOption(args, &i, &size_limit_in_bytes_, 1, std::numeric_limits<uint64_t>::max(),
+                         true)) {
         return false;
       }
-      if (!android::base::ParseInt(args[i].c_str(), &start_profiling_fd_, 0)) {
-        LOG(ERROR) << "Invalid start_profiling_fd: " << args[i];
+    } else if (args[i] == "--start_profiling_fd") {
+      if (!GetUintOption(args, &i, &start_profiling_fd_)) {
         return false;
       }
     } else if (args[i] == "--symfs") {
@@ -1051,6 +1036,11 @@
   if (ShouldOmitRecord(record)) {
     return true;
   }
+  if (size_limit_in_bytes_ > 0u) {
+    if (size_limit_in_bytes_ < record_file_writer_->GetDataSectionSize()) {
+      return event_selection_set_.GetIOEventLoop()->ExitLoop();
+    }
+  }
   last_record_timestamp_ = record->Timestamp();
   if (unwind_dwarf_callchain_) {
     if (post_unwind_) {
diff --git a/simpleperf/cmd_record_test.cpp b/simpleperf/cmd_record_test.cpp
index a01c6bc..cab1a38 100644
--- a/simpleperf/cmd_record_test.cpp
+++ b/simpleperf/cmd_record_test.cpp
@@ -570,3 +570,17 @@
   TemporaryFile tmpfile;
   ASSERT_TRUE(RecordCmd()->Run({"-o", tmpfile.path, "--", "sleep", "1"}));
 }
+
+TEST(record_cmd, size_limit_option) {
+  std::vector<std::unique_ptr<Workload>> workloads;
+  CreateProcesses(1, &workloads);
+  std::string pid = std::to_string(workloads[0]->GetPid());
+  TemporaryFile tmpfile;
+  ASSERT_TRUE(RecordCmd()->Run({"-o", tmpfile.path, "-p", pid, "--size-limit", "1k", "--duration",
+                                "1"}));
+  std::unique_ptr<RecordFileReader> reader = RecordFileReader::CreateInstance(tmpfile.path);
+  ASSERT_TRUE(reader);
+  ASSERT_GT(reader->FileHeader().data.size, 1000u);
+  ASSERT_LT(reader->FileHeader().data.size, 2000u);
+  ASSERT_FALSE(RunRecordCmd({"--size-limit", "0"}));
+}
diff --git a/simpleperf/cmd_report.cpp b/simpleperf/cmd_report.cpp
index 6a757f0..1a4cb4c 100644
--- a/simpleperf/cmd_report.cpp
+++ b/simpleperf/cmd_report.cpp
@@ -26,7 +26,6 @@
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
-#include <android-base/parsedouble.h>
 #include <android-base/parseint.h>
 #include <android-base/stringprintf.h>
 #include <android-base/strings.h>
@@ -561,11 +560,7 @@
       }
       Dso::SetKallsyms(kallsyms);
     } else if (args[i] == "--max-stack") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      if (!android::base::ParseUint(args[i].c_str(), &callgraph_max_stack_)) {
-        LOG(ERROR) << "invalid arg for --max-stack: " << args[i];
+      if (!GetUintOption(args, &i, &callgraph_max_stack_)) {
         return false;
       }
     } else if (args[i] == "-n") {
@@ -581,13 +576,9 @@
       }
       report_filename_ = args[i];
     } else if (args[i] == "--percent-limit") {
-      if (!NextArgumentOrError(args, &i)) {
+      if (!GetDoubleOption(args, &i, &callgraph_percent_limit_)) {
         return false;
       }
-      if (!android::base::ParseDouble(args[i].c_str(),
-                                      &callgraph_percent_limit_, 0.0)) {
-        LOG(ERROR) << "invalid arg for --percent-limit: " << args[i];
-      }
     } else if (args[i] == "--pids" || args[i] == "--tids") {
       const std::string& option = args[i];
       std::unordered_set<int>& filter =
diff --git a/simpleperf/cmd_stat.cpp b/simpleperf/cmd_stat.cpp
index 21c74ea..da5b603 100644
--- a/simpleperf/cmd_stat.cpp
+++ b/simpleperf/cmd_stat.cpp
@@ -28,7 +28,6 @@
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
-#include <android-base/parsedouble.h>
 #include <android-base/strings.h>
 
 #include "command.h"
@@ -531,21 +530,11 @@
     } else if (args[i] == "--csv") {
       csv_ = true;
     } else if (args[i] == "--duration") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      if (!android::base::ParseDouble(args[i].c_str(), &duration_in_sec_,
-                                      1e-9)) {
-        LOG(ERROR) << "Invalid duration: " << args[i].c_str();
+      if (!GetDoubleOption(args, &i, &duration_in_sec_, 1e-9)) {
         return false;
       }
     } else if (args[i] == "--interval") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      if (!android::base::ParseDouble(args[i].c_str(), &interval_in_ms_,
-                                      1e-9)) {
-        LOG(ERROR) << "Invalid interval: " << args[i].c_str();
+      if (!GetDoubleOption(args, &i, &interval_in_ms_, 1e-9)) {
         return false;
       }
     } else if (args[i] == "--interval-only-values") {
diff --git a/simpleperf/cmd_trace_sched.cpp b/simpleperf/cmd_trace_sched.cpp
index afef072..6d61d9e 100644
--- a/simpleperf/cmd_trace_sched.cpp
+++ b/simpleperf/cmd_trace_sched.cpp
@@ -21,7 +21,6 @@
 #include <vector>
 
 #include <android-base/logging.h>
-#include <android-base/parsedouble.h>
 #include <android-base/stringprintf.h>
 #include <android-base/test_utils.h>
 
@@ -140,27 +139,15 @@
   size_t i;
   for (i = 0; i < args.size(); ++i) {
     if (args[i] == "--duration") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      if (!android::base::ParseDouble(args[i].c_str(), &duration_in_sec_, 1e-9)) {
-        LOG(ERROR) << "Invalid duration for " << args[i-1];
+      if (!GetDoubleOption(args, &i, &duration_in_sec_, 1e-9)) {
         return false;
       }
     } else if (args[i] == "--check-spinloop") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      if (!android::base::ParseDouble(args[i].c_str(), &spinloop_check_period_in_sec_, 1e-9)) {
-        LOG(ERROR) << "Invalid check period for " << args[i-1];
+      if (!GetDoubleOption(args, &i, &spinloop_check_period_in_sec_, 1e-9)) {
         return false;
       }
     } else if (args[i] == "--spin-rate") {
-      if (!NextArgumentOrError(args, &i)) {
-        return false;
-      }
-      if (!android::base::ParseDouble(args[i].c_str(), &spinloop_check_rate_, 1e-9, 1.0)) {
-        LOG(ERROR) << "Invalid spin rate for " << args[i-1];
+      if (!GetDoubleOption(args, &i, &spinloop_check_rate_, 1e-9, 1.0)) {
         return false;
       }
     } else if (args[i] == "--show-threads") {
diff --git a/simpleperf/command.cpp b/simpleperf/command.cpp
index 63a7dfd..9caadaa 100644
--- a/simpleperf/command.cpp
+++ b/simpleperf/command.cpp
@@ -22,6 +22,8 @@
 #include <vector>
 
 #include <android-base/logging.h>
+#include <android-base/parsedouble.h>
+#include <android-base/parseint.h>
 #include <android-base/quick_exit.h>
 
 #include "utils.h"
@@ -36,6 +38,18 @@
   return true;
 }
 
+bool Command::GetDoubleOption(const std::vector<std::string>& args, size_t* pi, double* value,
+                              double min, double max) {
+  if (!NextArgumentOrError(args, pi)) {
+    return false;
+  }
+  if (!android::base::ParseDouble(args[*pi].c_str(), value, min, max)) {
+    LOG(ERROR) << "Invalid argument for option " << args[*pi - 1] << ": " << args[*pi];
+    return false;
+  }
+  return true;
+}
+
 void Command::ReportUnknownOption(const std::vector<std::string>& args, size_t i) {
   LOG(ERROR) << "Unknown option for " << name_ << " command: '" << args[i]
              << "'. Try `simpleperf help " << name_ << "`";
diff --git a/simpleperf/command.h b/simpleperf/command.h
index 6063dbf..2ce789e 100644
--- a/simpleperf/command.h
+++ b/simpleperf/command.h
@@ -19,10 +19,13 @@
 
 #include <functional>
 #include <memory>
+#include <limits>
 #include <string>
 #include <vector>
 
+#include <android-base/logging.h>
 #include <android-base/macros.h>
+#include <android-base/parseint.h>
 
 class Command {
  public:
@@ -48,6 +51,24 @@
 
   virtual bool Run(const std::vector<std::string>& args) = 0;
 
+  template <typename T>
+  bool GetUintOption(const std::vector<std::string>& args, size_t* pi, T* value, uint64_t min = 0,
+                     uint64_t max = std::numeric_limits<T>::max(), bool allow_suffixes = false) {
+    if (!NextArgumentOrError(args, pi)) {
+      return false;
+    }
+    uint64_t tmp_value;
+    if (!android::base::ParseUint(args[*pi], &tmp_value, max, allow_suffixes) || tmp_value < min) {
+      LOG(ERROR) << "Invalid argument for option " << args[*pi - 1] << ": " << args[*pi];
+      return false;
+    }
+    *value = static_cast<T>(tmp_value);
+    return true;
+  }
+
+  bool GetDoubleOption(const std::vector<std::string>& args, size_t* pi, double* value,
+                       double min = 0, double max = std::numeric_limits<double>::max());
+
  protected:
   bool NextArgumentOrError(const std::vector<std::string>& args, size_t* pi);
   void ReportUnknownOption(const std::vector<std::string>& args, size_t i);
diff --git a/simpleperf/command_test.cpp b/simpleperf/command_test.cpp
index 18cb569..29c745e 100644
--- a/simpleperf/command_test.cpp
+++ b/simpleperf/command_test.cpp
@@ -43,3 +43,30 @@
   UnRegisterCommand("mock1");
   ASSERT_EQ(command_count, GetAllCommandNames().size());
 }
+
+TEST(command, GetValueForOption) {
+  MockCommand command;
+  uint64_t value;
+  size_t i;
+  for (bool allow_suffixes : {true, false}) {
+    i = 0;
+    ASSERT_TRUE(command.GetUintOption({"-s", "156"}, &i, &value, 0,
+                                      std::numeric_limits<uint64_t>::max(), allow_suffixes));
+    ASSERT_EQ(i, 1u);
+    ASSERT_EQ(value, 156u);
+  }
+  i = 0;
+  ASSERT_TRUE(command.GetUintOption({"-s", "156k"}, &i, &value, 0,
+                                    std::numeric_limits<uint64_t>::max(), true));
+  ASSERT_EQ(value, 156 * (1ULL << 10));
+  i = 0;
+  ASSERT_FALSE(command.GetUintOption({"-s"}, &i, &value));
+  i = 0;
+  ASSERT_FALSE(command.GetUintOption({"-s", "0"}, &i, &value, 1));
+  i = 0;
+  ASSERT_FALSE(command.GetUintOption({"-s", "156"}, &i, &value, 0, 155));
+  i = 0;
+  double double_value;
+  ASSERT_TRUE(command.GetDoubleOption({"-s", "3.2"}, &i, &double_value, 0, 4));
+  ASSERT_DOUBLE_EQ(double_value, 3.2);
+}
diff --git a/simpleperf/record_file.h b/simpleperf/record_file.h
index c90b269..4ec6c31 100644
--- a/simpleperf/record_file.h
+++ b/simpleperf/record_file.h
@@ -47,6 +47,7 @@
   bool WriteAttrSection(const std::vector<EventAttrWithId>& attr_ids);
   bool WriteRecord(const Record& record);
 
+  uint64_t GetDataSectionSize() const { return data_section_size_; }
   bool ReadDataSection(const std::function<void(const Record*)>& callback);
 
   bool BeginWriteFeatures(size_t feature_count);
diff --git a/simpleperf/scripts/run_simpleperf_without_usb_connection.py b/simpleperf/scripts/run_simpleperf_without_usb_connection.py
index b26b7b9..a3524f6 100644
--- a/simpleperf/scripts/run_simpleperf_without_usb_connection.py
+++ b/simpleperf/scripts/run_simpleperf_without_usb_connection.py
@@ -45,6 +45,8 @@
     shell_cmd = 'cd /data/local/tmp && nohup ./simpleperf record ' + args.record_options
     if args.app:
         shell_cmd += ' --app ' + args.app
+    if args.size_limit:
+        shell_cmd += ' --size-limit ' + args.size_limit
     shell_cmd += ' >/data/local/tmp/simpleperf_output 2>&1'
     print('shell_cmd: %s' % shell_cmd)
     subproc = subprocess.Popen([adb.adb_path, 'shell', shell_cmd])
@@ -67,8 +69,8 @@
         print('Waiting for simpleperf process to finish...')
         while adb.run(['shell', 'pidof', 'simpleperf']):
             time.sleep(1)
-    adb.check_run(['pull', '/data/local/tmp/perf.data', args.perf_data_path])
     adb.run(['shell', 'cat', '/data/local/tmp/simpleperf_output'])
+    adb.check_run(['pull', '/data/local/tmp/perf.data', args.perf_data_path])
     print('The recording data has been collected in %s.' % args.perf_data_path)
 
 def main():
@@ -79,9 +81,12 @@
     start_parser.add_argument('-r', '--record_options',
                               default='-e task-clock:u -g',
                               help="""Set options for `simpleperf record` command.
-                                      Default is '-e task-clock:u -g'.""")
+                                      Default is `-e task-clock:u -g`.""")
     start_parser.add_argument('-p', '--app', help="""Profile an Android app, given the package
-                              name. Like -p com.example.android.myapp.""")
+                              name. Like `-p com.example.android.myapp`.""")
+    start_parser.add_argument('--size_limit', type=str,
+                              help="""Stop profiling when recording data reaches
+                                      [size_limit][K|M|G] bytes. Like `--size_limit 1M`.""")
     start_parser.set_defaults(func=start_recording)
     stop_parser = subparsers.add_parser('stop', help='Stop recording.')
     stop_parser.add_argument('-o', '--perf_data_path', default='perf.data', help="""The path to
diff --git a/simpleperf/scripts/test.py b/simpleperf/scripts/test.py
index 64bc4a4..b8aeeb6 100644
--- a/simpleperf/scripts/test.py
+++ b/simpleperf/scripts/test.py
@@ -451,7 +451,7 @@
     def test_run_simpleperf_without_usb_connection(self):
         self.adb.check_run(['shell', 'am', 'start', '-n', self.package_name + '/.MainActivity'])
         self.run_cmd(['run_simpleperf_without_usb_connection.py', 'start', '-p',
-                      self.package_name])
+                      self.package_name, '--size_limit', '1M'])
         self.adb.check_run(['kill-server'])
         time.sleep(3)
         self.run_cmd(['run_simpleperf_without_usb_connection.py', 'stop'])