[XRay][compiler-rt] xray::Array Freelist and Iterator Updates

Summary:
We found a bug while working on a benchmark for the profiling mode which
manifests as a segmentation fault in the profiling handler's
implementation. This change adds unit tests which replicate the
issues in isolation.

We've tracked this down as a bug in the implementation of the Freelist
in the `xray::Array` type. This happens when we trim the array by a
number of elements, where we've been incorrectly assigning pointers for
the links in the freelist of chunk nodes. We've taken the chance to add
more debug-only assertions to the code path and allow us to verify these
assumptions in debug builds.

In the process, we also took the opportunity to use iterators to
implement both `front()` and `back()` which exposes a bug in the
iterator decrement operation.  In particular, when we decrement past a
chunk size boundary, we end up moving too far back and reaching the
`SentinelChunk` prematurely.

This change unblocks us to allow for contributing the non-crashing
version of the benchmarks in the test-suite as well.

Reviewers: kpw

Subscribers: mgorny, llvm-commits

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

llvm-svn: 336644

3 files changed, 108 insertions, 22 deletions

diff --git a/compiler-rt/lib/xray/tests/CMakeLists.txt b/compiler-rt/lib/xray/tests/CMakeLists.txt
index eb17b095df6..ea5f2da4e10 100644
--- a/compiler-rt/lib/xray/tests/CMakeLists.txt
+++ b/compiler-rt/lib/xray/tests/CMakeLists.txt
@@ -68,7 +68,11 @@ function(get_xray_lib_for_arch arch lib)
 endfunction()
 
 set(XRAY_TEST_ARCH ${XRAY_SUPPORTED_ARCH})
-set(XRAY_UNITTEST_LINK_FLAGS ${CMAKE_THREAD_LIBS_INIT})
+set(XRAY_UNITTEST_LINK_FLAGS
+  ${CMAKE_THREAD_LIBS_INIT}
+  -fxray-instrument
+  -lstdc++  # TODO: Detect the correct standard library used!
+  )
 if (NOT APPLE)
   append_list_if(COMPILER_RT_HAS_LIBM -lm XRAY_UNITTEST_LINK_FLAGS)
   append_list_if(COMPILER_RT_HAS_LIBRT -lrt XRAY_UNITTEST_LINK_FLAGS)
@@ -94,7 +98,7 @@ macro(add_xray_unittest testname)
         RUNTIME "${XRAY_RUNTIME_LIBS}"
         DEPS gtest xray
         CFLAGS ${XRAY_UNITTEST_CFLAGS}
