[XRay] Improve FDR trace handling and error messaging

[dberris]

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.

[mboerger]

llvm/include/llvm/XRay/FDRRecords.h

129 ↗

(On Diff #172910)

This is a weird interface. You basically want to associate a type enum with a class type. But here you do it in a kind of inverse way. Since you don't really use it, I don't know what better to do though.

llvm/lib/XRay/FDRRecordProducer.cpp

80 ↗

(On Diff #172910)

Please reword this

84 ↗

(On Diff #172910)

one would expect the external 'OffsetPtr' to stay unaffected and instead of 'PreReadOffset' having a 'CurrentOffset' or so variable. Please clarify or change.

91 ↗

(On Diff #172910)

comment what this check does, or even better provide a static helper function that is used everywhere you do this

llvm/lib/XRay/FDRRecords.cpp

34 ↗

(On Diff #172910)

this screams for using a macro :-)

llvm/lib/XRay/FDRTraceWriter.cpp

46 ↗

(On Diff #172910)

so this should probably also get a helper function

[dberris]

Address some comments by @mboerger.

[dberris]

llvm/include/llvm/XRay/FDRRecords.h

129 ↗

(On Diff #172910)

We use it in LLVM's own RTTI implementation (see https://llvm.org/docs/HowToSetUpLLVMStyleRTTI.html).

llvm/lib/XRay/FDRRecordProducer.cpp

84 ↗

(On Diff #172910)

The API for the DataExtractor documents the semantics of these get*(...) functions, which says that the offset pointer is updated only if there were no errors encountered while reading the data. This is why we need to check against whether the offset we got before reading the data is the same, and if so deal with the error condition.

llvm/lib/XRay/FDRRecords.cpp

34 ↗

(On Diff #172910)

Haha -- or a language feature. :(

llvm/lib/XRay/FDRTraceWriter.cpp

46 ↗

(On Diff #172910)

This is only done here, where we're synthesising the record being serialised. There's no other place AFAICT where we're doing this particular operation in the LLVM side of the project.

In compiler-rt, we do this serialization but using bitfields.

compiler-rt/trunk/lib/xray/tests/unit/fdr_controller_test.cc

TEST_F(BufferManagementTest, HandlesOverflow) {
   ASSERT_THAT(BQ->finalize(), Eq(BufferQueue::ErrorCode::Ok));
 
   std::string Serialized = serialize(*BQ, 3);
   llvm::DataExtractor DE(Serialized, true, 8);
   auto TraceOrErr = llvm::xray::loadTrace(DE);
   EXPECT_THAT_EXPECTED(TraceOrErr, HasValue(SizeIs(kBuffers * 2)));
 }
 
+TEST_F(BufferManagementTest, HandlesOverflowWithArgs) {
+  uint64_t TSC = 1;
+  uint16_t CPU = 1;
+  uint64_t ARG = 1;
+  for (size_t I = 0; I < kBuffers + 1; ++I) {
+    ASSERT_TRUE(C->functionEnterArg(1, TSC++, CPU, ARG++));
+    ASSERT_TRUE(C->functionExit(1, TSC++, CPU));
+  }
+  ASSERT_TRUE(C->flush());
+  ASSERT_THAT(BQ->finalize(), Eq(BufferQueue::ErrorCode::Ok));
+
+  std::string Serialized = serialize(*BQ, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(TraceOrErr, HasValue(SizeIs(kBuffers)));
+}
+
 TEST_F(BufferManagementTest, HandlesOverflowWithCustomEvents) {
   uint64_t TSC = 1;
   uint16_t CPU = 1;
   int32_t D = 0x9009;
   for (size_t I = 0; I < kBuffers; ++I) {
     ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
     ASSERT_TRUE(C->functionExit(1, TSC++, CPU));
     ASSERT_TRUE(C->customEvent(TSC++, CPU, &D, sizeof(D)));

compiler-rt/trunk/lib/xray/xray_fdr_controller.h

if (finalized() ||
       return returnBuffer();
 
     auto Delta = TSC - LatestTSC;
     LatestTSC = TSC;
     LastFunctionEntryTSC = 0;
     UndoableFunctionEnters = 0;
     UndoableTailExits = 0;
 
-    W.writeFunction(FDRLogWriter::FunctionRecordKind::EnterArg, mask(FuncId),
-                    Delta);
-    return W.writeMetadata<MetadataRecord::RecordKinds::CallArgument>(Arg);
+    return W.writeFunctionWithArg(FDRLogWriter::FunctionRecordKind::EnterArg,
+                                  mask(FuncId), Delta, Arg);
   }
 
   bool functionExit(int32_t FuncId, uint64_t TSC,
                     uint16_t CPU) XRAY_NEVER_INSTRUMENT {
     if (finalized() ||
         !prepareBuffer(sizeof(MetadataRecord) + sizeof(FunctionRecord)))
       return returnBuffer();
 

compiler-rt/trunk/lib/xray/xray_fdr_log_writer.h

bool writeFunction(FunctionRecordKind Kind, int32_t FuncId, int32_t Delta) {
     R.Type = 0;
     R.RecordKind = uint8_t(Kind);
     R.FuncId = FuncId;
     R.TSCDelta = Delta;
     writeRecord(R);
     return true;
   }
 
+  bool writeFunctionWithArg(FunctionRecordKind Kind, int32_t FuncId,
+                            int32_t Delta, uint64_t Arg) {
+    // We need to write the function with arg into the buffer, and then
+    // atomically update the buffer extents. This ensures that any reads
+    // synchronised on the buffer extents record will always see the writes
+    // that happen before the atomic update.
+    FunctionRecord R;
+    R.Type = 0;
+    R.RecordKind = uint8_t(Kind);
+    R.FuncId = FuncId;
+    R.TSCDelta = Delta;
+    MetadataRecord A =
+        createMetadataRecord<MetadataRecord::RecordKinds::CallArgument>(Arg);
+    NextRecord = reinterpret_cast<char *>(internal_memcpy(
+                     NextRecord, reinterpret_cast<char *>(&R), sizeof(R))) +
+                 sizeof(R);
+    NextRecord = reinterpret_cast<char *>(internal_memcpy(
+                     NextRecord, reinterpret_cast<char *>(&A), sizeof(A))) +
+                 sizeof(A);
+    atomic_fetch_add(&Buffer.Extents, sizeof(R) + sizeof(A),
+                     memory_order_acq_rel);
+    return true;
+  }
+
   bool writeCustomEvent(int32_t Delta, const void *Event, int32_t EventSize) {
     // We write the metadata record and the custom event data into the buffer
     // first, before we atomically update the extents for the buffer. This
     // allows us to ensure that any threads reading the extents of the buffer
     // will only ever see the full metadata and custom event payload accounted
     // (no partial writes accounted).
     MetadataRecord R =
         createMetadataRecord<MetadataRecord::RecordKinds::CustomEventMarker>(

llvm/trunk/include/llvm/XRay/FDRRecordProducer.h

public:
   virtual Expected<std::unique_ptr<Record>> produce() = 0;
   virtual ~RecordProducer() = default;
 };
 
 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,
                           uint32_t &OP)
       : Header(FH), E(DE), OffsetPtr(OP) {}
 
   /// This producer encapsulates the logic for loading a File-backed
   /// RecordProducer hidden behind a DataExtractor.
   Expected<std::unique_ptr<Record>> produce() override;
 };
 
 } // namespace xray
 } // namespace llvm
 
 #endif // LLVM_INCLUDE_LLVM_XRAY_FDRRECORDPRODUCER_H_

llvm/trunk/include/llvm/XRay/FDRRecords.h

 //===- FDRRecords.h - XRay Flight Data Recorder Mode Records --------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // Define types and operations on these types that represent the different kinds
 // of records we encounter in XRay flight data recorder mode traces.
 //
 //===----------------------------------------------------------------------===//
 #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 {
 
 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.
   virtual Error apply(RecordVisitor &V) = 0;
 
   virtual ~Record() = default;
 };
 
 class MetadataRecord : public Record {
-protected:
-  static constexpr int kMetadataBodySize = 15;
-  friend class RecordInitializer;
-
 public:
   enum class MetadataType : unsigned {
     Unknown,
     BufferExtents,
     WallClockTime,
     NewCPUId,
     TSCWrap,
     CustomEvent,
     CallArg,
     PIDEntry,
     NewBuffer,
     EndOfBuffer,
     TypedEvent,
   };
 
-  Type type() const override { return Type::Metadata; }
+protected:
+  static constexpr int kMetadataBodySize = 15;
+  friend class RecordInitializer;
+
+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;
+  }
 
-  // All metadata records must know to provide the type of their open
-  // metadata record.
-  virtual MetadataType metadataType() const = 0;
+  MetadataType metadataType() const { return MT; }
 
   virtual ~MetadataRecord() = default;
 };
 
 // What follows are specific Metadata record types which encapsulate the
 // information associated with specific metadata record types in an FDR mode
 // log.
 class BufferExtents : public MetadataRecord {
   uint64_t Size = 0;
   friend class RecordInitializer;
 
 public:
-  BufferExtents() = default;
-  explicit BufferExtents(uint64_t S) : MetadataRecord(), Size(S) {}
-
-  MetadataType metadataType() const override {
-    return MetadataType::BufferExtents;
-  }
+  BufferExtents()
+      : MetadataRecord(RecordKind::RK_Metadata_BufferExtents,
+                       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 {
   uint64_t Seconds = 0;
   uint32_t Nanos = 0;
   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 {
   uint16_t CPUId = 0;
   uint64_t TSC = 0;
   friend class RecordInitializer;
 
 public:
-  NewCPUIDRecord() = default;
-  NewCPUIDRecord(uint16_t C, uint64_t T) : MetadataRecord(), CPUId(C), TSC(T) {}
-
-  MetadataType metadataType() const override { return MetadataType::NewCPUId; }
+  NewCPUIDRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_NewCPUId,
+                       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 {
   uint64_t BaseTSC = 0;
   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 {
   int32_t Size = 0;
   uint64_t TSC = 0;
   uint16_t CPU = 0;
   std::string Data{};
   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; }
   uint16_t cpu() const { return CPU; }
   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 {
   int32_t Size = 0;
   int32_t Delta = 0;
   std::string Data{};
   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 {
   int32_t Size = 0;
   int32_t Delta = 0;
   uint16_t EventType = 0;
   std::string Data{};
   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; }
   uint16_t eventType() const { return EventType; }
   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 {
   uint64_t Arg;
   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 {
   int32_t PID = 0;
   friend class RecordInitializer;
 
 public:
-  PIDRecord() = default;
-  explicit PIDRecord(int32_t P) : MetadataRecord(), PID(P) {}
-
-  MetadataType metadataType() const override { return MetadataType::PIDEntry; }
+  PIDRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_PIDEntry,
+                       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 {
   int32_t TID = 0;
   friend class RecordInitializer;
 
 public:
-  NewBufferRecord() = default;
-  explicit NewBufferRecord(int32_t T) : MetadataRecord(), TID(T) {}
-
-  MetadataType metadataType() const override { return MetadataType::NewBuffer; }
+  NewBufferRecord()
+      : MetadataRecord(RecordKind::RK_Metadata_NewBuffer,
+                       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 {
   RecordTypes Kind;
   int32_t FuncId;
   uint32_t Delta;
   friend class RecordInitializer;
 
   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.
   RecordTypes recordType() const { return Kind; }
   int32_t functionId() const { return FuncId; }
   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 {
 public:
   virtual ~RecordVisitor() = default;
 
   // Support all specific kinds of records:
   virtual Error visit(BufferExtents &) = 0;

llvm/trunk/lib/XRay/FDRRecordProducer.cpp

 //===- FDRRecordProducer.cpp - XRay FDR Mode Record Producer --------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 #include "llvm/XRay/FDRRecordProducer.h"
 #include "llvm/Support/DataExtractor.h"
 
+#include <cstdint>
+
 namespace llvm {
 namespace xray {
 
 namespace {
 
-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,
-  };
 
+Expected<std::unique_ptr<Record>>
+metadataRecordType(const XRayFileHeader &Header, uint8_t T) {
+
   if (T >= static_cast<uint8_t>(MetadataRecordKinds::EnumEndMarker))
     return createStringError(std::make_error_code(std::errc::invalid_argument),
                              "Invalid metadata record type: %d", T);
   switch (T) {
   case MetadataRecordKinds::NewBufferKind:
     return make_unique<NewBufferRecord>();
   case MetadataRecordKinds::EndOfBufferKind:
     if (Header.Version >= 2)
metadataRecordType(const XRayFileHeader &Header, uint8_t T) {
   case MetadataRecordKinds::PidKind:
     return make_unique<PIDRecord>();
   case MetadataRecordKinds::EnumEndMarker:
     llvm_unreachable("Invalid MetadataRecordKind");
   }
   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:
   //
   //   - offset 0: A '1' indicates a metadata record, a '0' indicates a function
   //     record.
   //   - offsets 1-7: For metadata records, this will indicate the kind of
   //     metadata record should be loaded.
   //
   // We read first byte, then create the appropriate type of record to consume
   // the rest of the bytes.
   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);
 
-  // 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)
       return joinErrors(
           MetadataRecordOrErr.takeError(),
           createStringError(
               std::make_error_code(std::errc::executable_format_error),
               "Encountered an unsupported metadata record (%d) at offset %d.",
               LoadedType, PreReadOffset));
     R = std::move(MetadataRecordOrErr.get());
   } else {
     R = llvm::make_unique<FunctionRecord>();
   }
   RecordInitializer RI(E, OffsetPtr);
 
   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);
 }
 
 } // namespace xray
 } // namespace llvm

llvm/trunk/lib/XRay/FDRRecords.cpp

 Error CallArgRecord::apply(RecordVisitor &V) { return V.visit(*this); }
 Error PIDRecord::apply(RecordVisitor &V) { return V.visit(*this); }
 Error NewBufferRecord::apply(RecordVisitor &V) { return V.visit(*this); }
 Error EndBufferRecord::apply(RecordVisitor &V) { return V.visit(*this); }
 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

llvm/trunk/lib/XRay/FDRTraceWriter.cpp

template <
           int>::type = 0>
   static size_t write(support::endian::Writer &OS, Tuple &&) {
     return 0;
   }
 };
 
 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);
   auto Bytes = IndexedWriter<0>::write(OS, T);
   assert(Bytes <= 15 && "Must only ever write at most 16 byte metadata!");
   // Pad out with appropriate numbers of zero's.
   for (; Bytes < 15; ++Bytes)
     OS.write('\0');
if (auto E = writeMetadata<5u>(OS, R.size(), R.delta()))
     return E;
   auto D = R.data();
   ArrayRef<char> Bytes(D.data(), D.size());
   OS.write(Bytes);
   return Error::success();
 }
 
 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());
   OS.write(Bytes);
   return Error::success();
 }
 
 Error FDRTraceWriter::visit(CallArgRecord &R) {

llvm/trunk/lib/XRay/RecordInitializer.cpp

Error RecordInitializer::visit(CustomEventRecord &R) {
   auto BeginOffset = OffsetPtr;
   auto PreReadOffset = OffsetPtr;
   R.Size = E.getSigned(&OffsetPtr, sizeof(int32_t));
   if (PreReadOffset == OffsetPtr)
     return createStringError(
         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)
     return createStringError(
         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
Error RecordInitializer::visit(CustomEventRecord &R) {
   if (!E.isValidOffsetForDataOfSize(OffsetPtr, R.Size))
     return createStringError(
         std::make_error_code(std::errc::bad_address),
         "Cannot read %d bytes of custom event data from offset %d.", R.Size,
         OffsetPtr);
 
   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();
 }
 
 Error RecordInitializer::visit(CustomEventRecordV5 &R) {
   if (!E.isValidOffsetForDataOfSize(OffsetPtr,
                                     MetadataRecord::kMetadataBodySize))
     return createStringError(std::make_error_code(std::errc::bad_address),
                              "Invalid offset for a custom event record (%d).",
                              OffsetPtr);
 
   auto BeginOffset = OffsetPtr;
   auto PreReadOffset = OffsetPtr;
 
   R.Size = E.getSigned(&OffsetPtr, sizeof(int32_t));
   if (PreReadOffset == OffsetPtr)
     return createStringError(
         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)
     return createStringError(
         std::make_error_code(std::errc::invalid_argument),
         "Cannot read a custom event record TSC delta 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.
   if (!E.isValidOffsetForDataOfSize(OffsetPtr, R.Size))
     return createStringError(
         std::make_error_code(std::errc::bad_address),
         "Cannot read %d bytes of custom event data from offset %d.", R.Size,
         OffsetPtr);
 
   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();
 }
 
 Error RecordInitializer::visit(TypedEventRecord &R) {
   if (!E.isValidOffsetForDataOfSize(OffsetPtr,
                                     MetadataRecord::kMetadataBodySize))
     return createStringError(std::make_error_code(std::errc::bad_address),
                              "Invalid offset for a typed event record (%d).",
                              OffsetPtr);
 
   auto BeginOffset = OffsetPtr;
   auto PreReadOffset = OffsetPtr;
 
   R.Size = E.getSigned(&OffsetPtr, sizeof(int32_t));
   if (PreReadOffset == OffsetPtr)
     return createStringError(
         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)
     return createStringError(
         std::make_error_code(std::errc::invalid_argument),
         "Cannot read a typed event record TSC delta field at offset %d.",
         OffsetPtr);
 
Error RecordInitializer::visit(TypedEventRecord &R) {
   if (!E.isValidOffsetForDataOfSize(OffsetPtr, R.Size))
     return createStringError(
         std::make_error_code(std::errc::bad_address),
         "Cannot read %d bytes of custom event data from offset %d.", R.Size,
         OffsetPtr);
 
   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();
 }
 
 Error RecordInitializer::visit(CallArgRecord &R) {
   if (!E.isValidOffsetForDataOfSize(OffsetPtr,
                                     MetadataRecord::kMetadataBodySize))
     return createStringError(std::make_error_code(std::errc::bad_address),
Error RecordInitializer::visit(FunctionRecord &R) {
 
   auto BeginOffset = OffsetPtr;
   auto PreReadOffset = BeginOffset;
   uint32_t Buffer = E.getU32(&OffsetPtr);
   if (PreReadOffset == OffsetPtr)
     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):
   case static_cast<unsigned>(RecordTypes::EXIT):
   case static_cast<unsigned>(RecordTypes::TAIL_EXIT):
     R.Kind = static_cast<RecordTypes>(FunctionType);
     break;
   default:

llvm/trunk/unittests/XRay/FDRProducerConsumerTest.cpp

 
 namespace llvm {
 namespace xray {
 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);
 }
 
 template <> std::unique_ptr<Record> MakeRecord<NewCPUIDRecord>() {
   return make_unique<NewCPUIDRecord>(1, 2);
 }
 

 template <> std::unique_ptr<Record> MakeRecord<PIDRecord>() {
   return make_unique<PIDRecord>(1);
 }
 
 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;
     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(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();
 
   DataExtractor DE(this->Data, sys::IsLittleEndianHost, 8);
   uint32_t OffsetPtr = 0;
   auto HeaderOrErr = readBinaryFormatHeader(DE, OffsetPtr);
   if (!HeaderOrErr)
TYPED_TEST_P(RoundTripTest, RoundTripsSingleValue) {
   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)));
-  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