[XRay] Fix enter function tracing for record unwriting

[dberris]

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.

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

TEST_F(FunctionSequenceTest, RewindingIntermediaryTailExits) {
   // Serialize the buffers then test to see we find that all the calls have been
   // unwound because all of them are under the cycle counter threshold.
   std::string Serialized = serialize(*BQ, 3);
   llvm::DataExtractor DE(Serialized, true, 8);
   auto TraceOrErr = llvm::xray::loadTrace(DE);
   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{};
   std::unique_ptr<BufferQueue> BQ;
   std::unique_ptr<FDRLogWriter> W;
   std::unique_ptr<FDRController<>> C;
 
   static constexpr size_t kBuffers = 10;

compiler-rt/trunk/lib/xray/xray_fdr_controller.h

bool functionEnter(int32_t FuncId, uint64_t TSC,
     if (finalized() ||
         !prepareBuffer(sizeof(MetadataRecord) + sizeof(FunctionRecord)))
       return returnBuffer();
 
     auto PreambleStatus = recordPreamble(TSC, CPU);
     if (PreambleStatus == PreambleResult::InvalidBuffer)
       return returnBuffer();
 
-    UndoableFunctionEnters = (PreambleStatus == PreambleResult::WroteMetadata)
-                                 ? 1
-                                 : UndoableFunctionEnters + 1;
+    if (PreambleStatus == PreambleResult::WroteMetadata) {
+      UndoableFunctionEnters = 1;
+      UndoableTailExits = 0;
+    } else {
+      ++UndoableFunctionEnters;
+    }
+
     auto Delta = TSC - LatestTSC;
     LastFunctionEntryTSC = TSC;
     LatestTSC = TSC;
     return W.writeFunction(FDRLogWriter::FunctionRecordKind::Enter,
                            mask(FuncId), Delta);
   }
 
   bool functionTailExit(int32_t FuncId, uint64_t TSC,