[XRay] Fix enter function tracing for record unwriting

Summary:
Before this change, we could run into a situation where we may try to
undo tail exit records after writing metadata records before a function
enter event. This change rectifies that by resetting the tail exit
counter after writing the metadata records.

Reviewers: mboerger

Subscribers: llvm-commits

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

llvm-svn: 346475

1 files changed, 44 insertions, 0 deletions

diff --git a/compiler-rt/lib/xray/tests/unit/fdr_controller_test.cc b/compiler-rt/lib/xray/tests/unit/fdr_controller_test.cc
index 075bf83127b..1f02173256e 100644
--- a/compiler-rt/lib/xray/tests/unit/fdr_controller_test.cc
+++ b/compiler-rt/lib/xray/tests/unit/fdr_controller_test.cc
@@ -222,6 +222,50 @@ TEST_F(FunctionSequenceTest, RewindingIntermediaryTailExits) {
   EXPECT_THAT_EXPECTED(TraceOrErr, HasValue(IsEmpty()));
 }
 
+TEST_F(FunctionSequenceTest, RewindingAfterMigration) {
+  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));
+  ASSERT_TRUE(C->functionEnter(2, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(3, TSC++, CPU));
+
+  // Next we tail-exit into a new function multiple times.
+  ASSERT_TRUE(C->functionTailExit(3, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(4, TSC++, CPU));
+  ASSERT_TRUE(C->functionTailExit(4, TSC++, CPU));
+
+  // But before we enter the next function, we migrate to a different CPU.
+  CPU = 2;
+  ASSERT_TRUE(C->functionEnter(5, TSC++, CPU));
+  ASSERT_TRUE(C->functionTailExit(5, TSC++, CPU));
+  ASSERT_TRUE(C->functionEnter(6, TSC++, CPU));
+
+  // Then we exit them one at a time, in reverse order of entry.
+  ASSERT_TRUE(C->functionExit(6, TSC++, CPU));
+  ASSERT_TRUE(C->functionExit(2, TSC++, CPU));
+  ASSERT_TRUE(C->functionExit(1, TSC++, CPU));
+
+  ASSERT_TRUE(C->flush());
+  ASSERT_EQ(BQ->finalize(), BufferQueue::ErrorCode::Ok);
+
+  // Serialize buffers then test that we can find all the events that span the
+  // CPU migration.
+  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)),
+          AllOf(FuncId(2), RecordType(llvm::xray::RecordTypes::ENTER)),
+          AllOf(FuncId(2), RecordType(llvm::xray::RecordTypes::EXIT)),
+          AllOf(FuncId(1), RecordType(llvm::xray::RecordTypes::EXIT)))));
+}
+
 class BufferManagementTest : public ::testing::Test {
 protected:
   BufferQueue::Buffer B{};