llvm/Support/ThreadSafeAllocator.h

0001 //===- ThreadSafeAllocator.h ------------------------------------*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008
0009 #ifndef LLVM_SUPPORT_THREADSAFEALLOCATOR_H
0010 #define LLVM_SUPPORT_THREADSAFEALLOCATOR_H
0011
0012 #include "llvm/ADT/STLExtras.h"
0013 #include "llvm/Support/Allocator.h"
0014 #include <atomic>
0015
0016 namespace llvm {
0017
0018 /// Thread-safe allocator adaptor. Uses a spin lock on the assumption that
0019 /// contention here is extremely rare.
0020 ///
0021 /// TODO: Using a spin lock on every allocation can be quite expensive when
0022 /// contention is high. Since this is mainly used for BumpPtrAllocator and
0023 /// SpecificBumpPtrAllocator, it'd be better to have a specific thread-safe
0024 /// BumpPtrAllocator implementation that only use a fair lock when allocating a
0025 /// new slab but otherwise using atomic and be lock-free.
0026 template <class AllocatorType> class ThreadSafeAllocator {
0027   struct LockGuard {
0028     LockGuard(std::atomic_flag &Flag) : Flag(Flag) {
0029       if (LLVM_UNLIKELY(Flag.test_and_set(std::memory_order_acquire)))
0030         while (Flag.test_and_set(std::memory_order_acquire)) {
0031         }
0032     }
0033     ~LockGuard() { Flag.clear(std::memory_order_release); }
0034     std::atomic_flag &Flag;
0035   };
0036
0037 public:
0038   auto Allocate(size_t N) {
0039     return applyLocked([N](AllocatorType &Alloc) { return Alloc.Allocate(N); });
0040   }
0041
0042   auto Allocate(size_t Size, size_t Align) {
0043     return applyLocked([Size, Align](AllocatorType &Alloc) {
0044       return Alloc.Allocate(Size, Align);
0045     });
0046   }
0047
0048   template <typename FnT,
0049             typename T = typename llvm::function_traits<FnT>::result_t>
0050   T applyLocked(FnT Fn) {
0051     LockGuard Lock(Flag);
0052     return Fn(Alloc);
0053   }
0054
0055 private:
0056   AllocatorType Alloc;
0057   std::atomic_flag Flag = ATOMIC_FLAG_INIT;
0058 };
0059
0060 } // namespace llvm
0061
0062 #endif // LLVM_SUPPORT_THREADSAFEALLOCATOR_H