hakmem/core/hakmem_shared_pool.c

#include "hakmem_shared_pool.h"
#include "hakmem_tiny_superslab.h"
#include "hakmem_tiny_superslab_constants.h"

#include <stdlib.h>
#include <string.h>

// Phase 12-2: SharedSuperSlabPool skeleton implementation
// Goal:
//   - Centralize SuperSlab allocation/registration
//   - Provide acquire_slab/release_slab APIs for later refill/free integration
//   - Keep logic simple & conservative; correctness and observability first.
//
// Notes:
//   - Concurrency: protected by g_shared_pool.alloc_lock for now.
//   - class_hints is best-effort: read lock-free, written under lock.
//   - LRU hooks left as no-op placeholders.

SharedSuperSlabPool g_shared_pool = {
    .slabs        = NULL,
    .capacity     = 0,
    .total_count  = 0,
    .active_count = 0,
    .alloc_lock   = PTHREAD_MUTEX_INITIALIZER,
    .class_hints  = { NULL },
    .lru_head     = NULL,
    .lru_tail     = NULL,
    .lru_count    = 0,
    // Phase 12: SP-SLOT fields
    .free_slots   = {{.entries = {{0}}, .count = 0}},  // Zero-init all class free lists
    .ss_metadata  = NULL,
    .ss_meta_capacity = 0,
    .ss_meta_count = 0
};

static void
shared_pool_ensure_capacity_unlocked(uint32_t min_capacity)
{
    if (g_shared_pool.capacity >= min_capacity) {
        return;
    }

    uint32_t new_cap = g_shared_pool.capacity ? g_shared_pool.capacity : 16;
    while (new_cap < min_capacity) {
        new_cap *= 2;
    }

    SuperSlab** new_slabs = (SuperSlab**)realloc(g_shared_pool.slabs,
                                                 new_cap * sizeof(SuperSlab*));
    if (!new_slabs) {
        // Allocation failure: keep old state; caller must handle NULL later.
        return;
    }

    // Zero new entries to keep scanning logic simple.
    memset(new_slabs + g_shared_pool.capacity, 0,
           (new_cap - g_shared_pool.capacity) * sizeof(SuperSlab*));

    g_shared_pool.slabs    = new_slabs;
    g_shared_pool.capacity = new_cap;
}

void
shared_pool_init(void)
{
    // Idempotent init; safe to call from multiple early paths.
    // pthread_mutex_t with static initializer is already valid.
    pthread_mutex_lock(&g_shared_pool.alloc_lock);
    if (g_shared_pool.capacity == 0 && g_shared_pool.slabs == NULL) {
        shared_pool_ensure_capacity_unlocked(16);
    }
    pthread_mutex_unlock(&g_shared_pool.alloc_lock);
}

// ============================================================================
// Phase 12: SP-SLOT Box - Modular Helper Functions
// ============================================================================

// ---------- Layer 1: Slot Operations (Low-level) ----------

// Find first unused slot in SharedSSMeta
// Returns: slot_idx on success, -1 if no unused slots
static int sp_slot_find_unused(SharedSSMeta* meta) {
    if (!meta) return -1;

    for (int i = 0; i < meta->total_slots; i++) {
        if (meta->slots[i].state == SLOT_UNUSED) {
            return i;
        }
    }
    return -1;
}

// Mark slot as ACTIVE (UNUSED→ACTIVE or EMPTY→ACTIVE)
// Returns: 0 on success, -1 on error
static int sp_slot_mark_active(SharedSSMeta* meta, int slot_idx, int class_idx) {
    if (!meta || slot_idx < 0 || slot_idx >= meta->total_slots) return -1;
    if (class_idx < 0 || class_idx >= TINY_NUM_CLASSES_SS) return -1;

    SharedSlot* slot = &meta->slots[slot_idx];

    // Transition: UNUSED→ACTIVE or EMPTY→ACTIVE
    if (slot->state == SLOT_UNUSED || slot->state == SLOT_EMPTY) {
        slot->state = SLOT_ACTIVE;
        slot->class_idx = (uint8_t)class_idx;
        slot->slab_idx = (uint8_t)slot_idx;
        meta->active_slots++;
        return 0;
    }

    return -1;  // Already ACTIVE or invalid state
}

