[XRay] Handle allocator exhaustion in segmented array

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.

Reviewers: mboerger, eizan

Subscribers: llvm-commits

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

llvm-svn: 344880

1 files changed, 8 insertions, 4 deletions

diff --git a/compiler-rt/lib/xray/xray_segmented_array.h b/compiler-rt/lib/xray/xray_segmented_array.h
index c723c7de083..42f53be1e01 100644
--- a/compiler-rt/lib/xray/xray_segmented_array.h
+++ b/compiler-rt/lib/xray/xray_segmented_array.h
@@ -78,6 +78,8 @@ public:
 
   static SegmentBase SentinelSegment;
 
+  using size_type = size_t;
+
 private:
   AllocatorType *Alloc;
   SegmentBase *Head = &SentinelSegment;
@@ -334,9 +336,8 @@ public:
     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);
@@ -346,8 +347,11 @@ public:
                                 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;