[XRay] Improve FDR trace handling and error messaging

Summary:
This change covers a number of things spanning LLVM and compiler-rt,
which are related in a non-trivial way.

In LLVM, we have a library that handles the FDR mode even log loading,
which uses C++'s runtime polymorphism feature to better faithfully
represent the events that are written down by the FDR mode runtime. We
do this by interpreting a trace that's serliased in a common format
agreed upon by both the trace loading library and the FDR mode runtime.
This library is under active development, which consists of features
allowing us to reconstitute a higher-level event log.

This event log is used by the conversion and visualisation tools we have
for interpreting XRay traces.

One of the tools we have is a diagnostic tool in llvm-xray called
`fdr-dump` which we've been using to debug our expectations of what the
FDR runtime should be writing and what the logical FDR event log
structures are. We use this fairly extensively to reason about why some
non-trivial traces we're generating with FDR mode runtimes fail to
convert or fail to parse correctly.

One of these failures we've found in manual debugging of some of the
traces we've seen involve an inconsistency between the buffer extents (a
record indicating how many bytes to follow are part of a logical
thread's event log) and the record of the bytes written into the log --
sometimes it turns out the data could be garbage, due to buffers being
recycled, but sometimes we're seeing the buffer extent indicating a log
is "shorter" than the actual records associated with the buffer. This
case happens particularly with function entry records with a call
argument.

This change for now updates the FDR mode runtime to write the bytes for
the function call and arg record before updating the buffer extents
atomically, allowing multiple threads to see a consistent view of the
data in the buffer using the atomic counter associated with a buffer.
What we're trying to prevent here is partial updates where we see the
intermediary updates to the buffer extents (function record size then
call argument record size) becoming observable from another thread, for
instance, one doing the serialization/flushing.

