TCMallocator.cpp

#include "TCMallocator.h"
#include <sys/mman.h>
#include <cstdlib>

TCMallocator::TCMallocator() {
    init_size_map();

    for (std::size_t i = 0; i < kNumClasses; ++i) {
        central_caches_[i].non_empty_list.next = &central_caches_[i].non_empty_list;
        central_caches_[i].non_empty_list.prev = &central_caches_[i].non_empty_list;
        central_caches_[i].empty_list.next = &central_caches_[i].empty_list;
        central_caches_[i].empty_list.prev = &central_caches_[i].empty_list;
    }

    for (std::size_t i = 0; i < kMaxCachedPages; ++i) {
        free_span_lists_[i].next = &free_span_lists_[i];
        free_span_lists_[i].prev = &free_span_lists_[i];
    }
}

TCMallocator::~TCMallocator() = default;

void TCMallocator::init_size_map() {
    std::size_t current_size = 8;
    for (std::size_t i = 0; i < kNumClasses; ++i) {
        if (i > 0) {
            if (current_size <= 128) current_size += 8;
            else if (current_size <= 1024) current_size += 32;
            else if (current_size <= 8192) current_size += 128;
            else current_size += 1024;
        }
        if (current_size > kMaxSmallSize) current_size = kMaxSmallSize;
        class_to_size_map_[i] = current_size;
    }
    
    size_to_class_map_[0] = 0;
    std::size_t sc = 0;
    for (std::size_t b = 1; b <= kMaxSmallSize; ++b) {
        while (sc < kNumClasses && class_to_size_map_[sc] < b) {
            sc++;
        }
        size_to_class_map_[b] = (sc < kNumClasses ? sc : kNumClasses - 1);
    }
}

// FastPageMap 实现
TCMallocator::FastPageMap::FastPageMap() {
    std::memset(roots_, 0, sizeof(roots_));
}

TCMallocator::FastPageMap::~FastPageMap() {
    for (int i = 0; i < (1 << kBits1); ++i) {
        if (roots_[i]) {
            delete roots_[i];
        }
    }
}

TCMallocator::Span* TCMallocator::FastPageMap::get(PageID number) const {
    std::size_t i1 = number >> kBits2;
    std::size_t i2 = number & (kLeafSize - 1);
    if (i1 >= (1 << kBits1)) return nullptr;
    Leaf* leaf = roots_[i1];
    if (!leaf) return nullptr;
    return leaf->values[i2];
}

void TCMallocator::FastPageMap::set(PageID number, Span* span) {
    std::size_t i1 = number >> kBits2;
    std::size_t i2 = number & (kLeafSize - 1);
    if (i1 >= (1 << kBits1)) return;

    Leaf* leaf = roots_[i1];
    if (!leaf) {
        std::lock_guard<std::mutex> lock(mtx_);
        leaf = roots_[i1];
        if (!leaf) {
            leaf = new Leaf();
            roots_[i1] = leaf;
        }
    }
    leaf->values[i2] = span;
}

void TCMallocator::FastPageMap::remove_range(PageID start, std::size_t num_pages) {
    for (std::size_t i = 0; i < num_pages; ++i) {
        std::size_t number = start + i;
        std::size_t i1 = number >> kBits2;
        std::size_t i2 = number & (kLeafSize - 1);
        if (i1 < (1 << kBits1) && roots_[i1]) {
            roots_[i1]->values[i2] = nullptr;
        }
    }
}

// 无锁化 ThreadCache 获取
TCMallocator::ThreadCache* TCMallocator::get_thread_cache() {
    thread_local ThreadCache* local_tc = nullptr;
    if (__builtin_expect(local_tc != nullptr, 1)) {
        return local_tc;
    }
    local_tc = new ThreadCache();
    std::memset(local_tc, 0, sizeof(ThreadCache));
    return local_tc;
}

