memory.cpp

#include "memory.hpp"

// 每个线程的自由列表头指针
thread_local FixedSizeMemory::Node* thread_local_free_list = nullptr;

// 每个线程当前缓存的块数量
thread_local std::size_t thread_local_count = 0;

// 控制是否强制绕过本地缓存，直接归还给 Page
thread_local bool tl_bypass_local_cache = false;

// 线程缓存看门人 (RAII)
// 作用：记录当前线程归属的内存池，并在线程退出时自动归还缓存
struct ThreadCacheJanitor {
    FixedSizeMemory* pool = nullptr;

    ~ThreadCacheJanitor() {
        // 当线程退出时，如果缓存不为空且关联了内存池，则进行归还
        if (thread_local_free_list && pool) {
            pool->flush_thread_cache();
        }
    }
};

// 每个线程持有一个看门人实例
thread_local ThreadCacheJanitor tl_janitor;

// 实现 flush_thread_cache
void FixedSizeMemory::flush_thread_cache() {
    if (!thread_local_free_list) return;

    // 直接复用 drain_local_list 将整个链表归还
    // 这里会将 thread_local_free_list 全部清空
    drain_local_list(thread_local_free_list, thread_local_count);

    // 重置本地缓存状态
    thread_local_free_list = nullptr;
    thread_local_count = 0;
}


// 构造
FixedSizeMemory::FixedSizeMemory(std::size_t blocks_size, std::size_t blocks_per_page)
    : blocks_size(blocks_size), blocks_per_page(blocks_per_page) {
    adjusted_block_sz_ = adjusted_block_size(blocks_size);
}

// 析构
FixedSizeMemory::~FixedSizeMemory() {
    shutdown_.store(true, std::memory_order_release);

    while (active_ops_.load(std::memory_order_acquire) > 0) {
        std::this_thread::yield();
    }
    // 对象销毁时，清理当前线程残留的本地缓存
    // 防止下一个内存池实例复用这些指向已释放内存的野指针
    if (thread_local_free_list != nullptr) {
        thread_local_free_list = nullptr;
        thread_local_count = 0;
    }
    pages_.clear(); // 自动析构所有 Page（释放 base）
}

// 分配块
// 分配块 (使用 TLFL)
void* FixedSizeMemory::allocate() {
    if (shutdown_.load(std::memory_order_acquire)) {
        return nullptr;
    }
    // 记录当前线程正在使用的内存池实例
    // 这样当线程退出时，Janitor 知道要把内存还给谁
    if (tl_janitor.pool == nullptr) {
        tl_janitor.pool = this;
    }
    // 1. 快速路径：从线程本地缓存获取 (无锁)
    if (thread_local_free_list != nullptr) {
        Node* p = thread_local_free_list;
        thread_local_free_list = thread_local_free_list->next;
        thread_local_count--;
        // 由于是从本地缓存分配，active_ops_的增减逻辑可以绕过
        // 但为了保持原有的统计逻辑，我们仍需维护 active_ops_
        active_ops_.fetch_add(1, std::memory_order_relaxed);
        struct DecrementGuard {
            std::atomic<size_t>& counter;
            ~DecrementGuard() { counter.fetch_sub(1, std::memory_order_release); }
        } guard{active_ops_};
        
        return p;
    }
    
    // 2. 慢速路径：本地缓存为空，需要批量获取
    active_ops_.fetch_add(1, std::memory_order_relaxed);
    struct DecrementGuard {
        std::atomic<size_t>& counter;
        ~DecrementGuard() { counter.fetch_sub(1, std::memory_order_release); }
    } guard{active_ops_};

    try {
        refill_local_list(REFILL_COUNT);
    } catch (const std::bad_alloc& e) {
        // refill 失败 (如 expand 失败)
        return nullptr;
    }

    // 3. 再次从本地缓存获取
    if (thread_local_free_list != nullptr) {
        Node* p = thread_local_free_list;
        thread_local_free_list = thread_local_free_list->next;
        thread_local_count--;
        return p;
    }
    
    return nullptr;
}

