memory_pool.hpp

#pragma once
#ifndef MEMORY_POOL_HPP
#define MEMORY_POOL_HPP

#include "os_memory.hpp"

#include <cstdio>
#include <cstdint>
#include <cstddef>
#include <cstdlib>
#include <cmath>
#include <cassert>

#include <iostream>
#include <algorithm>
#include <utility>
#include <type_traits>
#include <new>
#include <mutex>
#include <shared_mutex>
#include <memory>
#include <atomic>
#include <array>
#include <vector>

// ============================================================================
// Alignment constants  对齐常量（用于头部布局与内部对齐约束）
// ----------------------------------------------------------------------------
// CLASS_DEFAULT_ALIGNMENT: used for aligning header structures so that commonly
// accessed fields fall into separate cache lines and avoid false sharing under
// concurrent workloads.
// CLASS_DEFAULT_ALIGNMENT：用于对齐头部结构，使常访问字段尽量分布在不同缓存行，
// 从而在并发场景下降低伪共享的概率。
//
// DEFAULT_ALIGNMENT: platform default alignment for regular allocations.
// DEFAULT_ALIGNMENT：平台默认的分配对齐。
//
// MIN_ALLOWED_ALIGNMENT and MAX_ALLOWED_ALIGNMENT: the allocator validates
// requested alignment for sanity. A request larger than the default alignment
// takes the large‑alignment path implemented with AlignHeader below.
// MIN_ALLOWED_ALIGNMENT 和 MAX_ALLOWED_ALIGNMENT：用于校验来访对齐参数。当请求
// 对齐大于平台默认值时，将走下方带 AlignHeader 的大对齐路径。
// ============================================================================

static constexpr std::size_t CLASS_DEFAULT_ALIGNMENT = 64;
static constexpr std::size_t DEFAULT_ALIGNMENT       = alignof(void*);
static constexpr std::size_t MIN_ALLOWED_ALIGNMENT   = 2;
static constexpr std::size_t MAX_ALLOWED_ALIGNMENT   = 64 * 1024;

// ============================================================================
// Probe for one hundred twenty‑eight bit atomic compare and swap availability
// 探测硬件与编译器是否支持一百二十八位宽度的原子“比较并交换”操作
// ----------------------------------------------------------------------------
// Some platforms provide native one hundred twenty‑eight bit atomic operations
// which let us implement a Treiber stack head tagged with a monotonically
// increasing counter inside a single atomic variable. When this is not
// available we keep a lock‑guarded fallback. The public shape of the buckets
// does not change.
// 某些平台原生支持一百二十八位宽度的原子操作，这样可以把 Treiber 栈的指针与随
// 计数一起装入一个原子变量中。若平台不支持，则退化为配合互斥量的实现。对外“桶”
// 的形态保持不变。
// ============================================================================
#if (defined(__x86_64__) || defined(_M_X64) || defined(__aarch64__))
#  define SUPPORT_128BIT_CAS 1
#else
#  define SUPPORT_128BIT_CAS 0
#endif

// ============================================================================
// AlignHeader  大对齐辅助头部
// ----------------------------------------------------------------------------
// Purpose: record the original pointer and the total byte count when the user
// requests an alignment greater than the platform default. The aligned data
// pointer returned to the caller is advanced beyond this header, and this
// header allows correct release by passing the original pointer and size back
// to the operating system allocator.
// 目的：当用户请求的对齐大于平台默认对齐时，记录原始指针与实际申请的总字节数。
// 返回给调用者的对齐后数据指针位于该头部之后；释放时利用该头部中的原始指针与大小
// 正确归还给操作系统分配器。
// ============================================================================
static constexpr std::uint64_t ALIGN_SENTINEL = 0xDEADBEEFCAFEBABEull;
struct AlignHeader {
    std::uint64_t tag;  
    // Must equal ALIGN_SENTINEL. The value guards against mis‑interpretation
    // of nearby memory during deallocation.
    // 必须等于 ALIGN_SENTINEL，用于在释放时校验，避免把邻近内存误当作头部。

