From bb725223f8aa76efafb10fe8f29f21b161732037 Mon Sep 17 00:00:00 2001 From: Sam Li Date: Tue, 23 Jun 2026 14:35:05 -0400 Subject: [PATCH] feat(mem): CBM_MAX_MEMORY_MB explicit memory-budget override MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a CBM_MAX_MEMORY_MB env var that overrides the ram_fraction × total_RAM budget with an explicit cap in MiB. Lets RAM-constrained hosts (no cgroup) cap the in-memory graph, and lets containers pin a budget below the cgroup limit so headroom is left for sibling processes. Same precedence as CBM_WORKERS: explicit override > implicit detection. Budget math is extracted into a pure, testable cbm_mem_resolve_budget() since cbm_mem_init is one-shot per process. Logs source=env|ram_fraction on mem.init and warns on invalid/clamped values. Closes #580. Signed-off-by: Sam Li --- README.md | 1 + src/foundation/mem.c | 55 +++++++++++++++++++++++++++++++++++-- src/foundation/mem.h | 9 ++++++ tests/test_mem.c | 65 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 128 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index b48a297f0..07bbcfcfc 100644 --- a/README.md +++ b/README.md @@ -444,6 +444,7 @@ codebase-memory-mcp config reset auto_index # reset to default | `CBM_DIAGNOSTICS` | `false` | Set to `1` or `true` to enable periodic diagnostics output to `/tmp/cbm-diagnostics-.json`. | | `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_MAX_MEMORY_MB` | *(50% of detected RAM)* | Explicit memory budget in MiB, overriding the default `ram_fraction × total RAM`. Caps the in-memory graph budget on RAM-constrained hosts, and lets containers pin a budget *below* the detected cgroup limit to leave headroom for sibling processes. Clamped to physical/cgroup RAM; non-positive/invalid values are ignored with a warning. | | `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. | ```bash diff --git a/src/foundation/mem.c b/src/foundation/mem.c index 67ef4d14e..9f164bfde 100644 --- a/src/foundation/mem.c +++ b/src/foundation/mem.c @@ -17,6 +17,7 @@ #include #include #include +#include #ifdef _WIN32 #ifndef WIN32_LEAN_AND_MEAN @@ -106,6 +107,30 @@ static void check_pressure(size_t rss) { /* ── Public API ────────────────────────────────────────────────── */ +size_t cbm_mem_resolve_budget(size_t total_ram, double ram_fraction, const char *max_memory_mb) { + if (ram_fraction <= 0.0 || ram_fraction > MAX_RAM_FRACTION) { + ram_fraction = DEFAULT_RAM_FRACTION; + } + size_t budget = (size_t)((double)total_ram * ram_fraction); + + /* Explicit CBM_MAX_MEMORY_MB override (positive integer MiB) wins over the + * fraction-derived default. Clamp to total_ram when known so we never claim + * more than physical/cgroup RAM. Invalid / non-positive values are ignored + * (caller logs a warning). */ + if (max_memory_mb != NULL && max_memory_mb[0] != '\0') { + char *end = NULL; + long long want_mb = strtoll(max_memory_mb, &end, CBM_DECIMAL_BASE); + if (end != max_memory_mb && want_mb > 0) { + size_t want = (size_t)want_mb * MB_DIVISOR; + if (total_ram > 0 && want > total_ram) { + want = total_ram; + } + budget = want; + } + } + return budget; +} + void cbm_mem_init(double ram_fraction) { int expected = 0; if (!atomic_compare_exchange_strong(&g_initialized, &expected, 1)) { @@ -124,13 +149,39 @@ void cbm_mem_init(double ram_fraction) { mi_option_set(mi_option_purge_delay, 0); /* immediate purge, no 1s delay */ cbm_system_info_t info = cbm_system_info(); - g_budget = (size_t)((double)info.total_ram * ram_fraction); + + /* CBM_MAX_MEMORY_MB env override: an explicit memory budget in MiB that + * takes precedence over the ram_fraction-derived value. Lets RAM- + * constrained hosts cap the in-memory graph budget, and lets containers + * pin a budget *below* the detected cgroup limit so headroom is left for + * sibling processes (e.g. an MCP client/parent). Same precedence shape as + * the CBM_WORKERS override: explicit override > implicit detection. (#580) */ + char env_buf[CBM_SZ_32]; + const char *env = cbm_safe_getenv("CBM_MAX_MEMORY_MB", env_buf, sizeof(env_buf), NULL); + g_budget = cbm_mem_resolve_budget(info.total_ram, ram_fraction, env); + + const char *budget_source = "ram_fraction"; + if (env != NULL && env[0] != '\0') { + char *end = NULL; + long long want_mb = strtoll(env, &end, CBM_DECIMAL_BASE); + if (end != env && want_mb > 0) { + budget_source = "env"; + if (info.total_ram > 0 && (size_t)want_mb * MB_DIVISOR > info.total_ram) { + char cap_mb[CBM_SZ_32]; + snprintf(cap_mb, sizeof(cap_mb), "%zu", info.total_ram / MB_DIVISOR); + cbm_log_warn("mem.max.clamped", "requested_mb", env, "cap_mb", cap_mb); + } + } else { + cbm_log_warn("mem.max.invalid", "value", env, "fallback", "ram_fraction"); + } + } 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", + budget_source); } size_t cbm_mem_rss(void) { diff --git a/src/foundation/mem.h b/src/foundation/mem.h index 5001410a8..3613de0d8 100644 --- a/src/foundation/mem.h +++ b/src/foundation/mem.h @@ -12,10 +12,19 @@ #include /* Initialize memory budget = ram_fraction * total_physical_ram. + * The CBM_MAX_MEMORY_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); +/* Pure budget resolver shared by cbm_mem_init (exposed for testing). + * Returns ram_fraction * total_ram, unless `max_memory_mb` is a positive + * integer string (the CBM_MAX_MEMORY_MB override) — then it returns that many + * MiB, clamped to total_ram when total_ram > 0. Invalid / non-positive + * overrides fall back to the fraction-derived value. Reads no globals/env. */ +size_t cbm_mem_resolve_budget(size_t total_ram, double ram_fraction, const char *max_memory_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); diff --git a/tests/test_mem.c b/tests/test_mem.c index debb9b505..688a51073 100644 --- a/tests/test_mem.c +++ b/tests/test_mem.c @@ -290,6 +290,64 @@ TEST(mem_init_second_call_noop) { PASS(); } +/* ── CBM_MAX_MEMORY_MB budget override (pure resolver) ──────────── + * cbm_mem_init is one-shot per process, so the override logic lives in the + * pure cbm_mem_resolve_budget() helper which we can exercise directly. (#580) */ + +#define CBM_TEST_MB ((size_t)1024 * 1024) + +TEST(resolve_budget_no_override_uses_fraction) { + /* No env override → ram_fraction × total_ram. */ + size_t total = 8192 * CBM_TEST_MB; + ASSERT_EQ(cbm_mem_resolve_budget(total, 0.5, NULL), 4096 * CBM_TEST_MB); + ASSERT_EQ(cbm_mem_resolve_budget(total, 0.25, ""), 2048 * CBM_TEST_MB); + PASS(); +} + +TEST(resolve_budget_invalid_fraction_defaults) { + /* Out-of-range fractions fall back to the 0.5 default. */ + size_t total = 8192 * CBM_TEST_MB; + ASSERT_EQ(cbm_mem_resolve_budget(total, 0.0, NULL), 4096 * CBM_TEST_MB); + ASSERT_EQ(cbm_mem_resolve_budget(total, -1.0, NULL), 4096 * CBM_TEST_MB); + ASSERT_EQ(cbm_mem_resolve_budget(total, 1.5, NULL), 4096 * CBM_TEST_MB); + PASS(); +} + +TEST(resolve_budget_override_wins) { + /* The key use case: pin a budget *below* the fraction default. */ + size_t total = 8192 * CBM_TEST_MB; + ASSERT_EQ(cbm_mem_resolve_budget(total, 0.5, "2048"), 2048 * CBM_TEST_MB); + /* Override above the fraction default is also honored (up to total_ram). */ + ASSERT_EQ(cbm_mem_resolve_budget(total, 0.5, "6144"), 6144 * CBM_TEST_MB); + PASS(); +} + +TEST(resolve_budget_override_clamped_to_total) { + /* Override larger than physical/cgroup RAM clamps to total_ram. */ + size_t total = 1024 * CBM_TEST_MB; + ASSERT_EQ(cbm_mem_resolve_budget(total, 0.5, "100000"), total); + PASS(); +} + +TEST(resolve_budget_override_when_total_unknown) { + /* Detection failed (total_ram == 0): override still yields a usable budget + * and is not clamped to zero. */ + ASSERT_EQ(cbm_mem_resolve_budget(0, 0.5, "512"), 512 * CBM_TEST_MB); + PASS(); +} + +TEST(resolve_budget_invalid_override_falls_back) { + /* Non-numeric, zero, and negative overrides are ignored. */ + size_t total = 8192 * CBM_TEST_MB; + size_t fraction_budget = 4096 * CBM_TEST_MB; + ASSERT_EQ(cbm_mem_resolve_budget(total, 0.5, "abc"), fraction_budget); + ASSERT_EQ(cbm_mem_resolve_budget(total, 0.5, "0"), fraction_budget); + ASSERT_EQ(cbm_mem_resolve_budget(total, 0.5, "-512"), fraction_budget); + PASS(); +} + +#undef CBM_TEST_MB + /* ── Arena integration tests ──────────────────────────────────── */ TEST(arena_alloc_and_destroy) { @@ -653,6 +711,13 @@ SUITE(mem) { RUN_TEST(mem_init_negative_fraction); RUN_TEST(mem_init_over_one_fraction); RUN_TEST(mem_init_second_call_noop); + /* CBM_MAX_MEMORY_MB budget override */ + RUN_TEST(resolve_budget_no_override_uses_fraction); + RUN_TEST(resolve_budget_invalid_fraction_defaults); + RUN_TEST(resolve_budget_override_wins); + RUN_TEST(resolve_budget_override_clamped_to_total); + RUN_TEST(resolve_budget_override_when_total_unknown); + RUN_TEST(resolve_budget_invalid_override_falls_back); /* Arena integration */ RUN_TEST(arena_alloc_and_destroy); RUN_TEST(arena_grow_tracks_sizes);