To do both diagnose this issue properly, we need to be able to honour
the extents being set in the `BufferExtents` records marking the
beginning of the logical buffers when reading an FDR trace. Since LLVM
doesn't use C++'s RTTI mechanism, we instead follow the advice in the
documentation for LLVM Style RTTI
(https://llvm.org/docs/HowToSetUpLLVMStyleRTTI.html). We then rely on
this RTTI feature to ensure that our file-based record producer (our
streaming "deserializer") can honour the extents of individual buffers
as we interpret traces.

This also sets us up to be able to eventually do smart
skipping/continuation of FDR logs, seeking instead to find BufferExtents
records in cases where we find potentially recoverable errors. In the
meantime, we make this change to operate in a strict mode when reading
logical buffers with extent records.

Reviewers: mboerger

Subscribers: hiraditya, llvm-commits, jfb

Differential Revision: https://reviews.llvm.org/D54201

llvm-svn: 346473

include/llvm/XRay/FDRRecordProducer.h

@@ -29,6 +29,11 @@ class FileBasedRecordProducer : public RecordProducer {
   const XRayFileHeader &Header;
   DataExtractor &E;
   uint32_t &OffsetPtr;
+  uint32_t CurrentBufferBytes = 0;
+
+  // Helper function which gets the next record by speculatively reading through
+  // the log, finding a buffer extents record.
+  Expected<std::unique_ptr<Record>> findNextBufferExtent();
 
 public:
   FileBasedRecordProducer(const XRayFileHeader &FH, DataExtractor &DE,

include/llvm/XRay/FDRRecords.h

@@ -14,10 +14,14 @@
 #ifndef LLVM_LIB_XRAY_FDRRECORDS_H_
 #define LLVM_LIB_XRAY_FDRRECORDS_H_
 
+#include <cstdint>
+#include <string>
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/DataExtractor.h"
 #include "llvm/Support/Error.h"
 #include "llvm/XRay/XRayRecord.h"
-#include <cstdint>
 
 namespace llvm {
 namespace xray {
@@ -26,21 +30,37 @@ class RecordVisitor;
 class RecordInitializer;
 
 class Record {
-protected:
-  enum class Type {
-    Unknown,
-    Function,
-    Metadata,
+public:
+  enum class RecordKind {
+    RK_Metadata,
+    RK_Metadata_BufferExtents,
+    RK_Metadata_WallClockTime,
+    RK_Metadata_NewCPUId,
+    RK_Metadata_TSCWrap,
+    RK_Metadata_CustomEvent,
+    RK_Metadata_CustomEventV5,
+    RK_Metadata_CallArg,
+    RK_Metadata_PIDEntry,
+    RK_Metadata_NewBuffer,
+    RK_Metadata_EndOfBuffer,
+    RK_Metadata_TypedEvent,
+    RK_Metadata_LastMetadata,
+    RK_Function,
   };
 
+  static StringRef kindToString(RecordKind K);
+
+private:
+  const RecordKind T;
+
 public:
   Record(const Record &) = delete;
   Record(Record &&) = delete;
   Record &operator=(const Record &) = delete;
   Record &operator=(Record &&) = delete;
-  Record() = default;
+  explicit Record(RecordKind T) : T(T) {}
 
-  virtual Type type() const = 0;
+  RecordKind getRecordType() const { return T; }
 
   // Each Record should be able to apply an abstract visitor, and choose the
   // appropriate function in the visitor to invoke, given its own type.
@@ -50,10 +70,6 @@ public:
 };
 
 class MetadataRecord : public Record {
-protected:
-  static constexpr int kMetadataBodySize = 15;
-  friend class RecordInitializer;
-
 public:
   enum class MetadataType : unsigned {
     Unknown,
@@ -69,11 +85,22 @@ public:
     TypedEvent,
   };
 
-  Type type() const override { return Type::Metadata; }
+protected:
+  static constexpr int kMetadataBodySize = 15;
+  friend class RecordInitializer;
 
-  // All metadata records must know to provide the type of their open
-  // metadata record.
-  virtual MetadataType metadataType() const = 0;
+private:
+  const MetadataType MT;
+
+public:
+  explicit MetadataRecord(RecordKind T, MetadataType M) : Record(T), MT(M) {}
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() >= RecordKind::RK_Metadata &&
+           R->getRecordType() <= RecordKind::RK_Metadata_LastMetadata;
+  }
+
+  MetadataType metadataType() const { return MT; }
 
   virtual ~MetadataRecord() = default;
 };
@@ -86,16 +113,22 @@ class BufferExtents : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  BufferExtents() = default;
-  explicit BufferExtents(uint64_t S) : MetadataRecord(), Size(S) {}
+  BufferExtents()
+      : MetadataRecord(RecordKind::RK_Metadata_BufferExtents,
+                       MetadataType::BufferExtents) {}
 
-  MetadataType metadataType() const override {
-    return MetadataType::BufferExtents;
-  }
+  explicit BufferExtents(uint64_t S)
+      : MetadataRecord(RecordKind::RK_Metadata_BufferExtents,
+                       MetadataType::BufferExtents),
+        Size(S) {}
 
   uint64_t size() const { return Size; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_BufferExtents;
+  }
 };
 
 class WallclockRecord : public MetadataRecord {
@@ -104,18 +137,23 @@ class WallclockRecord : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  WallclockRecord() = default;
-  explicit WallclockRecord(uint64_t S, uint32_t N)
-      : MetadataRecord(), Seconds(S), Nanos(N) {}
+  WallclockRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_WallClockTime,
+                       MetadataType::WallClockTime) {}
 
-  MetadataType metadataType() const override {
-    return MetadataType::WallClockTime;
-  }
+  explicit WallclockRecord(uint64_t S, uint32_t N)
+      : MetadataRecord(RecordKind::RK_Metadata_WallClockTime,
+                       MetadataType::WallClockTime),
+        Seconds(S), Nanos(N) {}
 
   uint64_t seconds() const { return Seconds; }
   uint32_t nanos() const { return Nanos; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_WallClockTime;
+  }
 };
 
 class NewCPUIDRecord : public MetadataRecord {
@@ -124,16 +162,24 @@ class NewCPUIDRecord : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  NewCPUIDRecord() = default;
-  NewCPUIDRecord(uint16_t C, uint64_t T) : MetadataRecord(), CPUId(C), TSC(T) {}
+  NewCPUIDRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_NewCPUId,
+                       MetadataType::NewCPUId) {}
 
-  MetadataType metadataType() const override { return MetadataType::NewCPUId; }
+  NewCPUIDRecord(uint16_t C, uint64_t T)
+      : MetadataRecord(RecordKind::RK_Metadata_NewCPUId,
+                       MetadataType::NewCPUId),
+        CPUId(C), TSC(T) {}
 
   uint16_t cpuid() const { return CPUId; }
 
   uint64_t tsc() const { return TSC; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_NewCPUId;
+  }
 };
 
 class TSCWrapRecord : public MetadataRecord {
@@ -141,14 +187,21 @@ class TSCWrapRecord : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  TSCWrapRecord() = default;
-  explicit TSCWrapRecord(uint64_t B) : MetadataRecord(), BaseTSC(B) {}
+  TSCWrapRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_TSCWrap, MetadataType::TSCWrap) {
+  }
 
-  MetadataType metadataType() const override { return MetadataType::TSCWrap; }
+  explicit TSCWrapRecord(uint64_t B)
+      : MetadataRecord(RecordKind::RK_Metadata_TSCWrap, MetadataType::TSCWrap),
+        BaseTSC(B) {}
 
   uint64_t tsc() const { return BaseTSC; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_TSCWrap;
+  }
 };
 
 class CustomEventRecord : public MetadataRecord {
@@ -159,13 +212,14 @@ class CustomEventRecord : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  CustomEventRecord() = default;
-  explicit CustomEventRecord(uint64_t S, uint64_t T, uint16_t C, std::string D)
-      : MetadataRecord(), Size(S), TSC(T), CPU(C), Data(std::move(D)) {}
+  CustomEventRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_CustomEvent,
+                       MetadataType::CustomEvent) {}
 
-  MetadataType metadataType() const override {
-    return MetadataType::CustomEvent;
-  }
+  explicit CustomEventRecord(uint64_t S, uint64_t T, uint16_t C, std::string D)
+      : MetadataRecord(RecordKind::RK_Metadata_CustomEvent,
+                       MetadataType::CustomEvent),
+        Size(S), TSC(T), CPU(C), Data(std::move(D)) {}
 
   int32_t size() const { return Size; }
   uint64_t tsc() const { return TSC; }
@@ -173,6 +227,10 @@ public:
   StringRef data() const { return Data; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_CustomEvent;
+  }
 };
 
 class CustomEventRecordV5 : public MetadataRecord {
@@ -182,19 +240,24 @@ class CustomEventRecordV5 : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  CustomEventRecordV5() = default;
-  explicit CustomEventRecordV5(int32_t S, int32_t D, std::string P)
-      : MetadataRecord(), Size(S), Delta(D), Data(std::move(P)) {}
+  CustomEventRecordV5()
+      : MetadataRecord(RecordKind::RK_Metadata_CustomEventV5,
+                       MetadataType::CustomEvent) {}
 
-  MetadataType metadataType() const override {
-    return MetadataType::CustomEvent;
-  }
+  explicit CustomEventRecordV5(int32_t S, int32_t D, std::string P)
+      : MetadataRecord(RecordKind::RK_Metadata_CustomEventV5,
+                       MetadataType::CustomEvent),
+        Size(S), Delta(D), Data(std::move(P)) {}
 
   int32_t size() const { return Size; }
   int32_t delta() const { return Delta; }
   StringRef data() const { return Data; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_CustomEventV5;
+  }
 };
 
 class TypedEventRecord : public MetadataRecord {
@@ -205,13 +268,14 @@ class TypedEventRecord : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  TypedEventRecord() = default;
-  explicit TypedEventRecord(int32_t S, int32_t D, uint16_t E, std::string P)
-      : MetadataRecord(), Size(S), Delta(D), Data(std::move(P)) {}
+  TypedEventRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_TypedEvent,
+                       MetadataType::TypedEvent) {}
 
