Feature: limit_allocator

README.md

@@ -9,6 +9,7 @@ It contains:
 - `integral_template_variant`: A wrapper type for `std::variant` guarantees to only contain variants of the form `T<ix>` where $\texttt{ix}\in [\texttt{first},\texttt{last}]$ (inclusive).
 - `for_{types,values,range}`: Compile time for loops for types, values or ranges
 - `polymorphic_allocator`: Like `std::pmr::polymorphic_allocator` but with static dispatch
+- `limit_allocator`: Allocator wrapper that limits the amount of memory that is allowed to be allocated
 - `DICE_DEFER`/`DICE_DEFER_TO_SUCCES`/`DICE_DEFER_TO_FAIL`: On-the-fly RAII for types that do not support it natively (similar to go's `defer` keyword)
 - `overloaded`: Composition for `std::variant` visitor lambdas
 - `flex_array`: A combination of `std::array`, `std::span` and a `vector` with small buffer optimization
@@ -64,6 +65,10 @@ The problem with `mmap` allocations is that they will be placed at an arbitrary
 therefore absolute pointers will cause segfaults if the segment is reloaded.
 Which means: vtables will not work (because they use absolute pointers) and therefore you cannot use `std::pmr::polymorphic_allocator`.
 
+### `limit_allocator`
+Allocator wrapper that limits the amount of memory that can be allocated through the inner allocator.
+If the limit is exceeded it will throw `std::bad_alloc`.
+
 ### `DICE_DEFER`/`DICE_DEFER_TO_SUCCES`/`DICE_DEFER_TO_FAIL`
 A mechanism similar to go's `defer` keyword, which can be used to defer some action to scope exit.
 The primary use-case for this is on-the-fly RAII-like resource management for types that do not support RAII (for example C types).

examples/CMakeLists.txt

@@ -80,3 +80,10 @@ target_link_libraries(example_variant2
         dice-template-library::dice-template-library
 )
 
+add_executable(example_limit_allocator
+        example_limit_allocator.cpp)
+target_link_libraries(example_limit_allocator
+        PRIVATE
+        dice-template-library::dice-template-library
+)
+

examples/example_limit_allocator.cpp

@@ -0,0 +1,26 @@
+#include <dice/template-library/limit_allocator.hpp>
+
+#include <cassert>
+#include <vector>
+
+
+int main() {
+	std::vector<int, dice::template_library::limit_allocator<int>> vec{dice::template_library::limit_allocator<int>{3 * sizeof(int)}};
+	vec.push_back(1);
+	vec.push_back(2);
+
+	try {
+		vec.push_back(3);
+		assert(false);
+	} catch (...) {
+	}
+
+	vec.pop_back();
+	vec.push_back(4);
+
+	try {
+		vec.push_back(5);
+		assert(false);
+	} catch (...) {
+	}
+}

include/dice/template-library/limit_allocator.hpp

@@ -0,0 +1,176 @@
+#ifndef DICE_TEMPLATELIBRARY_LIMITALLOCATOR_HPP
+#define DICE_TEMPLATELIBRARY_LIMITALLOCATOR_HPP
+
+#include <atomic>
+#include <cstddef>
+#include <memory>
+#include <new>
+#include <type_traits>
+#include <utility>
+
+namespace dice::template_library {
+
+	/**
+	 * The synchronization policy of a limit_allocator
+	 */
+	enum struct limit_allocator_syncness : bool {
+		sync, ///< thread-safe (synchronized)
+		unsync///< not thread-safe (unsynchronized)
+	};
+
+	namespace detail_limit_allocator {
+		template<limit_allocator_syncness syncness>
+		struct limit_allocator_control_block;
+
+		template<>
+		struct limit_allocator_control_block<limit_allocator_syncness::sync> {
+			std::atomic<size_t> bytes_left = 0;
+
+			void allocate(size_t n_bytes) {
+				auto old = bytes_left.load(std::memory_order_relaxed);
+
+				do {
+					if (old < n_bytes) [[unlikely]] {
+						throw std::bad_alloc{};
+					}
+				} while (!bytes_left.compare_exchange_weak(old, old - n_bytes, std::memory_order_relaxed, std::memory_order_relaxed));
+			}
+
+			void deallocate(size_t n_bytes) noexcept {
+				bytes_left.fetch_add(n_bytes, std::memory_order_relaxed);
+			}
+		};
+
+		template<>
+		struct limit_allocator_control_block<limit_allocator_syncness::unsync> {
+			size_t bytes_left = 0;
+
+			void allocate(size_t n_bytes) {
+				if (bytes_left < n_bytes) [[unlikely]] {
+					throw std::bad_alloc{};
+				}
+				bytes_left -= n_bytes;
+			}
+
+			void deallocate(size_t n_bytes) noexcept {
+				bytes_left += n_bytes;
+			}
+		};
+	}// namespace detail_limit_allocator
+
+	/**
+	 * Allocator wrapper that limits the amount of memory its underlying allocator
+	 * is allowed to allocate.
+	 *
+	 * @tparam T value type of the allocator (the thing that it allocates)
+	 * @tparam Allocator the underlying allocator
+	 * @tparam syncness determines the synchronization of the limit
+	 */
+	template<typename T, template<typename> typename Allocator = std::allocator, limit_allocator_syncness syncness = limit_allocator_syncness::sync>
+	struct limit_allocator {
+		using control_block_type = detail_limit_allocator::limit_allocator_control_block<syncness>;
+		using value_type = T;
+		using upstream_allocator_type = Allocator<T>;
+		using pointer = typename std::allocator_traits<upstream_allocator_type>::pointer;
+		using const_pointer = typename std::allocator_traits<upstream_allocator_type>::const_pointer;
+		using void_pointer = typename std::allocator_traits<upstream_allocator_type>::void_pointer;
+		using const_void_pointer = typename std::allocator_traits<upstream_allocator_type>::const_void_pointer;
+		using size_type = typename std::allocator_traits<upstream_allocator_type>::size_type;
+		using difference_type = typename std::allocator_traits<upstream_allocator_type>::difference_type;
+
+		using propagate_on_container_copy_assignment = typename std::allocator_traits<upstream_allocator_type>::propagate_on_container_copy_assignment;
+		using propagate_on_container_move_assignment = typename std::allocator_traits<upstream_allocator_type>::propagate_on_container_move_assignment;
+		using propagate_on_container_swap = typename std::allocator_traits<upstream_allocator_type>::propagate_on_container_swap;
+		using is_always_equal = std::false_type;
+
+		template<typename U>
+		struct rebind {
+			using other = limit_allocator<U, Allocator, syncness>;
+		};
+
+	private:
+		template<typename, template<typename> typename, limit_allocator_syncness>
+		friend struct limit_allocator;
+
+		std::shared_ptr<control_block_type> control_block_;
+		[[no_unique_address]] upstream_allocator_type inner_;
+
+		constexpr limit_allocator(std::shared_ptr<control_block_type> const &control_block, upstream_allocator_type const &alloc)
+			requires(std::is_default_constructible_v<upstream_allocator_type>)
+			: control_block_{control_block},
+			  inner_{alloc} {
+		}
+
+	public:
+		explicit constexpr limit_allocator(size_t bytes_limit)
+			requires(std::is_default_constructible_v<upstream_allocator_type>)
+			: control_block_{std::make_shared<control_block_type>(bytes_limit)},
+			  inner_{} {
+		}
+
+		constexpr limit_allocator(limit_allocator const &other) noexcept(std::is_nothrow_move_constructible_v<upstream_allocator_type>) = default;
+		constexpr limit_allocator(limit_allocator &&other) noexcept(std::is_nothrow_copy_constructible_v<upstream_allocator_type>) = default;
+		constexpr limit_allocator &operator=(limit_allocator const &other) noexcept(std::is_nothrow_copy_assignable_v<upstream_allocator_type>) = default;
+		constexpr limit_allocator &operator=(limit_allocator &&other) noexcept(std::is_nothrow_move_assignable_v<upstream_allocator_type>) = default;
+		constexpr ~limit_allocator() = default;
+
+		template<typename U>
+		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 &>)
+			: control_block_{other.control_block_},
+			  inner_{other.inner_} {
+		}
+
+		constexpr limit_allocator(size_t bytes_limit, upstream_allocator_type const &upstream)
+			: control_block_{std::make_shared<control_block_type>(bytes_limit)},
+			  inner_{upstream} {
+		}
+
+		constexpr limit_allocator(size_t bytes_limit, upstream_allocator_type &&upstream)
+			: control_block_{std::make_shared<control_block_type>(bytes_limit)},
+			  inner_{std::move(upstream)} {
+		}
+
+		template<typename... Args>
+		explicit constexpr limit_allocator(size_t bytes_limit, std::in_place_t, Args &&...args)
+			: control_block_{std::make_shared<control_block_type>(bytes_limit)},
+			  inner_{std::forward<Args>(args)...} {
+		}
+
+		constexpr pointer allocate(size_t n) {
+			control_block_->allocate(n * sizeof(T));
+
+			try {
+				return std::allocator_traits<upstream_allocator_type>::allocate(inner_, n);
+			} catch (...) {
+				control_block_->deallocate(n * sizeof(T));
+				throw;
+			}
+		}
+
+		constexpr void deallocate(pointer ptr, size_t n) {
+			std::allocator_traits<upstream_allocator_type>::deallocate(inner_, ptr, n);
+			control_block_->deallocate(n * sizeof(T));
+		}
+
+		constexpr limit_allocator select_on_container_copy_construction() const {
+			return limit_allocator{control_block_, std::allocator_traits<upstream_allocator_type>::select_on_container_copy_construction(inner_)};
+		}
+
+		[[nodiscard]] upstream_allocator_type const &upstream_allocator() const noexcept {
+			return inner_;
+		}
+
+		friend constexpr void swap(limit_allocator &a, limit_allocator &b) noexcept(std::is_nothrow_swappable_v<upstream_allocator_type>)
+			requires(std::is_swappable_v<upstream_allocator_type>)
+		{
+			using std::swap;
+			swap(a.control_block_, b.control_block_);
+			swap(a.inner_, b.inner_);
+		}
+
+		bool operator==(limit_allocator const &other) const noexcept = default;
+		bool operator!=(limit_allocator const &other) const noexcept = default;
+	};
+}// namespace dice::template_library
+
+#endif// DICE_TEMPLATELIBRARY_LIMITALLOCATOR_HPP

