固定内存块尺寸的内存池原理及代码

Posted 2020-06-11

有些情形会需要申请大量的固定尺寸的内存块，若一个个都用malloc申请效率很低，这种情况非常适合使用内存池解决。

下面就是一个固定内存块尺寸的内存池的完整源码。注：其中的内存申请不是使用的malloc，而是自己定义的torch::HeapMalloc，简单修改下即可。

代码详情请见Github【点击】

先看代码再看讲解

/*
 * 固定尺寸的内存池
 * 说明：
 *  - 只适用于需要大量固定尺寸内存块的情况
 */
template<size_t _Size>
class FixedMemoryPool
{
public:
    enum { _CHUNK_COUNT = 1024/_Size };

    FixedMemoryPool();
    ~FixedMemoryPool();

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

    /*
     * 从内存池获得一个内存块
     * 注意：
     *  - 用完后最好在放回内存池中(有利于提高性能)
     *  - 不能使用::free()释放内存，内存池会统一释放池中的内存块
     */
    void* Alloc();
    
    /*
     * 将内存块放回内存池中
     * 注意：必须是从内存池申请的内存块才能放回去
     */
    void Free(void *mem);

    /*
     * 清空内存池
     * 注意：调用此接口后，通过Alloc申请的内存块会全部被释放
     */
    void Clear();

    /*
     * 获得内存池总大小(单位Byte)
     */
    size_t GetCapacity();

    /*
     * 获得内存块的大小(单位Byte)
     */
    size_t GetItemSize();

private:
    struct Chunk;
    Chunk* MakeChunk();

private:
    union ChunkItem {
        ChunkItem*  next;
        uint8_t     mem[_Size];
    };
    struct Chunk { ChunkItem itemArray[_CHUNK_COUNT]; };

    ChunkItem*          m_rootItem;
    std::list<Chunk *>  m_chunkList;
};

template<size_t _Size>
FixedMemoryPool<_Size>::FixedMemoryPool() {
    m_rootItem = nullptr;
    Chunk *chunk = this->MakeChunk();
    if (!chunk)
        return;
    m_chunkList.push_back(chunk);
    m_rootItem = chunk->itemArray;
}

template<size_t _Size>
FixedMemoryPool<_Size>::~FixedMemoryPool() {
    this->Clear();
}

template<size_t _Size>
FixedMemoryPool<_Size>::FixedMemoryPool(const FixedMemoryPool &) {
    m_rootItem = nullptr;
    Chunk *chunk = this->MakeChunk();
    if (!chunk)
        return;
    m_chunkList.push_back(chunk);
    m_rootItem = chunk->itemArray;
}

template<size_t _Size>
void* FixedMemoryPool<_Size>::Alloc() {
    if (!m_rootItem) {
        Chunk *chunk = this->MakeChunk();
        if (!chunk) return nullptr;
        m_chunkList.push_back(chunk);
        m_rootItem = chunk->itemArray;
    }
    void *item = m_rootItem;
    m_rootItem = m_rootItem->next;
    memset(item, 0, _Size);
    return item;
}

template<size_t _Size>
void FixedMemoryPool<_Size>::Free(void *mem) {
    if (!mem) return;
    ChunkItem *item = static_cast<ChunkItem*>(mem);
    item->next = m_rootItem;
    m_rootItem = item;
}

template<size_t _Size>
void FixedMemoryPool<_Size>::Clear() {
    for (Chunk *chunk : m_chunkList) {
        torch::HeapFree(chunk);
    }
    m_chunkList.clear();
}

template<size_t _Size>
size_t FixedMemoryPool<_Size>::GetCapacity() {
    return m_chunkList.size() * _CHUNK_COUNT * _Size;
}

template<size_t _Size>
size_t FixedMemoryPool<_Size>::GetItemSize() {
    return _Size;
}

template<size_t _Size>
typename FixedMemoryPool<_Size>::Chunk* FixedMemoryPool<_Size>::MakeChunk() {
    Chunk *chunk = (Chunk *)torch::HeapMalloc(sizeof(Chunk));
    if (!chunk)
        return nullptr;
    for (int i = 0; i < _CHUNK_COUNT - 1; i++) {
        chunk->itemArray[i].next = &(chunk->itemArray[i+1]);
    }
    chunk->itemArray[_CHUNK_COUNT-1].next = nullptr;
    return chunk;
}

 

首先要清除这个内存池的结构：

 