-  MetadataType metadataType() const override {
-    return MetadataType::TypedEvent;
-  }
+  explicit TypedEventRecord(int32_t S, int32_t D, uint16_t E, std::string P)
+      : MetadataRecord(RecordKind::RK_Metadata_TypedEvent,
+                       MetadataType::TypedEvent),
+        Size(S), Delta(D), Data(std::move(P)) {}
 
   int32_t size() const { return Size; }
   int32_t delta() const { return Delta; }
@@ -219,6 +283,10 @@ public:
   StringRef data() const { return Data; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_TypedEvent;
+  }
 };
 
 class CallArgRecord : public MetadataRecord {
@@ -226,14 +294,21 @@ class CallArgRecord : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  CallArgRecord() = default;
-  explicit CallArgRecord(uint64_t A) : MetadataRecord(), Arg(A) {}
+  CallArgRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_CallArg, MetadataType::CallArg) {
+  }
 
-  MetadataType metadataType() const override { return MetadataType::CallArg; }
+  explicit CallArgRecord(uint64_t A)
+      : MetadataRecord(RecordKind::RK_Metadata_CallArg, MetadataType::CallArg),
+        Arg(A) {}
 
   uint64_t arg() const { return Arg; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_CallArg;
+  }
 };
 
 class PIDRecord : public MetadataRecord {
@@ -241,14 +316,22 @@ class PIDRecord : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  PIDRecord() = default;
-  explicit PIDRecord(int32_t P) : MetadataRecord(), PID(P) {}
+  PIDRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_PIDEntry,
+                       MetadataType::PIDEntry) {}
 
