hakmem/core/hakmem_smallmid.c

/**
 * hakmem_smallmid.c - Small-Mid Allocator Front Box Implementation
 *
 * Phase 17-1: Front Box Only (No Dedicated SuperSlab Backend)
 *
 * Strategy (ChatGPT reviewed):
 * - Thin front layer with TLS freelist (256B/512B/1KB)
 * - Backend: Use existing Tiny SuperSlab/SharedPool APIs
 * - Goal: Measure performance impact before building dedicated backend
 * - A/B test: Does Small-Mid front improve 256-1KB performance?
 *
 * Architecture:
 * - 3 size classes: 256B/512B/1KB (reduced from 5)
 * - TLS freelist for fast alloc/free (static inline)
 * - Backend: Call Tiny allocator APIs (reuse existing infrastructure)
 * - ENV controlled (HAKMEM_SMALLMID_ENABLE=1)
 *
 * Created: 2025-11-16
 * Updated: 2025-11-16 (Phase 17-1 revision - Front Box only)
 */

#include "hakmem_smallmid.h"
#include "hakmem_build_flags.h"
#include "hakmem_smallmid_superslab.h"  // Phase 17-2: Dedicated backend
#include "tiny_region_id.h"             // For header writing
#include <string.h>
#include <pthread.h>

// ============================================================================
// TLS State
// ============================================================================

__thread void*    g_smallmid_tls_head[SMALLMID_NUM_CLASSES]  = {NULL};
__thread uint32_t g_smallmid_tls_count[SMALLMID_NUM_CLASSES] = {0};

// ============================================================================
// Size Class Table (Phase 17-1: 3 classes)
// ============================================================================

const size_t g_smallmid_class_sizes[SMALLMID_NUM_CLASSES] = {
    256,   // SM0: 256B
    512,   // SM1: 512B
    1024   // SM2: 1KB
};

// ============================================================================
// Global State
// ============================================================================

static pthread_mutex_t g_smallmid_init_lock = PTHREAD_MUTEX_INITIALIZER;
static int g_smallmid_initialized = 0;
static int g_smallmid_enabled = -1;  // -1 = not checked, 0 = disabled, 1 = enabled

// ============================================================================
// Statistics (Debug)
// ============================================================================

#ifdef HAKMEM_SMALLMID_STATS
SmallMidStats g_smallmid_stats = {0};

void smallmid_print_stats(void) {
    fprintf(stderr, "\n=== Small-Mid Allocator Statistics ===\n");
    fprintf(stderr, "Total allocs:        %lu\n", g_smallmid_stats.total_allocs);
    fprintf(stderr, "Total frees:         %lu\n", g_smallmid_stats.total_frees);
    fprintf(stderr, "TLS hits:            %lu\n", g_smallmid_stats.tls_hits);
    fprintf(stderr, "TLS misses:          %lu\n", g_smallmid_stats.tls_misses);
    fprintf(stderr, "SuperSlab refills:   %lu\n", g_smallmid_stats.superslab_refills);
    if (g_smallmid_stats.total_allocs > 0) {
        double hit_rate = (double)g_smallmid_stats.tls_hits / g_smallmid_stats.total_allocs * 100.0;
        fprintf(stderr, "TLS hit rate:        %.2f%%\n", hit_rate);
    }
    fprintf(stderr, "=======================================\n\n");
}
#endif

// ============================================================================
// ENV Control
// ============================================================================

bool smallmid_is_enabled(void) {
    if (__builtin_expect(g_smallmid_enabled == -1, 0)) {
        const char* env = getenv("HAKMEM_SMALLMID_ENABLE");
        g_smallmid_enabled = (env && atoi(env) == 1) ? 1 : 0;

        if (g_smallmid_enabled) {
            SMALLMID_LOG("Small-Mid allocator ENABLED (ENV: HAKMEM_SMALLMID_ENABLE=1)");
        } else {
            SMALLMID_LOG("Small-Mid allocator DISABLED (default, set HAKMEM_SMALLMID_ENABLE=1 to enable)");
        }
    }
    return (g_smallmid_enabled == 1);
}

// ============================================================================
// Initialization
// ============================================================================