    void*         raw;  
    // The original pointer returned by the underlying operating system call.
    // 操作系统底层分配调用返回的原始指针。

    std::size_t   size; 
    // The total number of bytes actually allocated for the aligned request.
    // 针对对齐请求实际申请的总字节数。
};
static constexpr std::size_t ALIGN_HEADER_BYTES = sizeof(AlignHeader);

// ============================================================================
// NotAlignHeader  默认对齐路径的通用头部
// ----------------------------------------------------------------------------
// Purpose: tag the owner tier and preserve a pointer to the real header of the
// chosen tier. The user pointer returned by the public allocate function is the
// address directly after this structure. During release we look back, read this
// record, and dispatch to the correct tier deallocation routine.
// 目的：标记“所属层级”并保存指向真实头部的指针。对外的 allocate 返回值位于该结构
// 体之后；释放时回溯本记录并路由到正确的层级释放函数。
// ============================================================================
struct NotAlignHeader {
    std::uint32_t owner_type;  
    // 1 = Small, 2 = Medium, 3 = Large, 4 = Huge.  
    // 取值一至四分别代表小块、中等块、大块、超大块层。

    void*         raw;  
    // Pointer to the real tier‑specific header.
    // 指向该层真实头部结构体的指针。
};
static constexpr std::size_t NOT_ALIGN_HEADER_BYTES = sizeof(NotAlignHeader);

// ============================================================================
// SmallMemoryHeader  小块头部
// ----------------------------------------------------------------------------
// Layout note: only two fields are appended at the end for slab back‑reference
// and slot index. The original field order and semantics stay untouched to keep
// binary compatibility for any tooling that inspects the header.
// Paper: <The Slab Allocator: An Object-Caching Kernel Memory Allocator>
// URL: https://people.eecs.berkeley.edu/~kubitron/cs194-24/hand-outs/bonwick_slab.pdf
// 结构说明：仅在末尾追加两个字段以便回溯到所在 slab 以及槽位索引；原有字段顺序与
// 语义保持不变，从而维持可能依赖该头部布局的工具的二进制兼容性。
// slab（对象缓存式内存分配中的“板块”，由若干页组成，用于存放固定大小对象的槽位）
// ============================================================================
struct alignas(CLASS_DEFAULT_ALIGNMENT) SmallMemoryHeader {
    static constexpr std::uint32_t MAGIC = 0x534D4853;  
    // Magic value used for basic integrity checking during release.
    // 用于在释放时执行基本一致性校验的魔术字。

    std::uint32_t                   magic {MAGIC};
    // The magic is initialized on construction and cleared before a block
    // enters the thread local cache to reduce accidental double free mistakes.
    // 构造时写入该魔术字，进入当前线程缓存前会清零，以降低重复释放的风险。

    std::uint32_t                   bucket_index {0};   
    // Index of the size class bucket this header belongs to.
    // 该头部所属的尺寸桶索引。

    std::uint32_t                   block_size {0};     
    // User requested size for the block, including the NotAlignHeader stored
    // just in front of the user pointer for default alignment.
    // 用户请求的字节数；在默认对齐路径中，包含位于用户指针之前的 NotAlignHeader。

    std::atomic<bool>               is_free {true};     
    // Whether the block is considered free from the header point of view.
    // 从头部视角标记该块是否空闲。

    SmallMemoryHeader*              next {nullptr};     
    // Next pointer for the per‑thread cache list or the global stack list.
    // 作为线程缓存链或全局栈链的下一指针。

    std::uint8_t                    in_tls {0};         
    // Whether the block currently resides in the thread local cache list of
    // the current thread.
    // 该块是否当前处在本线程的线程本地缓存链中。

    void*                           parent_slab {nullptr}; 
    // Back‑pointer to the owning slab; typed as void pointer to keep the header
    // independent from the SmallMemoryManager internal type layout.
    // 回指其所属 slab；以 void 指针形式保存，使头部不依赖 SmallMemoryManager 的内部
    // 类型布局。

