[XRay] Fix TSC and atomic custom/typed event accounting

[dberris]

Summary

This is a follow-on change to D53858 which turns out to have had a TSC
accounting bug when writing out function exit records in FDR mode.

This change adds a number of tests to ensure that:

• We are handling the delta between the exit TSC and the last TSC we've seen.

• We are writing the custom event and typed event records as a single update to the buffer extents.

• We are able to catch boundary conditions when loading FDR logs.

We introduce a TSC matcher to the test helpers, which we use in the
testing/verification of the TSC accounting change.

[dberris]

Write down correct delta.

[mboerger]

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

126

to see if...

compiler-rt/lib/xray/tests/unit/fdr_log_writer_test.cc

102

s/THen/Then/
where is first?

112

ASSERT_THAT(B.Data, IsNull());

compiler-rt/lib/xray/xray_fdr_log_writer.h

116

s/first before/first, before/

[dberris]

Address comments by @mboerger.

[mboerger]

LGTM

compiler-rt/lib/xray/tests/CMakeLists.txt

foreach(arch ${XRAY_TEST_ARCH})
         SOURCES ${TEST_SOURCES} ${COMPILER_RT_GTEST_SOURCE}
                 ${COMPILER_RT_GMOCK_SOURCE}
 
         # Note that any change in the implementations will cause all the unit
         # tests to be re-built. This is by design, but may be cumbersome during
         # the build/test cycle.
         COMPILE_DEPS ${TEST_SOURCES} ${COMPILER_RT_GTEST_SOURCE}
         ${XRAY_HEADERS} ${XRAY_ALL_SOURCE_FILES_ABS_PATHS}
+        "test_helpers.h"
         RUNTIME "${XRAY_RUNTIME_LIBS}"
         DEPS gtest xray llvm-xray LLVMXRay LLVMTestingSupport
         CFLAGS ${XRAY_UNITTEST_CFLAGS}
         LINK_FLAGS ${TARGET_LINK_FLAGS} ${XRAY_UNITTEST_LINK_FLAGS}
         )
       set_target_properties(XRayUnitTests
         PROPERTIES
         RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})

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

 
 namespace __xray {
 namespace {
 
 using ::llvm::HasValue;
 using ::llvm::xray::testing::FuncId;
 using ::llvm::xray::testing::HasArg;
 using ::llvm::xray::testing::RecordType;
+using ::llvm::xray::testing::TSCIs;
 using ::testing::AllOf;
 using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::Field;
+using ::testing::Gt;
 using ::testing::IsEmpty;
 using ::testing::SizeIs;
 
 class FunctionSequenceTest : public ::testing::Test {
 protected:
   BufferQueue::Buffer B{};
   std::unique_ptr<BufferQueue> BQ;
   std::unique_ptr<FDRLogWriter> W;
TEST_F(FunctionSequenceTest, ArgsAreHandledAndKept) {
   EXPECT_THAT_EXPECTED(
       TraceOrErr,
       HasValue(ElementsAre(
           AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::ENTER_ARG),
                 HasArg(4)),
           AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::EXIT)))));
 }
 
+TEST_F(FunctionSequenceTest, PreservedCallsHaveCorrectTSC) {
+  C = llvm::make_unique<FDRController<>>(BQ.get(), B, *W, clock_gettime, 1000);
+  uint64_t TSC = 1;
+  uint16_t CPU = 0;
+  ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(2, TSC++, CPU));
+  ASSERT_TRUE(C->functionExit(2, TSC++, CPU));
+  ASSERT_TRUE(C->functionExit(1, TSC += 1000, CPU));
+  ASSERT_TRUE(C->flush());
+  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+
+  // Serialize the buffers then test to see we find the *no* records, because
+  // the function entry-exit comes under the cycle threshold.
+  std::string Serialized = serialize(*BQ, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(
+      TraceOrErr,
+      HasValue(ElementsAre(
+          AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::ENTER),
+                TSCIs(Eq(1uL))),
+          AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::EXIT),
+                TSCIs(Gt(1000uL))))));
+}
+
 TEST_F(FunctionSequenceTest, RewindingMultipleCalls) {
   C = llvm::make_unique<FDRController<>>(BQ.get(), B, *W, clock_gettime, 1000);
 
   // First we construct an arbitrarily deep function enter/call stack.
   // We also ensure that we are in the same CPU.
   uint64_t TSC = 1;
   uint16_t CPU = 1;
   ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
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, 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)));
+  }
+  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, HandlesFinalizedBufferQueue) {
   uint64_t TSC = 1;
   uint16_t CPU = 1;
 
   // First write one function entry.
   ASSERT_TRUE(C->functionEnter(1, TSC++, CPU));
 
   // Then we finalize the buffer queue, simulating the case where the logging

compiler-rt/lib/xray/tests/unit/fdr_log_writer_test.cc

 namespace __xray {
 namespace {
 
 static constexpr size_t kSize = 4096;
 
 using ::llvm::HasValue;
 using ::llvm::xray::testing::FuncId;
 using ::llvm::xray::testing::RecordType;
-using ::testing::Eq;
 using ::testing::AllOf;
-using ::testing::IsEmpty;
 using ::testing::ElementsAre;
+using ::testing::Eq;
+using ::testing::IsEmpty;
 
 // Exercise the common code path where we initialize a buffer and are able to
 // write some records successfully.
 TEST(FdrLogWriterTest, WriteSomeRecords) {
   bool Success = false;
   BufferQueue Buffers(kSize, 1, Success);
   BufferQueue::Buffer B;
   ASSERT_EQ(Buffers.getBuffer(B), BufferQueue::ErrorCode::Ok);
TEST(FdrLogWriterTest, WriteSomeRecords) {
   auto TraceOrErr = llvm::xray::loadTrace(DE);
   EXPECT_THAT_EXPECTED(
       TraceOrErr,
       HasValue(ElementsAre(
           AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::ENTER)),
           AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::EXIT)))));
 }
 