ChunkList[
    Chunk[ChunkItem, ... ChunkItem],
    Chunk[ChunkItem, ... ChunkItem],
    ...
    Chunk[ChunkItem, ... ChunkItem]
]

 

ChunkItem非常重要，他是一个union，其中有两个成员：

    union ChunkItem {
        ChunkItem*  next; // 指向下一个ChunkItem
        uint8_t     mem[_Size]; // _Size为模板参数
    };
　　 struct Chunk { ChunkItem itemArray[_CHUNK_COUNT]; }; 

 

当ChunkItem在池中时，只会使用next域，使用next将一个Chunk中的所有ChunkItem串联成一个链表。

到这里有人发现了，Chunk里面不是一个ChunkItem的数组吗，干毛还用next串联成链表。

我想说的是，此中有深意呀。首先这样更加灵活，当然主要还和内存块的回收有关系，内存块回收回来时这个ChunkItem的位置不确定(在数组中的位置，处于哪一个Chunk)，所以直接将他串到链表头就可以了，没必要关心。所以这个Chunk只是刚初始化时是数组和链表等同，后面经过Alloc和Free后，链表和数组就不等同了。

mem字段比较简单了，由末班参数_Size控制其大小，一般也就是ChunkItem的大小(_Size若大于4)，ChunItem本身就是分配给用户使用的内存块，尺寸也是用户通过模板参数传入的_Size

 

---

 

    ChunkItem*          m_rootItem;
    std::list<Chunk *>  m_chunkList;

 

先说m_rootItem，其实就是指向可用的ChunkItem的指针，当Alloc申请时就可以非常简单的

    void *item = m_rootItem;
    m_rootItem = m_rootItem->next;
    memset(item, 0, _Size);
    return item;

 

m_chunkList，其实就是Chunk的记录list，主要用于内存释放，因为一个Chunk的内存块数固定，所以当内存块被申请完之后，就重新构造一个Chunk出来，然后将所有Chunk都保存到m_chunkList。这样就可以：

    for (Chunk *chunk : m_chunkList) {
        torch::HeapFree(chunk);
    }
    m_chunkList.clear();

 

 

下面看看比较关键的Alloc的实现：

template<size_t _Size>
void* FixedMemoryPool<_Size>::Alloc() {
    if (!m_rootItem) { // 当m_rootItem为空，代表没有可用的内存块了。 此时需要重新构造一个Chunk出来
        Chunk *chunk = this->MakeChunk(); // 构造一个Chunk出来，详情见下面
        if (!chunk) return nullptr;
        m_chunkList.push_back(chunk); // 将新构造的Chunk，记录到list中，否则会造成内存泄露
        m_rootItem = chunk->itemArray; // 让m_rootItem指向，新构造的Chunk的ItemArray中的第一个元素
    }
    // 直接将m_rootItem返回，然后将m_rootItem指向下一个
    void *item = m_rootItem;
    m_rootItem = m_rootItem->next;
    memset(item, 0, _Size);
    return item;
}

 

MakeChunk的细节：

template<size_t _Size>
typename FixedMemoryPool<_Size>::Chunk* FixedMemoryPool<_Size>::MakeChunk() {
    Chunk *chunk = (Chunk *)torch::HeapMalloc(sizeof(Chunk)); // 申请一个Chunk的内存块，struct Chunk {ChunkItem itemArray[_CHUNK_COUNT];};
    if (!chunk)
        return nullptr;
    for (int i = 0; i < _CHUNK_COUNT - 1; i++) { // 链表结构的初始化，将头一个指向下一个ChunkItem
        chunk->itemArray[i].next = &(chunk->itemArray[i+1]);
    }
    chunk->itemArray[_CHUNK_COUNT-1].next = nullptr; // 注意！！！最后一个一定要指向null，这样最后一个被申请后，m_rootItem就变为null了
    return chunk;
}

  

Free比较简单：

template<size_t _Size>
void FixedMemoryPool<_Size>::Free(void *mem) {
    if (!mem) return;
    ChunkItem *item = static_cast<ChunkItem*>(mem); 
    item->next = m_rootItem; // 直接将内存块放到链表头
    m_rootItem = item;
}

 

固定尺寸的内存池实现并不复杂，但是却非常有用，Alloc和Free都只是移动一个指针就可以完成内存的申请和释放(在不触发MakeChunk的情况下)，非常高效。

这里只介绍了内存池的核心算法，其他地方都比较简单。