mem_pool.h
代码语言:c复制#pragma once
#include <memory>
#include <mutex>
#include <vector>
#include <atomic>
// 利用自定义智能指针析构函数实现自动回收的内存池
class MemoryPool {
public:
MemoryPool()
{
}
explicit MemoryPool(size_t count)
{
// 池子最少1000个
maxCount_ = count > 1000 ? count : 1000;
}
~MemoryPool()
{
std::unique_lock<std::mutex> lock(lock_);
for (auto p : freeConn_) {
free(p);
}
freeConn_.clear();
}
std::shared_ptr<unsigned char> Malloc(size_t size)
{
unsigned char* ptr = nullptr;
// 如果大于1024字节就直接malloc内存
if (size > 1024) {
ptr = (unsigned char*)malloc(size);
if (ptr == nullptr) {
return nullptr;
}
// 自定义智能指针析构
return std::shared_ptr<unsigned char>(ptr, [this](unsigned char* p) {
free(p);
});
}
std::unique_lock<std::mutex> lock(lock_);
if (freeConn_.empty()) {
if (curCount_ < maxCount_) {
// 内存池大小为1024字节
ptr = (unsigned char*)malloc(1024);
if (ptr == nullptr) {
return nullptr;
}
curCount_;
} else {
return nullptr;
}
} else {
ptr = freeConn_.back();
freeConn_.pop_back();
}
// 自定义智能指针析构
return std::shared_ptr<unsigned char>(ptr, [this](unsigned char* p) {
// give back to pool
std::lock_guard<std::mutex> lock(lock_);
freeConn_.push_back(p);
});
}
size_t MaxCount() const
{
return maxCount_;
}
size_t CurCount() const
{
return curCount_;
}
size_t FreeCount() const
{
return freeConn_.size();
}
private:
size_t maxCount_ = 1000;
size_t curCount_ = 0;
std::vector<unsigned char*> freeConn_;
std::mutex lock_;
};
gtest测试代码
代码语言:c复制#include <gtest/gtest.h>
#include "mem_pool.h"
TEST(MemPoolTest, TestAlloc1)
{
MemoryPool pool;
for (int i=0;i<2000; i) {
pool.Malloc(100);
}
// 因为使用完立即就被回收了所以malloc 2000次只需要一个池子
EXPECT_EQ(1, pool.CurCount());
EXPECT_EQ(1, pool.FreeCount());
}
TEST(MemPoolTest, TestAlloc2)
{
MemoryPool pool;
std::vector<std::shared_ptr<unsigned char>> vecs;
for (int i=0;i<2000; i) {
auto p = pool.Malloc(100);
if (p != nullptr) {
vecs.emplace_back(p);
}
}
EXPECT_TRUE(nullptr == pool.Malloc(1024));
EXPECT_TRUE(nullptr != pool.Malloc(1025));
EXPECT_EQ(1000, pool.CurCount());
EXPECT_EQ(0, pool.FreeCount());
EXPECT_EQ(1000, vecs.size());
vecs.pop_back();
EXPECT_EQ(1, pool.FreeCount());
vecs.pop_back();
EXPECT_EQ(2, pool.FreeCount());
vecs.clear();
EXPECT_EQ(1000, pool.CurCount());
EXPECT_EQ(1000, pool.FreeCount());
}
TEST(MemPoolTest, TestAlloc3)
{
MemoryPool pool;
unsigned char* p = nullptr;
{
auto ptr = pool.Malloc(100);
p = ptr.get();
}
{
auto ptr = pool.Malloc(200);
EXPECT_TRUE(p == ptr.get());
}
EXPECT_EQ(1, pool.CurCount());
EXPECT_EQ(1, pool.FreeCount());
}