void smallmid_init(void) {
    if (g_smallmid_initialized) return;

    pthread_mutex_lock(&g_smallmid_init_lock);

    if (!g_smallmid_initialized) {
        SMALLMID_LOG("Initializing Small-Mid Front Box...");

        // Check ENV
        if (!smallmid_is_enabled()) {
            SMALLMID_LOG("Small-Mid allocator is disabled, skipping initialization");
            g_smallmid_initialized = 1;
            pthread_mutex_unlock(&g_smallmid_init_lock);
            return;
        }

        // Phase 17-1: No dedicated backend - use existing Tiny infrastructure
        // No additional initialization needed (TLS state is static)

        g_smallmid_initialized = 1;
        SMALLMID_LOG("Small-Mid Front Box initialized (3 classes: 256B/512B/1KB, backend=Tiny)");
    }

    pthread_mutex_unlock(&g_smallmid_init_lock);
}

// ============================================================================
// TLS Freelist Operations
// ============================================================================

/**
 * smallmid_tls_pop - Pop a block from TLS freelist
 *
 * @param class_idx  Size class index
 * @return           Block pointer (with header), or NULL if empty
 */
static inline void* smallmid_tls_pop(int class_idx) {
    void* head = g_smallmid_tls_head[class_idx];
    if (!head) return NULL;

    // Read next pointer (stored at offset 0 in user data, after 1-byte header)
    void* next = *(void**)((uint8_t*)head + 1);
    g_smallmid_tls_head[class_idx] = next;
    g_smallmid_tls_count[class_idx]--;

    #ifdef HAKMEM_SMALLMID_STATS
    __atomic_fetch_add(&g_smallmid_stats.tls_hits, 1, __ATOMIC_RELAXED);
    #endif

    return head;
}

/**
 * smallmid_tls_push - Push a block to TLS freelist
 *
 * @param class_idx  Size class index
 * @param ptr        Block pointer (with header)
 * @return           true on success, false if TLS full
 */
static inline bool smallmid_tls_push(int class_idx, void* ptr) {
    uint32_t capacity = smallmid_tls_capacity(class_idx);
    if (g_smallmid_tls_count[class_idx] >= capacity) {
        return false;  // TLS full
    }

    // Write next pointer (at offset 0 in user data, after 1-byte header)
    void* head = g_smallmid_tls_head[class_idx];
    *(void**)((uint8_t*)ptr + 1) = head;
    g_smallmid_tls_head[class_idx] = ptr;
    g_smallmid_tls_count[class_idx]++;

    return true;
}

// ============================================================================
// TLS Refill (Phase 17-2: Batch refill from dedicated SuperSlab)
// ============================================================================

/**
 * smallmid_tls_refill - Refill TLS freelist from SuperSlab
 *
 * @param class_idx  Size class index
 * @return           true on success, false on failure
 *
 * Strategy (Phase 17-2):
 * - Batch refill 8-16 blocks from dedicated SmallMid SuperSlab
 * - No Tiny delegation (completely separate backend)
 * - Amortizes SuperSlab lookup cost across multiple blocks
 * - Expected cost: ~1-2 instructions per block (amortized)
 */
static bool smallmid_tls_refill(int class_idx) {
    // Determine batch size based on size class
    const int batch_sizes[SMALLMID_NUM_CLASSES] = {
        SMALLMID_REFILL_BATCH_256B,   // 16 blocks
        SMALLMID_REFILL_BATCH_512B,   // 12 blocks
        SMALLMID_REFILL_BATCH_1KB     // 8 blocks
    };

    int batch_max = batch_sizes[class_idx];
    void* batch[16];  // Max batch size

    // Call SuperSlab batch refill
    int refilled = smallmid_refill_batch(class_idx, batch, batch_max);
    if (refilled == 0) {
        SMALLMID_LOG("smallmid_tls_refill: SuperSlab refill failed (class=%d)", class_idx);
        return false;
    }

    #ifdef HAKMEM_SMALLMID_STATS
    __atomic_fetch_add(&g_smallmid_stats.tls_misses, 1, __ATOMIC_RELAXED);
    __atomic_fetch_add(&g_smallmid_stats.superslab_refills, 1, __ATOMIC_RELAXED);
    #endif

    // Push blocks to TLS freelist (in reverse order for LIFO)
    for (int i = refilled - 1; i >= 0; i--) {
        void* user_ptr = batch[i];
        void* base = (uint8_t*)user_ptr - 1;

        if (!smallmid_tls_push(class_idx, base)) {
            // TLS full - should not happen with proper batch sizing
            SMALLMID_LOG("smallmid_tls_refill: TLS push failed (class=%d, i=%d)", class_idx, i);
            break;
        }
    }

    SMALLMID_LOG("smallmid_tls_refill: Refilled %d blocks (class=%d)", refilled, class_idx);
    return true;
}

// ============================================================================
// Allocation
// ============================================================================

