llvm/unittests/Support/ThreadSafeAllocatorTest.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-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(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
  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       [&](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
 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
 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