5 files changed, 27 insertions, 14 deletions

diff --git a/llvm/lib/XRay/FDRTraceWriter.cpp b/llvm/lib/XRay/FDRTraceWriter.cpp
index d0206e775a8..4f40593cba0 100644
--- a/llvm/lib/XRay/FDRTraceWriter.cpp
+++ b/llvm/lib/XRay/FDRTraceWriter.cpp
@@ -94,9 +94,10 @@ Error FDRTraceWriter::visit(TSCWrapRecord &R) {
 }
 
 Error FDRTraceWriter::visit(CustomEventRecord &R) {
-  if (auto E = writeMetadata<5u>(OS, R.size(), R.tsc()))
+  if (auto E = writeMetadata<5u>(OS, R.size(), R.tsc(), R.cpu()))
     return E;
-  ArrayRef<char> Bytes(R.data().data(), R.data().size());
+  auto D = R.data();
+  ArrayRef<char> Bytes(D.data(), D.size());
   OS.write(Bytes);
   return Error::success();
 }
@@ -127,7 +128,7 @@ Error FDRTraceWriter::visit(FunctionRecord &R) {
   OS.write(TypeRecordFuncId);
   OS.write(R.delta());
   return Error::success();
-} // namespace xray
+}
 
 } // namespace xray
 } // namespace llvm
diff --git a/llvm/lib/XRay/FileHeaderReader.cpp b/llvm/lib/XRay/FileHeaderReader.cpp
index 967e85f30d2..9dea217840b 100644
--- a/llvm/lib/XRay/FileHeaderReader.cpp
+++ b/llvm/lib/XRay/FileHeaderReader.cpp
@@ -63,8 +63,7 @@ Expected<XRayFileHeader> readBinaryFormatHeader(DataExtractor &HeaderExtractor,
   // Manually advance the offset pointer 16 bytes, after getting a raw memcpy
   // from the underlying data.
   OffsetPtr += 16;
-  if (FileHeader.Version != 1 && FileHeader.Version != 2 &&
-      FileHeader.Version != 3)
+  if (FileHeader.Version < 1 || FileHeader.Version > 4)
     return createStringError(std::make_error_code(std::errc::invalid_argument),
                              "Unsupported XRay file version: %d at offset %d",
                              FileHeader.Version, OffsetPtr);
diff --git a/llvm/lib/XRay/RecordInitializer.cpp b/llvm/lib/XRay/RecordInitializer.cpp
index fe76f7d79fb..2ebaa1cec26 100644
--- a/llvm/lib/XRay/RecordInitializer.cpp
+++ b/llvm/lib/XRay/RecordInitializer.cpp
@@ -118,6 +118,19 @@ Error RecordInitializer::visit(CustomEventRecord &R) {
         std::make_error_code(std::errc::invalid_argument),
         "Cannot read a custom event TSC field at offset %d.", OffsetPtr);
 
+  // For version 4 onwards, of the FDR log, we want to also capture the CPU ID
+  // of the custom event.
+  if (Version >= 4) {
+    PreReadOffset = OffsetPtr;
+    R.CPU = E.getU16(&OffsetPtr);
+    if (PreReadOffset == OffsetPtr)
+      return createStringError(
+          std::make_error_code(std::errc::invalid_argument),
+          "Missing CPU field at offset %d", OffsetPtr);
+  }
+
+  assert(OffsetPtr > BeginOffset &&
+         OffsetPtr - BeginOffset <= MetadataRecord::kMetadataBodySize);
   OffsetPtr += MetadataRecord::kMetadataBodySize - (OffsetPtr - BeginOffset);
 
   // Next we read in a fixed chunk of data from the given offset.
diff --git a/llvm/lib/XRay/RecordPrinter.cpp b/llvm/lib/XRay/RecordPrinter.cpp
index 09b25ddba25..81d77f67cc1 100644
--- a/llvm/lib/XRay/RecordPrinter.cpp
+++ b/llvm/lib/XRay/RecordPrinter.cpp
@@ -35,8 +35,9 @@ Error RecordPrinter::visit(TSCWrapRecord &R) {
 }
 
 Error RecordPrinter::visit(CustomEventRecord &R) {
-  OS << formatv("<Custom Event: tsc = {0}, size = {1}, data = '{2}'>", R.tsc(),
-                R.size(), R.data())
+  OS << formatv(
+            "<Custom Event: tsc = {0}, cpu = {1}, size = {2}, data = '{3}'>",
+            R.tsc(), R.cpu(), R.size(), R.data())
      << Delim;
   return Error::success();
 }
diff --git a/llvm/lib/XRay/Trace.cpp b/llvm/lib/XRay/Trace.cpp
index 1d7c723864d..e7b878cb83f 100644
--- a/llvm/lib/XRay/Trace.cpp
+++ b/llvm/lib/XRay/Trace.cpp
@@ -310,12 +310,11 @@ Error loadFDRLog(StringRef Data, bool IsLittleEndian,
   {
     for (auto &PTB : Index) {
       auto &Blocks = PTB.second;
-      llvm::sort(
-          Blocks,
-          [](const BlockIndexer::Block &L, const BlockIndexer::Block &R) {
-            return (L.WallclockTime->seconds() < R.WallclockTime->seconds() &&
-                    L.WallclockTime->nanos() < R.WallclockTime->nanos());
-          });
+      llvm::sort(Blocks, [](const BlockIndexer::Block &L,
+                            const BlockIndexer::Block &R) {
+        return (L.WallclockTime->seconds() < R.WallclockTime->seconds() &&
+                L.WallclockTime->nanos() < R.WallclockTime->nanos());
+      });
       auto Adder = [&](const XRayRecord &R) { Records.push_back(R); };
       TraceExpander Expander(Adder, FileHeader.Version);
       for (auto &B : Blocks) {
@@ -435,7 +434,7 @@ Expected<Trace> llvm::xray::loadTrace(const DataExtractor &DE, bool Sort) {
     }
     break;
   case FLIGHT_DATA_RECORDER_FORMAT:
-    if (Version == 1 || Version == 2 || Version == 3) {
+    if (Version >= 1 && Version <= 4) {
       if (auto E = loadFDRLog(DE.getData(), DE.isLittleEndian(), T.FileHeader,
                               T.Records))
         return std::move(E);