[BOLT][NFC] Simplify preprocessProfile

[Amir]

Summary

Move out prepareToParse lambda, generalize it to handle mem events perf process.

[rafauler]

bolt/lib/Profile/DataAggregator.cpp

525–528

I don't think this will be called, because we're calling "exit(1)" directly

[Amir]

Revert deleteTempFiles change

[Amir]

bolt/lib/Profile/DataAggregator.cpp

525–528

You're right. Call to exit doesn't unwind stack and destructors are not called.

[rafauler]

LGTM

bolt/include/bolt/Profile/DataAggregator.h

private:
 
   /// Consume the entire line
   void consumeRestOfLine();
 
   /// True if the next token in the parsing buffer is a new line, but don't
   /// consume it (peek only).
   bool checkNewLine();
 
+  using PerfProcessErrorCallbackTy = std::function<void(int, StringRef)>;
+  /// Prepare to parse data from a given perf script invocation.
+  /// Returns an invocation exit code.
+  int prepareToParse(StringRef Name, PerfProcessInfo &Process,
+                     PerfProcessErrorCallbackTy Callback);
+
   /// Parse a single LBR entry as output by perf script -Fbrstack
   ErrorOr<LBREntry> parseLBREntry();
 
   /// Parse and pre-aggregate branch events.
   std::error_code parseBranchEvents();
 
   /// Process all branch events.
   void processBranchEvents();

bolt/lib/Profile/DataAggregator.cpp

if (MMapInfoIter != BinaryMMapInfo.end()) {
       if (errs().has_colors())
         errs().resetColor();
 
       exit(1);
     }
   }
 }
 
