hakmem/core/box/ss_budget_box.c

// ss_budget_box.c - Superslab Budget Box
// Box Theory: Budget/limit guard for Superslab growth.
// - ENV:
//     HAKMEM_SS_BUDGET_GLOBAL : global cap (0 = unlimited, default varies)
//     HAKMEM_SS_BUDGET_C0..C7 : per-class cap override (0 = unlimited)
//     HAKMEM_SS_BUDGET_C7     : shorthand most often used
// - Profile hint:
//     HAKMEM_TINY_PROFILE=larson_guard → stricter defaults.

#include "ss_budget_box.h"

#include <stdatomic.h>
#include <stdlib.h>
#include <strings.h>
#include <stdio.h>

#include "ss_stats_box.h"

static _Atomic int g_budget_init = 0;
static int g_ss_budget_global = 0;
static int g_ss_budget_per_class[8] = {0};

static int ss_budget_parse_env(const char* name, int fallback) {
    const char* e = getenv(name);
    if (e && *e) {
        int v = atoi(e);
        if (v < 0) v = 0;
        return v;
    }
    return fallback;
}

static void ss_budget_init_once(void) {
    if (atomic_load_explicit(&g_budget_init, memory_order_acquire)) {
        return;
    }

    // Profile hint: larson_guard uses tighter defaults to cap RSS.
    const char* profile = getenv("HAKMEM_TINY_PROFILE");
    int is_larson_guard = (profile && strcasecmp(profile, "larson_guard") == 0);

    // Defaults: unlimited unless larson_guard
    int default_global = is_larson_guard ? 512 : 0;
    g_ss_budget_global = ss_budget_parse_env("HAKMEM_SS_BUDGET_GLOBAL", default_global);

    for (int i = 0; i < 8; i++) {
        int def = 0;
        if (is_larson_guard) {
            // Larson guard: modest per-class caps, C7 is a bit looser.
            def = (i == 7) ? 192 : 96;
        }
        g_ss_budget_per_class[i] = def;
    }

    // Per-class overrides: HAKMEM_SS_BUDGET_C7 or HAKMEM_SS_BUDGET_C{idx}
    for (int i = 0; i < 8; i++) {
        char buf[32];
        snprintf(buf, sizeof(buf), "HAKMEM_SS_BUDGET_C%d", i);
        int override = ss_budget_parse_env(buf, g_ss_budget_per_class[i]);
        g_ss_budget_per_class[i] = override;
    }
    // Support the legacy shorthand HAKMEM_SS_BUDGET_C7
    g_ss_budget_per_class[7] =
        ss_budget_parse_env("HAKMEM_SS_BUDGET_C7", g_ss_budget_per_class[7]);

    atomic_store_explicit(&g_budget_init, 1, memory_order_release);
}

static inline uint64_t ss_budget_global_live_sum(void) {
    uint64_t sum = 0;
    for (int i = 0; i < 8; i++) {
        sum += atomic_load_explicit(&g_ss_live_by_class[i], memory_order_relaxed);
    }
    return sum;
}

bool ss_budget_on_alloc(int class_idx) {
    ss_budget_init_once();

    if (class_idx < 0 || class_idx >= 8) {
        return true;  // outside Tiny; do not gate here
    }

    uint64_t live_cls = atomic_load_explicit(&g_ss_live_by_class[class_idx],
                                             memory_order_relaxed);
    int class_cap = g_ss_budget_per_class[class_idx];
    if (class_cap > 0 && live_cls >= (uint64_t)class_cap) {
        static _Atomic uint32_t log_once = 0;
        if (atomic_fetch_add_explicit(&log_once, 1, memory_order_relaxed) < 4) {
            fprintf(stderr,
                    "[SS_BUDGET_DENY] class=%d live=%llu cap=%d\n",
                    class_idx,
                    (unsigned long long)live_cls,
                    class_cap);
        }
        return false;
    }

    int global_cap = g_ss_budget_global;
    if (global_cap > 0) {
        uint64_t live_total = ss_budget_global_live_sum();
        if (live_total >= (uint64_t)global_cap) {
            static _Atomic uint32_t g_log_once = 0;
            if (atomic_fetch_add_explicit(&g_log_once, 1, memory_order_relaxed) < 4) {
                fprintf(stderr,
                        "[SS_BUDGET_DENY_GLOBAL] live_total=%llu cap=%d class=%d\n",
                        (unsigned long long)live_total,
                        global_cap,
                        class_idx);
            }
            return false;
        }
    }

    return true;
}

void ss_budget_on_free(int class_idx) {
    (void)class_idx;
    ss_budget_init_once();
    // We currently rely on ss_stats_on_ss_free_class() to update live counters.
}