Index: lib/xray/tests/unit/segmented_array_test.cc =================================================================== --- lib/xray/tests/unit/segmented_array_test.cc +++ lib/xray/tests/unit/segmented_array_test.cc @@ -1,9 +1,13 @@ +#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; @@ -12,6 +16,10 @@ 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::AllocatorType; AllocatorType A(1 << 4); @@ -161,6 +169,23 @@ EXPECT_EQ(Data.size(), SegmentX2); } +TEST(SegmentedArrayTest, HandleExhaustedAllocator) { + using AllocatorType = typename Array::AllocatorType; + constexpr auto Segment = Array::SegmentSize; + constexpr auto MaxElements = Array::ElementsPerSegment; + AllocatorType A(Segment); + Array Data(A); + for (auto i = MaxElements; i > 0u; --i) + EXPECT_NE(Data.AppendEmplace(static_cast(i), static_cast(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; Index: lib/xray/tests/unit/test_helpers.h =================================================================== --- lib/xray/tests/unit/test_helpers.h +++ lib/xray/tests/unit/test_helpers.h @@ -14,8 +14,9 @@ #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/... @@ -54,6 +55,19 @@ std::string serialize(BufferQueue &Buffers, int32_t Version); +template void PrintTo(const Array &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_ Index: lib/xray/xray_segmented_array.h =================================================================== --- lib/xray/xray_segmented_array.h +++ lib/xray/xray_segmented_array.h @@ -78,6 +78,8 @@ static SegmentBase SentinelSegment; + using size_type = size_t; + private: AllocatorType *Alloc; SegmentBase *Head = &SentinelSegment; @@ -334,9 +336,8 @@ 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 @@ 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;