-Error DataAggregator::preprocessProfile(BinaryContext &BC) {
-  this->BC = &BC;
-
-  if (opts::ReadPreAggregated) {
-    parsePreAggregated();
-    return Error::success();
-  }
-
-  if (std::optional<StringRef> FileBuildID = BC.getFileBuildID()) {
-    outs() << "BOLT-INFO: binary build-id is:     " << *FileBuildID << "\n";
-    processFileBuildID(*FileBuildID);
-  } else {
-    errs() << "BOLT-WARNING: build-id will not be checked because we could "
-              "not read one from input binary\n";
-  }
-
-  auto prepareToParse = [&](StringRef Name, PerfProcessInfo &Process) {
-    std::string Error;
-    outs() << "PERF2BOLT: waiting for perf " << Name
-           << " collection to finish...\n";
-    sys::ProcessInfo PI = sys::Wait(Process.PI, std::nullopt, &Error);
+int DataAggregator::prepareToParse(StringRef Name, PerfProcessInfo &Process,
+                                   PerfProcessErrorCallbackTy Callback) {
 
-    if (!Error.empty()) {
-      errs() << "PERF-ERROR: " << PerfPath << ": " << Error << "\n";
-      deleteTempFiles();
-      exit(1);
-    }
 
-    if (PI.ReturnCode != 0) {
-      ErrorOr<std::unique_ptr<MemoryBuffer>> ErrorMB =
-          MemoryBuffer::getFileOrSTDIN(Process.StderrPath.data());
-      StringRef ErrBuf = (*ErrorMB)->getBuffer();
 
-      errs() << "PERF-ERROR: return code " << PI.ReturnCode << "\n";
-      errs() << ErrBuf;
-      deleteTempFiles();
-      exit(1);
-    }
+    Callback(PI.ReturnCode, ErrBuf);
+    return PI.ReturnCode;
 
-    ErrorOr<std::unique_ptr<MemoryBuffer>> MB =
-        MemoryBuffer::getFileOrSTDIN(Process.StdoutPath.data());
-    if (std::error_code EC = MB.getError()) {
-      errs() << "Cannot open " << Process.StdoutPath.data() << ": "
-             << EC.message() << "\n";
-      deleteTempFiles();
-      exit(1);
-    }
 
-    FileBuf = std::move(*MB);
-    ParsingBuf = FileBuf->getBuffer();
-    Col = 0;
-    Line = 1;
+  return PI.ReturnCode;
+}
+
+Error DataAggregator::preprocessProfile(BinaryContext &BC) {
+  this->BC = &BC;
+
+  if (opts::ReadPreAggregated) {
+    parsePreAggregated();
+    return Error::success();
+  }
+
+  if (std::optional<StringRef> FileBuildID = BC.getFileBuildID()) {
+    outs() << "BOLT-INFO: binary build-id is:     " << *FileBuildID << "\n";
+    processFileBuildID(*FileBuildID);
+  } else {
+    errs() << "BOLT-WARNING: build-id will not be checked because we could "
+              "not read one from input binary\n";
+  }
+
+  auto ErrorCallback = [](int ReturnCode, StringRef ErrBuf) {
+    errs() << "PERF-ERROR: return code " << ReturnCode << "\n" << ErrBuf;
+    exit(1);
+  };
+
+  auto MemEventsErrorCallback = [&](int ReturnCode, StringRef ErrBuf) {
+    Regex NoData("Samples for '.*' event do not have ADDR attribute set. "
+                 "Cannot print 'addr' field.");
+    if (!NoData.match(ErrBuf))
+      ErrorCallback(ReturnCode, ErrBuf);
   };
 
   if (opts::LinuxKernelMode) {
     // Current MMap parsing logic does not work with linux kernel.
     // MMap entries for linux kernel uses PERF_RECORD_MMAP
     // format instead of typical PERF_RECORD_MMAP2 format.
     // Since linux kernel address mapping is absolute (same as
     // in the ELF file), we avoid parsing MMap in linux kernel mode.
     // While generating optimized linux kernel binary, we may need
     // to parse MMap entries.
 
     // In linux kernel mode, we analyze and optimize
     // all linux kernel binary instructions, irrespective
     // of whether they are due to system calls or due to
     // interrupts. Therefore, we cannot ignore interrupt
     // in Linux kernel mode.
     opts::IgnoreInterruptLBR = false;
   } else {
-    prepareToParse("mmap events", MMapEventsPPI);
+    prepareToParse("mmap events", MMapEventsPPI, ErrorCallback);
     if (parseMMapEvents())
       errs() << "PERF2BOLT: failed to parse mmap events\n";
   }
 
-  prepareToParse("task events", TaskEventsPPI);
+  prepareToParse("task events", TaskEventsPPI, ErrorCallback);
   if (parseTaskEvents())
     errs() << "PERF2BOLT: failed to parse task events\n";
 
   filterBinaryMMapInfo();
-  prepareToParse("events", MainEventsPPI);
+  prepareToParse("events", MainEventsPPI, ErrorCallback);
 
   if (opts::HeatmapMode) {
     if (std::error_code EC = printLBRHeatMap()) {
       errs() << "ERROR: failed to print heat map: " << EC.message() << '\n';
       exit(1);
     }
     exit(0);
   }
 
   if ((!opts::BasicAggregation && parseBranchEvents()) ||
       (opts::BasicAggregation && parseBasicEvents()))
     errs() << "PERF2BOLT: failed to parse samples\n";
 
   // We can finish early if the goal is just to generate data for autofdo
   if (opts::WriteAutoFDOData) {
     if (std::error_code EC = writeAutoFDOData(opts::OutputFilename))
       errs() << "Error writing autofdo data to file: " << EC.message() << "\n";
 
     deleteTempFiles();
     exit(0);
   }
 
   // Special handling for memory events
-  std::string Error;
-  sys::ProcessInfo PI = sys::Wait(MemEventsPPI.PI, std::nullopt, &Error);
-  if (PI.ReturnCode != 0) {
-    ErrorOr<std::unique_ptr<MemoryBuffer>> MB =
-        MemoryBuffer::getFileOrSTDIN(MemEventsPPI.StderrPath.data());
-    StringRef ErrBuf = (*MB)->getBuffer();
-
-    deleteTempFiles();
-
-    Regex NoData("Samples for '.*' event do not have ADDR attribute set. "
-                 "Cannot print 'addr' field.");
-    if (!NoData.match(ErrBuf)) {
-      errs() << "PERF-ERROR: return code " << PI.ReturnCode << "\n";
-      errs() << ErrBuf;
-      exit(1);
-    }
+  if (prepareToParse("mem events", MemEventsPPI, MemEventsErrorCallback))
     return Error::success();
-  }
-
-  ErrorOr<std::unique_ptr<MemoryBuffer>> MB =
-      MemoryBuffer::getFileOrSTDIN(MemEventsPPI.StdoutPath.data());
-  if (std::error_code EC = MB.getError()) {
-    errs() << "Cannot open " << MemEventsPPI.StdoutPath.data() << ": "
-           << EC.message() << "\n";
-    deleteTempFiles();
-    exit(1);
-  }
 
-  FileBuf = std::move(*MB);
-  ParsingBuf = FileBuf->getBuffer();
-  Col = 0;
-  Line = 1;
   if (const std::error_code EC = parseMemEvents())
     errs() << "PERF2BOLT: failed to parse memory events: " << EC.message()
            << '\n';
 
   deleteTempFiles();
 
   return Error::success();
 }