    std::uint32_t                   slot_index {0};        
    // Index of the slot inside the slab bitmap, from zero to sixty‑three.
    // 在 slab 位图内的槽位索引，范围零至六十三。

    void* data() { return reinterpret_cast<char*>(this) + sizeof(SmallMemoryHeader); }
    // Return the user data pointer which is located immediately after the
    // header.  返回紧随头部之后的用户数据指针。
};

// ============================================================================
// MediumMemoryHeader  中等块头部
// ============================================================================
struct alignas(CLASS_DEFAULT_ALIGNMENT) MediumMemoryHeader {
    static constexpr std::uint32_t MAGIC = 0x4D4D4853;  
    // Magic value for integrity checking during release.
    // 用于释放时的完整性校验。

    std::uint32_t                   magic {MAGIC};
    std::size_t                     block_size {0};
    std::atomic<bool>               is_free {true};
    MediumMemoryHeader*             next {nullptr};

    void* data() { return reinterpret_cast<char*>(this) + sizeof(MediumMemoryHeader); }
};

// ============================================================================
// LargeMemoryHeader and HugeMemoryHeader  大块与超大块头部
// ============================================================================
struct alignas(CLASS_DEFAULT_ALIGNMENT) LargeMemoryHeader {
    static constexpr std::uint32_t MAGIC = 0x4C4D4853;  
    std::uint32_t                   magic {MAGIC};
    std::size_t                     block_size {0};
    void* data() { return reinterpret_cast<char*>(this) + sizeof(LargeMemoryHeader); }
};

struct alignas(CLASS_DEFAULT_ALIGNMENT) HugeMemoryHeader {
    static constexpr std::uint32_t MAGIC = 0x484D4853;  
    std::uint32_t                   magic {MAGIC};
    std::size_t                     block_size {0};
    void* data() { return reinterpret_cast<char*>(this) + sizeof(HugeMemoryHeader); }
};

/*
================================================================================
SmallMemoryManager  小块内存管理器
--------------------------------------------------------------------------------
Public shape is preserved: each size class keeps a Treiber stack of header
pointers globally and a per‑thread cache list. Internally, slabs are used to
materialize many headers for the same bucket size. A slab contains up to sixty
four slots that correspond one to one with headers. The slab structure stores a
sixty‑four bit bitmap; bit value one means occupied, zero means free. Slot state
transitions are flipped with atomic compare and swap while global push and pop
remain the same Treiber operations on header pointers.

对外形态保持不变：每个尺寸类别在全局保留一个 Treiber 栈（头部指针）和一个
线程本地缓存链。内部使用 slab 在一次操作系统分配中生成同尺寸的多个对象头。
每个 slab 包含最多六十四个槽位，与头部一一对应；slab 保存一个六十四位的位图，
其中一表示占用，零表示空闲。槽位状态通过原子“比较并交换”翻转；而全局的入栈与
出栈仍是对头部指针执行 Treiber 算法，不改变外观与调用方式。
================================================================================
*/
struct alignas(CLASS_DEFAULT_ALIGNMENT) SmallMemoryManager {
    // ---------------------------------------------------------------------
    // Bucket sizing table  尺寸类别表
    // The first thirty‑two entries progress almost linearly for small sizes.
    // The latter thirty‑two grow roughly geometrically until one mebibyte.
    // 前三十二项近似线性地覆盖较小尺寸；后三十二项以几何方式增长直到一兆字节。
    // ---------------------------------------------------------------------
    static constexpr std::size_t BUCKET_COUNT = 64;
    static constexpr std::array<std::size_t, BUCKET_COUNT> BUCKET_SIZES = {
        8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128,
        136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256,
        336, 432, 560, 728, 944, 1224, 1584, 2048, 2656, 3448, 4472, 5800, 7520, 9744, 12640, 16384,
        21248, 27560, 35736, 46344, 60104, 77936, 101072, 131072, 169984, 220440, 285872, 370728, 480776, 623488, 808568, 1048576
    };

