include/paimon/memory/memory_pool.h - paimon-cpp - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *   http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 #pragma once

 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <functional>
 #include <memory>
 #include <new>
 #include <utility>

 #include "paimon/macros.h"
 #include "paimon/visibility.h"

 namespace paimon {

 class MemoryPool;

 /// Create a default implementation of memory pool.
 /// @return Unique pointer to a newly created `MemoryPool` instance.
 PAIMON_EXPORT std::unique_ptr<MemoryPool> GetMemoryPool();

 /// Get a system-wide singleton memory pool.
 /// @return Shared pointer to the singleton `MemoryPool` instance.
 PAIMON_EXPORT std::shared_ptr<MemoryPool> GetDefaultPool();

 /// Abstract base class for memory pool implementations that provides controlled memory management.
 class PAIMON_EXPORT MemoryPool {
  public:
     virtual ~MemoryPool() = default;
     MemoryPool& operator=(const MemoryPool& other) = delete;
     MemoryPool& operator=(MemoryPool& other) = delete;
     MemoryPool& operator=(MemoryPool&& other) = delete;

     /// Allocate memory from the pool.
     ///
     /// Allocates a block of memory with the specified size and alignment.
     /// The returned memory is uninitialized and must be properly constructed
     /// before use.
     ///
     /// @param size Number of bytes to allocate.
     /// @param alignment Memory alignment requirement (0 for default alignment).
     /// @return Pointer to allocated memory, or nullptr on failure.
     virtual void* Malloc(uint64_t size, uint64_t alignment = 0) = 0;

     /// Reallocate memory to a new size.
     ///
     /// Changes the size of a previously allocated memory block. The contents
     /// of the memory block are preserved up to the minimum of the old and new sizes.
     ///
     /// @param p Pointer to the memory block to reallocate.
     /// @param old_size Current size of the memory block.
     /// @param new_size New desired size of the memory block.
     /// @param alignment Memory alignment requirement (0 for default alignment).
     /// @return Pointer to the reallocated memory block.
     virtual void* Realloc(void* p, size_t old_size, size_t new_size, uint64_t alignment = 0) = 0;

     /// Deallocate memory back to the pool.
     ///
     /// Releases a previously allocated memory block back to the pool.
     /// The size must match the size used during allocation.
     /// The alignment used during allocation is not provided, subclass should store it
     /// by themselves if needed.
     ///
     /// @param p Pointer to the memory block to deallocate.
     /// @param size Size of the memory block (must match allocation size).
     virtual void Free(void* p, uint64_t size) = 0;

     /// Deallocate memory back to the pool with specified alignment.
     ///
     /// Releases a previously allocated memory block back to the pool.
     /// The size and alignment must match the values used during allocation.
     /// Subclass can override this method to optimize deallocation based on alignment.
     ///
     /// @param p Pointer to the memory block to deallocate.
     /// @param size Size of the memory block (must match allocation size).
     /// @param alignment Alignment of the memory block (must match allocation alignment).
     virtual void Free(void* p, uint64_t size, uint64_t alignment) {
         Free(p, size);
     }

     /// Get current memory usage.
     ///
     /// Returns the amount of memory currently allocated from this pool.
     /// This includes all outstanding allocations that have not been freed.
     ///
     /// @return Current memory usage in bytes.
     virtual uint64_t CurrentUsage() const = 0;

     /// Get peak memory usage.
     ///
     /// Returns the maximum amount of memory that has been allocated from
     /// this pool at any point in time since its creation.
     ///
     /// @return Peak memory usage in bytes.
     virtual uint64_t MaxMemoryUsage() const = 0;

     /// Custom deleter for use with std::unique_ptr that integrates with memory pools.
     ///
     /// AllocatorDelete provides automatic memory deallocation through the memory pool
     /// when used with std::unique_ptr. It stores both the pool reference and the
     /// allocation size to ensure proper cleanup.
     ///
     /// @tparam T Type of object being managed.
     template <typename T>
     class AllocatorDelete {
      private:
         std::reference_wrapper<MemoryPool> pool;
         size_t size{0};

      public:
         AllocatorDelete() : pool(*GetMemoryPool()) {}
         AllocatorDelete(MemoryPool& _pool, size_t _size) : pool(_pool), size(_size) {}
         AllocatorDelete(AllocatorDelete const&) = default;

         template <typename U>
         AllocatorDelete& operator=(const AllocatorDelete<U>& other) {
             pool = other.GetPool();
             size = other.GetSize();
             return *this;
         }

         AllocatorDelete& operator=(AllocatorDelete const& other) {
             pool = other.GetPool();
             size = other.GetSize();
             return *this;
         }

         AllocatorDelete& operator=(AllocatorDelete&& other) {
             pool = other.GetPool();
             size = other.GetSize();
             return *this;
         }

         template <typename U>
         AllocatorDelete(const AllocatorDelete<U>& other)  // NOLINT(google-explicit-constructor)
             : pool(other.GetPool()), size(other.GetSize()) {}

         MemoryPool& GetPool() const {
             return pool;
         }

         size_t GetSize() const {
             return size;
         }

         void operator()(T* p) const {
             if (p) {
                 if (PAIMON_UNLIKELY(size == 0)) {
                     // This indicates a programming error where the deleter was not properly
                     // initialized with the correct allocation size. In debug builds, this
                     // will trigger an assertion. In release builds, we cannot safely
                     // deallocate without knowing the size, which may lead to undefined behavior.
                     assert(false &&
                            "Mismatched pointer allocation detected - deleter size is zero");
                     return;  // Cannot safely deallocate without size information
                 }
                 p->~T();
                 pool.get().Free(reinterpret_cast<void*>(p), size);
             }
         }
     };

     /// Allocate an object on the pool and return a unique_ptr with custom deleter.
     ///
     /// This method provides a convenient way to allocate objects that are automatically
     /// managed by the memory pool. The returned unique_ptr will automatically deallocate
     /// the memory through the pool when it goes out of scope.
     ///
     /// @tparam T Type of object to allocate.
     /// @tparam Args Types of constructor arguments.
     /// @param args Constructor arguments to forward to T's constructor.
     /// @return unique_ptr managing the allocated object with pool-aware deleter.
     template <typename T, typename... Args>
     std::unique_ptr<T, AllocatorDelete<T>> AllocateUnique(Args&&... args) {
         struct DeferCondDeallocate {
             bool& cond;
             MemoryPool& pool;
             void* p;
             ~DeferCondDeallocate() {
                 if (PAIMON_UNLIKELY(!cond)) {
                     pool.Free(p, sizeof(T));
                 }
             }
         };
         auto p = Malloc(sizeof(T));
         {
             bool constructed = false;
             DeferCondDeallocate handler{constructed, *this, p};
             new (p) T(std::forward<Args>(args)...);
             constructed = true;
         }
         return std::unique_ptr<T, AllocatorDelete<T>>(reinterpret_cast<T*>(p),
                                                       AllocatorDelete<T>(*this, sizeof(T)));
     }
 };

 template <class T>
 using pooled_unique_ptr = std::unique_ptr<T, MemoryPool::AllocatorDelete<T>>;

 #define PAIMON_UNIQUE_PTR pooled_unique_ptr

 }  // namespace paimon
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing,
	* software distributed under the License is distributed on an
	* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	#pragma once

	#include <cassert>
	#include <cstddef>
	#include <cstdint>
	#include <functional>
	#include <memory>
	#include <new>
	#include <utility>

	#include "paimon/macros.h"
	#include "paimon/visibility.h"

	namespace paimon {

	class MemoryPool;

	/// Create a default implementation of memory pool.
	/// @return Unique pointer to a newly created `MemoryPool` instance.
	PAIMON_EXPORT std::unique_ptr<MemoryPool> GetMemoryPool();

	/// Get a system-wide singleton memory pool.
	/// @return Shared pointer to the singleton `MemoryPool` instance.
	PAIMON_EXPORT std::shared_ptr<MemoryPool> GetDefaultPool();

	/// Abstract base class for memory pool implementations that provides controlled memory management.
	class PAIMON_EXPORT MemoryPool {
	public:
	virtual ~MemoryPool() = default;
	MemoryPool& operator=(const MemoryPool& other) = delete;
	MemoryPool& operator=(MemoryPool& other) = delete;
	MemoryPool& operator=(MemoryPool&& other) = delete;

	/// Allocate memory from the pool.
	///
	/// Allocates a block of memory with the specified size and alignment.
	/// The returned memory is uninitialized and must be properly constructed
	/// before use.
	///
	/// @param size Number of bytes to allocate.
	/// @param alignment Memory alignment requirement (0 for default alignment).
	/// @return Pointer to allocated memory, or nullptr on failure.
	virtual void* Malloc(uint64_t size, uint64_t alignment = 0) = 0;

	/// Reallocate memory to a new size.
	///
	/// Changes the size of a previously allocated memory block. The contents
	/// of the memory block are preserved up to the minimum of the old and new sizes.
	///
	/// @param p Pointer to the memory block to reallocate.
	/// @param old_size Current size of the memory block.
	/// @param new_size New desired size of the memory block.
	/// @param alignment Memory alignment requirement (0 for default alignment).
	/// @return Pointer to the reallocated memory block.
	virtual void* Realloc(void* p, size_t old_size, size_t new_size, uint64_t alignment = 0) = 0;

	/// Deallocate memory back to the pool.
	///
	/// Releases a previously allocated memory block back to the pool.
	/// The size must match the size used during allocation.
	/// The alignment used during allocation is not provided, subclass should store it
	/// by themselves if needed.
	///
	/// @param p Pointer to the memory block to deallocate.
	/// @param size Size of the memory block (must match allocation size).
	virtual void Free(void* p, uint64_t size) = 0;

	/// Deallocate memory back to the pool with specified alignment.
	///
	/// Releases a previously allocated memory block back to the pool.
	/// The size and alignment must match the values used during allocation.
	/// Subclass can override this method to optimize deallocation based on alignment.
	///
	/// @param p Pointer to the memory block to deallocate.
	/// @param size Size of the memory block (must match allocation size).
	/// @param alignment Alignment of the memory block (must match allocation alignment).
	virtual void Free(void* p, uint64_t size, uint64_t alignment) {
	Free(p, size);
	}

	/// Get current memory usage.
	///
	/// Returns the amount of memory currently allocated from this pool.
	/// This includes all outstanding allocations that have not been freed.
	///
	/// @return Current memory usage in bytes.
	virtual uint64_t CurrentUsage() const = 0;

	/// Get peak memory usage.
	///
	/// Returns the maximum amount of memory that has been allocated from
	/// this pool at any point in time since its creation.
	///
	/// @return Peak memory usage in bytes.
	virtual uint64_t MaxMemoryUsage() const = 0;

	/// Custom deleter for use with std::unique_ptr that integrates with memory pools.
	///
	/// AllocatorDelete provides automatic memory deallocation through the memory pool
	/// when used with std::unique_ptr. It stores both the pool reference and the
	/// allocation size to ensure proper cleanup.
	///
	/// @tparam T Type of object being managed.
	template <typename T>
	class AllocatorDelete {
	private:
	std::reference_wrapper<MemoryPool> pool;
	size_t size{0};

	public:
	AllocatorDelete() : pool(*GetMemoryPool()) {}
	AllocatorDelete(MemoryPool& _pool, size_t _size) : pool(_pool), size(_size) {}
	AllocatorDelete(AllocatorDelete const&) = default;

	template <typename U>
	AllocatorDelete& operator=(const AllocatorDelete<U>& other) {
	pool = other.GetPool();
	size = other.GetSize();
	return *this;
	}

	AllocatorDelete& operator=(AllocatorDelete const& other) {
	pool = other.GetPool();
	size = other.GetSize();
	return *this;
	}

	AllocatorDelete& operator=(AllocatorDelete&& other) {
	pool = other.GetPool();
	size = other.GetSize();
	return *this;
	}

	template <typename U>
	AllocatorDelete(const AllocatorDelete<U>& other) // NOLINT(google-explicit-constructor)
	: pool(other.GetPool()), size(other.GetSize()) {}

	MemoryPool& GetPool() const {
	return pool;
	}

	size_t GetSize() const {
	return size;
	}

	void operator()(T* p) const {
	if (p) {
	if (PAIMON_UNLIKELY(size == 0)) {
	// This indicates a programming error where the deleter was not properly
	// initialized with the correct allocation size. In debug builds, this
	// will trigger an assertion. In release builds, we cannot safely
	// deallocate without knowing the size, which may lead to undefined behavior.
	assert(false &&
	"Mismatched pointer allocation detected - deleter size is zero");
	return; // Cannot safely deallocate without size information
	}
	p->~T();
	pool.get().Free(reinterpret_cast<void*>(p), size);
	}
	}
	};

	/// Allocate an object on the pool and return a unique_ptr with custom deleter.
	///
	/// This method provides a convenient way to allocate objects that are automatically
	/// managed by the memory pool. The returned unique_ptr will automatically deallocate
	/// the memory through the pool when it goes out of scope.
	///
	/// @tparam T Type of object to allocate.
	/// @tparam Args Types of constructor arguments.
	/// @param args Constructor arguments to forward to T's constructor.
	/// @return unique_ptr managing the allocated object with pool-aware deleter.
	template <typename T, typename... Args>
	std::unique_ptr<T, AllocatorDelete<T>> AllocateUnique(Args&&... args) {
	struct DeferCondDeallocate {
	bool& cond;
	MemoryPool& pool;
	void* p;
	~DeferCondDeallocate() {
	if (PAIMON_UNLIKELY(!cond)) {
	pool.Free(p, sizeof(T));
	}
	}
	};
	auto p = Malloc(sizeof(T));
	{
	bool constructed = false;
	DeferCondDeallocate handler{constructed, *this, p};
	new (p) T(std::forward<Args>(args)...);
	constructed = true;
	}
	return std::unique_ptr<T, AllocatorDelete<T>>(reinterpret_cast<T*>(p),
	AllocatorDelete<T>(*this, sizeof(T)));
	}
	};

	template <class T>
	using pooled_unique_ptr = std::unique_ptr<T, MemoryPool::AllocatorDelete<T>>;

	#define PAIMON_UNIQUE_PTR pooled_unique_ptr

	} // namespace paimon