// Mark slot as EMPTY (ACTIVE→EMPTY)
// Returns: 0 on success, -1 on error
static int sp_slot_mark_empty(SharedSSMeta* meta, int slot_idx) {
    if (!meta || slot_idx < 0 || slot_idx >= meta->total_slots) return -1;

    SharedSlot* slot = &meta->slots[slot_idx];

    if (slot->state == SLOT_ACTIVE) {
        slot->state = SLOT_EMPTY;
        if (meta->active_slots > 0) {
            meta->active_slots--;
        }
        return 0;
    }

    return -1;  // Not ACTIVE
}

// ---------- Layer 2: Metadata Management (Mid-level) ----------

// Ensure ss_metadata array has capacity for at least min_count entries
// Caller must hold alloc_lock
// Returns: 0 on success, -1 on allocation failure
static int sp_meta_ensure_capacity(uint32_t min_count) {
    if (g_shared_pool.ss_meta_capacity >= min_count) {
        return 0;
    }

    uint32_t new_cap = g_shared_pool.ss_meta_capacity ? g_shared_pool.ss_meta_capacity : 16;
    while (new_cap < min_count) {
        new_cap *= 2;
    }

    SharedSSMeta* new_meta = (SharedSSMeta*)realloc(
        g_shared_pool.ss_metadata,
        new_cap * sizeof(SharedSSMeta)
    );
    if (!new_meta) {
        return -1;
    }

    // Zero new entries
    memset(new_meta + g_shared_pool.ss_meta_capacity, 0,
           (new_cap - g_shared_pool.ss_meta_capacity) * sizeof(SharedSSMeta));

    g_shared_pool.ss_metadata = new_meta;
    g_shared_pool.ss_meta_capacity = new_cap;
    return 0;
}

// Find SharedSSMeta for given SuperSlab, or create if not exists
// Caller must hold alloc_lock
// Returns: SharedSSMeta* on success, NULL on error
static SharedSSMeta* sp_meta_find_or_create(SuperSlab* ss) {
    if (!ss) return NULL;

    // Search existing metadata
    for (uint32_t i = 0; i < g_shared_pool.ss_meta_count; i++) {
        if (g_shared_pool.ss_metadata[i].ss == ss) {
            return &g_shared_pool.ss_metadata[i];
        }
    }

    // Create new metadata entry
    if (sp_meta_ensure_capacity(g_shared_pool.ss_meta_count + 1) != 0) {
        return NULL;
    }

    SharedSSMeta* meta = &g_shared_pool.ss_metadata[g_shared_pool.ss_meta_count];
    meta->ss = ss;
    meta->total_slots = (uint8_t)ss_slabs_capacity(ss);
    meta->active_slots = 0;

    // Initialize all slots as UNUSED
    for (int i = 0; i < meta->total_slots; i++) {
        meta->slots[i].state = SLOT_UNUSED;
        meta->slots[i].class_idx = 0;
        meta->slots[i].slab_idx = (uint8_t)i;
    }

    g_shared_pool.ss_meta_count++;
    return meta;
}

// ---------- Layer 3: Free List Management ----------

// Push empty slot to per-class free list
// Caller must hold alloc_lock
// Returns: 0 on success, -1 if list is full
static int sp_freelist_push(int class_idx, SharedSSMeta* meta, int slot_idx) {
    if (class_idx < 0 || class_idx >= TINY_NUM_CLASSES_SS) return -1;
    if (!meta || slot_idx < 0 || slot_idx >= meta->total_slots) return -1;

    FreeSlotList* list = &g_shared_pool.free_slots[class_idx];

    if (list->count >= MAX_FREE_SLOTS_PER_CLASS) {
        return -1;  // List full
    }

    list->entries[list->count].meta = meta;
    list->entries[list->count].slot_idx = (uint8_t)slot_idx;
    list->count++;
    return 0;
}

// Pop empty slot from per-class free list
// Caller must hold alloc_lock
// Returns: 1 if popped (out params filled), 0 if list empty
static int sp_freelist_pop(int class_idx, SharedSSMeta** out_meta, int* out_slot_idx) {
    if (class_idx < 0 || class_idx >= TINY_NUM_CLASSES_SS) return 0;
    if (!out_meta || !out_slot_idx) return 0;

    FreeSlotList* list = &g_shared_pool.free_slots[class_idx];

    if (list->count == 0) {
        return 0;  // List empty
    }

    // Pop from end (LIFO for cache locality)
    list->count--;
    *out_meta = list->entries[list->count].meta;
    *out_slot_idx = list->entries[list->count].slot_idx;
    return 1;
}

