139 lines · cpp
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-exception6//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(LBusy, [&]() { 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(LEnd, [&]() { 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} // namespace72 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 [&](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 105TEST(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 124TEST(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#endif139