[XRay] Handle allocator exhaustion in segmented array

[dberris]

Summary

This change allows us to handle allocator exhaustion properly in the
segmented array implementation. Before this change, we relied on the
caller of the trim function to provide a valid number of elements to
trim. This change allows us to do the right thing in case the elements
to trim is greater than the size of the container.

[mboerger]

LGTM

lib/xray/tests/unit/segmented_array_test.cc

+#include "test_helpers.h"
 #include "xray_segmented_array.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace __xray {
 namespace {
 
+using ::testing::SizeIs;
+
 struct TestData {
   s64 First;
   s64 Second;
 
   // Need a constructor for emplace operations.
   TestData(s64 F, s64 S) : First(F), Second(S) {}
 };
 
+void PrintTo(const TestData &D, std::ostream *OS) {
+  *OS << "{ " << D.First << ", " << D.Second << " }";
+}
+
 TEST(SegmentedArrayTest, ConstructWithAllocators) {
   using AllocatorType = typename Array<TestData>::AllocatorType;
   AllocatorType A(1 << 4);
   Array<TestData> Data(A);
   (void)Data;
 }
 
 TEST(SegmentedArrayTest, ConstructAndPopulate) {
TEST(SegmentedArrayTest, IteratorTrimBehaviour) {
 
   // Then we ensure that adding elements back works just fine.
   for (auto i = SegmentX2; i > 0u; --i) {
     Data.AppendEmplace(static_cast<s64>(i), static_cast<s64>(i));
   }
   EXPECT_EQ(Data.size(), SegmentX2);
 }
 
+TEST(SegmentedArrayTest, HandleExhaustedAllocator) {
+  using AllocatorType = typename Array<TestData>::AllocatorType;
+  constexpr auto Segment = Array<TestData>::SegmentSize;
+  constexpr auto MaxElements = Array<TestData>::ElementsPerSegment;
+  AllocatorType A(Segment);
+  Array<TestData> Data(A);
+  for (auto i = MaxElements; i > 0u; --i)
+    EXPECT_NE(Data.AppendEmplace(static_cast<s64>(i), static_cast<s64>(i)),
+              nullptr);
+  EXPECT_EQ(Data.AppendEmplace(0, 0), nullptr);
+  EXPECT_THAT(Data, SizeIs(MaxElements));
+
+  // Trimming more elements than there are in the container should be fine.
+  Data.trim(MaxElements + 1);
+  EXPECT_THAT(Data, SizeIs(0u));
+}
+
 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;

lib/xray/tests/unit/test_helpers.h

 //===-- test_helpers.h ----------------------------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is a part of XRay, a function call tracing system.
 //
 //===----------------------------------------------------------------------===//
 #ifndef COMPILER_RT_LIB_XRAY_TESTS_TEST_HELPERS_H_
 #define COMPILER_RT_LIB_XRAY_TESTS_TEST_HELPERS_H_
 
 #include "xray_buffer_queue.h"
-#include "llvm/XRay/XRayRecord.h"
+#include "xray_segmented_array.h"
 #include "llvm/XRay/Trace.h"
+#include "llvm/XRay/XRayRecord.h"
 #include "gmock/gmock.h"
 
 // TODO: Move these to llvm/include/Testing/XRay/...
 namespace llvm {
 namespace xray {
 
 std::string RecordTypeAsString(RecordTypes T);
 void PrintTo(RecordTypes T, std::ostream *OS);

 } // namespace testing
 } // namespace xray
 } // namespace llvm
 
 namespace __xray {
 
 std::string serialize(BufferQueue &Buffers, int32_t Version);
 
+template <class T> void PrintTo(const Array<T> &A, std::ostream *OS) {
+  *OS << "[";
+  bool first = true;
+  for (const auto &E : A) {
+    if (!first) {
+      *OS << ", ";
+    }
+    PrintTo(E, OS);
+    first = false;
+  }
+  *OS << "]";
+}
+
 } // namespace __xray
 
 #endif // COMPILER_RT_LIB_XRAY_TESTS_TEST_HELPERS_H_

lib/xray/xray_segmented_array.h

public:
   static constexpr size_t ElementsPerSegment =
       (SegmentSize - sizeof(Segment)) / next_pow2(sizeof(T));
 
   static_assert(ElementsPerSegment > 0,
                 "Must have at least 1 element per segment.");
 
   static SegmentBase SentinelSegment;
 
+  using size_type = size_t;
+
 private:
   AllocatorType *Alloc;
   SegmentBase *Head = &SentinelSegment;
   SegmentBase *Tail = &SentinelSegment;
   size_t Size = 0;
 
   // Here we keep track of segments in the freelist, to allow us to re-use
   // segments when elements are trimmed off the end.
public:
   }
 
   /// Remove N Elements from the end. This leaves the blocks behind, and not
   /// require allocation of new blocks for new elements added after trimming.
   void trim(size_t Elements) XRAY_NEVER_INSTRUMENT {
     if (Elements == 0)
       return;
 
-    DCHECK_LE(Elements, Size);
-    DCHECK_GT(Size, 0);
     auto OldSize = Size;
+    Elements = Elements >= Size ? Size : Elements;
     Size -= Elements;
 
     DCHECK_NE(Head, &SentinelSegment);
     DCHECK_NE(Tail, &SentinelSegment);
 
     for (auto SegmentsToTrim = (nearest_boundary(OldSize, ElementsPerSegment) -
                                 nearest_boundary(Size, ElementsPerSegment)) /
                                ElementsPerSegment;
          SegmentsToTrim > 0; --SegmentsToTrim) {
-      DCHECK_NE(Head, &SentinelSegment);
-      DCHECK_NE(Tail, &SentinelSegment);
+
+      // We want to short-circuit if the trace is already empty.
+      if (Head == &SentinelSegment && Head == Tail)
+        return;
+
       // Put the tail into the Freelist.
       auto *FreeSegment = Tail;
       Tail = Tail->Prev;
       if (Tail == &SentinelSegment)
         Head = Tail;
       else
         Tail->Next = &SentinelSegment;