    // ---------------------------------------------------------------------
    // Thread local cache  线程本地缓存（行为不变）
    // Each bucket index has one singly linked list in the current thread.
    // 当前线程为每个尺寸桶维持一个单链表作为缓存。
    // ---------------------------------------------------------------------
    struct ThreadLocalCache {
        SmallMemoryHeader* buckets[BUCKET_COUNT] = { nullptr };
        std::size_t        deallocation_counter = 0;
    };
    static thread_local ThreadLocalCache thread_local_cache;

    // ---------------------------------------------------------------------
    // Global bucket head with Treiber stack  全局桶头（Treiber 栈）
    // To reduce the classic pointer value reappearance problem often called the
    // A‑B‑A problem in literature, a pointer tag is used on platforms that
    // support one hundred twenty‑eight bit atomic operations. When the wide
    // atomic form is not available, a lock‑guarded fallback path is used.
    // 为降低文献中常提及的指针值再现问题（常被称作 A‑B‑A 问题），在支持一百二十八位
    // 原子操作的平台上使用“带标签的指针”。若平台不支持宽原子，则退化为带互斥量的路径。
    // ---------------------------------------------------------------------
    struct alignas(CLASS_DEFAULT_ALIGNMENT) PointerTag { SmallMemoryHeader* pointer; std::uint64_t tag; };
    struct alignas(CLASS_DEFAULT_ALIGNMENT) GlobalBucket {
    #if SUPPORT_128BIT_CAS
        std::atomic<PointerTag> head;
    #else
        std::atomic<SmallMemoryHeader*> head;
        std::mutex                      mutex;
    #endif
    };
    std::array<GlobalBucket, BUCKET_COUNT> global_buckets;

    // ---------------------------------------------------------------------
    // SmallSlab  slab 描述结构（仅内部使用）
    // Holds the occupancy bitmap and derived sizing for iterating to the
    // header at a slot index.  保存位图与尺寸信息，以便按槽位索引定位到头部。
    // ---------------------------------------------------------------------
    struct alignas(CLASS_DEFAULT_ALIGNMENT) SmallSlab {
        std::atomic<std::uint64_t>  bitmap {0};   // one means occupied, zero means free  一为占用，零为空闲
        std::atomic<std::uint32_t>  used {0};     // number of occupied slots  已占用槽数量
        std::uint32_t               capacity {0}; // maximum slot count, up to sixty‑four  槽位容量（至多六十四）
        std::uint32_t               bucket_index {0};    // owning bucket index  所属尺寸桶索引
        std::size_t                 bucket_bytes {0};    // bytes per slot including default header  单槽字节（含默认头）
        std::size_t                 block_bytes {0};     // header plus payload bytes  头部加负载大小
        SmallMemoryHeader* header_at(std::uint32_t i) {
            char* base = reinterpret_cast<char*>(this) + sizeof(SmallSlab);
            return reinterpret_cast<SmallMemoryHeader*>(base + static_cast<std::size_t>(i) * block_bytes);
        }
    };

    // ---------------------------------------------------------------------
    // Tracking for operating system allocations that back slabs  记录 slab 的底层分配
    // These records are used for unified release during manager destruction.
    // 在管理器析构时统一释放，避免碎片与泄漏。
    // ---------------------------------------------------------------------
    std::mutex                                   chunk_mutex;
    std::vector<std::pair<void*, std::size_t>>   allocated_chunks;

    // ---------------------------------------------------------------------
    // Map a requested size to a bucket index  将请求尺寸映射为桶索引
    // A binary search over the constant table amortizes to a tiny constant time.
    // 通过对常量表执行二分查找实现，摊销时间开销极小。
    // ---------------------------------------------------------------------
    [[nodiscard]] static std::size_t calculate_bucket_index(std::size_t bytes) {
        std::size_t low = 0, high = BUCKET_COUNT - 1;
        while (low < high) {
            const std::size_t middle = (low + high) >> 1;
            if (bytes <= BUCKET_SIZES[middle]) high = middle; else low = middle + 1;
        }
        return low;
    }

