Mooncake/mooncake-store/include/aligned_client_buffer.hpp

62 lines
2.0 KiB
C++

#pragma once
#include "client_buffer.hpp"
namespace mooncake {
/**
* AlignedClientBufferAllocator extends ClientBufferAllocator to provide
* 4096-byte aligned memory allocation, which is required for O_DIRECT I/O
* operations with io_uring.
*
* This allocator ensures:
* - Base memory address is aligned to 4096 bytes (page size)
* - Can be registered with UringFile for zero-copy I/O
* - Maintains all features of the parent ClientBufferAllocator
*/
class AlignedClientBufferAllocator : public ClientBufferAllocator {
public:
// Alignment requirement for O_DIRECT I/O
static constexpr size_t kDirectIOAlignment = 4096;
/**
* Create an AlignedClientBufferAllocator with aligned memory
* @param size Total size of the buffer to allocate
* @param protocol Optional protocol string (unused, for compatibility)
* @param use_hugepage Whether to use huge pages for allocation
* @return Shared pointer to the allocator, or nullptr on failure
*/
static std::shared_ptr<AlignedClientBufferAllocator> create(
size_t size, const std::string& protocol = "",
bool use_hugepage = false);
/**
* Get the base pointer of the aligned buffer
* @return Base address of the allocated buffer
*/
[[nodiscard]] void* get_base_pointer() const { return getBase(); }
/**
* Get the total size of the aligned buffer
* @return Total size in bytes
*/
[[nodiscard]] size_t get_total_size() const { return size(); }
/**
* Destructor - properly frees the aligned memory
*/
~AlignedClientBufferAllocator();
private:
// Private constructor - use create() factory method
AlignedClientBufferAllocator(void* aligned_buffer, size_t size,
const std::string& protocol,
bool use_hugepage);
// Store whether we own the memory (for cleanup)
bool owns_memory_;
size_t allocated_size_; // Store the actual allocated size for cleanup
};
} // namespace mooncake