1 //===- llvm/unittest/Support/ThreadSafeAllocatorTest.cpp ------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8
9 #include "llvm/Support/ThreadSafeAllocator.h"
10 #include "llvm/Config/llvm-config.h"
11 #include "llvm/Support/ThreadPool.h"
12 #include "gtest/gtest.h"
13 #include <atomic>
14 #include <thread>
15
16 using namespace llvm;
17
18 namespace {
19
20 struct AllocCondition {
21 std::mutex BusyLock, EndLock;
22 std::condition_variable Busy, End;
23 bool IsBusy = false, IsEnd = false;
24 std::atomic<unsigned> BytesAllocated = 0;
25
startAllocation__anon300027430111::AllocCondition26 void startAllocation() {
27 {
28 std::lock_guard<std::mutex> Lock(BusyLock);
29 IsBusy = true;
30 }
31 Busy.notify_all();
32 }
waitAllocationStarted__anon300027430111::AllocCondition33 void waitAllocationStarted() {
34 std::unique_lock<std::mutex> LBusy(BusyLock);
35 Busy.wait(LBusy, [&]() { return IsBusy; });
36 IsBusy = false;
37 }
finishAllocation__anon300027430111::AllocCondition38 void finishAllocation() {
39 {
40 std::lock_guard<std::mutex> Lock(EndLock);
41 IsEnd = true;
42 }
43 End.notify_all();
44 }
waitAllocationFinished__anon300027430111::AllocCondition45 void waitAllocationFinished() {
46 std::unique_lock<std::mutex> LEnd(EndLock);
47 // Wait for end state.
48 End.wait(LEnd, [&]() { return IsEnd; });
49 IsEnd = false;
50 }
51 };
52
53 class MockAllocator : public AllocatorBase<MockAllocator> {
54 public:
55 MockAllocator() = default;
56
Allocate(size_t Size,size_t Alignment)57 void *Allocate(size_t Size, size_t Alignment) {
58 C.startAllocation();
59 C.waitAllocationFinished();
60 C.BytesAllocated += Size;
61 return Reserved;
62 }
63
getAllocCondition()64 AllocCondition &getAllocCondition() { return C; }
65
66 private:
67 char Reserved[16];
68 AllocCondition C;
69 };
70
71 } // namespace
72
73 #if (LLVM_ENABLE_THREADS)
TEST(ThreadSafeAllocatorTest,AllocWait)74 TEST(ThreadSafeAllocatorTest, AllocWait) {
75 ThreadSafeAllocator<MockAllocator> Alloc;
76 AllocCondition *C;
77 // Get the allocation from the allocator first since this requires a lock.
78 Alloc.applyLocked(
79 [&](MockAllocator &Alloc) { C = &Alloc.getAllocCondition(); });
80 DefaultThreadPool Threads;
81 // First allocation of 1 byte.
82 Threads.async([&Alloc]() {
83 char *P = (char *)Alloc.Allocate(1, alignof(char));
84 P[0] = 0;
85 });
86 // No allocation yet.
87 EXPECT_EQ(C->BytesAllocated, 0u);
88 C->waitAllocationStarted(); // wait till 1st alloocation starts.
89 // Second allocation of 2 bytes.
90 Threads.async([&Alloc]() {
91 char *P = (char *)Alloc.Allocate(2, alignof(char));
92 P[1] = 0;
93 });
94 C->finishAllocation(); // finish 1st allocation.
95
96 C->waitAllocationStarted(); // wait till 2nd allocation starts.
97 // still 1 byte allocated since 2nd allocation is not finished yet.
98 EXPECT_EQ(C->BytesAllocated, 1u);
99 C->finishAllocation(); // finish 2nd allocation.
100
101 Threads.wait(); // all allocations done.
102 EXPECT_EQ(C->BytesAllocated, 3u);
103 }
104
TEST(ThreadSafeAllocatorTest,AllocWithAlign)105 TEST(ThreadSafeAllocatorTest, AllocWithAlign) {
106 ThreadSafeAllocator<BumpPtrAllocator> Alloc;
107 DefaultThreadPool Threads;
108
109 for (unsigned Index = 1; Index < 100; ++Index)
110 Threads.async(
111 [&Alloc](unsigned I) {
112 int *P = (int *)Alloc.Allocate(sizeof(int) * I, alignof(int));
113 P[I - 1] = I;
114 },
115 Index);
116
117 Threads.wait();
118
119 Alloc.applyLocked([](BumpPtrAllocator &Alloc) {
120 EXPECT_EQ(4950U * sizeof(int), Alloc.getBytesAllocated());
121 });
122 }
123
TEST(ThreadSafeAllocatorTest,SpecificBumpPtrAllocator)124 TEST(ThreadSafeAllocatorTest, SpecificBumpPtrAllocator) {
125 ThreadSafeAllocator<SpecificBumpPtrAllocator<int>> Alloc;
126 DefaultThreadPool Threads;
127
128 for (unsigned Index = 1; Index < 100; ++Index)
129 Threads.async(
130 [&Alloc](unsigned I) {
131 int *P = Alloc.Allocate(I);
132 P[I - 1] = I;
133 },
134 Index);
135
136 Threads.wait();
137 }
138 #endif
139