Merge pull request DeusData#954 from yangsec888/feat/mem-budget-resolver-clamp

feat(mem): pure resolver + strict parse + clamp for CBM_MEM_BUDGET_MB

Author: DeusData

README.md

@@ -489,6 +489,7 @@ codebase-memory-mcp config reset auto_index              # reset to default
 | `CBM_DOWNLOAD_URL` | *(GitHub releases)* | Override the download URL for updates. Used for testing or self-hosted deployments. |
 | `CBM_LOG_LEVEL` | `info` | Set the minimum log level. Accepted values (case-insensitive): `debug`, `info`, `warn`, `error`, `none` — or their numeric equivalents `0`–`4` matching the internal enum. Logs go to stderr; stdout is reserved for MCP JSON-RPC. |
 | `CBM_WORKERS` | *(detected)* | Override the parallel-indexing worker count returned by `cbm_default_worker_count`. Useful inside containers where `sysconf(_SC_NPROCESSORS_ONLN)` reports host CPUs rather than the cgroup's effective quota. Range 1–256; invalid values are ignored with a warning. |
+| `CBM_MEM_BUDGET_MB` | *(detected)* | Override the in-memory graph budget with an explicit cap in MiB, taking precedence over the `ram_fraction × total_RAM` default. Useful on bare-metal hosts without a cgroup limit, or to pin a budget *below* the cgroup limit so headroom is left for sibling processes. Must be a positive integer; it is clamped to detected total RAM (logged as `mem.budget.clamped`), and non-numeric or non-positive values are ignored with a warning (`mem.budget.env.invalid`). |
 | `CBM_DUMP_VERIFY_MIN_RATIO` | `0.5` | After indexing, compare persisted SQLite node count to the in-memory dump count. When persisted nodes fall below this fraction of committed nodes (and committed > 50), `index_repository` returns `status:"degraded"` instead of silent `indexed`. Range 0–1; set `0` to disable. Invalid values are ignored with a warning. |
 
 ```bash

src/foundation/mem.c

@@ -14,9 +14,12 @@
 
 #define MAX_RAM_FRACTION 1.0
 #define DEFAULT_RAM_FRACTION 0.5
+#include <errno.h>
 #include <mimalloc.h>
 #include <stdatomic.h>
+#include <stdint.h>
 #include <stdio.h>
+#include <stdlib.h>
 
 #ifdef _WIN32
 #ifndef WIN32_LEAN_AND_MEAN
@@ -121,6 +124,55 @@ double cbm_mem_ram_fraction_for_total(size_t total_ram_bytes) {
     return RAM_FRACTION_DEFAULT;
 }
 
+cbm_mem_budget_t cbm_mem_resolve_budget(size_t total_ram, double ram_fraction,
+                                        const char *budget_mb) {
+    if (ram_fraction <= 0.0 || ram_fraction > MAX_RAM_FRACTION) {
+        ram_fraction = DEFAULT_RAM_FRACTION;
+    }
+    cbm_mem_budget_t result = {
+        .budget = (size_t)((double)total_ram * ram_fraction),
+        .source = "ram_fraction",
+        .clamped = false,
+        .invalid = false,
+    };
+
+    if (budget_mb == NULL || budget_mb[0] == '\0') {
+        return result; /* no override → fraction-derived budget */
+    }
+
+    /* Strict parse, matching the src/foundation/limits.c convention: reject
+     * trailing garbage (`*end`), overflow (errno==ERANGE), and non-positive
+     * values. This turns a fat-fingered value (e.g. "8GB", or a 20-digit typo)
+     * into a clean fallback-with-warning rather than a silently wrong budget. */
+    errno = 0;
+    char *end = NULL;
+    long long want_mb = strtoll(budget_mb, &end, CBM_DECIMAL_BASE);
+    if (errno != 0 || end == budget_mb || *end != '\0' || want_mb <= 0) {
+        result.invalid = true; /* keep the fraction-derived budget */
+        return result;
+    }
+
+    result.source = "CBM_MEM_BUDGET_MB";
+    size_t want = (size_t)want_mb;
+    if (total_ram > 0) {
+        /* Compare in MiB space so a valid-but-huge request (e.g. 2^44 MiB, which
+         * would overflow the ×MiB byte multiply) clamps cleanly to total_ram. */
+        if (want > total_ram / MB_DIVISOR) {
+            result.budget = total_ram;
+            result.clamped = true;
+        } else {
+            result.budget = want * MB_DIVISOR;
+        }
+    } else if (want > SIZE_MAX / MB_DIVISOR) {
+        /* RAM detection failed (no clamp target) and the request is
+         * astronomically large — cap at SIZE_MAX rather than wrap. */
+        result.budget = SIZE_MAX;
+    } else {
+        result.budget = want * MB_DIVISOR;
+    }
+    return result;
+}
+
 void cbm_mem_init(double ram_fraction) {
     int expected = 0;
     if (!atomic_compare_exchange_strong(&g_initialized, &expected, 1)) {
@@ -150,28 +202,33 @@ void cbm_mem_init(double ram_fraction) {
     /* CBM_MEM_BUDGET_MB env override (memory analogue of CBM_WORKERS).
      * Lets users cap the budget directly without an enclosing cgroup —
      * useful on bare-metal hosts where cgroup memory limits are absent
-     * (#363). Explicit override > implicit RAM/cgroup detection. */
+     * (#363). Explicit override > implicit RAM/cgroup detection. The budget
+     * math (fraction default, override, clamp-to-total) lives in the pure,
+     * testable cbm_mem_resolve_budget(); this path only reads the env and
+     * emits the log/warn lines. */
+    cbm_system_info_t info = cbm_system_info();
+
     char env_buf[CBM_SZ_32];
-    if (cbm_safe_getenv("CBM_MEM_BUDGET_MB", env_buf, sizeof(env_buf), NULL) != NULL) {
-        long mb = strtol(env_buf, NULL, CBM_DECIMAL_BASE);
-        if (mb > 0) {
-            g_budget = (size_t)mb * MB_DIVISOR;
-            char ovr_mb[CBM_SZ_32];
-            snprintf(ovr_mb, sizeof(ovr_mb), "%ld", mb);
-            cbm_log_info("mem.init", "budget_mb", ovr_mb, "source", "CBM_MEM_BUDGET_MB");
-            return;
-        }
-        cbm_log_warn("mem.budget.env.invalid", "value", env_buf, "fallback", "ram_fraction");
-    }
+    const char *env = cbm_safe_getenv("CBM_MEM_BUDGET_MB", env_buf, sizeof(env_buf), NULL);
+    cbm_mem_budget_t resolved = cbm_mem_resolve_budget(info.total_ram, ram_fraction, env);
+    g_budget = resolved.budget;
 
-    cbm_system_info_t info = cbm_system_info();
-    g_budget = (size_t)((double)info.total_ram * ram_fraction);
+    /* The resolver is the single source of truth for the parse + clamp; this
+     * path only surfaces its outcome as log lines. */
+    if (resolved.invalid) {
+        cbm_log_warn("mem.budget.env.invalid", "value", env, "fallback", "ram_fraction");
+    } else if (resolved.clamped) {
+        char cap_mb[CBM_SZ_32];
+        snprintf(cap_mb, sizeof(cap_mb), "%zu", info.total_ram / MB_DIVISOR);
+        cbm_log_warn("mem.budget.clamped", "requested_mb", env, "cap_mb", cap_mb);
+    }
 
     char budget_mb[CBM_SZ_32];
     char ram_mb[CBM_SZ_32];
     snprintf(budget_mb, sizeof(budget_mb), "%zu", g_budget / MB_DIVISOR);
     snprintf(ram_mb, sizeof(ram_mb), "%zu", info.total_ram / MB_DIVISOR);
-    cbm_log_info("mem.init", "budget_mb", budget_mb, "total_ram_mb", ram_mb);
+    cbm_log_info("mem.init", "budget_mb", budget_mb, "total_ram_mb", ram_mb, "source",
+                 resolved.source);
 }
 
 size_t cbm_mem_rss(void) {

src/foundation/mem.h

@@ -15,10 +15,31 @@
 double cbm_mem_ram_fraction_for_total(size_t total_ram_bytes);
 
 /* Initialize memory budget = ram_fraction * total_physical_ram.
+ * The CBM_MEM_BUDGET_MB env var, when set to a positive integer, overrides
+ * this with an explicit budget in MiB (clamped to physical/cgroup RAM).
  * Thread-safe: only the first call takes effect.
  * Configures mimalloc options for reduced upfront memory. */
 void cbm_mem_init(double ram_fraction);
 
+/* Result of cbm_mem_resolve_budget: the resolved budget plus the metadata
+ * cbm_mem_init logs — so the parse/clamp logic lives in exactly ONE place and
+ * the caller never re-parses the env string. */
+typedef struct {
+    size_t budget;      /* resolved budget in bytes */
+    const char *source; /* log token: "ram_fraction" | "CBM_MEM_BUDGET_MB" */
+    bool clamped;       /* override was valid but exceeded total_ram → clamped down */
+    bool invalid;       /* override was present but unparseable / out-of-range / ≤0 */
+} cbm_mem_budget_t;
+
+/* Pure budget resolver shared by cbm_mem_init (exposed for testing).
+ * Returns ram_fraction * total_ram, unless `budget_mb` is a STRICTLY valid
+ * positive integer string (the CBM_MEM_BUDGET_MB override) — then it returns
+ * that many MiB, clamped to total_ram when total_ram > 0. Trailing garbage,
+ * overflow (ERANGE), and non-positive values are rejected (invalid=true) and
+ * fall back to the fraction-derived value. Reads no globals/env. */
+cbm_mem_budget_t cbm_mem_resolve_budget(size_t total_ram, double ram_fraction,
+                                        const char *budget_mb);
+
 /* Current RSS in bytes via mi_process_info().
  * Falls back to OS-specific queries when MI_OVERRIDE=0 (ASan builds). */
 size_t cbm_mem_rss(void);

tests/repro/repro_issue363.c

@@ -1,74 +1,38 @@
 /*
- * repro_issue363.c — Reproduce-first case for OPEN bug #363.
+ * repro_issue363.c — Regression guard for bug #363 (both axes now FIXED).
  *
  * Issue: #363 — "Linux: cbm_system_info / cbm_default_worker_count don't
  *               respect cgroup CPU/memory limits"
  *
- * ROOT CAUSE (two distinct axes):
+ * Both axes of #363 have shipped; this case is now a permanent GREEN guard:
  *
  *   CPU axis — FIXED in v0.8.0 (commit a5a3d1d).
  *     cbm_detect_cgroup_cpus() reads /sys/fs/cgroup/cpu.max (v2) or
- *     .../cpu/cpu.cfs_quota_us + .../cpu/cpu.cfs_period_us (v1) and the
- *     result is used by detect_system_linux() in system_info.c:226.
- *     cbm_default_worker_count() also honours the CBM_WORKERS env override
- *     (commit d952238).  Both are thoroughly tested in test_platform.c.
+ *     .../cpu/cpu.cfs_quota_us + .../cpu/cpu.cfs_period_us (v1); the result
+ *     feeds detect_system_linux(). cbm_default_worker_count() also honours the
+ *     CBM_WORKERS env override (commit d952238). Both tested in test_platform.c.
  *
- *   Memory axis — STILL OPEN (confirmed by reporter @mayurpise in the last
- *     open comment on #363, 2026-06-25).
- *     cbm_detect_cgroup_mem() similarly reads /sys/fs/cgroup/memory.max (v2)
- *     or .../memory/memory.limit_in_bytes (v1), and detect_system_linux()
- *     uses it (system_info.c:229).  BUT: there is NO env-override knob on
- *     the memory axis.  The CPU axis has CBM_WORKERS; the memory side has
- *     nothing.  On a bare-metal host with no enclosing cgroup, users cannot
- *     cap cbm_mem_init's budget without wrapping the process in a cgroup
- *     scope (as @mayurpise's workaround shows).
+ *   Memory axis — FIXED.
+ *     cbm_detect_cgroup_mem() reads /sys/fs/cgroup/memory.max (v2) or
+ *     .../memory/memory.limit_in_bytes (v1); detect_system_linux() applies
+ *     min(cgroup, host). The missing env knob — the "EXACT OPEN GAP" this
+ *     repro was filed for — now exists: cbm_mem_init() reads CBM_MEM_BUDGET_MB
+ *     and routes it through the pure cbm_mem_resolve_budget() (explicit
+ *     override > implicit detection, mirroring CBM_WORKERS).
  *
- * EXACT OPEN GAP:
- *   A CBM_MEM_BUDGET_MB environment variable (analogous to CBM_WORKERS) that
- *   cbm_mem_init() checks before computing g_budget from info.total_ram.
- *   If set to a valid integer N, cbm_mem_init() should set
- *   g_budget = N * 1024 * 1024, honouring it regardless of cgroup or host RAM.
- *
- * WHY THIS TEST IS RED:
- *   cbm_mem_init() (src/foundation/mem.c) reads cbm_system_info().total_ram
- *   and multiplies by ram_fraction.  It does NOT call cbm_safe_getenv for
- *   CBM_MEM_BUDGET_MB — the override path does not exist.  Setting
- *   CBM_MEM_BUDGET_MB=4096 has no effect; cbm_mem_budget() returns a value
- *   derived from host RAM (or cgroup RAM when inside a container), not from
- *   the env var.  The assertion ASSERT_EQ(cbm_mem_budget(), 4096*1024*1024)
- *   therefore fails on any host whose cgroup or physical RAM != exactly 4 GiB.
- *
- * ROOT CAUSE LOCATION:
- *   src/foundation/mem.c, cbm_mem_init(), after the mimalloc option block
- *   (currently around line 126):
- *     cbm_system_info_t info = cbm_system_info();
- *     g_budget = (size_t)((double)info.total_ram * ram_fraction);
- *   The fix is to insert a cbm_safe_getenv("CBM_MEM_BUDGET_MB", ...) lookup
- *   BEFORE this line and, if valid, set g_budget directly without involving
- *   info.total_ram — mirroring the CBM_WORKERS pattern in
- *   cbm_default_worker_count() (system_info.c:290).
- *
- * INTENDED FIX:
- *   1. In cbm_mem_init(): read CBM_MEM_BUDGET_MB; if set to a valid positive
- *      integer, use that value (in bytes) as g_budget and log it.
- *   2. Test: set CBM_MEM_BUDGET_MB=4096, call cbm_mem_init(0.5), assert
- *      cbm_mem_budget() == 4096 * 1024 * 1024.  This test goes GREEN when
- *      the override is wired.
- *   3. Complementary: on Linux, confirm cbm_system_info().total_ram is capped
- *      by the cgroup memory limit when present — already covered in
- *      test_platform.c via cbm_detect_cgroup_mem() unit tests, but an
- *      integration path via cbm_system_info() is untestable without a seam
- *      that lets callers override the hardcoded "/sys/fs/cgroup" root in
- *      detect_system_linux() (system_info.c:229).
+ * IMPORTANT — clamp semantics:
+ *   cbm_mem_resolve_budget() clamps the override to effective (cgroup/host) RAM
+ *   so it can never claim more than the box has. So a *large* override is NOT
+ *   honoured verbatim on a small host — it clamps to total_ram. This guard
+ *   therefore uses a small budget (REPRO363_BUDGET_MB) that is at/below any
+ *   realistic host's RAM, so it is honoured exactly and the assertion is stable
+ *   on every runner. The clamp behaviour itself is unit-tested separately in
+ *   tests/test_mem.c (resolve_budget_override_clamped_to_total).
  *
  * NOTE on cbm_mem_init() caching:
- *   g_budget is initialised once via atomic_compare_exchange_strong.
- *   The test must run in a process where cbm_mem_init() has NOT been called
- *   yet, OR the test must reset g_initialized — neither is supported today.
- *   The repro works as written because the repro runner does not call
- *   cbm_mem_init() before this suite.  If the initialisation guard is an
- *   issue, the fix also needs a cbm_mem_reset_for_test() hook (test-only,
- *   guarded by CBM_TEST_HOOKS or similar).
+ *   g_budget is initialised once via atomic_compare_exchange_strong, so this
+ *   guard relies on running before any suite that calls cbm_mem_init(); the
+ *   repro runner does not init the budget before this suite.
  */
 
 #include "test_framework.h"
@@ -77,39 +41,27 @@
 #include <stdint.h>
 #include <stdlib.h>
 
-#define REPRO363_BUDGET_MB 4096UL
+/* Deliberately small so the resolver honours it exactly (never clamps) on any
+ * host — see the clamp-semantics note above. */
+#define REPRO363_BUDGET_MB 128UL
 #define REPRO363_BUDGET_BYTES (REPRO363_BUDGET_MB * 1024UL * 1024UL)
 
 /*
  * repro_issue363_mem_budget_env_override
  *
- * Precondition: CBM_MEM_BUDGET_MB=4096 is set before cbm_mem_init() is
- * called.  The budget should be 4096 MiB regardless of host RAM or cgroup.
- *
- * RED condition (current code):
- *   cbm_mem_init() ignores CBM_MEM_BUDGET_MB entirely; cbm_mem_budget()
- *   returns host-RAM * fraction, not 4 GiB.  The assertion fires unless the
- *   test runner happens to be on a machine whose effective RAM is exactly
- *   8 GiB with fraction=0.5 — essentially never.
- *
- * GREEN condition (after fix):
- *   cbm_mem_init() reads CBM_MEM_BUDGET_MB, finds "4096", sets
- *   g_budget = 4096 * 1024 * 1024.  The assertion passes on any machine.
+ * Precondition: CBM_MEM_BUDGET_MB is set before the process's first
+ * cbm_mem_init(). The budget must equal the override (honoured exactly, since
+ * it is below any host's effective RAM), proving the env knob is wired.
  */
 TEST(repro_issue363_mem_budget_env_override) {
-    cbm_setenv("CBM_MEM_BUDGET_MB", "4096", 1);
+    cbm_setenv("CBM_MEM_BUDGET_MB", "128", 1);
 
     cbm_mem_init(0.5);
 
     size_t budget = cbm_mem_budget();
 
     cbm_unsetenv("CBM_MEM_BUDGET_MB");
 
-    /*
-     * RED on current code: budget derives from host/cgroup RAM, not the env
-     * var.  On any machine where effective RAM != 8192 MiB this fails.
-     * GREEN once CBM_MEM_BUDGET_MB is wired in cbm_mem_init().
-     */
     ASSERT_EQ((long long)budget, (long long)REPRO363_BUDGET_BYTES);
 
     PASS();