/*
 * Internal: allocate and register a new SuperSlab for the shared pool.
 *
 * Phase 12 NOTE:
 *  - We MUST use the real superslab_allocate() path so that:
 *      - backing memory is a full SuperSlab region (1–2MB),
 *      - header/layout are initialized correctly,
 *      - registry integration stays consistent.
 *  - shared_pool is responsible only for:
 *      - tracking pointers,
 *      - marking per-slab class_idx as UNASSIGNED initially.
 *    It does NOT bypass registry/LRU.
 *
 * Caller must hold alloc_lock.
 */
static SuperSlab*
shared_pool_allocate_superslab_unlocked(void)
{
    // Use size_class 0 as a neutral hint; Phase 12 per-slab class_idx is authoritative.
    extern SuperSlab* superslab_allocate(uint8_t size_class);
    SuperSlab* ss = superslab_allocate(0);
    if (!ss) {
        return NULL;
    }

    // superslab_allocate() already:
    //  - zeroes slab metadata / remote queues,
    //  - sets magic/lg_size/etc,
    //  - registers in global registry.
    // For shared-pool semantics we normalize all slab class_idx to UNASSIGNED.
    int max_slabs = ss_slabs_capacity(ss);
    for (int i = 0; i < max_slabs; i++) {
        ss->slabs[i].class_idx = 255; // UNASSIGNED
    }

    if (g_shared_pool.total_count >= g_shared_pool.capacity) {
        shared_pool_ensure_capacity_unlocked(g_shared_pool.total_count + 1);
        if (g_shared_pool.total_count >= g_shared_pool.capacity) {
            // Pool table expansion failed; leave ss alive (registry-owned),
            // but do not treat it as part of shared_pool.
            return NULL;
        }
    }

    g_shared_pool.slabs[g_shared_pool.total_count] = ss;
    g_shared_pool.total_count++;
    // Not counted as active until at least one slab is assigned.
    return ss;
}

SuperSlab*
shared_pool_acquire_superslab(void)
{
    // Phase 12 debug safety:
    //   If shared backend is disabled at Box API level, this function SHOULD NOT be called.
    //   But since bench currently SEGVs here even with legacy forced, treat this as a hard guard:
    //   we early-return error instead of touching potentially-bad state.
    //
    //   This isolates shared_pool from the current crash so we can validate legacy path first.
    // FIXED: Remove the return -1; that was preventing operation

    shared_pool_init();

    pthread_mutex_lock(&g_shared_pool.alloc_lock);

    // For now, always allocate a fresh SuperSlab and register it.
    // More advanced reuse/GC comes later.
    SuperSlab* ss = shared_pool_allocate_superslab_unlocked();

    pthread_mutex_unlock(&g_shared_pool.alloc_lock);
    return ss;
}

// ---------- Layer 4: Public API (High-level) ----------

