Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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-<pid>.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
Expand Down
55 changes: 53 additions & 2 deletions src/foundation/mem.c
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#include <mimalloc.h>
#include <stdatomic.h>
#include <stdio.h>
#include <stdlib.h>

#ifdef _WIN32
#ifndef WIN32_LEAN_AND_MEAN
Expand Down Expand Up @@ -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)) {
Expand All @@ -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) {
Expand Down
9 changes: 9 additions & 0 deletions src/foundation/mem.h
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,19 @@
#include <stddef.h>

/* 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);
Expand Down
65 changes: 65 additions & 0 deletions tests/test_mem.c
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
Expand Down
Loading