    // ---------------------------------------------------------------------
    // Public operations  对外操作（名称与签名保持不变）
    // ---------------------------------------------------------------------
    void* allocate(std::size_t bytes, std::size_t alignment);
    void  deallocate(SmallMemoryHeader* header);

    // Flush the current thread cache back to the global buckets.  将本线程缓存批量并回全局栈。
    void  flush_thread_local_cache();

    // Release all slabs and clear global stacks.  释放所有 slab 并清空全局栈。
    void  release_resources();

private:
    // ---------------------------------------------------------------------
    // Internal helpers  内部辅助方法（仅小块层使用）
    // ---------------------------------------------------------------------
    SmallSlab*        make_new_slab(std::size_t bucket_index, std::size_t bucket_bytes, std::size_t alignment);
    static inline bool try_mark_slot_acquired(SmallSlab* slab, std::uint32_t slot_index);
    static inline bool try_mark_slot_released(SmallSlab* slab, std::uint32_t slot_index);
};

// ============================================================================
// MediumMemoryManager  中等块管理器（接口保持不变）
// ----------------------------------------------------------------------------
// Public interface remains the same. Internally a simple free list hierarchy is
// provided here; the complete body is placed in the source file.  对外接口不变，
// 具体实现见源文件。
// ============================================================================
struct MediumMemoryManager {
    static constexpr std::size_t MIN_BUCKET_BYTES_UNIT = 1 << 20;  // one mebibyte  一兆字节
    static constexpr int         LEVEL_COUNT = 10;                  // number of levels  等级数量
    static constexpr std::size_t CHUNK_SIZE  = 64ull << 20;         // default request size  默认请求大小

    struct alignas(CLASS_DEFAULT_ALIGNMENT) PointerTag { MediumMemoryHeader* pointer; std::uint64_t tag; };
    struct alignas(CLASS_DEFAULT_ALIGNMENT) FreeList {
        std::atomic<PointerTag> head;
        std::atomic_size_t      count {0};
    };

    std::atomic<uint16_t> free_list_level_mask {0};

    struct alignas(CLASS_DEFAULT_ALIGNMENT) MergeRequest { MediumMemoryHeader* block; int order; };
    static constexpr size_t MERGE_QUEUE_SIZE = 128;
    struct alignas(CLASS_DEFAULT_ALIGNMENT) MergeQueue {
        std::atomic<MergeRequest> requests[MERGE_QUEUE_SIZE];
        std::atomic<std::size_t>  head {0};
        std::atomic<std::size_t>  tail {0};
    };

    MergeQueue        merge_queue;
    std::atomic<bool> merge_worker_active {false};

    std::array<FreeList, LEVEL_COUNT>         free_lists;
    std::mutex                                 chunk_mutex;
    std::vector<std::pair<void*, std::size_t>> allocated_chunks;

    void* allocate(std::size_t bytes, std::size_t alignment);
    void  deallocate(MediumMemoryHeader* header);
    void  release_resources();

private:
    static int         order_from_size(std::size_t actual_byte_size) {
        std::size_t need = actual_byte_size;
        if (need < MIN_BUCKET_BYTES_UNIT) need = MIN_BUCKET_BYTES_UNIT;
        int order = 0; std::size_t cap = MIN_BUCKET_BYTES_UNIT;
        while (cap < need && order < LEVEL_COUNT - 1) { cap <<= 1; ++order; }
        return order;
    }
    static std::size_t size_from_order(int order) { return MIN_BUCKET_BYTES_UNIT << order; }

    void                prepare_block(MediumMemoryHeader* block, int order);
    void                push_block(MediumMemoryHeader* header, int order);
    MediumMemoryHeader* pop_block(int order);
    void                process_merge_queue();
    bool                try_remove_from_freelist(MediumMemoryHeader* header, int order);
    void                try_merge_buddy(MediumMemoryHeader* block, int order);
    MediumMemoryHeader* split_to_order(MediumMemoryHeader* block, int from_order, int to_order);
    MediumMemoryHeader* request_new_chunk(int min_order, std::size_t alignment);
};