TCMallocator::Span* TCMallocator::map_from_system(std::size_t num_pages) {
    std::size_t bytes = num_pages * kPageSize;
    void* ptr = mmap(nullptr, bytes, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (ptr == MAP_FAILED) return nullptr;

    uintptr_t addr = reinterpret_cast<uintptr_t>(ptr);
    Span* span = new Span();
    span->start_address = addr;
    span->num_pages = num_pages;
    span->size_class = 0;
    span->free_list = nullptr;
    span->refcount = 0;
    span->is_used = true;
    return span;
}

// 核心优化：Page Heap 内存池化缓存复用
TCMallocator::Span* TCMallocator::allocate_pages(std::size_t num_pages) {
    std::lock_guard<std::mutex> lock(page_heap_mtx_);

    // 1. 尝试从 Page Heap 的空闲链表中寻找合适的 Span
    if (num_pages < kMaxCachedPages) {
        Span* list_head = &free_span_lists_[num_pages];
        if (list_head->next != list_head) {
            Span* span = list_head->next;
            // 从链表中摘除
            span->prev->next = span->next;
            span->next->prev = span->prev;
            
            PageID start_page_id = span->start_address >> kPageShift;
            for (std::size_t i = 0; i < num_pages; ++i) {
                metadata_map_.set(start_page_id + i, span);
            }
            return span;
        }
    }

    // 2. 缓存没有，向系统申请
    Span* span = map_from_system(std::max(num_pages, std::size_t(32))); // 预分配大块
    if (!span) return nullptr;

    PageID start_page_id = span->start_address >> kPageShift;
    for (std::size_t i = 0; i < num_pages; ++i) {
        metadata_map_.set(start_page_id + i, span);
    }
    return span;
}

void TCMallocator::deallocate_pages(Span* span) {
    std::lock_guard<std::mutex> lock(page_heap_mtx_);
    PageID start_page_id = span->start_address >> kPageShift;
    metadata_map_.remove_range(start_page_id, span->num_pages);

    // 缓存回 Page Heap，而不是直接 munmap
    if (span->num_pages < kMaxCachedPages) {
        Span* list_head = &free_span_lists_[span->num_pages];
        span->next = list_head->next;
        span->prev = list_head;
        list_head->next->prev = span;
        list_head->next = span;
    } else {
        munmap(reinterpret_cast<void*>(span->start_address), span->num_pages * kPageSize);
        delete span;
    }
}

TCMallocator::Span* TCMallocator::get_one_span(std::size_t size_class) {
    Span* span = allocate_pages(1);
    if (!span) return nullptr;

    std::size_t obj_size = class_to_size(size_class);
    char* start_ptr = reinterpret_cast<char*>(span->start_address);
    char* end_ptr = start_ptr + kPageSize;

    Node* head = nullptr;
    char* curr = start_ptr;
    while (curr + obj_size <= end_ptr) {
        Node* node = reinterpret_cast<Node*>(curr);
        node->next = head;
        head = node;
        curr += obj_size;
    }
    span->free_list = head;
    span->size_class = size_class;
    return span;
}

TCMallocator::Node* TCMallocator::fetch_from_central_cache(std::size_t size_class, std::size_t& batch_size) {
    CentralCacheBucket& bucket = central_caches_[size_class];
    std::lock_guard<std::mutex> lock(bucket.mtx);

    Span* span = bucket.non_empty_list.next;
    if (span == &bucket.non_empty_list) {
        span = get_one_span(size_class);
        if (!span) return nullptr;
        span->next = bucket.non_empty_list.next;
        span->prev = &bucket.non_empty_list;
        bucket.non_empty_list.next->prev = span;
        bucket.non_empty_list.next = span;
    }

    Node* head = span->free_list;
    Node* tail = head;
    std::size_t count = 1;
    while (tail && tail->next && count < batch_size) {
        tail = tail->next;
        count++;
    }
    
    if (tail) {
        span->free_list = tail->next;
        tail->next = nullptr;
    }
    batch_size = count;
    span->refcount += count;

    if (span->free_list == nullptr) {
        span->prev->next = span->next;
        span->next->prev = span->prev;

        span->next = bucket.empty_list.next;
        span->prev = &bucket.empty_list;
        bucket.empty_list.next->prev = span;
        bucket.empty_list.next = span;
    }

    return head;
}

void TCMallocator::release_to_central_cache(Node* start, Node* end, std::size_t count, std::size_t size_class) {
    CentralCacheBucket& bucket = central_caches_[size_class];
    std::lock_guard<std::mutex> lock(bucket.mtx);

    PageID page_id = reinterpret_cast<std::uintptr_t>(start) >> kPageShift;
    Span* span = metadata_map_.get(page_id);
    assert(span != nullptr);

    bool was_empty = (span->free_list == nullptr);

    end->next = span->free_list;
    span->free_list = start;
    span->refcount -= count;

    if (span->refcount == 0) {
        span->prev->next = span->next;
        span->next->prev = span->prev;
        deallocate_pages(span);
    } 
    else if (was_empty) {
        span->prev->next = span->next;
        span->next->prev = span->prev;

        span->next = bucket.non_empty_list.next;
        span->prev = &bucket.non_empty_list;
        bucket.non_empty_list.next->prev = span;
        bucket.non_empty_list.next = span;
    }
}

void* TCMallocator::allocate(std::size_t size) {
    if (size > kMaxSmallSize) {
        std::size_t num_pages = (size + kPageSize - 1) / kPageSize;
        Span* span = allocate_pages(num_pages);
        if (span) {
            span->size_class = kLargeSizeClass;
        }
        return span ? reinterpret_cast<void*>(span->start_address) : nullptr;
    }

    std::size_t sc = size_to_class(size);
    assert(sc < kNumClasses);
    
    ThreadCache* tc = get_thread_cache();
    ThreadCache::FreeList& fl = tc->lists[sc];

    if (fl.list == nullptr) {
        std::size_t batch_size = 8;
        Node* fetched = fetch_from_central_cache(sc, batch_size);
        if (!fetched) return nullptr;
        fl.list = fetched;
        fl.length += batch_size;
    }

    Node* node = fl.list;
    fl.list = node->next;
    fl.length--;
    return node;
}

void TCMallocator::deallocate(void* ptr) {
    if (!ptr) return;

    PageID page_id = reinterpret_cast<std::uintptr_t>(ptr) >> kPageShift;
    Span* span = metadata_map_.get(page_id);
    if (!span) return;

    if (span->size_class == kLargeSizeClass) {
        deallocate_pages(span);
        return;
    }

    std::size_t sc = span->size_class;
    ThreadCache* tc = get_thread_cache();
    ThreadCache::FreeList& fl = tc->lists[sc];

    Node* node = reinterpret_cast<Node*>(ptr);
    node->next = fl.list;
    fl.list = node;
    fl.length++;
}