[trace][intel pt] Add a nice parser for the trace size

rnofenko · Walter Erquinigo · commit db73a52d7b19 · 2022-07-13T10:53:14.000-07:00
Thanks to rnofenko@fb.com for coming up with these changes. This diff adds support for passing units in the trace size inputs. For example, it's now possible to specify 64KB as the trace size, instead of the problematic 65536. This makes the user experience a bit friendlier. Differential Revision: https://reviews.llvm.org/D129613
diff --git a/lldb/source/Plugins/Trace/intel-pt/CommandObjectTraceStartIntelPT.cpp b/lldb/source/Plugins/Trace/intel-pt/CommandObjectTraceStartIntelPT.cpp
@@ -32,13 +32,12 @@ Status CommandObjectThreadTraceStartIntelPT::CommandOptions::SetOptionValue(
 
   switch (short_option) {
   case 's': {
-    int64_t ipt_trace_size;
-    if (option_arg.empty() || option_arg.getAsInteger(0, ipt_trace_size) ||
-        ipt_trace_size < 0)
-      error.SetErrorStringWithFormat("invalid integer value for option '%s'",
-                                     option_arg.str().c_str());
+    if (Optional<uint64_t> bytes =
+            ParsingUtils::ParseUserFriendlySizeExpression(option_arg))
+      m_ipt_trace_size = *bytes;
     else
-      m_ipt_trace_size = ipt_trace_size;
+      error.SetErrorStringWithFormat("invalid bytes expression for '%s'",
+                                     option_arg.str().c_str());
     break;
   }
   case 't': {
@@ -98,24 +97,21 @@ Status CommandObjectProcessTraceStartIntelPT::CommandOptions::SetOptionValue(
 
   switch (short_option) {
   case 's': {
-    int64_t ipt_trace_size;
-    if (option_arg.empty() || option_arg.getAsInteger(0, ipt_trace_size) ||
-        ipt_trace_size < 0)
-      error.SetErrorStringWithFormat("invalid integer value for option '%s'",
-                                     option_arg.str().c_str());
+    if (Optional<uint64_t> bytes =
+            ParsingUtils::ParseUserFriendlySizeExpression(option_arg))
+      m_ipt_trace_size = *bytes;
     else
-      m_ipt_trace_size = ipt_trace_size;
+      error.SetErrorStringWithFormat("invalid bytes expression for '%s'",
+                                     option_arg.str().c_str());
     break;
   }
   case 'l': {
-    int64_t process_buffer_size_limit;
-    if (option_arg.empty() ||
-        option_arg.getAsInteger(0, process_buffer_size_limit) ||
-        process_buffer_size_limit < 0)
-      error.SetErrorStringWithFormat("invalid integer value for option '%s'",
-                                     option_arg.str().c_str());
+    if (Optional<uint64_t> bytes =
+            ParsingUtils::ParseUserFriendlySizeExpression(option_arg))
+      m_process_buffer_size_limit = *bytes;
     else
-      m_process_buffer_size_limit = process_buffer_size_limit;
+      error.SetErrorStringWithFormat("invalid bytes expression for '%s'",
+                                     option_arg.str().c_str());
     break;
   }
   case 't': {
@@ -168,3 +164,45 @@ bool CommandObjectProcessTraceStartIntelPT::DoExecute(
 
   return result.Succeeded();
 }
+
+Optional<uint64_t>
+ParsingUtils::ParseUserFriendlySizeExpression(llvm::StringRef size_expression) {
+  if (size_expression.empty()) {
+    return llvm::None;
+  }
+  const uint64_t kBytesMultiplier = 1;
+  const uint64_t kKibiBytesMultiplier = 1024;
+  const uint64_t kMebiBytesMultiplier = 1024 * 1024;
+
+  DenseMap<StringRef, uint64_t> multipliers = {
+      {"mib", kMebiBytesMultiplier}, {"mb", kMebiBytesMultiplier},
+      {"m", kMebiBytesMultiplier},   {"kib", kKibiBytesMultiplier},
+      {"kb", kKibiBytesMultiplier},  {"k", kKibiBytesMultiplier},
+      {"b", kBytesMultiplier},       {"", kBytesMultiplier}};
+
+  const auto non_digit_index = size_expression.find_first_not_of("0123456789");
+  if (non_digit_index == 0) { // expression starts from from non-digit char.
+    return llvm::None;
+  }
+
+  const llvm::StringRef number_part =
+      non_digit_index == llvm::StringRef::npos
+          ? size_expression
+          : size_expression.substr(0, non_digit_index);
+  uint64_t parsed_number;
+  if (number_part.getAsInteger(10, parsed_number)) {
+    return llvm::None;
+  }
+
+  if (non_digit_index != llvm::StringRef::npos) { // if expression has units.
+    const auto multiplier = size_expression.substr(non_digit_index).lower();
+
+    auto it = multipliers.find(multiplier);
+    if (it == multipliers.end())
+      return llvm::None;
+
+    return parsed_number * it->second;
+  } else {
+    return parsed_number;
+  }
+}
diff --git a/lldb/source/Plugins/Trace/intel-pt/CommandObjectTraceStartIntelPT.h b/lldb/source/Plugins/Trace/intel-pt/CommandObjectTraceStartIntelPT.h
@@ -109,6 +109,23 @@ class CommandObjectProcessTraceStartIntelPT : public CommandObjectParsed {
   CommandOptions m_options;
 };
 
+namespace ParsingUtils {
+/// Convert an integral size expression like 12KiB or 4MB into bytes. The units
+/// are taken loosely to help users input sizes into LLDB, e.g. KiB and KB are
+/// considered the same (2^20 bytes) for simplicity.
+///
+/// \param[in] size_expression
+///     String expression which is an integral number plus a unit that can be
+///     lower or upper case. Supported units: K, KB and KiB for 2^10 bytes; M,
+///     MB and MiB for 2^20 bytes; and B for bytes. A single integral number is
+///     considered bytes.
+/// \return
+///   The converted number of bytes or \a llvm::None if the expression is
+///   invalid.
+llvm::Optional<uint64_t>
+ParseUserFriendlySizeExpression(llvm::StringRef size_expression);
+} // namespace ParsingUtils
+
 } // namespace trace_intel_pt
 } // namespace lldb_private
 
diff --git a/lldb/source/Plugins/Trace/intel-pt/TraceIntelPT.cpp b/lldb/source/Plugins/Trace/intel-pt/TraceIntelPT.cpp
@@ -476,7 +476,20 @@ Error TraceIntelPT::Start(llvm::ArrayRef<lldb::tid_t> tids,
 
   if (configuration) {
     if (StructuredData::Dictionary *dict = configuration->GetAsDictionary()) {
-      dict->GetValueForKeyAsInteger("iptTraceSize", ipt_trace_size);
+      llvm::StringRef ipt_trace_size_not_parsed;
+      if (dict->GetValueForKeyAsString("iptTraceSize",
+                                       ipt_trace_size_not_parsed)) {
+        if (Optional<uint64_t> bytes =
+                ParsingUtils::ParseUserFriendlySizeExpression(
+                    ipt_trace_size_not_parsed))
+          ipt_trace_size = *bytes;
+        else
+          return createStringError(inconvertibleErrorCode(),
+                                   "iptTraceSize is wrong bytes expression");
+      } else {
+        dict->GetValueForKeyAsInteger("iptTraceSize", ipt_trace_size);
+      }
+
       dict->GetValueForKeyAsBoolean("enableTsc", enable_tsc);
       dict->GetValueForKeyAsInteger("psbPeriod", psb_period);
     } else {
diff --git a/lldb/source/Plugins/Trace/intel-pt/TraceIntelPTOptions.td b/lldb/source/Plugins/Trace/intel-pt/TraceIntelPTOptions.td
@@ -11,7 +11,9 @@ let Command = "thread trace start intel pt" in {
     Arg<"Value">,
     Desc<"Trace size in bytes per thread. It must be a power of 2 greater "
          "than or equal to 4096 (2^12). The trace is circular keeping "
-         "the most recent data. Defaults to 4096 bytes.">;
+         "the most recent data. Defaults to 4096 bytes. It's possible to "
+         "specify size using multiples of unit bytes, e.g., 4KB, 1MB, 1MiB, "
+         "where 1K is 1024 bytes and 1M is 1048576 bytes.">;
   def thread_trace_start_intel_pt_tsc: Option<"tsc", "t">,
     Group<1>,
     Desc<"Enable the use of TSC timestamps. This is supported on all devices "
@@ -40,7 +42,8 @@ let Command = "process trace start intel pt" in {
     Arg<"Value">,
     Desc<"Size in bytes used by each individual per-thread or per-cpu trace "
          "buffer. It must be a power of 2 greater than or equal to 4096 (2^12) "
-         "bytes.">;
+         "bytes. It's possible to specify a unit for these bytes, like 4KB, "
+         "16KiB or 1MB. Lower case units are allowed for convenience.">;
   def process_trace_start_intel_pt_per_cpu_tracing:
     Option<"per-cpu-tracing", "c">,
     Group<1>,
@@ -53,7 +56,8 @@ let Command = "process trace start intel pt" in {
          "option forces the capture of TSC timestamps (see --tsc). Also, this "
          "option can't be used simulatenously with any other trace sessions "
          "because of its system-wide nature.">;
-  def process_trace_start_intel_pt_process_size_limit: Option<"total-size-limit", "l">,
+  def process_trace_start_intel_pt_process_size_limit:
+    Option<"total-size-limit", "l">,
     Group<1>,
     Arg<"Value">,
     Desc<"Maximum total trace size per process in bytes. This limit applies to "
@@ -62,7 +66,9 @@ let Command = "process trace start intel pt" in {
          "Whenever a thread is attempted to be traced due to this command and "
          "the limit would be reached, the process is stopped with a "
          "\"processor trace\" reason, so that the user can retrace the process "
-         "if needed. Defaults to 500MB.">;
+         "if needed. Defaults to 500MB. It's possible to specify a unit for "
+         "these bytes, like 4KB, 16KiB or 1MB. Lower case units are allowed "
+         "for convenience.">;
   def process_trace_start_intel_pt_tsc: Option<"tsc", "t">,
     Group<1>,
     Desc<"Enable the use of TSC timestamps. This is supported on all devices "
diff --git a/lldb/test/API/commands/trace/TestTraceStartStop.py b/lldb/test/API/commands/trace/TestTraceStartStop.py
@@ -46,6 +46,62 @@ def testStartSessionWithWrongSize(self):
 
         self.traceStartThread(iptTraceSize=1048576)
 
+    @testSBAPIAndCommands
+    def testStartSessionWithSizeDeclarationInUnits(self):
+        self.expect("file " + os.path.join(self.getSourceDir(), "intelpt-trace", "a.out"))
+        self.expect("b main")
+        self.expect("r")
+
+        self.traceStartThread(
+            error=True, iptTraceSize="abc",
+            substrs=["invalid bytes expression for 'abc'"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="123.12",
+            substrs=["invalid bytes expression for '123.12'"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="\"\"",
+            substrs=["invalid bytes expression for ''"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="2000B",
+            substrs=["The intel pt trace size must be a power of 2 greater than or equal to 4096 (2^12) bytes. It was 2000"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="3MB",
+            substrs=["The intel pt trace size must be a power of 2 greater than or equal to 4096 (2^12) bytes. It was 3145728"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="3MiB",
+            substrs=["The intel pt trace size must be a power of 2 greater than or equal to 4096 (2^12) bytes. It was 3145728"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="3mib",
+            substrs=["The intel pt trace size must be a power of 2 greater than or equal to 4096 (2^12) bytes. It was 3145728"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="3M",
+            substrs=["The intel pt trace size must be a power of 2 greater than or equal to 4096 (2^12) bytes. It was 3145728"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="3KB",
+            substrs=["The intel pt trace size must be a power of 2 greater than or equal to 4096 (2^12) bytes. It was 3072"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="3KiB",
+            substrs=["The intel pt trace size must be a power of 2 greater than or equal to 4096 (2^12) bytes. It was 3072"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="3K",
+            substrs=["The intel pt trace size must be a power of 2 greater than or equal to 4096 (2^12) bytes. It was 3072"])
+
+        self.traceStartThread(
+            error=True, iptTraceSize="3MS",
+            substrs=["invalid bytes expression for '3MS'"])
+
+        self.traceStartThread(iptTraceSize="1048576")
+
     @skipIf(oslist=no_match(['linux']), archs=no_match(['i386', 'x86_64']))
     def testSBAPIHelp(self):
         self.expect("file " + os.path.join(self.getSourceDir(), "intelpt-trace", "a.out"))
