Memory_Pool.cpp

#include <iostream>
#include <vector>
#include <array>
#include <cstddef>
#include <cassert>
#include <atomic>
#include <mutex>
#include <new>
#include <string>
#include <optional>
#include <functional>

// ============================================================================
// 1. 对齐工具
// ============================================================================
inline std::size_t align_up(std::size_t size, std::size_t alignment) noexcept {
    return (size + alignment - 1) & ~(alignment - 1);
}

// ============================================================================
// 2. LockFreeStack — 无锁 Treiber Stack (节点放内存块头部)
// ============================================================================
struct FreeNode {
    FreeNode* next;
};

class LockFreeStack {
public:
    void Push(void* ptr) noexcept {
        auto* node = static_cast<FreeNode*>(ptr);
        FreeNode* head = head_.load(std::memory_order_relaxed);
        do {
            node->next = head;
        } while (!head_.compare_exchange_weak(
            head, node, std::memory_order_release, std::memory_order_relaxed));
    }

    void* Pop() noexcept {
        FreeNode* head = head_.load(std::memory_order_acquire);
        while (head) {
            FreeNode* next = head->next;
            if (head_.compare_exchange_weak(
                    head, next, std::memory_order_acquire, std::memory_order_relaxed)) {
                return head;
            }
        }
        return nullptr;
    }

    bool Empty() const noexcept { return head_.load(std::memory_order_acquire) == nullptr; }

private:
    std::atomic<FreeNode*> head_{nullptr};
};

// ============================================================================
// 3. FixedSizeMemoryPool — 单尺寸内存池
// ============================================================================
class FixedSizeMemoryPool {
public:
    FixedSizeMemoryPool(std::size_t block_size, std::size_t blocks_per_chunk = 256)
        : block_size_(std::max(block_size, sizeof(FreeNode)))
        , blocks_per_chunk_(blocks_per_chunk) {}

    ~FixedSizeMemoryPool() {
        for (char* page : chunks_) { ::operator delete[](page); }
    }

    void* Allocate() {
        void* ptr = free_list_.Pop();
        if (!ptr) {
            Expand();
            ptr = free_list_.Pop();
        }
        return ptr;
    }

    void Deallocate(void* ptr) noexcept {
        if (ptr) free_list_.Push(ptr);
    }

    std::size_t BlockSize() const noexcept { return block_size_; }
    std::size_t ChunkCount() const noexcept { return chunks_.size(); }

private:
    void Expand() {
        std::size_t bytes = block_size_ * blocks_per_chunk_;
        char* page = static_cast<char*>(::operator new[](bytes));
        {
            std::lock_guard lock(mtx_);
            chunks_.push_back(page);
        }
        for (std::size_t i = 0; i < blocks_per_chunk_; ++i) {
            free_list_.Push(page + i * block_size_);
        }
    }

    LockFreeStack       free_list_;
    std::vector<char*>  chunks_;
    std::size_t         block_size_;
    std::size_t         blocks_per_chunk_;
    std::mutex          mtx_;
};

// ============================================================================
// 4. SlabAllocator — 多尺寸 Slab 分配器 (libumem 风格)
// ============================================================================
class SlabAllocator {
    struct Slab {
        std::size_t         block_size;
        std::size_t         blocks_per_chunk;
        std::vector<char*>  chunks;
        LockFreeStack       free_list;
        std::mutex          mtx;
    };

    static constexpr std::size_t kMaxSlabSize = 4096;
    static constexpr std::size_t kSlabCount   = 8;

    static constexpr std::array<std::size_t, kSlabCount> kSizes = {
        16, 32, 64, 128, 256, 512, 1024, 4096
    };

public:
    SlabAllocator() {
        for (std::size_t i = 0; i < kSlabCount; ++i) {
            slabs_[i].block_size       = kSizes[i];
            slabs_[i].blocks_per_chunk = std::max<std::size_t>(1, kMaxSlabSize / kSizes[i]);
        }
    }

    void* Allocate(std::size_t size) noexcept {
        auto idx = IndexFor(size);
        if (!idx) return nullptr;
        auto& slab = slabs_[*idx];

        void* ptr = slab.free_list.Pop();
        if (!ptr) {
            std::lock_guard lock(slab.mtx);
            // 双重检查
            ptr = slab.free_list.Pop();
            if (!ptr) {
                std::size_t bytes = slab.block_size * slab.blocks_per_chunk;
                char* page = static_cast<char*>(::operator new[](bytes));
                slab.chunks.push_back(page);
                for (std::size_t i = 0; i < slab.blocks_per_chunk; ++i) {
                    slab.free_list.Push(page + i * slab.block_size);
                }
                ptr = slab.free_list.Pop();
            }
        }
        return ptr;
    }

