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_tiny.h"      // For backend: hak_tiny_alloc / hak_tiny_free
#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;
}

// ============================================================================
// Backend: Use Tiny Allocator APIs (Phase 17-1)
// ============================================================================

/**
 * smallmid_backend_alloc - Allocate from Tiny backend
 *
 * @param size  Allocation size
 * @return      Allocated pointer (user pointer, no Small-Mid header)
 *
 * Phase 17-1: Delegate to existing Tiny allocator infrastructure
 * This reuses Tiny's SuperSlab/SharedPool without building dedicated backend
 */
static inline void* smallmid_backend_alloc(size_t size) {
    #ifdef HAKMEM_SMALLMID_STATS
    __atomic_fetch_add(&g_smallmid_stats.tls_misses, 1, __ATOMIC_RELAXED);
    #endif

    // Call Tiny allocator (reuses existing SuperSlab/SharedPool)
    void* ptr = hak_tiny_alloc(size);
    SMALLMID_LOG("smallmid_backend_alloc(%zu) = %p (via Tiny)", size, ptr);
    return ptr;
}

/**
 * smallmid_backend_free - Free to Tiny backend
 *
 * @param ptr   User pointer (no Small-Mid header)
 * @param size  Allocation size
 *
 * Phase 17-1: Delegate to existing Tiny allocator infrastructure
 */
static inline void smallmid_backend_free(void* ptr, size_t size) {
    (void)size;  // Unused: Tiny free reads header, doesn't need size
    SMALLMID_LOG("smallmid_backend_free(%p) (via Tiny)", ptr);
    hak_tiny_free(ptr);
}

// ============================================================================
// 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();
    }

    // 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)
    }

    // Slow path: Allocate from Tiny backend (no refill, direct delegation)
    // Phase 17-1: Simplified - no TLS refill, just pass through to Tiny
    void* backend_ptr = smallmid_backend_alloc(size);
    if (!backend_ptr) {
        SMALLMID_LOG("smallmid_alloc(%zu) = NULL (backend alloc failed)", size);
        return NULL;
    }

    SMALLMID_LOG("smallmid_alloc(%zu) = %p (backend alloc, class=%d)", size, backend_ptr, class_idx);
    return backend_ptr;  // Backend returns user pointer directly
}

// ============================================================================
// 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-1: Read header to identify if this is a Small-Mid TLS allocation
    // or a backend (Tiny) allocation
    uint8_t* base = (uint8_t*)ptr - 1;
    uint8_t header = *base;

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

    if (magic == 0xb0 && class_idx >= 0 && class_idx < SMALLMID_NUM_CLASSES) {
        // This is a Small-Mid TLS allocation, 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: Delegate to Tiny backend
        SMALLMID_LOG("smallmid_free(%p): TLS full, delegating to backend", ptr);
        // Fall through to backend free
    }

    // This is a backend (Tiny) allocation, or TLS full - delegate to Tiny
    // Tiny will handle the free based on its own header (0xa0)
    size_t size = 0;  // Tiny free doesn't need size, it reads header
    smallmid_backend_free(ptr, size);
}

// ============================================================================
// 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;
    }
}