memorypool.h

#ifndef MEMORYPOOL_H
#define MEMORYPOOL_H

#include <vector>
#include <unordered_map>
#include <mutex>
#include <atomic>
#include "memory_env.h"

// Device memory pool
class MemoryPool {
public:
    explicit MemoryPool(size_t poolSize);
    ~MemoryPool();

    // Allocate a memory block.
    void AllocateBlock(uint32_t size, int &blockId);

    // Release a memory block.
    void FreeBlock(int blockId);

    // Obtain the physical address of a memory block.
    void GetBlockPtr(int blockId, void *&addr);

private:
    // Generate a memory block index.
    uint64_t GenerateBlocksId();

    std::atomic<uint64_t> id_ = 0;
    std::mutex blockMutex_;
    void *baseMemPtr_ = nullptr;
    void *curMemPtr_ = nullptr;
    int64_t remainSize_ = 0;
    std::unordered_map<int, MemoryBlock> freeBlocks_;
    std::unordered_map<int, MemoryBlock> usedBlocks_;
};

#endif