-  MetadataType metadataType() const override { return MetadataType::PIDEntry; }
+  explicit PIDRecord(int32_t P)
+      : MetadataRecord(RecordKind::RK_Metadata_PIDEntry,
+                       MetadataType::PIDEntry),
+        PID(P) {}
 
   int32_t pid() const { return PID; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_PIDEntry;
+  }
 };
 
 class NewBufferRecord : public MetadataRecord {
@@ -256,25 +339,35 @@ class NewBufferRecord : public MetadataRecord {
   friend class RecordInitializer;
 
 public:
-  NewBufferRecord() = default;
-  explicit NewBufferRecord(int32_t T) : MetadataRecord(), TID(T) {}
+  NewBufferRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_NewBuffer,
+                       MetadataType::NewBuffer) {}
 
-  MetadataType metadataType() const override { return MetadataType::NewBuffer; }
+  explicit NewBufferRecord(int32_t T)
+      : MetadataRecord(RecordKind::RK_Metadata_NewBuffer,
+                       MetadataType::NewBuffer),
+        TID(T) {}
 
   int32_t tid() const { return TID; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_NewBuffer;
+  }
 };
 
 class EndBufferRecord : public MetadataRecord {
 public:
-  EndBufferRecord() = default;
-
-  MetadataType metadataType() const override {
-    return MetadataType::EndOfBuffer;
-  }
+  EndBufferRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_EndOfBuffer,
+                       MetadataType::EndOfBuffer) {}
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Metadata_EndOfBuffer;
+  }
 };
 
 class FunctionRecord : public Record {
@@ -286,11 +379,10 @@ class FunctionRecord : public Record {
   static constexpr unsigned kFunctionRecordSize = 8;
 
 public:
-  FunctionRecord() = default;
-  explicit FunctionRecord(RecordTypes K, int32_t F, uint32_t D)
-      : Record(), Kind(K), FuncId(F), Delta(D) {}
+  FunctionRecord() : Record(RecordKind::RK_Function) {}
 
-  Type type() const override { return Type::Function; }
+  explicit FunctionRecord(RecordTypes K, int32_t F, uint32_t D)
+      : Record(RecordKind::RK_Function), Kind(K), FuncId(F), Delta(D) {}
 
   // A function record is a concrete record type which has a number of common
   // properties.
@@ -299,6 +391,10 @@ public:
   uint32_t delta() const { return Delta; }
 
   Error apply(RecordVisitor &V) override;
+
+  static bool classof(const Record *R) {
+    return R->getRecordType() == RecordKind::RK_Function;
+  }
 };
 
 class RecordVisitor {

lib/XRay/FDRRecordProducer.cpp

@@ -9,29 +9,32 @@
 #include "llvm/XRay/FDRRecordProducer.h"
 #include "llvm/Support/DataExtractor.h"
 
+#include <cstdint>
+
 namespace llvm {
 namespace xray {
 
 namespace {
 
+// Keep this in sync with the values written in the XRay FDR mode runtime in
+// compiler-rt.
+enum MetadataRecordKinds : uint8_t {
+  NewBufferKind,
+  EndOfBufferKind,
+  NewCPUIdKind,
+  TSCWrapKind,
+  WalltimeMarkerKind,
+  CustomEventMarkerKind,
+  CallArgumentKind,
+  BufferExtentsKind,
+  TypedEventMarkerKind,
+  PidKind,
+  // This is an end marker, used to identify the upper bound for this enum.
+  EnumEndMarker,
+};
+
 Expected<std::unique_ptr<Record>>
 metadataRecordType(const XRayFileHeader &Header, uint8_t T) {
-  // Keep this in sync with the values written in the XRay FDR mode runtime in
-  // compiler-rt.
-  enum MetadataRecordKinds : uint8_t {
-    NewBufferKind,
-    EndOfBufferKind,
-    NewCPUIdKind,
-    TSCWrapKind,
-    WalltimeMarkerKind,
-    CustomEventMarkerKind,
-    CallArgumentKind,
-    BufferExtentsKind,
-    TypedEventMarkerKind,
-    PidKind,
-    // This is an end marker, used to identify the upper bound for this enum.
-    EnumEndMarker,
-  };
 
   if (T >= static_cast<uint8_t>(MetadataRecordKinds::EnumEndMarker))
     return createStringError(std::make_error_code(std::errc::invalid_argument),
@@ -70,9 +73,70 @@ metadataRecordType(const XRayFileHeader &Header, uint8_t T) {
   llvm_unreachable("Unhandled MetadataRecordKinds enum value");
 }
 
+constexpr bool isMetadataIntroducer(uint8_t FirstByte) {
+  return FirstByte & 0x01u;
+}
+
 } // namespace
 
+Expected<std::unique_ptr<Record>>
+FileBasedRecordProducer::findNextBufferExtent() {
+  // We seek one byte at a time until we find a suitable buffer extents metadata
+  // record introducer.
+  std::unique_ptr<Record> R;
+  while (!R) {
+    auto PreReadOffset = OffsetPtr;
+    uint8_t FirstByte = E.getU8(&OffsetPtr);
+    if (OffsetPtr == PreReadOffset)
+      return createStringError(
+          std::make_error_code(std::errc::executable_format_error),
+          "Failed reading one byte from offset %d.", OffsetPtr);
+
+    if (isMetadataIntroducer(FirstByte)) {
+      auto LoadedType = FirstByte >> 1;
+      if (LoadedType == MetadataRecordKinds::BufferExtentsKind) {
+        auto MetadataRecordOrErr = metadataRecordType(Header, LoadedType);
+        if (!MetadataRecordOrErr)
+          return MetadataRecordOrErr.takeError();
+
+        R = std::move(MetadataRecordOrErr.get());
+        RecordInitializer RI(E, OffsetPtr);
+        if (auto Err = R->apply(RI))
+          return std::move(Err);
+        return std::move(R);
+      }
+    }
+  }
+  llvm_unreachable("Must always terminate with either an error or a record.");
+}
+
 Expected<std::unique_ptr<Record>> FileBasedRecordProducer::produce() {
+  // First, we set up our result record.
+  std::unique_ptr<Record> R;
+
+  // Before we do any further reading, we should check whether we're at the end
+  // of the current buffer we're been consuming. In FDR logs version >= 3, we
+  // rely on the buffer extents record to determine how many bytes we should be
+  // considering as valid records.
+  if (Header.Version >= 3 && CurrentBufferBytes == 0) {
+    // Find the next buffer extents record.
+    auto BufferExtentsOrError = findNextBufferExtent();
+    if (!BufferExtentsOrError)
+      return joinErrors(
+          BufferExtentsOrError.takeError(),
+          createStringError(
+              std::make_error_code(std::errc::executable_format_error),
+              "Failed to find the next BufferExtents record."));
+
+    R = std::move(BufferExtentsOrError.get());
+    assert(R != nullptr);
+    assert(isa<BufferExtents>(R.get()));
+    auto BE = dyn_cast<BufferExtents>(R.get());
+    CurrentBufferBytes = BE->size();
+    return std::move(R);
+  }
+
+  //
   // At the top level, we read one byte to determine the type of the record to
   // create. This byte will comprise of the following bits:
   //
@@ -90,11 +154,8 @@ Expected<std::unique_ptr<Record>> FileBasedRecordProducer::produce() {
         std::make_error_code(std::errc::executable_format_error),
         "Failed reading one byte from offset %d.", OffsetPtr);
 
-  // Set up our result record.
-  std::unique_ptr<Record> R;
-
   // For metadata records, handle especially here.
-  if (FirstByte & 0x01) {
+  if (isMetadataIntroducer(FirstByte)) {
     auto LoadedType = FirstByte >> 1;
     auto MetadataRecordOrErr = metadataRecordType(Header, LoadedType);
     if (!MetadataRecordOrErr)
@@ -113,6 +174,22 @@ Expected<std::unique_ptr<Record>> FileBasedRecordProducer::produce() {
   if (auto Err = R->apply(RI))
     return std::move(Err);
 
+  // If we encountered a BufferExtents record, we should record the remaining
+  // bytes for the current buffer, to determine when we should start ignoring
+  // potentially malformed data and looking for buffer extents records.
+  if (auto BE = dyn_cast<BufferExtents>(R.get())) {
+    CurrentBufferBytes = BE->size();
+  } else if (Header.Version >= 3) {
+    if (OffsetPtr - PreReadOffset > CurrentBufferBytes)
+      return createStringError(
+          std::make_error_code(std::errc::executable_format_error),
+          "Buffer over-read at offset %d (over-read by %d bytes); Record Type "
+          "= %s.",
+          OffsetPtr, (OffsetPtr - PreReadOffset) - CurrentBufferBytes,
+          Record::kindToString(R->getRecordType()).data());
+
+    CurrentBufferBytes -= OffsetPtr - PreReadOffset;
+  }
   assert(R != nullptr);
   return std::move(R);
 }

lib/XRay/FDRRecords.cpp

@@ -29,5 +29,39 @@ Error FunctionRecord::apply(RecordVisitor &V) { return V.visit(*this); }
 Error CustomEventRecordV5::apply(RecordVisitor &V) { return V.visit(*this); }
 Error TypedEventRecord::apply(RecordVisitor &V) { return V.visit(*this); }
 
+StringRef Record::kindToString(RecordKind K) {
+  switch (K) {
+  case RecordKind::RK_Metadata:
+    return "Metadata";
+  case RecordKind::RK_Metadata_BufferExtents:
+    return "Metadata:BufferExtents";
+  case RecordKind::RK_Metadata_WallClockTime:
+    return "Metadata:WallClockTime";
+  case RecordKind::RK_Metadata_NewCPUId:
+    return "Metadata:NewCPUId";
+  case RecordKind::RK_Metadata_TSCWrap:
+    return "Metadata:TSCWrap";
+  case RecordKind::RK_Metadata_CustomEvent:
+    return "Metadata:CustomEvent";
+  case RecordKind::RK_Metadata_CustomEventV5:
+    return "Metadata:CustomEventV5";
+  case RecordKind::RK_Metadata_CallArg:
+    return "Metadata:CallArg";
+  case RecordKind::RK_Metadata_PIDEntry:
+    return "Metadata:PIDEntry";
+  case RecordKind::RK_Metadata_NewBuffer:
+    return "Metadata:NewBuffer";
+  case RecordKind::RK_Metadata_EndOfBuffer:
+    return "Metadata:EndOfBuffer";
+  case RecordKind::RK_Metadata_TypedEvent:
+    return "Metadata:TypedEvent";
+  case RecordKind::RK_Metadata_LastMetadata:
+    return "Metadata:LastMetadata";
+  case RecordKind::RK_Function:
+    return "Function";
+  }
+  return "Unknown";
+}
+
 } // namespace xray
 } // namespace llvm

lib/XRay/FDRTraceWriter.cpp

@@ -43,7 +43,9 @@ template <size_t Index> struct IndexedWriter {
 
 template <uint8_t Kind, class... Values>
 Error writeMetadata(support::endian::Writer &OS, Values &&... Ds) {
-  uint8_t FirstByte = (Kind << 1) | uint8_t{0x01};
+  // The first bit in the first byte of metadata records is always set to 1, so
+  // we ensure this is the case when we write out the first byte of the record.
+  uint8_t FirstByte = (static_cast<uint8_t>(Kind) << 1) | uint8_t{0x01u};
   auto T = std::make_tuple(std::forward<Values>(std::move(Ds))...);
   // Write in field order.
   OS.write(FirstByte);
@@ -112,7 +114,7 @@ Error FDRTraceWriter::visit(CustomEventRecordV5 &R) {
 }
 
 Error FDRTraceWriter::visit(TypedEventRecord &R) {
-  if (auto E = writeMetadata<7u>(OS, R.size(), R.delta(), R.eventType()))
+  if (auto E = writeMetadata<8u>(OS, R.size(), R.delta(), R.eventType()))
     return E;
   auto D = R.data();
   ArrayRef<char> Bytes(D.data(), D.size());

lib/XRay/RecordInitializer.cpp

@@ -111,6 +111,12 @@ Error RecordInitializer::visit(CustomEventRecord &R) {
         std::make_error_code(std::errc::invalid_argument),
         "Cannot read a custom event record size field offset %d.", OffsetPtr);
 
+  if (R.Size <= 0)
+    return createStringError(
+        std::make_error_code(std::errc::bad_address),
+        "Invalid size for custom event (size = %d) at offset %d.", R.Size,
+        OffsetPtr);
+
   PreReadOffset = OffsetPtr;
   R.TSC = E.getU64(&OffsetPtr);
   if (PreReadOffset == OffsetPtr)
@@ -142,11 +148,21 @@ Error RecordInitializer::visit(CustomEventRecord &R) {
 
   std::vector<uint8_t> Buffer;
   Buffer.resize(R.Size);
+  PreReadOffset = OffsetPtr;
   if (E.getU8(&OffsetPtr, Buffer.data(), R.Size) != Buffer.data())
     return createStringError(
         std::make_error_code(std::errc::invalid_argument),
         "Failed reading data into buffer of size %d at offset %d.", R.Size,
         OffsetPtr);
+
+  assert(OffsetPtr >= PreReadOffset);
+  if (OffsetPtr - PreReadOffset != static_cast<uint32_t>(R.Size))
+    return createStringError(
+        std::make_error_code(std::errc::invalid_argument),
+        "Failed reading enough bytes for the custom event payload -- read %d "
+        "expecting %d bytes at offset %d.",
+        OffsetPtr - PreReadOffset, R.Size, PreReadOffset);
+
   R.Data.assign(Buffer.begin(), Buffer.end());
   return Error::success();
 }
@@ -167,6 +183,12 @@ Error RecordInitializer::visit(CustomEventRecordV5 &R) {
         std::make_error_code(std::errc::invalid_argument),
         "Cannot read a custom event record size field offset %d.", OffsetPtr);
 
+  if (R.Size <= 0)
+    return createStringError(
+        std::make_error_code(std::errc::bad_address),
+        "Invalid size for custom event (size = %d) at offset %d.", R.Size,
+        OffsetPtr);
+
   PreReadOffset = OffsetPtr;
   R.Delta = E.getSigned(&OffsetPtr, sizeof(int32_t));
   if (PreReadOffset == OffsetPtr)
@@ -188,11 +210,21 @@ Error RecordInitializer::visit(CustomEventRecordV5 &R) {
 
   std::vector<uint8_t> Buffer;
   Buffer.resize(R.Size);
+  PreReadOffset = OffsetPtr;
   if (E.getU8(&OffsetPtr, Buffer.data(), R.Size) != Buffer.data())
     return createStringError(
         std::make_error_code(std::errc::invalid_argument),
         "Failed reading data into buffer of size %d at offset %d.", R.Size,
         OffsetPtr);
+
+  assert(OffsetPtr >= PreReadOffset);
+  if (OffsetPtr - PreReadOffset != static_cast<uint32_t>(R.Size))
+    return createStringError(
+        std::make_error_code(std::errc::invalid_argument),
+        "Failed reading enough bytes for the custom event payload -- read %d "
+        "expecting %d bytes at offset %d.",
+        OffsetPtr - PreReadOffset, R.Size, PreReadOffset);
+
   R.Data.assign(Buffer.begin(), Buffer.end());
   return Error::success();
 }
@@ -213,6 +245,12 @@ Error RecordInitializer::visit(TypedEventRecord &R) {
         std::make_error_code(std::errc::invalid_argument),
         "Cannot read a typed event record size field offset %d.", OffsetPtr);
 
+  if (R.Size <= 0)
+    return createStringError(
+        std::make_error_code(std::errc::bad_address),
+        "Invalid size for typed event (size = %d) at offset %d.", R.Size,
+        OffsetPtr);
+
   PreReadOffset = OffsetPtr;
   R.Delta = E.getSigned(&OffsetPtr, sizeof(int32_t));
   if (PreReadOffset == OffsetPtr)
@@ -241,11 +279,21 @@ Error RecordInitializer::visit(TypedEventRecord &R) {
 
   std::vector<uint8_t> Buffer;
   Buffer.resize(R.Size);
+  PreReadOffset = OffsetPtr;
   if (E.getU8(&OffsetPtr, Buffer.data(), R.Size) != Buffer.data())
     return createStringError(
         std::make_error_code(std::errc::invalid_argument),
         "Failed reading data into buffer of size %d at offset %d.", R.Size,
         OffsetPtr);
+
+  assert(OffsetPtr >= PreReadOffset);
+  if (OffsetPtr - PreReadOffset != static_cast<uint32_t>(R.Size))
+    return createStringError(
+        std::make_error_code(std::errc::invalid_argument),
+        "Failed reading enough bytes for the typed event payload -- read %d "
+        "expecting %d bytes at offset %d.",
+        OffsetPtr - PreReadOffset, R.Size, PreReadOffset);
+
   R.Data.assign(Buffer.begin(), Buffer.end());
   return Error::success();
 }
@@ -337,7 +385,11 @@ Error RecordInitializer::visit(FunctionRecord &R) {
     return createStringError(std::make_error_code(std::errc::bad_address),
                              "Cannot read function id field from offset %d.",
                              OffsetPtr);
-  unsigned FunctionType = (Buffer >> 1) & 0x07;
+
+  // To get the function record type, we shift the buffer one to the right
+  // (truncating the function record indicator) then take the three bits
+  // (0b0111) to get the record type as an unsigned value.
+  unsigned FunctionType = (Buffer >> 1) & 0x07u;
   switch (FunctionType) {
   case static_cast<unsigned>(RecordTypes::ENTER):
   case static_cast<unsigned>(RecordTypes::ENTER_ARG):

unittests/XRay/FDRProducerConsumerTest.cpp

@@ -30,13 +30,10 @@ namespace {
 using ::testing::Eq;
 using ::testing::IsEmpty;
 using ::testing::Not;
+using ::testing::SizeIs;
 
 template <class RecordType> std::unique_ptr<Record> MakeRecord();
 
-template <> std::unique_ptr<Record> MakeRecord<BufferExtents>() {
-  return make_unique<BufferExtents>(1);
-}
-
 template <> std::unique_ptr<Record> MakeRecord<NewBufferRecord>() {
   return make_unique<NewBufferRecord>(1);
 }
@@ -69,10 +66,18 @@ template <> std::unique_ptr<Record> MakeRecord<FunctionRecord>() {
   return make_unique<FunctionRecord>(RecordTypes::ENTER, 1, 2);
 }
 
+template <> std::unique_ptr<Record> MakeRecord<CustomEventRecordV5>() {
+  return make_unique<CustomEventRecordV5>(4, 1, "data");
+}
+
+template <> std::unique_ptr<Record> MakeRecord<TypedEventRecord>() {
+  return make_unique<TypedEventRecord>(4, 1, 2, "data");
+}
+
 template <class T> class RoundTripTest : public ::testing::Test {
 public:
   RoundTripTest() : Data(), OS(Data) {
-    H.Version = 3;
+    H.Version = 4;
     H.Type = 1;
     H.ConstantTSC = true;
     H.NonstopTSC = true;
@@ -92,9 +97,36 @@ protected:
 
 TYPED_TEST_CASE_P(RoundTripTest);
 
+template <class T> class RoundTripTestV5 : public ::testing::Test {
+public:
+  RoundTripTestV5() : Data(), OS(Data) {
+    H.Version = 5;
+    H.Type = 1;
+    H.ConstantTSC = true;
+    H.NonstopTSC = true;
+    H.CycleFrequency = 3e9;
+
+    Writer = make_unique<FDRTraceWriter>(OS, H);
+    Rec = MakeRecord<T>();
+  }
+
+protected:
+  std::string Data;
+  raw_string_ostream OS;
+  XRayFileHeader H;
+  std::unique_ptr<FDRTraceWriter> Writer;
+  std::unique_ptr<Record> Rec;
+};
+
+TYPED_TEST_CASE_P(RoundTripTestV5);
+
 // This test ensures that the writing and reading implementations are in sync --
 // that given write(read(write(R))) == R.
 TYPED_TEST_P(RoundTripTest, RoundTripsSingleValue) {
+  // Always write a buffer extents record which will cover the correct size of
+  // the record, for version 3 and up.
+  BufferExtents BE(200);
+  ASSERT_FALSE(errorToBool(BE.apply(*this->Writer)));
   auto &R = this->Rec;
   ASSERT_FALSE(errorToBool(R->apply(*this->Writer)));
   this->OS.flush();
@@ -125,17 +157,67 @@ TYPED_TEST_P(RoundTripTest, RoundTripsSingleValue) {
 
   EXPECT_EQ(Data2.substr(sizeof(XRayFileHeader)),
             this->Data.substr(sizeof(XRayFileHeader)));
-  EXPECT_THAT(Records[0]->type(), Eq(R->type()));
+  ASSERT_THAT(Records, SizeIs(2));
+  EXPECT_THAT(Records[1]->getRecordType(), Eq(R->getRecordType()));
 }
 
 REGISTER_TYPED_TEST_CASE_P(RoundTripTest, RoundTripsSingleValue);
 
+// We duplicate the above case for the V5 version using different types and
+// encodings.
+TYPED_TEST_P(RoundTripTestV5, RoundTripsSingleValue) {
+  BufferExtents BE(200);
+  ASSERT_FALSE(errorToBool(BE.apply(*this->Writer)));
+  auto &R = this->Rec;
+  ASSERT_FALSE(errorToBool(R->apply(*this->Writer)));
+  this->OS.flush();
+
+  DataExtractor DE(this->Data, sys::IsLittleEndianHost, 8);
+  uint32_t OffsetPtr = 0;
+  auto HeaderOrErr = readBinaryFormatHeader(DE, OffsetPtr);
+  if (!HeaderOrErr)
+    FAIL() << HeaderOrErr.takeError();
+
+  FileBasedRecordProducer P(HeaderOrErr.get(), DE, OffsetPtr);
+  std::vector<std::unique_ptr<Record>> Records;
+  LogBuilderConsumer C(Records);
+  while (DE.isValidOffsetForDataOfSize(OffsetPtr, 1)) {
+    auto R = P.produce();
+    if (!R)
+      FAIL() << R.takeError();
+    if (auto E = C.consume(std::move(R.get())))
+      FAIL() << E;
+  }
+  ASSERT_THAT(Records, Not(IsEmpty()));
+  std::string Data2;
+  raw_string_ostream OS2(Data2);
+  FDRTraceWriter Writer2(OS2, this->H);
+  for (auto &P : Records)
+    ASSERT_FALSE(errorToBool(P->apply(Writer2)));
+  OS2.flush();
+
+  EXPECT_EQ(Data2.substr(sizeof(XRayFileHeader)),
+            this->Data.substr(sizeof(XRayFileHeader)));
+  ASSERT_THAT(Records, SizeIs(2));
+  EXPECT_THAT(Records[1]->getRecordType(), Eq(R->getRecordType()));
+}
+
+REGISTER_TYPED_TEST_CASE_P(RoundTripTestV5, RoundTripsSingleValue);
+
+// These are the record types we support for v4 and below.
 using RecordTypes =
-    ::testing::Types<BufferExtents, NewBufferRecord, NewCPUIDRecord,
-                     TSCWrapRecord, WallclockRecord, CustomEventRecord,
-                     CallArgRecord, BufferExtents, PIDRecord, FunctionRecord>;
+    ::testing::Types<NewBufferRecord, NewCPUIDRecord, TSCWrapRecord,
+                     WallclockRecord, CustomEventRecord, CallArgRecord,
+                     PIDRecord, FunctionRecord>;
 INSTANTIATE_TYPED_TEST_CASE_P(Records, RoundTripTest, RecordTypes);
 
+// For V5, we have two new types we're supporting.
+using RecordTypesV5 =
+    ::testing::Types<NewBufferRecord, NewCPUIDRecord, TSCWrapRecord,
+                     WallclockRecord, CustomEventRecordV5, TypedEventRecord,
+                     CallArgRecord, PIDRecord, FunctionRecord>;
+INSTANTIATE_TYPED_TEST_CASE_P(Records, RoundTripTestV5, RecordTypesV5);
+
 } // namespace
 } // namespace xray
 } // namespace llvm