1 files changed, 45 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 30276c61f56..84850c182ff 100644
--- a/compiler-rt/lib/xray/tests/unit/fdr_controller_test.cc
+++ b/compiler-rt/lib/xray/tests/unit/fdr_controller_test.cc
@@ -33,10 +33,12 @@ 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;
 
@@ -110,6 +112,31 @@ TEST_F(FunctionSequenceTest, ArgsAreHandledAndKept) {
           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 if we find the remaining 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);
 
@@ -211,6 +238,24 @@ TEST_F(BufferManagementTest, HandlesOverflow) {
   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;