[XRay] Bug fixes for FDR custom event and arg-logging

[dberris]

Summary

This change has a number of fixes for FDR mode in compiler-rt along with
changes to the tooling handling the traces in llvm.

In the runtime, we do the following:

• Advance the "last record" pointer appropriately when writing the custom event data in the log.

• Add XRAY_NEVER_INSTRUMENT in the rewinding routine.

• When collecting the argument of functions appropriately marked, we should not attempt to rewind them (and reset the counts of functions that can be re-wound).

In the tooling, we do the following:

• Remove the state logic in BlockIndexer and instead rely on the presence/absence of records to indicate blocks.

• Move the verifier into a loop associated with each block.

[dberris]

Retitle, update.

[mboerger]

compiler-rt/lib/xray/xray_fdr_logging.cc

293

Put the formula on one line and maybe state that the number is 0x0FFFFFFFF

623

what about the 'second'

[dberris]

Address comments by @mboerger.

[dberris]

0xfff... -> 0x0fff...

[mboerger]

LGTM

compiler-rt/lib/xray/xray_fdr_logging.cc

static void writeCallArgumentMetadata(uint64_t A) XRAY_NEVER_INSTRUMENT {
   CallArg.RecordKind = uint8_t(MetadataRecord::RecordKinds::CallArgument);
 
   internal_memcpy(CallArg.Data, &A, sizeof(A));
   internal_memcpy(TLD.RecordPtr, &CallArg, sizeof(MetadataRecord));
   TLD.RecordPtr += sizeof(MetadataRecord);
   incrementExtents(sizeof(MetadataRecord));
 }
 
