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
16using namespace llvm;
17
18namespace {
19
20struct 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
53class MockAllocator : public AllocatorBase<MockAllocator> {
54public:
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
66private:
67 char Reserved[16];
68 AllocCondition C;
69};
70
71} // namespace
72
73#if (LLVM_ENABLE_THREADS)
74TEST(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
105TEST(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
124TEST(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

source code of llvm/unittests/Support/ThreadSafeAllocatorTest.cpp