// 取回块
// 取回块 (使用 TLFL)
void FixedSizeMemory::deallocate(void* p) noexcept {
    if (!p) return;

    // 同上，确保 Janitor 知道归还目标
    if (tl_janitor.pool == nullptr) {
        tl_janitor.pool = this;
    }

    Node* node = static_cast<Node*>(p);

    // 1. 快速路径：归还给线程本地缓存 (无锁)
    if (thread_local_count < LOCAL_CACHE_SIZE) {
        node->next = thread_local_free_list;
        thread_local_free_list = node;
        thread_local_count++;
        
        // 同样，active_ops_的增减逻辑可以绕过
        active_ops_.fetch_add(1, std::memory_order_relaxed);
        struct DecrementGuard {
            std::atomic<size_t>& counter;
            ~DecrementGuard() { counter.fetch_sub(1, std::memory_order_release); }
        } op_guard{active_ops_};
        
        return;
    }

    // 2. 慢速路径：本地缓存已满，需要批量归还
    active_ops_.fetch_add(1, std::memory_order_relaxed);
    struct DecrementGuard {
        std::atomic<size_t>& counter;
        ~DecrementGuard() { counter.fetch_sub(1, std::memory_order_release); }
    } op_guard{active_ops_};

    // 将待回收的 p 暂存到链表头部，此时 thread_local_count == LOCAL_CACHE_SIZE + 1
    node->next = thread_local_free_list;
    thread_local_free_list = node;
    thread_local_count++; 

    const std::size_t drain_num = LOCAL_CACHE_SIZE / 2; 

    // 找到要归还的块的尾部
    Node* current = thread_local_free_list;
    for (std::size_t i = 0; i < drain_num - 1; ++i) {
        if (current->next == nullptr) break; 
        current = current->next;
    }
    
    // 分割链表
    Node* to_drain_head = thread_local_free_list;
    Node* new_local_head = current->next;
    current->next = nullptr; // 结束归还的子链表

    // 更新本地缓存
    thread_local_free_list = new_local_head;
    thread_local_count -= drain_num;

    // 3. 批量归还到 Page
    drain_local_list(to_drain_head, drain_num);
}

// 获取内存池信息
std::size_t FixedSizeMemory::get_blocks_size() const {
    return blocks_size;
}

std::size_t FixedSizeMemory::get_blocks_per_page() const {
    return blocks_per_page;
}

std::size_t FixedSizeMemory::page_count() const noexcept {
    std::lock_guard<std::mutex> g(big_mtx_);
    return pages_.size();
}

std::size_t FixedSizeMemory::block_size() const noexcept {
    return adjusted_block_sz_;
}

std::size_t FixedSizeMemory::free_count() const noexcept {
    std::size_t total = 0;
    std::lock_guard<std::mutex> container_guard(big_mtx_);
    for (const auto& spg : pages_) {
        Page* pg = spg.get();
        std::lock_guard<std::mutex> page_lock(pg->mtx);
        for (Node* cur = pg->free_list; cur; cur = cur->next) ++total;
    }
    return total;
}

std::size_t FixedSizeMemory::allocated_count() const {
    return allocated.load(std::memory_order_relaxed);
}

// Page 析构函数
FixedSizeMemory::Page::~Page() {
    if (base) {
        ::operator delete(base);
    }
}

// 惰性页回收
void FixedSizeMemory::decommission_pages() {
    std::lock_guard<std::mutex> guard(big_mtx_);

    const std::size_t total_pages = pages_.size();
    const std::size_t empty_count = empty_pages_.size();

    if (empty_count == 0 || (double)empty_count / total_pages < DECOMMISSION_THRESHOLD) {
        return;
    }

    const std::size_t pages_to_keep = 1;
    const std::size_t pages_to_remove = empty_count - std::min(empty_count, pages_to_keep);
    if (pages_to_remove == 0) return;

    for (std::size_t i = 0; i < pages_to_remove; ++i) {
        Page* pg_to_remove = empty_pages_.back();
        empty_pages_.pop_back();

        auto it = std::remove_if(pages_.begin(), pages_.end(),
            [&](const std::shared_ptr<Page>& spg) {
                return spg.get() == pg_to_remove;
            });

        assert(it != pages_.end() && "Decommissioning page not found in pages_ container!");
        pages_.erase(it, pages_.end());

        pg_to_remove->is_empty_candidate = false;
    }
}

// 当内存页都满的时候，扩一页
FixedSizeMemory::Page* FixedSizeMemory::expand() {
    const std::size_t block_sz = adjusted_block_sz_;
    const std::size_t page_sz = block_sz * blocks_per_page;
    void* raw = ::operator new(page_sz);
    Page* pg = nullptr;

    try {
        std::lock_guard<std::mutex> g(big_mtx_);
        pages_.emplace_back(std::make_shared<Page>());
        pg = pages_.back().get();
        pg->base = static_cast<char*>(raw);
        pg->capacity = blocks_per_page;
        pg->inuse = 0;
        pg->free_list = nullptr;
    } catch (...) {
        ::operator delete(raw);
        throw;
    }

    Node* head = nullptr;
    char* base = pg->base;
    for (std::size_t i = 0; i < blocks_per_page; ++i) {
        Node* node = reinterpret_cast<Node*>(base + i * block_sz);
        node->next = head;
        head = node;
    }

    std::lock_guard<std::mutex> page_lock(pg->mtx);
    pg->free_list = head;
    return pg;
}


