1+ #ifndef DICE_TEMPLATELIBRARY_LIMITALLOCATOR_HPP
2+ #define DICE_TEMPLATELIBRARY_LIMITALLOCATOR_HPP
3+
4+ #include < atomic>
5+ #include < cstddef>
6+ #include < memory>
7+ #include < new>
8+ #include < type_traits>
9+ #include < utility>
10+
11+ namespace dice ::template_library {
12+
13+ /* *
14+ * The synchronization policy of a limit_allocator
15+ */
16+ enum struct limit_allocator_syncness : bool {
17+ sync, // /< thread-safe (synchronized)
18+ unsync, // /< not thread-safe (unsynchronized)
19+ };
20+
21+ namespace detail_limit_allocator {
22+ template <limit_allocator_syncness syncness>
23+ struct limit_allocator_control_block ;
24+
25+ template <>
26+ struct limit_allocator_control_block <limit_allocator_syncness::sync> {
27+ std::atomic<size_t > bytes_left = 0 ;
28+
29+ void allocate (size_t n_bytes) {
30+ auto old = bytes_left.load (std::memory_order_relaxed);
31+
32+ do {
33+ if (old < n_bytes) [[unlikely]] {
34+ throw std::bad_alloc{};
35+ }
36+ } while (!bytes_left.compare_exchange_weak (old, old - n_bytes, std::memory_order_relaxed, std::memory_order_relaxed));
37+ }
38+
39+ void deallocate (size_t n_bytes) noexcept {
40+ bytes_left.fetch_add (n_bytes, std::memory_order_relaxed);
41+ }
42+ };
43+
44+ template <>
45+ struct limit_allocator_control_block <limit_allocator_syncness::unsync> {
46+ size_t bytes_left = 0 ;
47+
48+ void allocate (size_t n_bytes) {
49+ if (bytes_left < n_bytes) [[unlikely]] {
50+ throw std::bad_alloc{};
51+ }
52+ bytes_left -= n_bytes;
53+ }
54+
55+ void deallocate (size_t n_bytes) noexcept {
56+ bytes_left += n_bytes;
57+ }
58+ };
59+ }// namespace detail_limit_allocator
60+
61+ /* *
62+ * Allocator wrapper that limits the amount of memory its underlying allocator
63+ * is allowed to allocate.
64+ *
65+ * @tparam T value type of the allocator (the thing that it allocates)
66+ * @tparam Allocator the underlying allocator
67+ * @tparam syncness determines the synchronization of the limit
68+ */
69+ template <typename T, template <typename > typename Allocator = std::allocator, limit_allocator_syncness syncness = limit_allocator_syncness::sync>
70+ struct limit_allocator {
71+ using control_block_type = detail_limit_allocator::limit_allocator_control_block<syncness>;
72+ using value_type = T;
73+ using upstream_allocator_type = Allocator<T>;
74+ using pointer = typename std::allocator_traits<upstream_allocator_type>::pointer;
75+ using const_pointer = typename std::allocator_traits<upstream_allocator_type>::const_pointer;
76+ using void_pointer = typename std::allocator_traits<upstream_allocator_type>::void_pointer;
77+ using const_void_pointer = typename std::allocator_traits<upstream_allocator_type>::const_void_pointer;
78+ using size_type = typename std::allocator_traits<upstream_allocator_type>::size_type;
79+ using difference_type = typename std::allocator_traits<upstream_allocator_type>::difference_type;
80+
81+ using propagate_on_container_copy_assignment = typename std::allocator_traits<upstream_allocator_type>::propagate_on_container_copy_assignment;
82+ using propagate_on_container_move_assignment = typename std::allocator_traits<upstream_allocator_type>::propagate_on_container_move_assignment;
83+ using propagate_on_container_swap = typename std::allocator_traits<upstream_allocator_type>::propagate_on_container_swap;
84+ using is_always_equal = std::false_type;
85+
86+ template <typename U>
87+ struct rebind {
88+ using other = limit_allocator<U, Allocator, syncness>;
89+ };
90+
91+ private:
92+ template <typename , template <typename > typename , limit_allocator_syncness>
93+ friend struct limit_allocator ;
94+
95+ std::shared_ptr<control_block_type> control_block_;
96+ [[no_unique_address]] upstream_allocator_type inner_;
97+
98+ constexpr limit_allocator (std::shared_ptr<control_block_type> const &control_block, upstream_allocator_type const &alloc)
99+ requires(std::is_default_constructible_v<upstream_allocator_type>)
100+ : control_block_{control_block},
101+ inner_{alloc} {
102+ }
103+
104+ public:
105+ explicit constexpr limit_allocator (size_t bytes_limit)
106+ requires(std::is_default_constructible_v<upstream_allocator_type>)
107+ : control_block_{std::make_shared<control_block_type>(bytes_limit)},
108+ inner_{} {
109+ }
110+
111+ constexpr limit_allocator (limit_allocator const &other) noexcept (std::is_nothrow_move_constructible_v<upstream_allocator_type>) = default;
112+ constexpr limit_allocator (limit_allocator &&other) noexcept (std::is_nothrow_copy_constructible_v<upstream_allocator_type>) = default;
113+ constexpr limit_allocator &operator =(limit_allocator const &other) noexcept (std::is_nothrow_copy_assignable_v<upstream_allocator_type>) = default ;
114+ constexpr limit_allocator &operator =(limit_allocator &&other) noexcept (std::is_nothrow_move_assignable_v<upstream_allocator_type>) = default ;
115+ constexpr ~limit_allocator () = default ;
116+
117+ template <typename U>
118+ constexpr limit_allocator (limit_allocator<U, Allocator> const &other) noexcept (std::is_nothrow_constructible_v<upstream_allocator_type, typename limit_allocator<U, Allocator>::upstream_allocator_type const &>)
119+ : control_block_{other.control_block_ },
120+ inner_{other.inner_ } {
121+ }
122+
123+ constexpr limit_allocator (size_t bytes_limit, upstream_allocator_type const &upstream)
124+ : control_block_{std::make_shared<control_block_type>(bytes_limit)},
125+ inner_{upstream} {
126+ }
127+
128+ constexpr limit_allocator (size_t bytes_limit, upstream_allocator_type &&upstream)
129+ : control_block_{std::make_shared<control_block_type>(bytes_limit)},
130+ inner_{std::move (upstream)} {
131+ }
132+
133+ template <typename ... Args>
134+ explicit constexpr limit_allocator (size_t bytes_limit, std::in_place_t , Args &&...args)
135+ : control_block_{std::make_shared<control_block_type>(bytes_limit)},
136+ inner_{std::forward<Args>(args)...} {
137+ }
138+
139+ constexpr pointer allocate (size_t n) {
140+ control_block_->allocate (n * sizeof (T));
141+
142+ try {
143+ return std::allocator_traits<upstream_allocator_type>::allocate (inner_, n);
144+ } catch (...) {
145+ control_block_->deallocate (n * sizeof (T));
146+ throw ;
147+ }
148+ }
149+
150+ constexpr void deallocate (pointer ptr, size_t n) {
151+ std::allocator_traits<upstream_allocator_type>::deallocate (inner_, ptr, n);
152+ control_block_->deallocate (n * sizeof (T));
153+ }
154+
155+ constexpr limit_allocator select_on_container_copy_construction () const {
156+ return limit_allocator{control_block_, std::allocator_traits<upstream_allocator_type>::select_on_container_copy_construction (inner_)};
157+ }
158+
159+ [[nodiscard]] upstream_allocator_type const &upstream_allocator () const noexcept {
160+ return inner_;
161+ }
162+
163+ friend constexpr void swap (limit_allocator &a, limit_allocator &b) noexcept (std::is_nothrow_swappable_v<upstream_allocator_type>)
164+ requires(std::is_swappable_v<upstream_allocator_type>)
165+ {
166+ using std::swap;
167+ swap (a.control_block_ , b.control_block_ );
168+ swap (a.inner_ , b.inner_ );
169+ }
170+
171+ bool operator ==(limit_allocator const &other) const noexcept = default ;
172+ bool operator !=(limit_allocator const &other) const noexcept = default ;
173+ };
174+ }// namespace dice::template_library
175+
176+ #endif // DICE_TEMPLATELIBRARY_LIMITALLOCATOR_HPP
0 commit comments