[XRay][tools] Support tail-call exits before we write them in the runtime

[dberris]

Summary

This change adds support for explicit tail-exit records to be written by
the XRay runtime. This lets us differentiate the tail exit
records/events in the log, and allows us to treat those exit events
especially in the future. For now we allow printing those out in YAML
(and reading them in).

[pelikan]

Nitpick: maybe update docs/XRayFDRFormat.rst in the same diff.

[dberris]

• fixup: update to make tests clean; typo in docs

include/llvm/XRay/XRayRecord.h

struct XRayFileHeader {
   // buffer.
   char FreeFormData[16];
 };
 
 /// Determines the supported types of records that could be seen in XRay traces.
 /// This may or may not correspond to actual record types in the raw trace (as
 /// the loader implementation may synthesize this information in the process of
 /// of loading).
-enum class RecordTypes { ENTER, EXIT };
+enum class RecordTypes { ENTER, EXIT, TAIL_EXIT };
 
 struct XRayRecord {
   /// The type of record.
   uint16_t RecordType;
 
   /// The CPU where the thread is running. We assume number of CPUs <= 65536.
   uint16_t CPU;
 

include/llvm/XRay/YAMLXRayRecord.h

 namespace yaml {
 
 // YAML Traits
 // -----------
 template <> struct ScalarEnumerationTraits<xray::RecordTypes> {
   static void enumeration(IO &IO, xray::RecordTypes &Type) {
     IO.enumCase(Type, "function-enter", xray::RecordTypes::ENTER);
     IO.enumCase(Type, "function-exit", xray::RecordTypes::EXIT);
+    IO.enumCase(Type, "function-tail-exit", xray::RecordTypes::TAIL_EXIT);
   }
 };
 
 template <> struct MappingTraits<xray::YAMLXRayFileHeader> {
   static void mapping(IO &IO, xray::YAMLXRayFileHeader &Header) {
     IO.mapRequired("version", Header.Version);
     IO.mapRequired("type", Header.Type);
     IO.mapRequired("constant-tsc", Header.ConstantTSC);

lib/XRay/Trace.cpp

Error readBinaryFormatHeader(StringRef Data, XRayFileHeader &FileHeader) {
   uint32_t OffsetPtr = 0;
   FileHeader.Version = HeaderExtractor.getU16(&OffsetPtr);
   FileHeader.Type = HeaderExtractor.getU16(&OffsetPtr);
   uint32_t Bitfield = HeaderExtractor.getU32(&OffsetPtr);
   FileHeader.ConstantTSC = Bitfield & 1uL;
   FileHeader.NonstopTSC = Bitfield & 1uL << 1;
   FileHeader.CycleFrequency = HeaderExtractor.getU64(&OffsetPtr);
   std::memcpy(&FileHeader.FreeFormData, Data.bytes_begin() + OffsetPtr, 16);
-  if (FileHeader.Version != 1)
+  if (FileHeader.Version != 1 && FileHeader.Version != 2)
     return make_error<StringError>(
         Twine("Unsupported XRay file version: ") + Twine(FileHeader.Version),
         std::make_error_code(std::errc::invalid_argument));
   return Error::success();
 }
 
 Error loadNaiveFormatLog(StringRef Data, XRayFileHeader &FileHeader,
                          std::vector<XRayRecord> &Records) {
for (auto S = Data.drop_front(32); !S.empty(); S = S.drop_front(32)) {
     auto Type = RecordExtractor.getU8(&OffsetPtr);
     switch (Type) {
     case 0:
       Record.Type = RecordTypes::ENTER;
       break;
     case 1:
       Record.Type = RecordTypes::EXIT;
       break;
+    case 2:
+      if (FileHeader.Version == 2)
+        Record.Type = RecordTypes::TAIL_EXIT;
+      else
+        return make_error<StringError>(
+            Twine("Unsupported record type: ") + Twine(int{Type}) +
+                "; Tail exits only supported in Version 2 of the log.",
+            std::make_error_code(std::errc::executable_format_error));
     default:
       return make_error<StringError>(
           Twine("Unknown record type '") + Twine(int{Type}) + "'",
           std::make_error_code(std::errc::executable_format_error));
     }
     Record.FuncId = RecordExtractor.getSigned(&OffsetPtr, sizeof(int32_t));
     Record.TSC = RecordExtractor.getU64(&OffsetPtr);
     Record.TId = RecordExtractor.getU32(&OffsetPtr);
default:
     Record.RecordType = 0; // Record is type NORMAL.
     // Strip off record type bit and use the next three bits.
     uint8_t RecordType = (RecordFirstByte >> 1) & 0x07;
     switch (RecordType) {
     case static_cast<uint8_t>(RecordTypes::ENTER):
       Record.Type = RecordTypes::ENTER;
       break;
     case static_cast<uint8_t>(RecordTypes::EXIT):
-    case 2: // TAIL_EXIT is not yet defined in RecordTypes.
-      Record.Type = RecordTypes::EXIT;
+    case 2:
+      Record.Type = RecordTypes::TAIL_EXIT;
       break;
     default:
       // Cast to an unsigned integer to not interpret the record type as a char.
       return make_error<StringError>(
           Twine("Illegal function record type: ")
               .concat(Twine(static_cast<unsigned>(RecordType))),
           std::make_error_code(std::errc::executable_format_error));
     }
if (isMetadataRecord) {
         return E;
     }
     State.CurrentBufferConsumed += RecordSize;
   }
 
   // Having iterated over everything we've been given, we've either consumed
   // everything and ended up in the end state, or were told to skip the rest.
   bool Finished = State.Expects == FDRState::Token::SCAN_TO_END_OF_THREAD_BUF &&
                   State.CurrentBufferSize == State.CurrentBufferConsumed;
   if (State.Expects != FDRState::Token::NEW_BUFFER_RECORD_OR_EOF && !Finished)
     return make_error<StringError>(
         Twine("Encountered EOF with unexpected state expectation ") +
             fdrStateToTwine(State.Expects) +
             ". Remaining expected bytes in thread buffer total " +
             Twine(State.CurrentBufferSize - State.CurrentBufferConsumed),
         std::make_error_code(std::errc::executable_format_error));
 
Expected<Trace> llvm::xray::loadTraceFile(StringRef Filename, bool Sort) {
   uint32_t OffsetPtr = 0;
   uint16_t Version = HeaderExtractor.getU16(&OffsetPtr);
   uint16_t Type = HeaderExtractor.getU16(&OffsetPtr);
 
   enum BinaryFormatType { NAIVE_FORMAT = 0, FLIGHT_DATA_RECORDER_FORMAT = 1 };
 
   Trace T;
   if (Version == 1 && Type == NAIVE_FORMAT) {
-    if (auto E =
-            loadNaiveFormatLog(Data, T.FileHeader, T.Records))
+    if (auto E = loadNaiveFormatLog(Data, T.FileHeader, T.Records))
       return std::move(E);
   } else if (Version == 1 && Type == FLIGHT_DATA_RECORDER_FORMAT) {
     if (auto E = loadFDRLog(Data, T.FileHeader, T.Records))
       return std::move(E);
   } else {
     if (auto E = loadYAMLLog(Data, T.FileHeader, T.Records))
       return std::move(E);
   }

tools/llvm-xray/xray-account.cc

if (Record.TSC < CurrentMaxTSC)
     return false;
 
   auto &ThreadStack = PerThreadFunctionStack[Record.TId];
   switch (Record.Type) {
   case RecordTypes::ENTER: {
     ThreadStack.emplace_back(Record.FuncId, Record.TSC);
     break;
   }
-  case RecordTypes::EXIT: {
+  case RecordTypes::EXIT:
+  case RecordTypes::TAIL_EXIT: {
     if (ThreadStack.empty())
       return false;
 
     if (ThreadStack.back().first == Record.FuncId) {
       const auto &Top = ThreadStack.back();
       recordLatency(Top.first, diff(Top.second, Record.TSC));
       ThreadStack.pop_back();
       break;
static void format(const llvm::xray::RecordTypes &T, raw_ostream &Stream,
                      StringRef Style) {
     switch(T) {
       case RecordTypes::ENTER:
         Stream << "enter";
         break;
       case RecordTypes::EXIT:
         Stream << "exit";
         break;
+      case RecordTypes::TAIL_EXIT:
+        Stream << "tail-exit";
+        break;
     }
   }
 };
 } // namespace llvm
 
 static CommandRegistration Unused(&Account, []() -> Error {
   InstrumentationMap Map;
   if (!AccountInstrMap.empty()) {

tools/llvm-xray/xray-converter.cc

for (const auto &R : Records) {
     Writer.write(static_cast<uint8_t>(R.CPU));
     switch (R.Type) {
     case RecordTypes::ENTER:
       Writer.write(uint8_t{0});
       break;
     case RecordTypes::EXIT:
       Writer.write(uint8_t{1});
       break;
+    case RecordTypes::TAIL_EXIT:
+      Writer.write(uint8_t{2});
     }
     Writer.write(R.FuncId);
     Writer.write(R.TSC);
     Writer.write(R.TId);
     Writer.write(Padding4B);
     Writer.write(Padding4B);
     Writer.write(Padding4B);
   }

tools/llvm-xray/xray-graph.cc

Error GraphRenderer::accountRecord(const XRayRecord &Record) {
   auto &ThreadStack = PerThreadFunctionStack[Record.TId];
   switch (Record.Type) {
   case RecordTypes::ENTER: {
     if (Record.FuncId != 0 && G.count(Record.FuncId) == 0)
       G[Record.FuncId].SymbolName = FuncIdHelper.SymbolOrNumber(Record.FuncId);
     ThreadStack.push_back({Record.FuncId, Record.TSC});
     break;
   }
-  case RecordTypes::EXIT: {
+  case RecordTypes::EXIT:
+  case RecordTypes::TAIL_EXIT: {
     // FIXME: Refactor this and the account subcommand to reduce code
     // duplication
     if (ThreadStack.size() == 0 || ThreadStack.back().FuncId != Record.FuncId) {
       if (!DeduceSiblingCalls)
         return make_error<StringError>("No matching ENTRY record",
                                        make_error_code(errc::invalid_argument));
       auto Parent = std::find_if(
           ThreadStack.rbegin(), ThreadStack.rend(),

tools/llvm-xray/xray-stacks.cc

case RecordTypes::ENTER: {
         TS.emplace_back(N, R.TSC);
       } else {
         // We found the callee in the stack trie, so we'll use that pointer
         // instead, add it to the stack associated with the TSC.
         TS.emplace_back(*I, R.TSC);
       }
       return AccountRecordStatus::OK;
     }
-    case RecordTypes::EXIT: {
+    case RecordTypes::EXIT:
+    case RecordTypes::TAIL_EXIT: {
       bool wasLastRecordExit = state->wasLastRecordExit;
       state->wasLastRecordExit = true;
       // The exit case is more interesting, since we want to be able to deduce
       // missing exit records. To do that properly, we need to look up the stack
       // and see whether the exit record matches any of the entry records. If it
       // does match, we attempt to record the durations as we pop the stack to
       // where we see the parent.
       if (TS.empty()) {