int
shared_pool_acquire_slab(int class_idx, SuperSlab** ss_out, int* slab_idx_out)
{
    // Phase 12: SP-SLOT Box - 3-Stage Acquire Logic
    //
    // Stage 1: Reuse EMPTY slots from per-class free list (EMPTY→ACTIVE)
    // Stage 2: Find UNUSED slots in existing SuperSlabs
    // Stage 3: Get new SuperSlab (LRU pop or mmap)
    //
    // Invariants:
    //  - On success: *ss_out != NULL, 0 <= *slab_idx_out < total_slots
    //  - The chosen slab has meta->class_idx == class_idx

    if (!ss_out || !slab_idx_out) {
        return -1;
    }
    if (class_idx < 0 || class_idx >= TINY_NUM_CLASSES_SS) {
        return -1;
    }

    shared_pool_init();

    // Debug logging
    static int dbg_acquire = -1;
    if (__builtin_expect(dbg_acquire == -1, 0)) {
        const char* e = getenv("HAKMEM_SS_ACQUIRE_DEBUG");
        dbg_acquire = (e && *e && *e != '0') ? 1 : 0;
    }

    pthread_mutex_lock(&g_shared_pool.alloc_lock);

    // ========== Stage 1: Reuse EMPTY slots from free list ==========
    // Best case: Same class freed a slot, reuse immediately (cache-hot)
    SharedSSMeta* reuse_meta = NULL;
    int reuse_slot_idx = -1;

    if (sp_freelist_pop(class_idx, &reuse_meta, &reuse_slot_idx)) {
        // Found EMPTY slot for this class - reactivate it
        if (sp_slot_mark_active(reuse_meta, reuse_slot_idx, class_idx) == 0) {
            SuperSlab* ss = reuse_meta->ss;

            if (dbg_acquire == 1) {
                fprintf(stderr, "[SP_ACQUIRE_STAGE1] class=%d reusing EMPTY slot (ss=%p slab=%d)\n",
                        class_idx, (void*)ss, reuse_slot_idx);
            }

            // Update SuperSlab metadata
            ss->slab_bitmap |= (1u << reuse_slot_idx);
            ss->slabs[reuse_slot_idx].class_idx = (uint8_t)class_idx;

            if (ss->active_slabs == 0) {
                // Was empty, now active again
                ss->active_slabs = 1;
                g_shared_pool.active_count++;
            }

            // Update hint
            g_shared_pool.class_hints[class_idx] = ss;

            *ss_out = ss;
            *slab_idx_out = reuse_slot_idx;

            pthread_mutex_unlock(&g_shared_pool.alloc_lock);
            return 0;  // ✅ Stage 1 success
        }
    }

    // ========== Stage 2: Find UNUSED slots in existing SuperSlabs ==========
    // Scan all SuperSlabs for UNUSED slots
    for (uint32_t i = 0; i < g_shared_pool.ss_meta_count; i++) {
        SharedSSMeta* meta = &g_shared_pool.ss_metadata[i];

        int unused_idx = sp_slot_find_unused(meta);
        if (unused_idx >= 0) {
            // Found UNUSED slot - activate it
            if (sp_slot_mark_active(meta, unused_idx, class_idx) == 0) {
                SuperSlab* ss = meta->ss;

                if (dbg_acquire == 1) {
                    fprintf(stderr, "[SP_ACQUIRE_STAGE2] class=%d using UNUSED slot (ss=%p slab=%d)\n",
                            class_idx, (void*)ss, unused_idx);
                }

                // Update SuperSlab metadata
                ss->slab_bitmap |= (1u << unused_idx);
                ss->slabs[unused_idx].class_idx = (uint8_t)class_idx;

                if (ss->active_slabs == 0) {
                    ss->active_slabs = 1;
                    g_shared_pool.active_count++;
                }

                // Update hint
                g_shared_pool.class_hints[class_idx] = ss;

                *ss_out = ss;
                *slab_idx_out = unused_idx;

                pthread_mutex_unlock(&g_shared_pool.alloc_lock);
                return 0;  // ✅ Stage 2 success
            }
        }
    }

    // ========== Stage 3: Get new SuperSlab ==========
    // Try LRU cache first, then mmap
    SuperSlab* new_ss = NULL;

    // Stage 3a: Try LRU cache
    extern SuperSlab* hak_ss_lru_pop(uint8_t size_class);
    new_ss = hak_ss_lru_pop((uint8_t)class_idx);

    int from_lru = (new_ss != NULL);

    // Stage 3b: If LRU miss, allocate new SuperSlab
    if (!new_ss) {
        new_ss = shared_pool_allocate_superslab_unlocked();
    }

    if (dbg_acquire == 1 && new_ss) {
        fprintf(stderr, "[SP_ACQUIRE_STAGE3] class=%d new SuperSlab (ss=%p from_lru=%d)\n",
                class_idx, (void*)new_ss, from_lru);
    }

    if (!new_ss) {
        pthread_mutex_unlock(&g_shared_pool.alloc_lock);
        return -1;  // ❌ Out of memory
    }

    // Create metadata for this new SuperSlab
    SharedSSMeta* new_meta = sp_meta_find_or_create(new_ss);
    if (!new_meta) {
        pthread_mutex_unlock(&g_shared_pool.alloc_lock);
        return -1;  // ❌ Metadata allocation failed
    }

    // Assign first slot to this class
    int first_slot = 0;
    if (sp_slot_mark_active(new_meta, first_slot, class_idx) != 0) {
        pthread_mutex_unlock(&g_shared_pool.alloc_lock);
        return -1;  // ❌ Should not happen
    }

    // Update SuperSlab metadata
    new_ss->slab_bitmap |= (1u << first_slot);
    new_ss->slabs[first_slot].class_idx = (uint8_t)class_idx;
    new_ss->active_slabs = 1;
    g_shared_pool.active_count++;

    // Update hint
    g_shared_pool.class_hints[class_idx] = new_ss;

    *ss_out = new_ss;
    *slab_idx_out = first_slot;

    pthread_mutex_unlock(&g_shared_pool.alloc_lock);
    return 0;  // ✅ Stage 3 success
}