+// Ensure that we can handle buffer re-use.
+TEST(FdrLogWriterTest, ReuseBuffers) {
+  bool Success = false;
+  BufferQueue Buffers(kSize, 1, Success);
+  BufferQueue::Buffer B;
+  ASSERT_EQ(Buffers.getBuffer(B), BufferQueue::ErrorCode::Ok);
+
+  FDRLogWriter Writer(B);
+  MetadataRecord Preamble[] = {
+      createMetadataRecord<MetadataRecord::RecordKinds::NewBuffer>(int32_t{1}),
+      createMetadataRecord<MetadataRecord::RecordKinds::WalltimeMarker>(
+          int64_t{1}, int32_t{2}),
+      createMetadataRecord<MetadataRecord::RecordKinds::Pid>(int32_t{1}),
+  };
+  ASSERT_THAT(Writer.writeMetadataRecords(Preamble),
+              Eq(sizeof(MetadataRecord) * 3));
+  ASSERT_TRUE(Writer.writeMetadata<MetadataRecord::RecordKinds::NewCPUId>(
+      uint16_t{1}, uint64_t{1}));
+  uint64_t TSC = 1;
+  ASSERT_TRUE(
+      Writer.writeFunction(FDRLogWriter::FunctionRecordKind::Enter, 1, TSC++));
+  ASSERT_TRUE(
+      Writer.writeFunction(FDRLogWriter::FunctionRecordKind::Exit, 1, TSC++));
+  ASSERT_EQ(Buffers.releaseBuffer(B), BufferQueue::ErrorCode::Ok);
+  ASSERT_EQ(B.Data, nullptr);
+
+  // THen we re-use the buffer, but only write one record.
+  ASSERT_EQ(Buffers.getBuffer(B), BufferQueue::ErrorCode::Ok);
+  Writer.resetRecord();
+  ASSERT_THAT(Writer.writeMetadataRecords(Preamble),
+              Eq(sizeof(MetadataRecord) * 3));
+  ASSERT_TRUE(Writer.writeMetadata<MetadataRecord::RecordKinds::NewCPUId>(
+      uint16_t{1}, uint64_t{1}));
+  ASSERT_TRUE(
+      Writer.writeFunction(FDRLogWriter::FunctionRecordKind::Enter, 1, TSC++));
+  ASSERT_EQ(Buffers.releaseBuffer(B), BufferQueue::ErrorCode::Ok);
+  ASSERT_EQ(B.Data, nullptr);
+  ASSERT_EQ(Buffers.finalize(), BufferQueue::ErrorCode::Ok);
+
+  // Then we validate that we only see the single enter record.
+  std::string Serialized = serialize(Buffers, 3);
+  llvm::DataExtractor DE(Serialized, true, 8);
+  auto TraceOrErr = llvm::xray::loadTrace(DE);
+  EXPECT_THAT_EXPECTED(
+      TraceOrErr, HasValue(ElementsAre(AllOf(
+                      FuncId(1), RecordType(llvm::xray::RecordTypes::ENTER)))));
+}
+
 TEST(FdrLogWriterTest, UnwriteRecords) {
   bool Success = false;
   BufferQueue Buffers(kSize, 1, Success);
   BufferQueue::Buffer B;
   ASSERT_EQ(Buffers.getBuffer(B), BufferQueue::ErrorCode::Ok);
 
   FDRLogWriter Writer(B);
   MetadataRecord Preamble[] = {

compiler-rt/lib/xray/tests/unit/test_helpers.h

 
 MATCHER_P(HasArg, A, "") {
   *result_listener << "where args contains " << A;
   return !arg.CallArgs.empty() &&
          std::any_of(arg.CallArgs.begin(), arg.CallArgs.end(),
                      [this](decltype(A) V) { return V == A; });
 }
 
+MATCHER_P(TSCIs, M, std::string("TSC is ") + ::testing::PrintToString(M)) {
+  return ::testing::Matcher<decltype(arg.TSC)>(M).MatchAndExplain(
+      arg.TSC, result_listener);
+}
+
 } // namespace testing
 } // namespace xray
 } // namespace llvm
 
 namespace __xray {
 
 std::string serialize(BufferQueue &Buffers, int32_t Version);
 

compiler-rt/lib/xray/xray_fdr_controller.h

bool functionExit(int32_t FuncId, uint64_t TSC,
     if (finalized() ||
         !prepareBuffer(sizeof(MetadataRecord) + sizeof(FunctionRecord)))
       return returnBuffer();
 
     auto PreambleStatus = functionPreamble(TSC, CPU);
     if (PreambleStatus == PreambleResult::InvalidBuffer)
       return returnBuffer();
 
+    auto Delta = TSC - LastFunctionEntryTSC;
     if (PreambleStatus == PreambleResult::NoChange &&
-        UndoableFunctionEnters != 0 &&
-        TSC - LastFunctionEntryTSC < CycleThreshold)
+        UndoableFunctionEnters != 0 && Delta < CycleThreshold)
       return rewindRecords(FuncId, TSC, CPU);
 
     LatestTSC = TSC;
     UndoableFunctionEnters = 0;
     UndoableTailExits = 0;
     return W.writeFunction(FDRLogWriter::FunctionRecordKind::Exit, mask(FuncId),
-                           TSC - LatestTSC);
+                           Delta);
   }
 
   bool customEvent(uint64_t TSC, uint16_t CPU, const void *Event,
                    int32_t EventSize) XRAY_NEVER_INSTRUMENT {
     if (finalized() ||
         !prepareBuffer((2 * sizeof(MetadataRecord)) + EventSize) ||
         functionPreamble(TSC, CPU) == PreambleResult::InvalidBuffer)
       return returnBuffer();

compiler-rt/lib/xray/xray_fdr_log_writer.h

bool writeFunction(FunctionRecordKind Kind, int32_t FuncId, int32_t Delta) {
     R.FuncId = FuncId;
     R.TSCDelta = Delta;
     writeRecord(R);
     return true;
   }
 
   bool writeCustomEvent(uint64_t TSC, uint16_t CPU, const void *Event,
                         int32_t EventSize) {
-    writeMetadata<MetadataRecord::RecordKinds::CustomEventMarker>(EventSize,
-                                                                  TSC, CPU);
-    internal_memcpy(NextRecord, Event, EventSize);
-    NextRecord += EventSize;
-    atomic_fetch_add(&Buffer.Extents, EventSize, memory_order_acq_rel);
+    // 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>(
+            EventSize, TSC, CPU);
+    NextRecord = reinterpret_cast<char *>(internal_memcpy(
+                     NextRecord, reinterpret_cast<char *>(&R), sizeof(R))) +
+                 sizeof(R);
+    NextRecord = reinterpret_cast<char *>(
+                     internal_memcpy(NextRecord, Event, EventSize)) +
+                 EventSize;
+    atomic_fetch_add(&Buffer.Extents, sizeof(R) + EventSize,
+                     memory_order_acq_rel);
     return true;
   }
 
   bool writeTypedEvent(uint64_t TSC, uint16_t EventType, const void *Event,
                        int32_t EventSize) {
-    writeMetadata<MetadataRecord::RecordKinds::TypedEventMarker>(EventSize, TSC,
-                                                                 EventType);
-    internal_memcpy(NextRecord, Event, EventSize);
-    NextRecord += EventSize;
+    // We do something similar when writing out typed events, see
+    // writeCustomEvent(...) above for details.
+    MetadataRecord R =
+        createMetadataRecord<MetadataRecord::RecordKinds::TypedEventMarker>(
+            EventSize, TSC, EventType);
+    NextRecord = reinterpret_cast<char *>(internal_memcpy(
+                     NextRecord, reinterpret_cast<char *>(&R), sizeof(R))) +
+                 sizeof(R);
+    NextRecord = reinterpret_cast<char *>(
+                     internal_memcpy(NextRecord, Event, EventSize)) +
+                 EventSize;
     atomic_fetch_add(&Buffer.Extents, EventSize, memory_order_acq_rel);
     return true;
   }
 
   char *getNextRecord() const { return NextRecord; }
 
   void resetRecord() {
     NextRecord = reinterpret_cast<char *>(Buffer.Data);

llvm/lib/XRay/FDRRecordProducer.cpp

Expected<std::unique_ptr<Record>> FileBasedRecordProducer::produce() {
   //     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) {
     auto LoadedType = FirstByte >> 1;
     auto MetadataRecordOrErr = metadataRecordType(Header, LoadedType);
     if (!MetadataRecordOrErr)
       return joinErrors(