feat(PP,ByteStream): new counting memory allocator

utils/common/countingallocator.h

@@ -0,0 +1,91 @@
+/* Copyright (C) 2024 MariaDB Corporation
+
+   This program is free software; you can redistribute it and/or
+   modify it under the terms of the GNU General Public License
+   as published by the Free Software Foundation; version 2 of
+   the License.
+
+   This program is distributed in the hope that it will be useful,
+   but WITHOUT ANY WARRANTY; without even the implied warranty of
+   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+   GNU General Public License for more details.
+
+   You should have received a copy of the GNU General Public License
+   along with this program; if not, write to the Free Software
+   Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
+   MA 02110-1301, USA. */
+
+#pragma once
+
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <atomic>
+#include <cstddef>
+#include <iostream>
+#include <utility>
+
+namespace allocators 
+{
+
+// const constexpr std::uint64_t CounterUpdateUnitSize = 4 * 1024 * 1024; 
+const constexpr std::int64_t MemoryLimitLowerBound = 100 * 1024 * 1024;  // WIP
+
+// Custom Allocator that tracks allocated memory using an atomic counter
+template <typename T>
+class CountingAllocator {
+public:
+    using value_type = T;
+
+    // Constructor accepting a reference to an atomic counter
+    explicit CountingAllocator(std::atomic<int64_t>& memoryLimit, const uint64_t lowerBound = MemoryLimitLowerBound) noexcept
+        : memoryLimitRef_(memoryLimit), memoryLimitLowerBound(lowerBound) {}
+
+    // Copy constructor (template to allow conversion between different types)
+    template <typename U>
+    CountingAllocator(const CountingAllocator<U>& other) noexcept
+        : memoryLimitRef_(other.memoryLimitRef_) {}
+
+    // Allocate memory for n objects of type T
+    T* allocate(std::size_t n) {
+        auto memCounted = memoryLimitRef_.fetch_sub(n * sizeof(T), std::memory_order_relaxed);
+        if (memCounted < memoryLimitLowerBound) {
+            memoryLimitRef_.fetch_add(n * sizeof(T), std::memory_order_relaxed);
+            throw std::bad_alloc();
+        }
+        
+        T* ptr = static_cast<T*>(::operator new(n * sizeof(T)));
+        // std::cout << "[Allocate] " << n * sizeof(T) << " bytes at " << static_cast<void*>(ptr)
+        //           << ". current timit: " << std::dec << memoryLimitRef_.load() << std::hex << " bytes.\n";
+        return ptr;
+    }
+
+    // Deallocate memory for n objects of type T
+    void deallocate(T* ptr, std::size_t n) noexcept {
+        ::operator delete(ptr);
+        memoryLimitRef_.fetch_add(n * sizeof(T), std::memory_order_relaxed);
+        // std::cout << "[Deallocate] " << n * sizeof(T) << " bytes from " << static_cast<void*>(ptr)
+        //           << ". current timit: " << std::dec << memoryLimitRef_.load() << std::hex << " bytes.\n";
+    }
+
+    // Equality operators (allocators are equal if they share the same counter)
+    template <typename U>
+    bool operator==(const CountingAllocator<U>& other) const noexcept {
+        return &memoryLimitRef_ == &other.memoryLimitRef_;
+    }
+
+    template <typename U>
+    bool operator!=(const CountingAllocator<U>& other) const noexcept {
+        return !(*this == other);
+    }
+
+private:
+    std::atomic<int64_t>& memoryLimitRef_;
+    int64_t memoryLimitLowerBound = 0;
+
+    // Grant access to other instances of CountingAllocator with different types
+    template <typename U>
+    friend class CountingAllocator;
+};
+
+}  // namespace allocators