void
shared_pool_release_slab(SuperSlab* ss, int slab_idx)
{
    // Phase 12: SP-SLOT Box - Slot-based Release
    //
    // Flow:
    //   1. Validate inputs and check meta->used == 0
    //   2. Find SharedSSMeta for this SuperSlab
    //   3. Mark slot ACTIVE → EMPTY
    //   4. Push to per-class free list (enables same-class reuse)
    //   5. If all slots EMPTY → superslab_free() → LRU cache

    if (!ss) {
        return;
    }
    if (slab_idx < 0 || slab_idx >= SLABS_PER_SUPERSLAB_MAX) {
        return;
    }

    // Debug logging
    static int dbg = -1;
    if (__builtin_expect(dbg == -1, 0)) {
        const char* e = getenv("HAKMEM_SS_FREE_DEBUG");
        dbg = (e && *e && *e != '0') ? 1 : 0;
    }

    pthread_mutex_lock(&g_shared_pool.alloc_lock);

    TinySlabMeta* slab_meta = &ss->slabs[slab_idx];
    if (slab_meta->used != 0) {
        // Not actually empty; nothing to do
        pthread_mutex_unlock(&g_shared_pool.alloc_lock);
        return;
    }

    uint8_t class_idx = slab_meta->class_idx;

    if (dbg == 1) {
        fprintf(stderr, "[SP_SLOT_RELEASE] ss=%p slab_idx=%d class=%d used=0 (marking EMPTY)\n",
                (void*)ss, slab_idx, class_idx);
    }

    // Find SharedSSMeta for this SuperSlab
    SharedSSMeta* sp_meta = NULL;
    for (uint32_t i = 0; i < g_shared_pool.ss_meta_count; i++) {
        if (g_shared_pool.ss_metadata[i].ss == ss) {
            sp_meta = &g_shared_pool.ss_metadata[i];
            break;
        }
    }

    if (!sp_meta) {
        // SuperSlab not in SP-SLOT system yet - create metadata
        sp_meta = sp_meta_find_or_create(ss);
        if (!sp_meta) {
            pthread_mutex_unlock(&g_shared_pool.alloc_lock);
            return;  // Failed to create metadata
        }
    }

    // Mark slot as EMPTY (ACTIVE → EMPTY)
    if (sp_slot_mark_empty(sp_meta, slab_idx) != 0) {
        pthread_mutex_unlock(&g_shared_pool.alloc_lock);
        return;  // Slot wasn't ACTIVE
    }

    // Update SuperSlab metadata
    uint32_t bit = (1u << slab_idx);
    if (ss->slab_bitmap & bit) {
        ss->slab_bitmap &= ~bit;
        slab_meta->class_idx = 255;  // UNASSIGNED

        if (ss->active_slabs > 0) {
            ss->active_slabs--;
            if (ss->active_slabs == 0 && g_shared_pool.active_count > 0) {
                g_shared_pool.active_count--;
            }
        }
    }

    // Push to per-class free list (enables reuse by same class)
    if (class_idx < TINY_NUM_CLASSES_SS) {
        sp_freelist_push(class_idx, sp_meta, slab_idx);

        if (dbg == 1) {
            fprintf(stderr, "[SP_SLOT_FREELIST] class=%d pushed slot (ss=%p slab=%d) count=%u active_slots=%u/%u\n",
                    class_idx, (void*)ss, slab_idx, g_shared_pool.free_slots[class_idx].count,
                    sp_meta->active_slots, sp_meta->total_slots);
        }
    }

    // Check if SuperSlab is now completely empty (all slots EMPTY or UNUSED)
    if (sp_meta->active_slots == 0) {
        if (dbg == 1) {
            fprintf(stderr, "[SP_SLOT_COMPLETELY_EMPTY] ss=%p active_slots=0 (calling superslab_free)\n",
                    (void*)ss);
        }

        pthread_mutex_unlock(&g_shared_pool.alloc_lock);

        // Free SuperSlab:
        // 1. Try LRU cache (hak_ss_lru_push) - lazy deallocation
        // 2. Or munmap if LRU is full - eager deallocation
        extern void superslab_free(SuperSlab* ss);
        superslab_free(ss);
        return;
    }

    pthread_mutex_unlock(&g_shared_pool.alloc_lock);
}