include/dice/template-library/polymorphic_allocator.hpp

@@ -337,6 +337,11 @@ namespace dice::template_library {
 		using propagate_on_container_swap = typename std::allocator_traits<upstream_allocator_type>::propagate_on_container_swap;
 		using is_always_equal = typename std::allocator_traits<upstream_allocator_type>::is_always_equal;
 
+		template<typename U>
+		struct rebind {
+			using other = offset_ptr_stl_allocator<U, Allocator>;
+		};
+
 	private:
 		template<typename, template<typename> typename>
 		friend struct offset_ptr_stl_allocator;

tests/CMakeLists.txt

@@ -63,3 +63,6 @@ custom_add_test(tests_variant2)
 
 add_executable(tests_type_traits tests_type_traits.cpp)
 custom_add_test(tests_type_traits)
+
+add_executable(tests_limit_allocator tests_limit_allocator.cpp)
+custom_add_test(tests_limit_allocator)

tests/tests_limit_allocator.cpp

@@ -0,0 +1,39 @@
+#define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN
+#include <doctest/doctest.h>
+
+#include <dice/template-library/limit_allocator.hpp>
+
+TEST_SUITE("limit_allocator sanity check") {
+	using namespace dice::template_library;
+
+	template<limit_allocator_syncness syncness>
+	void run_tests() {
+		limit_allocator<int> alloc{12 * sizeof(int)};
+
+		auto a = alloc.allocate(5);
+		auto b = alloc.allocate(5);
+		auto c = alloc.allocate(1);
+		alloc.deallocate(c, 1);
+		alloc.deallocate(b, 5);
+
+		CHECK_THROWS(alloc.allocate(8));
+
+		auto d = alloc.allocate(7);
+
+		alloc.deallocate(a, 5);
+		alloc.deallocate(d, 7);
+
+		auto e = alloc.allocate(12);
+		alloc.deallocate(e, 12);
+
+		CHECK_THROWS(alloc.allocate(13));
+	}
+
+	TEST_CASE("sync") {
+		run_tests<limit_allocator_syncness::sync>();
+	}
+
+	TEST_CASE("unsync") {
+		run_tests<limit_allocator_syncness::unsync>();
+	}
+}