-        LINK_FLAGS ${TARGET_LINK_FLAGS} ${XRAY_UNITTEST_LINK_FLAGS} -lstdc++)
+        LINK_FLAGS ${TARGET_LINK_FLAGS} ${XRAY_UNITTEST_LINK_FLAGS})
       set_target_properties(XRayUnitTests
         PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
     endforeach()
diff --git a/compiler-rt/lib/xray/tests/unit/function_call_trie_test.cc b/compiler-rt/lib/xray/tests/unit/function_call_trie_test.cc
index d4a79d9028c..6ad51aa148c 100644
--- a/compiler-rt/lib/xray/tests/unit/function_call_trie_test.cc
+++ b/compiler-rt/lib/xray/tests/unit/function_call_trie_test.cc
@@ -18,12 +18,6 @@ namespace __xray {
 
 namespace {
 
-TEST(FunctionCallTrieTest, Construction) {
-  // We want to make sure that we can create one of these without the set of
-  // allocators we need. This will by default use the global allocators.
-  FunctionCallTrie Trie;
-}
-
 TEST(FunctionCallTrieTest, ConstructWithTLSAllocators) {
   // FIXME: Support passing in configuration for allocators in the allocator
   // constructors.
@@ -61,6 +55,17 @@ TEST(FunctionCallTrieTest, MissingFunctionEntry) {
   ASSERT_TRUE(R.empty());
 }
 
+TEST(FunctionCallTrieTest, NoMatchingEntersForExit) {
+  auto A = FunctionCallTrie::InitAllocators();
+  FunctionCallTrie Trie(A);
+  Trie.enterFunction(2, 1);
+  Trie.enterFunction(3, 3);
+  Trie.exitFunction(1, 5);
+  const auto &R = Trie.getRoots();
+
+  ASSERT_TRUE(R.empty());
+}
+
 TEST(FunctionCallTrieTest, MissingFunctionExit) {
   auto A = FunctionCallTrie::InitAllocators();
   FunctionCallTrie Trie(A);
@@ -153,6 +158,31 @@ TEST(FunctionCallTrieTest, MissingIntermediaryExit) {
   EXPECT_EQ(F1.CumulativeLocalTime, 100);
 }
 
+TEST(FunctionCallTrieTest, DeepCallStack) {
+  // Simulate a relatively deep call stack (32 levels) and ensure that we can
+  // properly pop all the way up the stack.
+  profilingFlags()->setDefaults();
+  auto A = FunctionCallTrie::InitAllocators();
+  FunctionCallTrie Trie(A);
+  for (int i = 0; i < 32; ++i)
+    Trie.enterFunction(i + 1, i);
+  Trie.exitFunction(1, 33);
+
+  // Here, validate that we have a 32-level deep function call path from the
+  // root (1) down to the leaf (33).
+  const auto &R = Trie.getRoots();
+  ASSERT_EQ(R.size(), 1u);
+  auto F = R[0];
+  for (int i = 0; i < 32; ++i) {
+    EXPECT_EQ(F->FId, i + 1);
+    EXPECT_EQ(F->CallCount, 1);
+    if (F->Callees.empty() && i != 31)
+      FAIL() << "Empty callees for FId " << F->FId;
+    if (i != 31)
+      F = F->Callees[0].NodePtr;
+  }
+}
+
 // TODO: Test that we can handle cross-CPU migrations, where TSCs are not
 // guaranteed to be synchronised.
 TEST(FunctionCallTrieTest, DeepCopy) {
diff --git a/compiler-rt/lib/xray/tests/unit/segmented_array_test.cc b/compiler-rt/lib/xray/tests/unit/segmented_array_test.cc
index dd07e7f137b..539162d7d2f 100644
--- a/compiler-rt/lib/xray/tests/unit/segmented_array_test.cc
+++ b/compiler-rt/lib/xray/tests/unit/segmented_array_test.cc
@@ -107,32 +107,84 @@ TEST(SegmentedArrayTest, IteratorRetreat) {
 }
 
 TEST(SegmentedArrayTest, IteratorTrimBehaviour) {
-  Array<TestData> data;
-  ASSERT_TRUE(data.empty());
-  auto I0Begin = data.begin(), I0End = data.end();
-  // Add enough elements in data to have more than one chunk.
-  constexpr auto ChunkX2 = Array<TestData>::ChunkSize * 2;
+  using AllocatorType = typename Array<TestData>::AllocatorType;
+  AllocatorType A(1 << 10, 0);
+  Array<TestData> Data(A);
+  ASSERT_TRUE(Data.empty());
+  auto I0Begin = Data.begin(), I0End = Data.end();
+  // Add enough elements in Data to have more than one chunk.
+  constexpr auto Chunk = Array<TestData>::ChunkSize;
+  constexpr auto ChunkX2 = Chunk * 2;
   for (auto i = ChunkX2; i > 0u; --i) {
-    data.Append({static_cast<s64>(i), static_cast<s64>(i)});
+    Data.AppendEmplace(static_cast<s64>(i), static_cast<s64>(i));
   }
-  ASSERT_EQ(data.size(), ChunkX2);
+  ASSERT_EQ(Data.size(), ChunkX2);
+  {
+    auto &Back = Data.back();
+    ASSERT_EQ(Back.First, 1);
+    ASSERT_EQ(Back.Second, 1);
+  }
+
   // Trim one chunk's elements worth.
-  data.trim(ChunkX2 / 2);
-  ASSERT_EQ(data.size(), ChunkX2 / 2);
+  Data.trim(Chunk);
+  ASSERT_EQ(Data.size(), Chunk);
+
+  // Check that we are still able to access 'back' properly.
+  {
+    auto &Back = Data.back();
+    ASSERT_EQ(Back.First, static_cast<s64>(Chunk + 1));
+    ASSERT_EQ(Back.Second, static_cast<s64>(Chunk + 1));
+  }
+
   // Then trim until it's empty.
-  data.trim(ChunkX2 / 2);
-  ASSERT_TRUE(data.empty());
+  Data.trim(Chunk);
+  ASSERT_TRUE(Data.empty());
 
   // Here our iterators should be the same.
-  auto I1Begin = data.begin(), I1End = data.end();
+  auto I1Begin = Data.begin(), I1End = Data.end();
   EXPECT_EQ(I0Begin, I1Begin);
   EXPECT_EQ(I0End, I1End);
 
   // Then we ensure that adding elements back works just fine.
   for (auto i = ChunkX2; i > 0u; --i) {
-    data.Append({static_cast<s64>(i), static_cast<s64>(i)});
+    Data.AppendEmplace(static_cast<s64>(i), static_cast<s64>(i));
+  }
+  EXPECT_EQ(Data.size(), ChunkX2);
+}
+
+struct ShadowStackEntry {
+  uint64_t EntryTSC = 0;
+  uint64_t *NodePtr = nullptr;
+  ShadowStackEntry(uint64_t T, uint64_t *N) : EntryTSC(T), NodePtr(N) {}
+};
+
+TEST(SegmentedArrayTest, SimulateStackBehaviour) {
+  using AllocatorType = typename Array<ShadowStackEntry>::AllocatorType;
+  AllocatorType A(1 << 10, 0);
+  Array<ShadowStackEntry> Data(A);
+  static uint64_t Dummy = 0;
+  constexpr uint64_t Max = 9;
+
+  for (uint64_t i = 0; i < Max; ++i) {
+    auto P = Data.Append({i, &Dummy});
+    ASSERT_NE(P, nullptr);
+    ASSERT_EQ(P->NodePtr, &Dummy);
+    auto &Back = Data.back();
+    ASSERT_EQ(Back.NodePtr, &Dummy);
+    ASSERT_EQ(Back.EntryTSC, i);
+  }
+
+  // Simulate a stack by checking the data from the end as we're trimming.
+  auto Counter = Max;
+  ASSERT_EQ(Data.size(), size_t(Max));
+  while (!Data.empty()) {
+    const auto &Top = Data.back();
+    uint64_t *TopNode = Top.NodePtr;
+    EXPECT_EQ(TopNode, &Dummy) << "Counter = " << Counter;
+    Data.trim(1);
+    --Counter;
+    ASSERT_EQ(Data.size(), size_t(Counter));
   }
-  EXPECT_EQ(data.size(), ChunkX2);
 }
 
 } // namespace