-static void writeFunctionRecord(int FuncId, uint32_t TSCDelta,
+static void writeFunctionRecord(int32_t FuncId, uint32_t TSCDelta,
                                 XRayEntryType EntryType) XRAY_NEVER_INSTRUMENT {
+  constexpr int32_t MaxFuncId = (1 << 29) - 1;
+  if (UNLIKELY(FuncId > MaxFuncId)) {
+    if (Verbosity())
+      Report("Warning: Function ID '%d' > max function id: '%d'", FuncId,
+             MaxFuncId);
+    return;
+  }
+
   FunctionRecord FuncRecord;
   FuncRecord.Type = uint8_t(RecordType::Function);
+
   // Only take 28 bits of the function id.
-  FuncRecord.FuncId = FuncId & ~(0x0F << 28);
+  //
+  // We need to be careful about the sign bit and the bitwise operations being
+  // performed here. In effect, we want to truncate the value of the function id
+  // to the first 28 bits -- this means we need to mask the function id with (2
+  // ^ 28) - 1, or 0xfffffff.
+  //
+  FuncRecord.FuncId = FuncId & MaxFuncId;
   FuncRecord.TSCDelta = TSCDelta;
 
   auto &TLD = getThreadLocalData();
   switch (EntryType) {
   case XRayEntryType::ENTRY:
     ++TLD.NumConsecutiveFnEnters;
     FuncRecord.RecordKind = uint8_t(FunctionRecord::RecordKinds::FunctionEnter);
     break;

 }
 
 static atomic_uint64_t TicksPerSec{0};
 static atomic_uint64_t ThresholdTicks{0};
 
 // Re-point the thread local pointer into this thread's Buffer before the recent
 // "Function Entry" record and any "Tail Call Exit" records after that.
 static void rewindRecentCall(uint64_t TSC, uint64_t &LastTSC,
-                             uint64_t &LastFunctionEntryTSC, int32_t FuncId) {
+                             uint64_t &LastFunctionEntryTSC,
+                             int32_t FuncId) XRAY_NEVER_INSTRUMENT {
   auto &TLD = getThreadLocalData();
   TLD.RecordPtr -= FunctionRecSize;
   decrementExtents(FunctionRecSize);
   FunctionRecord FuncRecord;
   internal_memcpy(&FuncRecord, TLD.RecordPtr, FunctionRecSize);
   DCHECK(FuncRecord.RecordKind ==
              uint8_t(FunctionRecord::RecordKinds::FunctionEnter) &&
          "Expected to find function entry recording when rewinding.");

 //     the correct TSC delta.
 static uint32_t writeCurrentCPUTSC(ThreadLocalData &TLD, uint64_t TSC,
                                    uint8_t CPU) {
   if (CPU != TLD.CurrentCPU) {
     // We've moved to a new CPU.
     writeNewCPUIdMetadata(CPU, TSC);
     return 0;
   }
+
   // If the delta is greater than the range for a uint32_t, then we write out
   // the TSC wrap metadata entry with the full TSC, and the TSC for the
   // function record be 0.
   uint64_t Delta = TSC - TLD.LastTSC;
   if (Delta <= std::numeric_limits<uint32_t>::max())
     return Delta;
 
   writeTSCWrapMetadata(TSC);
static void processFunctionHook(int32_t FuncId, XRayEntryType Entry,
   // that the record we are about to write is going to fit into the buffer,
   // without overflowing the buffer.
   //
   // To do this properly, we use the following assumptions:
   //
   //   - The least number of bytes we will ever write is 8
   //     (sizeof(FunctionRecord)) only if the delta between the previous entry
   //     and this entry is within 32 bits.
-  //   - The most number of bytes we will ever write is 8 + 16 + 16 = 40.
+  //   - The most number of bytes we will ever write is 8 + 16 + 16 + 16 = 56.
   //     This is computed by:
   //
-  //       MaxSize = sizeof(FunctionRecord) + 2 * sizeof(MetadataRecord)
+  //       MaxSize = sizeof(FunctionRecord) + 3 * sizeof(MetadataRecord)
   //
   //     These arise in the following cases:
   //
   //       1. When the delta between the TSC we get and the previous TSC for the
   //          same CPU is outside of the uint32_t range, we end up having to
   //          write a MetadataRecord to indicate a "tsc wrap" before the actual
   //          FunctionRecord.
   //       2. When we learn that we've moved CPUs, we need to write a
   //          MetadataRecord to indicate a "cpu change", and thus write out the
   //          current TSC for that CPU before writing out the actual
   //          FunctionRecord.
   //       3. When we learn about a new CPU ID, we need to write down a "new cpu
   //          id" MetadataRecord before writing out the actual FunctionRecord.
-  //       4. The second MetadataRecord is the optional function call argument.
+  //       4. The third MetadataRecord is the optional function call argument.
   //
   // So the math we need to do is to determine whether writing 40 bytes past the
   // current pointer exceeds the buffer's maximum size. If we don't have enough
   // space to write 40 bytes in the buffer, we need get a new Buffer, set it up
   // properly before doing any further writing.
-  size_t MaxSize = FunctionRecSize + 2 * MetadataRecSize;
+  size_t MaxSize = FunctionRecSize + 3 * MetadataRecSize;
   if (!prepareBuffer(TSC, CPU, wall_clock_reader, MaxSize)) {
     TLD.BQ = nullptr;
     return;
   }
 
-  // By this point, we are now ready to write up to 40 bytes (explained above).
-  DCHECK((TLD.RecordPtr + MaxSize) - static_cast<char *>(TLD.Buffer.Data) >=
-             static_cast<ptrdiff_t>(MetadataRecSize) &&
-         "Misconfigured BufferQueue provided; Buffer size not large enough.");
-
   auto RecordTSCDelta = writeCurrentCPUTSC(TLD, TSC, CPU);
   TLD.LastTSC = TSC;
   TLD.CurrentCPU = CPU;
   switch (Entry) {
   case XRayEntryType::ENTRY:
-  case XRayEntryType::LOG_ARGS_ENTRY:
     // Update the thread local state for the next invocation.
     TLD.LastFunctionEntryTSC = TSC;
     break;
+  case XRayEntryType::LOG_ARGS_ENTRY:
+    // Update the thread local state for the next invocation, but also prevent
+    // rewinding when we have arguments logged.
+    TLD.LastFunctionEntryTSC = TSC;
+    TLD.NumConsecutiveFnEnters = 0;
+    TLD.NumTailCalls = 0;
+    break;
   case XRayEntryType::TAIL:
   case XRayEntryType::EXIT:
     // Break out and write the exit record if we can't erase any functions.
     if (TLD.NumConsecutiveFnEnters == 0 ||
         (TSC - TLD.LastFunctionEntryTSC) >=
             atomic_load_relaxed(&ThresholdTicks))
       break;
     rewindRecentCall(TSC, TLD.LastTSC, TLD.LastFunctionEntryTSC, FuncId);
void fdrLoggingHandleCustomEvent(void *Event,
   auto &TLD = getThreadLocalData();
   if (!isLogInitializedAndReady(TLD.BQ, TSC, CPU, clock_gettime))
     return;
 
   // Here we need to prepare the log to handle:
   //   - The metadata record we're going to write. (16 bytes)
   //   - The additional data we're going to write. Currently, that's the size
   //   of the event we're going to dump into the log as free-form bytes.
-  if (!prepareBuffer(TSC, CPU, clock_gettime, MetadataRecSize + EventSize)) {
+  if (!prepareBuffer(TSC, CPU, clock_gettime,
+                     MetadataRecSize + ReducedEventSize)) {
     TLD.BQ = nullptr;
     return;
   }
 
+  // We need to reset the counts for the number of functions we're able to
+  // rewind.
+  TLD.NumConsecutiveFnEnters = 0;
+  TLD.NumTailCalls = 0;
+
   // Write the custom event metadata record, which consists of the following
   // information:
   //   - 8 bytes (64-bits) for the full TSC when the event started.
   //   - 4 bytes (32-bits) for the length of the data.
   MetadataRecord CustomEvent;
   CustomEvent.Type = uint8_t(RecordType::Metadata);
   CustomEvent.RecordKind =
       uint8_t(MetadataRecord::RecordKinds::CustomEventMarker);
   constexpr auto TSCSize = sizeof(TC.TSC);
   internal_memcpy(&CustomEvent.Data, &ReducedEventSize, sizeof(int32_t));
-  internal_memcpy(&CustomEvent.Data[sizeof(int32_t)], &TSC, TSCSize);
+  internal_memcpy(&CustomEvent.Data + sizeof(int32_t), &TSC, TSCSize);
   internal_memcpy(TLD.RecordPtr, &CustomEvent, sizeof(CustomEvent));
   TLD.RecordPtr += sizeof(CustomEvent);
   internal_memcpy(TLD.RecordPtr, Event, ReducedEventSize);
-  incrementExtents(MetadataRecSize + EventSize);
+  TLD.RecordPtr += ReducedEventSize;
+  incrementExtents(MetadataRecSize + ReducedEventSize);
   endBufferIfFull();
 }
 
 void fdrLoggingHandleTypedEvent(
     uint16_t EventType, const void *Event,
     std::size_t EventSize) noexcept XRAY_NEVER_INSTRUMENT {
   auto TC = getTimestamp();
   auto &TSC = TC.TSC;
void fdrLoggingHandleTypedEvent(
   auto &TLD = getThreadLocalData();
   if (!isLogInitializedAndReady(TLD.BQ, TSC, CPU, clock_gettime))
     return;
 
   // Here we need to prepare the log to handle:
   //   - The metadata record we're going to write. (16 bytes)
   //   - The additional data we're going to write. Currently, that's the size
   //   of the event we're going to dump into the log as free-form bytes.
-  if (!prepareBuffer(TSC, CPU, clock_gettime, MetadataRecSize + EventSize)) {
+  if (!prepareBuffer(TSC, CPU, clock_gettime,
+                     MetadataRecSize + ReducedEventSize)) {
     TLD.BQ = nullptr;
     return;
   }
   // Write the custom event metadata record, which consists of the following
   // information:
   //   - 8 bytes (64-bits) for the full TSC when the event started.
   //   - 4 bytes (32-bits) for the length of the data.
   //   - 2 bytes (16-bits) for the event type. 3 bytes remain since one of the
   //       bytes has the record type (Metadata Record) and kind (TypedEvent).
   //       We'll log the error if the event type is greater than 2 bytes.
   //       Event types are generated sequentially, so 2^16 is enough.
   MetadataRecord TypedEvent;
   TypedEvent.Type = uint8_t(RecordType::Metadata);
   TypedEvent.RecordKind =
       uint8_t(MetadataRecord::RecordKinds::TypedEventMarker);
   constexpr auto TSCSize = sizeof(TC.TSC);
   internal_memcpy(&TypedEvent.Data, &ReducedEventSize, sizeof(int32_t));
   internal_memcpy(&TypedEvent.Data[sizeof(int32_t)], &TSC, TSCSize);
   internal_memcpy(&TypedEvent.Data[sizeof(int32_t) + TSCSize], &EventType,
                   sizeof(EventType));
   internal_memcpy(TLD.RecordPtr, &TypedEvent, sizeof(TypedEvent));
 
   TLD.RecordPtr += sizeof(TypedEvent);
   internal_memcpy(TLD.RecordPtr, Event, ReducedEventSize);
+  TLD.RecordPtr += ReducedEventSize;
   incrementExtents(MetadataRecSize + EventSize);
   endBufferIfFull();
 }
 
 XRayLogInitStatus fdrLoggingInit(size_t BufferSize, size_t BufferMax,
                                  void *Options,
                                  size_t OptionsSize) XRAY_NEVER_INSTRUMENT {
   if (Options == nullptr)

llvm/include/llvm/XRay/BlockIndexer.h

public:
   };
 
   // This maps the process + thread combination to a sequence of blocks.
   using Index = DenseMap<std::pair<uint64_t, int32_t>, std::vector<Block>>;
 
 private:
   Index &Indices;
 
-  enum class State : unsigned { SeekExtents, ExtentsFound, ThreadIDFound };
-
-  State CurrentState = State::SeekExtents;
   Block CurrentBlock{0, 0, nullptr, {}};
 
 public:
   explicit BlockIndexer(Index &I) : RecordVisitor(), Indices(I) {}
 
   Error visit(BufferExtents &) override;
   Error visit(WallclockRecord &) override;
   Error visit(NewCPUIDRecord &) override;

llvm/lib/XRay/BlockIndexer.cpp

 // thread and a range of records representing a block.
 //
 //===----------------------------------------------------------------------===//
 #include "llvm/XRay/BlockIndexer.h"
 
 namespace llvm {
 namespace xray {
 
-Error BlockIndexer::visit(BufferExtents &) {
-  if (CurrentState == State::ThreadIDFound) {
-    Index::iterator It;
-    std::tie(It, std::ignore) =
-        Indices.insert({{CurrentBlock.ProcessID, CurrentBlock.ThreadID}, {}});
-    It->second.push_back({CurrentBlock.ProcessID, CurrentBlock.ThreadID,
-                          CurrentBlock.WallclockTime,
-                          std::move(CurrentBlock.Records)});
-    CurrentBlock.ProcessID = 0;
-    CurrentBlock.ThreadID = 0;
-    CurrentBlock.WallclockTime = nullptr;
-    CurrentBlock.Records = {};
-  }
-  CurrentState = State::ExtentsFound;
-  return Error::success();
-}
+Error BlockIndexer::visit(BufferExtents &) { return Error::success(); }
 
 Error BlockIndexer::visit(WallclockRecord &R) {
   CurrentBlock.Records.push_back(&R);
   CurrentBlock.WallclockTime = &R;
   return Error::success();
 }
 
 Error BlockIndexer::visit(NewCPUIDRecord &R) {

 
 Error BlockIndexer::visit(PIDRecord &R) {
   CurrentBlock.ProcessID = R.pid();
   CurrentBlock.Records.push_back(&R);
   return Error::success();
 }
 
 Error BlockIndexer::visit(NewBufferRecord &R) {
-  CurrentState = State::ThreadIDFound;
+  if (!CurrentBlock.Records.empty())
+    if (auto E = flush())
+      return E;
+
   CurrentBlock.ThreadID = R.tid();
   CurrentBlock.Records.push_back(&R);
   return Error::success();
 }
 
 Error BlockIndexer::visit(EndBufferRecord &R) {
-  CurrentState = State::SeekExtents;
   CurrentBlock.Records.push_back(&R);
   return Error::success();
 }
 
 Error BlockIndexer::visit(FunctionRecord &R) {
   CurrentBlock.Records.push_back(&R);
   return Error::success();
 }
 
 Error BlockIndexer::flush() {
-  CurrentState = State::SeekExtents;
   Index::iterator It;
   std::tie(It, std::ignore) =
       Indices.insert({{CurrentBlock.ProcessID, CurrentBlock.ThreadID}, {}});
   It->second.push_back({CurrentBlock.ProcessID, CurrentBlock.ThreadID,
                         CurrentBlock.WallclockTime,
                         std::move(CurrentBlock.Records)});
   CurrentBlock.ProcessID = 0;
   CurrentBlock.ThreadID = 0;
   CurrentBlock.Records = {};
+  CurrentBlock.WallclockTime = nullptr;
   return Error::success();
 }
 
 } // namespace xray
 } // namespace llvm

llvm/lib/XRay/Trace.cpp

for (auto &R : FDRRecords)
       if (auto E = R->apply(Indexer))
         return E;
     if (auto E = Indexer.flush())
       return E;
   }
 
   // Then we verify the consistency of the blocks.
   {
-    BlockVerifier Verifier;
     for (auto &PTB : Index) {
       auto &Blocks = PTB.second;
       for (auto &B : Blocks) {
+        BlockVerifier Verifier;
         for (auto *R : B.Records)
           if (auto E = R->apply(Verifier))
             return E;
         if (auto E = Verifier.verify())
           return E;
-        Verifier.reset();
       }
     }
   }
 
   // This is now the meat of the algorithm. Here we sort the blocks according to
   // the Walltime record in each of the blocks for the same thread. This allows
   // us to more consistently recreate the execution trace in temporal order.
   // After the sort, we then reconstitute `Trace` records using a stateful