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 | |
26 | void startAllocation() { |
27 | { |
28 | std::lock_guard<std::mutex> Lock(BusyLock); |
29 | IsBusy = true; |
30 | } |
31 | Busy.notify_all(); |
32 | } |
33 | void waitAllocationStarted() { |
34 | std::unique_lock<std::mutex> LBusy(BusyLock); |
35 | Busy.wait(lock&: LBusy, p: [&]() { return IsBusy; }); |
36 | IsBusy = false; |
37 | } |
38 | void finishAllocation() { |
39 | { |
40 | std::lock_guard<std::mutex> Lock(EndLock); |
41 | IsEnd = true; |
42 | } |
43 | End.notify_all(); |
44 | } |
45 | void waitAllocationFinished() { |
46 | std::unique_lock<std::mutex> LEnd(EndLock); |
47 | // Wait for end state. |
48 | End.wait(lock&: LEnd, p: [&]() { return IsEnd; }); |
49 | IsEnd = false; |
50 | } |
51 | }; |
52 | |
53 | class MockAllocator : public AllocatorBase<MockAllocator> { |
54 | public: |
55 | MockAllocator() = default; |
56 | |
57 | void *Allocate(size_t Size, size_t Alignment) { |
58 | C.startAllocation(); |
59 | C.waitAllocationFinished(); |
60 | C.BytesAllocated += Size; |
61 | return Reserved; |
62 | } |
63 | |
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) |
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 | Fn: [&](MockAllocator &Alloc) { C = &Alloc.getAllocCondition(); }); |
80 | DefaultThreadPool Threads; |
81 | // First allocation of 1 byte. |
82 | Threads.async(F: [&Alloc]() { |
83 | char *P = (char *)Alloc.Allocate(Size: 1, Align: 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(F: [&Alloc]() { |
91 | char *P = (char *)Alloc.Allocate(Size: 2, Align: 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 | |
105 | TEST(ThreadSafeAllocatorTest, AllocWithAlign) { |
106 | ThreadSafeAllocator<BumpPtrAllocator> Alloc; |
107 | DefaultThreadPool Threads; |
108 | |
109 | for (unsigned Index = 1; Index < 100; ++Index) |
110 | Threads.async( |
111 | F: [&Alloc](unsigned I) { |
112 | int *P = (int *)Alloc.Allocate(Size: sizeof(int) * I, Align: alignof(int)); |
113 | P[I - 1] = I; |
114 | }, |
115 | ArgList&: Index); |
116 | |
117 | Threads.wait(); |
118 | |
119 | Alloc.applyLocked(Fn: [](BumpPtrAllocator &Alloc) { |
120 | EXPECT_EQ(4950U * sizeof(int), Alloc.getBytesAllocated()); |
121 | }); |
122 | } |
123 | |
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 | F: [&Alloc](unsigned I) { |
131 | int *P = Alloc.Allocate(N: I); |
132 | P[I - 1] = I; |
133 | }, |
134 | ArgList&: Index); |
135 | |
136 | Threads.wait(); |
137 | } |
138 | #endif |
139 | |