accelerator/include/memory_pool.hpp

//
// SPDX-FileCopyrightText: Hadad <hadad@linuxmail.org>
// SPDX-License-Identifier: Apache-2.0
//

#ifndef POCKET_TTS_MEMORY_POOL_HPP
#define POCKET_TTS_MEMORY_POOL_HPP

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <mutex>
#include <unordered_map>
#include <vector>

namespace pocket_tts_accelerator {

struct MemoryBlock {
    std::unique_ptr<std::uint8_t[]> data;
    std::size_t block_size;
    bool is_in_use;
    std::uint64_t last_access_timestamp;
};

class MemoryPool {
public:
    explicit MemoryPool(std::size_t initial_pool_size_bytes = 64 * 1024 * 1024);
    ~MemoryPool();

    MemoryPool(const MemoryPool&) = delete;
    MemoryPool& operator=(const MemoryPool&) = delete;
    MemoryPool(MemoryPool&&) = delete;
    MemoryPool& operator=(MemoryPool&&) = delete;

    std::uint8_t* allocate(std::size_t requested_size_bytes);
    void deallocate(std::uint8_t* pointer);
    void clear_unused_blocks();
    void reset_pool();

    std::size_t get_total_allocated_bytes() const;
    std::size_t get_total_used_bytes() const;
    std::size_t get_block_count() const;

private:
    std::size_t find_suitable_block_index(std::size_t requested_size) const;
    void create_new_block(std::size_t block_size);
    std::uint64_t get_current_timestamp() const;

    std::vector<MemoryBlock> memory_blocks;
    std::unordered_map<std::uint8_t*, std::size_t> pointer_to_block_index;
    mutable std::mutex pool_mutex;
    std::size_t total_allocated_bytes;
    std::size_t total_used_bytes;
    std::size_t maximum_pool_size_bytes;
};

class ScopedMemoryAllocation {
public:
    ScopedMemoryAllocation(MemoryPool& pool, std::size_t size);
    ~ScopedMemoryAllocation();

    ScopedMemoryAllocation(const ScopedMemoryAllocation&) = delete;
    ScopedMemoryAllocation& operator=(const ScopedMemoryAllocation&) = delete;
    ScopedMemoryAllocation(ScopedMemoryAllocation&& other) noexcept;
    ScopedMemoryAllocation& operator=(ScopedMemoryAllocation&& other) noexcept;

    std::uint8_t* get() const;
    std::size_t size() const;

private:
    MemoryPool* memory_pool_pointer;
    std::uint8_t* allocated_pointer;
    std::size_t allocation_size;
};

}

#endif