// 将线程本地缓存中 N 个块批量归还给 Page
void FixedSizeMemory::drain_local_list(Node* head, std::size_t count) {
    if (count == 0 || head == nullptr) return;

    const std::size_t block_sz = adjusted_block_sz_;
    std::shared_ptr<Page> target_spg = nullptr; // ⬅️ 新增/修改

    // 1. 查找块所属的 Page (需要 big_mtx_ 保护 pages_)
    {
        std::lock_guard<std::mutex> guard(big_mtx_);
        char* ptr = reinterpret_cast<char*>(head);
        
        for (auto& spg : pages_) {
            Page* pg = spg.get();
            // 检查块地址是否在 Page 的范围内
            if (ptr >= pg->base && ptr < pg->base + pg->capacity * block_sz) {
                target_spg = spg; // ⬅️ 获取 shared_ptr 的所有权，防止 Page 被回收
                break;
            }
        }
    } // 释放 big_mtx_

    if (!target_spg) { // ⬅️ 检查 target_spg
        std::cerr << "错误: 尝试释放不属于此内存池的内存块 (TLFL)!" << std::endl;
        assert(false && "pointer not from this pool");
        return;
    }
    
    Page* target_pg = target_spg.get(); // ⬅️ 现在安全地获取裸指针

    // 2. 批量归还块 (锁定 Page)
    bool page_became_empty = false;
    {
        std::lock_guard<std::mutex> page_lock(target_pg->mtx);
        
        // 将归还的链表 (head) 加到 Page 的自由列表头部 (LIFO)
        Node* tail = head;
        while (tail->next != nullptr) {
            tail = tail->next;
        }
        tail->next = target_pg->free_list;
        target_pg->free_list = head;

        // 3. 更新 Page 状态
        target_pg->inuse -= count;
        
        // 更新全局统计
        allocated.fetch_sub(count, std::memory_order_relaxed);
        
        // 4. 检查是否满足惰性回收条件
        if (target_pg->inuse == 0 && !target_pg->is_empty_candidate) {
            target_pg->is_empty_candidate = true;
            page_became_empty = true;
        }
    } // 释放 page_lock

    // 5. 将空闲页加入空闲列表 (需要 big_mtx_)
    if (page_became_empty) {
        std::lock_guard<std::mutex> guard(big_mtx_);
        // 再次检查，防止在释放 big_mtx_ 期间，其他线程已将此页回收或标记
        if (target_pg->is_empty_candidate) {
            empty_pages_.push_back(target_pg);
        }
    }
}


// 从一个Page中批量获取 N 个块给线程本地缓存
void FixedSizeMemory::refill_local_list(std::size_t count) {
    // 使用 while(true) 循环进行重试
    // 如果发生竞态条件（扩容后被别人抢了），不能直接 return，必须重试
    while (true) {
        std::shared_ptr<Page> target_page = nullptr;
        bool locked_by_search = false;
        
        // 1. 查找有空闲块的 Page (需要 big_mtx_ 保护 pages_ 数组)
        {
            std::lock_guard<std::mutex> container_guard(big_mtx_);

            for (size_t i = 0; i < pages_.size(); ++i) {
                Page* pg = pages_[i].get();
                // 尝试锁定 Page
                if (pg->mtx.try_lock()) {
                    // 如果锁定成功，且 free_list 不为空
                    if (pg->free_list != nullptr) {
                        target_page = pages_[i];
                        locked_by_search = true; // 已经持有锁了
                        break; 
                    }
                    pg->mtx.unlock(); // 不符合条件，解锁
                }
            }
        } // 释放 big_mtx_

        // 2. 如果没有找到，则扩展页
        if (target_page == nullptr) {
            // expand() 可能会抛出 bad_alloc，这符合预期
            expand(); 
            
            // [关键逻辑变革]：
            // 扩容后，不要试图在这里手动去拿 pages_.back()。
            // 因为在释放 big_mtx_ 后，pages_.back() 可能已经被其他线程拿走了。
            // 最安全的做法是：直接 continue，进入下一次循环。
            // 下一次循环的 "步骤1" 会自动发现这个新创建的空闲页（或者其他页）。
            continue; 
        }
        
        // 3. 接管锁 (Locked Page)
        // 此时我们一定持有了 target_page 的锁 (由步骤1的 try_lock 获得)
        std::unique_lock<std::mutex> page_lock(target_page->mtx, std::adopt_lock);
        
        // 4. [再次检查]：竞态条件防御
        // 虽然在步骤1检查了 free_list != nullptr，但为了代码的健壮性（以及防止未来逻辑变更引入bug）
        // 我们再次检查。如果此时为空，说明发生了极其罕见的情况，或者逻辑错误。
        if (target_page->free_list == nullptr) {
            // 这种情况应该通过 continue 解决，而不是崩溃或返回失败
            continue; 
        }

        // 5. 开始传输节点
        Node* head = target_page->free_list;
        Node* current = head;
        std::size_t actual_count = 1;

        // 安全遍历：确保不会超出链表长度
        for (std::size_t i = 0; i < count - 1; ++i) {
            if (current->next == nullptr) break; 
            current = current->next;
            actual_count++;
        }

        // 分割链表
        Node* new_page_list = current->next;
        current->next = nullptr; // 切断，current 变为本地链表的尾部
        
        // 更新 Page 状态
        target_page->free_list = new_page_list;
        target_page->inuse += actual_count;
        
        // 更新全局统计
        allocated.fetch_add(actual_count, std::memory_order_relaxed);

        // 6. 更新线程本地状态 (无锁操作)
        current->next = thread_local_free_list;
        thread_local_free_list = head;
        thread_local_count += actual_count;

        // 成功填充，退出函数
        return;
    }
}