1 files changed, 223 insertions, 0 deletions
diff --git a/contrib/llvm-project/compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp b/contrib/llvm-project/compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp
new file mode 100644
index 000000000000..d9f08b1e7efe
--- /dev/null
+++ b/contrib/llvm-project/compiler-rt/lib/ctx_profile/tests/CtxInstrProfilingTest.cpp
@@ -0,0 +1,223 @@
+#include "../CtxInstrProfiling.h"
+#include "gtest/gtest.h"
+#include <thread>
+
+using namespace __ctx_profile;
+
+class ContextTest : public ::testing::Test {
+  void SetUp() override { memset(&Root, 0, sizeof(ContextRoot)); }
+  void TearDown() override { __llvm_ctx_profile_free(); }
+
+public:
+  ContextRoot Root;
+};
+
+TEST(ArenaTest, ZeroInit) {
+  char Buffer[1024];
+  memset(Buffer, 1, 1024);
+  Arena *A = new (Buffer) Arena(10);
+  for (auto I = 0U; I < A->size(); ++I)
+    EXPECT_EQ(A->pos()[I], static_cast<char>(0));
+  EXPECT_EQ(A->size(), 10U);
+}
+
+TEST(ArenaTest, Basic) {
+  Arena *A = Arena::allocateNewArena(1024);
+  EXPECT_EQ(A->size(), 1024U);
+  EXPECT_EQ(A->next(), nullptr);
+
+  auto *M1 = A->tryBumpAllocate(1020);
+  EXPECT_NE(M1, nullptr);
+  auto *M2 = A->tryBumpAllocate(4);
+  EXPECT_NE(M2, nullptr);
+  EXPECT_EQ(M1 + 1020, M2);
+  EXPECT_EQ(A->tryBumpAllocate(1), nullptr);
+  Arena *A2 = Arena::allocateNewArena(2024, A);
+  EXPECT_EQ(A->next(), A2);
+  EXPECT_EQ(A2->next(), nullptr);
+  Arena::freeArenaList(A);
+  EXPECT_EQ(A, nullptr);
+}
+
+TEST_F(ContextTest, Basic) {
+  auto *Ctx = __llvm_ctx_profile_start_context(&Root, 1, 10, 4);
+  ASSERT_NE(Ctx, nullptr);
+  EXPECT_NE(Root.CurrentMem, nullptr);
+  EXPECT_EQ(Root.FirstMemBlock, Root.CurrentMem);
+  EXPECT_EQ(Ctx->size(), sizeof(ContextNode) + 10 * sizeof(uint64_t) +
+                             4 * sizeof(ContextNode *));
+  EXPECT_EQ(Ctx->counters_size(), 10U);
+  EXPECT_EQ(Ctx->callsites_size(), 4U);
+  EXPECT_EQ(__llvm_ctx_profile_current_context_root, &Root);
+  Root.Taken.CheckLocked();
+  EXPECT_FALSE(Root.Taken.TryLock());
+  __llvm_ctx_profile_release_context(&Root);
+  EXPECT_EQ(__llvm_ctx_profile_current_context_root, nullptr);
+  EXPECT_TRUE(Root.Taken.TryLock());
+  Root.Taken.Unlock();
+}
+
+TEST_F(ContextTest, Callsite) {
+  auto *Ctx = __llvm_ctx_profile_start_context(&Root, 1, 10, 4);
+  int FakeCalleeAddress = 0;
+  const bool IsScratch = isScratch(Ctx);
+  EXPECT_FALSE(IsScratch);
+  // This is the sequence the caller performs - it's the lowering of the
+  // instrumentation of the callsite "2". "2" is arbitrary here.
+  __llvm_ctx_profile_expected_callee[0] = &FakeCalleeAddress;
+  __llvm_ctx_profile_callsite[0] = &Ctx->subContexts()[2];
+  // This is what the callee does
+  auto *Subctx = __llvm_ctx_profile_get_context(&FakeCalleeAddress, 2, 3, 1);
+  // We expect the subcontext to be appropriately placed and dimensioned
+  EXPECT_EQ(Ctx->subContexts()[2], Subctx);
+  EXPECT_EQ(Subctx->counters_size(), 3U);
+  EXPECT_EQ(Subctx->callsites_size(), 1U);
+  // We reset these in _get_context.
+  EXPECT_EQ(__llvm_ctx_profile_expected_callee[0], nullptr);
+  EXPECT_EQ(__llvm_ctx_profile_callsite[0], nullptr);
+
+  EXPECT_EQ(Subctx->size(), sizeof(ContextNode) + 3 * sizeof(uint64_t) +
+                                1 * sizeof(ContextNode *));
+  __llvm_ctx_profile_release_context(&Root);
+}
+
+TEST_F(ContextTest, ScratchNoCollection) {
+  EXPECT_EQ(__llvm_ctx_profile_current_context_root, nullptr);
+  int FakeCalleeAddress = 0;
+  // this would be the very first function executing this. the TLS is empty,
+  // too.
+  auto *Ctx = __llvm_ctx_profile_get_context(&FakeCalleeAddress, 2, 3, 1);
+  // We never entered a context (_start_context was never called) - so the
+  // returned context must be scratch.
+  EXPECT_TRUE(isScratch(Ctx));
+}
+
+TEST_F(ContextTest, ScratchDuringCollection) {
+  auto *Ctx = __llvm_ctx_profile_start_context(&Root, 1, 10, 4);
+  int FakeCalleeAddress = 0;
+  int OtherFakeCalleeAddress = 0;
+  __llvm_ctx_profile_expected_callee[0] = &FakeCalleeAddress;
+  __llvm_ctx_profile_callsite[0] = &Ctx->subContexts()[2];
+  auto *Subctx =
+      __llvm_ctx_profile_get_context(&OtherFakeCalleeAddress, 2, 3, 1);
+  // We expected a different callee - so return scratch. It mimics what happens
+  // in the case of a signal handler - in this case, OtherFakeCalleeAddress is
+  // the signal handler.
+  EXPECT_TRUE(isScratch(Subctx));
+  EXPECT_EQ(__llvm_ctx_profile_expected_callee[0], nullptr);
+  EXPECT_EQ(__llvm_ctx_profile_callsite[0], nullptr);
+
+  int ThirdFakeCalleeAddress = 0;
+  __llvm_ctx_profile_expected_callee[1] = &ThirdFakeCalleeAddress;
+  __llvm_ctx_profile_callsite[1] = &Subctx->subContexts()[0];
+
+  auto *Subctx2 =
+      __llvm_ctx_profile_get_context(&ThirdFakeCalleeAddress, 3, 0, 0);
+  // We again expect scratch because the '0' position is where the runtime
+  // looks, so it doesn't matter the '1' position is populated correctly.
+  EXPECT_TRUE(isScratch(Subctx2));
+
+  __llvm_ctx_profile_expected_callee[0] = &ThirdFakeCalleeAddress;
+  __llvm_ctx_profile_callsite[0] = &Subctx->subContexts()[0];
+  auto *Subctx3 =
+      __llvm_ctx_profile_get_context(&ThirdFakeCalleeAddress, 3, 0, 0);
+  // We expect scratch here, too, because the value placed in
+  // __llvm_ctx_profile_callsite is scratch
+  EXPECT_TRUE(isScratch(Subctx3));
+
+  __llvm_ctx_profile_release_context(&Root);
+}
+
+TEST_F(ContextTest, NeedMoreMemory) {
+  auto *Ctx = __llvm_ctx_profile_start_context(&Root, 1, 10, 4);
+  int FakeCalleeAddress = 0;
+  const bool IsScratch = isScratch(Ctx);
+  EXPECT_FALSE(IsScratch);
+  const auto *CurrentMem = Root.CurrentMem;
+  __llvm_ctx_profile_expected_callee[0] = &FakeCalleeAddress;
+  __llvm_ctx_profile_callsite[0] = &Ctx->subContexts()[2];
+  // Allocate a massive subcontext to force new arena allocation
+  auto *Subctx =
+      __llvm_ctx_profile_get_context(&FakeCalleeAddress, 3, 1 << 20, 1);
+  EXPECT_EQ(Ctx->subContexts()[2], Subctx);
+  EXPECT_NE(CurrentMem, Root.CurrentMem);
+  EXPECT_NE(Root.CurrentMem, nullptr);
+}
+
+TEST_F(ContextTest, ConcurrentRootCollection) {
+  std::atomic<int> NonScratch = 0;
+  std::atomic<int> Executions = 0;
+
+  __sanitizer::Semaphore GotCtx;
+
+  auto Entrypoint = [&]() {
+    ++Executions;
+    auto *Ctx = __llvm_ctx_profile_start_context(&Root, 1, 10, 4);
+    GotCtx.Post();
+    const bool IS = isScratch(Ctx);
+    NonScratch += (!IS);
+    if (!IS) {
+      GotCtx.Wait();
+      GotCtx.Wait();
+    }
+    __llvm_ctx_profile_release_context(&Root);
+  };
+  std::thread T1(Entrypoint);
+  std::thread T2(Entrypoint);
+  T1.join();
+  T2.join();
+  EXPECT_EQ(NonScratch, 1);
+  EXPECT_EQ(Executions, 2);
+}
+
+TEST_F(ContextTest, Dump) {
+  auto *Ctx = __llvm_ctx_profile_start_context(&Root, 1, 10, 4);
+  int FakeCalleeAddress = 0;
+  __llvm_ctx_profile_expected_callee[0] = &FakeCalleeAddress;
+  __llvm_ctx_profile_callsite[0] = &Ctx->subContexts()[2];
+  auto *Subctx = __llvm_ctx_profile_get_context(&FakeCalleeAddress, 2, 3, 1);
+  (void)Subctx;
+  __llvm_ctx_profile_release_context(&Root);
+
+  struct Writer {
+    ContextRoot *const Root;
+    const size_t Entries;
+    bool State = false;
+    Writer(ContextRoot *Root, size_t Entries) : Root(Root), Entries(Entries) {}
+
+    bool write(const ContextNode &Node) {
+      EXPECT_FALSE(Root->Taken.TryLock());
+      EXPECT_EQ(Node.guid(), 1U);
+      EXPECT_EQ(Node.counters()[0], Entries);
+      EXPECT_EQ(Node.counters_size(), 10U);
+      EXPECT_EQ(Node.callsites_size(), 4U);
+      EXPECT_EQ(Node.subContexts()[0], nullptr);
+      EXPECT_EQ(Node.subContexts()[1], nullptr);
+      EXPECT_NE(Node.subContexts()[2], nullptr);
+      EXPECT_EQ(Node.subContexts()[3], nullptr);
+      const auto &SN = *Node.subContexts()[2];
+      EXPECT_EQ(SN.guid(), 2U);
+      EXPECT_EQ(SN.counters()[0], Entries);
+      EXPECT_EQ(SN.counters_size(), 3U);
+      EXPECT_EQ(SN.callsites_size(), 1U);
+      EXPECT_EQ(SN.subContexts()[0], nullptr);
+      State = true;
+      return true;
+    }
+  };
+  Writer W(&Root, 1);
+  EXPECT_FALSE(W.State);
+  __llvm_ctx_profile_fetch(&W, [](void *W, const ContextNode &Node) -> bool {
+    return reinterpret_cast<Writer *>(W)->write(Node);
+  });
+  EXPECT_TRUE(W.State);
+
+  // this resets all counters but not the internal structure.
+  __llvm_ctx_profile_start_collection();
+  Writer W2(&Root, 0);
+  EXPECT_FALSE(W2.State);
+  __llvm_ctx_profile_fetch(&W2, [](void *W, const ContextNode &Node) -> bool {
+    return reinterpret_cast<Writer *>(W)->write(Node);
+  });
+  EXPECT_TRUE(W2.State);
+}