// ============================================================================
// LargeMemoryManager  大块管理器（接口保持不变）
// ============================================================================
struct LargeMemoryManager {
    std::mutex                      tracking_mutex;      // protects active_blocks  保护 active_blocks
    std::vector<LargeMemoryHeader*> active_blocks;       // used only for leak checking  仅用于泄漏检测

    void* allocate(std::size_t bytes, std::size_t alignment);
    void  deallocate(LargeMemoryHeader* header);
    void  release_resources();
};

// ============================================================================
// HugeMemoryManager  超大块管理器（接口保持不变）
// ============================================================================
struct HugeMemoryManager {
    std::mutex                                 tracking_mutex;  // protects active_blocks  保护 active_blocks
    std::vector<std::pair<void*, std::size_t>> active_blocks;   // used for leak checking  用于泄漏检测

    void* allocate(std::size_t bytes, std::size_t alignment);
    void  deallocate(HugeMemoryHeader* header);
    void  release_resources();
};

// ============================================================================
// MemoryPool  统一入口：将请求路由到四个层级
// ----------------------------------------------------------------------------
// The surface area and naming do not change. A small request goes to the small
// tier, then medium, then large, and finally huge when the request crosses the
// huge threshold. Default alignment requests go through the NotAlignHeader; a
// request with a larger alignment is served by the AlignHeader path.
// 对外形态与命名不变：小请求走小块层，然后是中等层、再到大块层；当跨过超大块阈值
// 时进入超大层。默认对齐的请求经过 NotAlignHeader；大于默认对齐的请求经过 AlignHeader。
// ============================================================================
class MemoryPool {
private:
    static constexpr std::size_t SMALL_BLOCK_MAX_SIZE  = 1 * 1024 * 1024;
    static constexpr std::size_t MEDIUM_BLOCK_MAX_SIZE = 512 * 1024 * 1024;
    static constexpr std::size_t HUGE_BLOCK_THRESHOLD  = 1 * 1024 * 1024 * 1024;

    SmallMemoryManager  small_manager;
    MediumMemoryManager medium_manager;
    LargeMemoryManager  large_manager;
    HugeMemoryManager   huge_manager;

    std::atomic<bool>        is_destructing {false};
    static std::atomic<bool> construction_warning_shown;
    static std::atomic<int>  live_pools;  // only the last instance prints global counters  仅最后一个实例打印全局统计

public:
    MemoryPool();
    ~MemoryPool();

    void* allocate(std::size_t bytes, std::size_t alignment = MIN_ALLOWED_ALIGNMENT,
                   const char* source_file = nullptr, std::uint32_t source_line = 0, bool nothrow = false);
    // Allocate memory for the caller. When alignment exceeds the default,
    // the AlignHeader path is taken; otherwise the NotAlignHeader path is used.
    // 为调用者分配内存。若对齐超过默认值，走 AlignHeader 路径；否则走 NotAlignHeader 路径。

    void  deallocate(void* pointer);
    // Release a previously allocated pointer. The function looks back at the
    // preceding header to determine the owner tier before dispatching to the
    // correct release routine.  释放先前分配的指针；通过回溯前置头部确定所属层级并调用
    // 对应释放函数。

    void  flush_current_thread_cache();
    // Push the current thread cache back to the global bucket stacks in batch.
    // 将当前线程的小块缓存批量并回全局桶栈。
};

namespace os_memory::memory_pool {
    // Utility: return whether a value is a non‑zero power of two.  工具：判断是否为非零的二的幂。
    constexpr inline bool is_power_of_two(std::size_t value) noexcept {
#if !defined(_DEBUG)
        if (value != 1 || (value & (value - 1)) != 0) { return false; }
#else
        assert(value != 1 && (value & (value - 1)) == 0 && "ensure_power_of_two: value must be non-zero power of two");
#endif
        return true;
    }
}

// ============================================================================
// End of header  结束
// ============================================================================
#endif // MEMORY_POOL_HPP