void* smallmid_alloc(size_t size) {
    // Check if enabled
    if (!smallmid_is_enabled()) {
        return NULL;  // Disabled, fall through to Mid or other allocators
    }

    // Initialize if needed
    if (__builtin_expect(!g_smallmid_initialized, 0)) {
        smallmid_init();
        smallmid_superslab_init();  // Phase 17-2: Initialize SuperSlab backend
    }

    // Validate size range
    if (__builtin_expect(!smallmid_is_in_range(size), 0)) {
        SMALLMID_LOG("smallmid_alloc: size %zu out of range [%d-%d]",
                     size, SMALLMID_MIN_SIZE, SMALLMID_MAX_SIZE);
        return NULL;
    }

    // Get size class
    int class_idx = smallmid_size_to_class(size);
    if (__builtin_expect(class_idx < 0, 0)) {
        SMALLMID_LOG("smallmid_alloc: invalid class for size %zu", size);
        return NULL;
    }

    #ifdef HAKMEM_SMALLMID_STATS
    __atomic_fetch_add(&g_smallmid_stats.total_allocs, 1, __ATOMIC_RELAXED);
    #endif

    // Fast path: Pop from TLS freelist
    void* ptr = smallmid_tls_pop(class_idx);
    if (ptr) {
        SMALLMID_LOG("smallmid_alloc(%zu) = %p (TLS hit, class=%d)", size, ptr, class_idx);
        return (uint8_t*)ptr + 1;  // Return user pointer (skip header)
    }

    // TLS miss: Refill from SuperSlab (Phase 17-2: Batch refill)
    if (!smallmid_tls_refill(class_idx)) {
        SMALLMID_LOG("smallmid_alloc(%zu) = NULL (refill failed)", size);
        return NULL;
    }

    // Retry TLS pop after refill
    ptr = smallmid_tls_pop(class_idx);
    if (!ptr) {
        SMALLMID_LOG("smallmid_alloc(%zu) = NULL (TLS pop failed after refill)", size);
        return NULL;
    }

    SMALLMID_LOG("smallmid_alloc(%zu) = %p (TLS refill, class=%d)", size, ptr, class_idx);
    return (uint8_t*)ptr + 1;  // Return user pointer (skip header)
}

// ============================================================================
// Free
// ============================================================================

void smallmid_free(void* ptr) {
    if (!ptr) return;

    // Check if enabled
    if (!smallmid_is_enabled()) {
        return;  // Disabled, should not be called
    }

    #ifdef HAKMEM_SMALLMID_STATS
    __atomic_fetch_add(&g_smallmid_stats.total_frees, 1, __ATOMIC_RELAXED);
    #endif

    // Phase 17-2: Read header to identify size class
    uint8_t* base = (uint8_t*)ptr - 1;
    uint8_t header = *base;

    // Small-Mid allocations have magic 0xb0
    uint8_t magic = header & 0xf0;
    int class_idx = header & 0x0f;

    if (magic != 0xb0 || class_idx < 0 || class_idx >= SMALLMID_NUM_CLASSES) {
        // Invalid header - should not happen
        SMALLMID_LOG("smallmid_free(%p): Invalid header 0x%02x", ptr, header);
        return;
    }

    // Fast path: Push to TLS freelist
    if (smallmid_tls_push(class_idx, base)) {
        SMALLMID_LOG("smallmid_free(%p): pushed to TLS (class=%d)", ptr, class_idx);
        return;
    }

    // TLS full: Push to SuperSlab freelist (slow path)
    // TODO Phase 17-2.1: Implement SuperSlab freelist push
    // For now, just log and leak (will be fixed in next commit)
    SMALLMID_LOG("smallmid_free(%p): TLS full, SuperSlab freelist not yet implemented", ptr);

    // Placeholder: Write next pointer to freelist (unsafe without SuperSlab lookup)
    // This will be properly implemented with smallmid_superslab_lookup() in Phase 17-2.1
}

// ============================================================================
// Thread Cleanup
// ============================================================================

void smallmid_thread_exit(void) {
    if (!smallmid_is_enabled()) return;

    SMALLMID_LOG("smallmid_thread_exit: cleaning up TLS state");

    // Phase 17-1: Return TLS blocks to Tiny backend
    for (int i = 0; i < SMALLMID_NUM_CLASSES; i++) {
        void* head = g_smallmid_tls_head[i];
        while (head) {
            void* next = *(void**)((uint8_t*)head + 1);
            void* user_ptr = (uint8_t*)head + 1;
            smallmid_backend_free(user_ptr, 0);
            head = next;
        }
        g_smallmid_tls_head[i] = NULL;
        g_smallmid_tls_count[i] = 0;
    }
}