    void Deallocate(void* ptr, std::size_t size) noexcept {
        if (!ptr) return;
        auto idx = IndexFor(size);
        if (!idx) return;
        slabs_[*idx].free_list.Push(ptr);
    }

private:
    std::optional<std::size_t> IndexFor(std::size_t size) const noexcept {
        for (std::size_t i = 0; i < kSlabCount; ++i) {
            if (size <= kSizes[i]) return i;
        }
        return std::nullopt;
    }

    std::array<Slab, kSlabCount> slabs_;
};

// ============================================================================
// 5. 线程局部缓存 (TLS batch cache) — 减少跨核 CAS 竞争
// ============================================================================
class TlsCache {
    static constexpr std::size_t kBatchSize = 16;
public:
    // parent_pool 是 "尺寸已经选好" 的具体池，如 FixedSizeMemoryPool 或 SlabAllocator
    using AllocFn = std::function<void*()>;
    using FreeFn  = std::function<void(void*)>;

    TlsCache(AllocFn alloc, FreeFn free)
        : alloc_(std::move(alloc)), free_(std::move(free)) {}

    void* Get() {
        if (cache_.empty()) Refill();
        void* p = cache_.back();
        cache_.pop_back();
        return p;
    }

    void Put(void* ptr) {
        if (cache_.size() >= kBatchSize) Flush();
        cache_.push_back(ptr);
    }

    ~TlsCache() { Flush(); }

private:
    void Refill() {
        for (std::size_t i = 0; i < kBatchSize; ++i) {
            void* p = alloc_();
            if (!p) break;
            cache_.push_back(p);
        }
    }

    void Flush() {
        for (void* p : cache_) free_(p);
        cache_.clear();
    }

    AllocFn alloc_;
    FreeFn  free_;
    std::vector<void*> cache_;
};

// ============================================================================
// 6. 演示
// ============================================================================
struct Particle {
    float x, y, z;
    int   life;
    void  update() { ++life; }
};

// 全局内存池用于 Particle
static FixedSizeMemoryPool gParticlePool(sizeof(Particle), 2000);

// 仅作为演示，不建议在生产中全局替换 operator new
// 实际项目中请使用工厂函数或自定义 allocator
// 工厂函数代替全局 operator new，更安全清晰
inline Particle* MakeParticle(float x, float y, float z, int life) {
    auto* ptr = static_cast<Particle*>(gParticlePool.Allocate());
    new (ptr) Particle{x, y, z, life};
    return ptr;
}

inline void ReleaseParticle(Particle* p) {
    if (p) {
        p->~Particle();
        gParticlePool.Deallocate(p);
    }
}

int main() {
    // ---- 演示 FixedSizeMemoryPool ----
    {
        std::cout << "=== FixedSizeMemoryPool (Particle) ===" << std::endl;
        std::vector<Particle*> particles;
        for (int i = 0; i < 100'000; ++i) {
            particles.push_back(MakeParticle(1.0f, 2.0f, 3.0f, i));
        }
        for (auto* p : particles) ReleaseParticle(p);
        std::cout << "  Block size: " << gParticlePool.BlockSize()
                  << " bytes, chunks: " << gParticlePool.ChunkCount() << std::endl;
    }

    // ---- 演示 SlabAllocator ----
    {
        std::cout << "\n=== SlabAllocator ===" << std::endl;
        SlabAllocator slab;
        std::vector<std::pair<void*, std::size_t>> allocs;
        allocs.emplace_back(slab.Allocate(12), 12);   // → 16B slab
        allocs.emplace_back(slab.Allocate(120), 120); // → 128B slab
        allocs.emplace_back(slab.Allocate(512), 512); // → 512B slab
        allocs.emplace_back(slab.Allocate(1024), 1024);//→ 1024B slab
        allocs.emplace_back(slab.Allocate(4000), 4000);//→ 4096B slab
        for (auto [p, s] : allocs) {
            slab.Deallocate(p, s);
        }
        std::cout << "  SlabAllocator alloc/free OK" << std::endl;
    }

    // ---- 演示 TlsCache ----
    {
        std::cout << "\n=== TlsCache (batch) ===" << std::endl;
        FixedSizeMemoryPool smallPool(32, 64);
        TlsCache cache(
            [&]() { return smallPool.Allocate(); },
            [&](void* p) { smallPool.Deallocate(p); });

        std::vector<void*> ptrs;
        for (int i = 0; i < 100; ++i) ptrs.push_back(cache.Get());
        for (void* p : ptrs) cache.Put(p);
        std::cout << "  TlsCache batch alloc/free OK" << std::endl;
    }

    std::cout << "\nAll memory pools OK." << std::endl;
    return 0;
}