Skip to content

Commit 1e15c63

Browse files
feat: add memory-aware global index scheduling
Signed-off-by: Zhiyu <zhiyuzhang001@gmail.com>
1 parent dd3f4e8 commit 1e15c63

15 files changed

Lines changed: 1305 additions & 171 deletions

src/daemon/application.c

Lines changed: 415 additions & 30 deletions
Large diffs are not rendered by default.

src/daemon/application.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ typedef struct {
3636
cbm_index_worker_poll_t (*poll)(void *context, cbm_daemon_application_worker_t worker,
3737
const cbm_index_worker_result_t **result_out);
3838
bool (*cancel)(void *context, cbm_daemon_application_worker_t worker);
39+
bool (*observed_rss)(void *context, cbm_daemon_application_worker_t worker, size_t *bytes_out);
3940
const char *(*log_path)(void *context, cbm_daemon_application_worker_t worker);
4041
void (*destroy)(void *context, cbm_daemon_application_worker_t worker);
4142
} cbm_daemon_application_worker_ops_t;

src/daemon/application_internal.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,9 @@ int cbm_daemon_application_background_initializes_for_test(void);
3737
* that a request QUEUED rather than erroring or starting. */
3838
int cbm_daemon_application_busy_queue_waits_for_test(void);
3939

40+
/* Sum of sampled process-group RSS high-water marks for active index jobs. */
41+
size_t cbm_daemon_application_observed_index_rss_for_test(cbm_daemon_application_t *application);
42+
4043
/* Build the JSON-RPC error substituted for a reply too large to frame (#1375).
4144
* Exposed because the alternative — driving a real >10 MiB reply — needs a
4245
* ~20k-node index, a fixture cost the unit suite should not carry. The

src/discover/discover.c

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -413,6 +413,8 @@ typedef struct {
413413
int count;
414414
int capacity;
415415
int max_files;
416+
size_t max_total_bytes;
417+
size_t total_bytes;
416418
uint64_t deadline_ms;
417419
bool count_only;
418420
bool collect_excluded;
@@ -504,8 +506,15 @@ static void fl_add(file_list_t *fl, const char *abs_path, const char *rel_path,
504506
fl->limit_exceeded = true;
505507
return;
506508
}
509+
size_t measured_size = size > 0 ? (size_t)size : 0;
510+
if (fl->count_only && (fl->total_bytes > fl->max_total_bytes ||
511+
measured_size > fl->max_total_bytes - fl->total_bytes)) {
512+
fl->limit_exceeded = true;
513+
return;
514+
}
507515
if (fl->count_only) {
508516
fl->count++;
517+
fl->total_bytes += measured_size;
509518
return;
510519
}
511520
if (fl->count >= fl->capacity) {
@@ -1117,7 +1126,8 @@ static cbm_discover_status_t discover_impl(const char *repo_path, const cbm_disc
11171126
int *excluded_count_out,
11181127
cbm_ignored_file_t **ignored_out, int *ignored_count_out,
11191128
int *ignored_total_out, bool count_only, int max_files,
1120-
uint64_t deadline_ms) {
1129+
size_t max_total_bytes, uint64_t deadline_ms,
1130+
size_t *total_bytes_out) {
11211131
if (excluded_out) {
11221132
*excluded_out = NULL;
11231133
}
@@ -1133,6 +1143,9 @@ static cbm_discover_status_t discover_impl(const char *repo_path, const cbm_disc
11331143
if (ignored_total_out) {
11341144
*ignored_total_out = 0;
11351145
}
1146+
if (total_bytes_out) {
1147+
*total_bytes_out = 0;
1148+
}
11361149
if (!repo_path || !out || !count || (count_only && max_files < 0)) {
11371150
return CBM_DISCOVER_ERROR;
11381151
}
@@ -1206,6 +1219,7 @@ static cbm_discover_status_t discover_impl(const char *repo_path, const cbm_disc
12061219
/* Walk */
12071220
file_list_t fl = {
12081221
.max_files = count_only ? max_files : -1,
1222+
.max_total_bytes = count_only ? max_total_bytes : SIZE_MAX,
12091223
.deadline_ms = count_only ? deadline_ms : 0,
12101224
.count_only = count_only,
12111225
.collect_excluded = !count_only && excluded_out != NULL,
@@ -1224,6 +1238,9 @@ static cbm_discover_status_t discover_impl(const char *repo_path, const cbm_disc
12241238
cbm_discover_free_excluded(fl.excluded, fl.excluded_count);
12251239
cbm_discover_free_ignored(fl.ignored, fl.ignored_count);
12261240
*count = fl.count;
1241+
if (total_bytes_out) {
1242+
*total_bytes_out = fl.total_bytes;
1243+
}
12271244
if (fl.failed) {
12281245
return CBM_DISCOVER_ERROR;
12291246
}
@@ -1269,7 +1286,7 @@ int cbm_discover_ex2(const char *repo_path, const cbm_discover_opts_t *opts, cbm
12691286
cbm_ignored_file_t **ignored_out, int *ignored_count_out,
12701287
int *ignored_total_out) {
12711288
return discover_impl(repo_path, opts, out, count, excluded_out, excluded_count_out, ignored_out,
1272-
ignored_count_out, ignored_total_out, false, 0, 0);
1289+
ignored_count_out, ignored_total_out, false, 0, SIZE_MAX, 0, NULL);
12731290
}
12741291

12751292
cbm_discover_status_t cbm_discover_count_bounded(const char *repo_path,
@@ -1283,10 +1300,36 @@ cbm_discover_status_t cbm_discover_count_bounded(const char *repo_path,
12831300
}
12841301
cbm_file_info_t *files = NULL;
12851302
int count = 0;
1286-
cbm_discover_status_t status = discover_impl(repo_path, opts, &files, &count, NULL, NULL, NULL,
1287-
NULL, NULL, true, max_files, deadline_ms);
1303+
cbm_discover_status_t status =
1304+
discover_impl(repo_path, opts, &files, &count, NULL, NULL, NULL, NULL, NULL, true,
1305+
max_files, SIZE_MAX, deadline_ms, NULL);
1306+
cbm_discover_free(files, count);
1307+
*count_out = status == CBM_DISCOVER_ERROR ? -1 : count;
1308+
return status;
1309+
}
1310+
1311+
cbm_discover_status_t cbm_discover_measure_bounded(const char *repo_path,
1312+
const cbm_discover_opts_t *opts, int max_files,
1313+
size_t max_total_bytes, uint64_t deadline_ms,
1314+
int *count_out, size_t *total_bytes_out) {
1315+
if (count_out) {
1316+
*count_out = -1;
1317+
}
1318+
if (total_bytes_out) {
1319+
*total_bytes_out = 0;
1320+
}
1321+
if (!repo_path || !count_out || !total_bytes_out || max_files < 0) {
1322+
return CBM_DISCOVER_ERROR;
1323+
}
1324+
cbm_file_info_t *files = NULL;
1325+
int count = 0;
1326+
size_t total_bytes = 0;
1327+
cbm_discover_status_t status =
1328+
discover_impl(repo_path, opts, &files, &count, NULL, NULL, NULL, NULL, NULL, true,
1329+
max_files, max_total_bytes, deadline_ms, &total_bytes);
12881330
cbm_discover_free(files, count);
12891331
*count_out = status == CBM_DISCOVER_ERROR ? -1 : count;
1332+
*total_bytes_out = total_bytes;
12901333
return status;
12911334
}
12921335

src/discover/discover.h

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,15 @@ cbm_discover_status_t cbm_discover_count_bounded(const char *repo_path,
164164
const cbm_discover_opts_t *opts, int max_files,
165165
uint64_t deadline_ms, int *count_out);
166166

167+
/* Count indexable files and their aggregate source bytes without retaining a
168+
* file array. Stops when either limit would be exceeded. This is the daemon's
169+
* pre-spawn resource classifier, so a large job is identified before it can
170+
* consume a daily-worker slot. */
171+
cbm_discover_status_t cbm_discover_measure_bounded(const char *repo_path,
172+
const cbm_discover_opts_t *opts, int max_files,
173+
size_t max_total_bytes, uint64_t deadline_ms,
174+
int *count_out, size_t *total_bytes_out);
175+
167176
/* Like cbm_discover(), but also reports the directory subtrees that were
168177
* skipped during the walk (hardcoded ALWAYS_SKIP/FAST_SKIP dirs + gitignore
169178
* matches), so callers can surface which subtrees were dropped (#411).

src/foundation/subprocess.c

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,11 +22,17 @@
2222
* non-ASCII repo path survives CreateProcess (#423/#20) */
2323
#include <stdlib.h> /* free */
2424
#else
25+
#ifndef __APPLE__
26+
#include <dirent.h>
27+
#endif
2528
#include <errno.h>
2629
#include <fcntl.h>
2730
#include <signal.h>
2831
#ifdef __APPLE__
32+
#include <libproc.h>
2933
#include <spawn.h>
34+
#include <sys/proc_info.h>
35+
#include <sys/sysctl.h>
3036
extern char **environ;
3137
#endif
3238
#include <sys/stat.h>
@@ -423,6 +429,7 @@ struct cbm_subprocess {
423429
bool root_reaped;
424430
uint64_t force_started_ms;
425431
bool containment_failed;
432+
size_t observed_tree_rss;
426433
cbm_proc_result_t result;
427434

428435
#ifdef _WIN32
@@ -436,6 +443,116 @@ struct cbm_subprocess {
436443
#endif
437444
};
438445

446+
static bool cbm_subprocess_tree_rss_sample(cbm_subprocess_t *process, size_t *bytes_out) {
447+
*bytes_out = 0;
448+
#ifdef _WIN32
449+
JOBOBJECT_EXTENDED_LIMIT_INFORMATION information;
450+
memset(&information, 0, sizeof(information));
451+
if (!process->job || !QueryInformationJobObject(process->job, JobObjectExtendedLimitInformation,
452+
&information, sizeof(information), NULL)) {
453+
return false;
454+
}
455+
*bytes_out = (size_t)information.PeakJobMemoryUsed;
456+
return true;
457+
#elif defined(__APPLE__)
458+
int query[4] = {CTL_KERN, KERN_PROC, KERN_PROC_ALL, 0};
459+
size_t length = 0;
460+
if (sysctl(query, 3, NULL, &length, NULL, 0) != 0 || length == 0) {
461+
return false;
462+
}
463+
struct kinfo_proc *processes = malloc(length);
464+
if (!processes || sysctl(query, 3, processes, &length, NULL, 0) != 0) {
465+
free(processes);
466+
return false;
467+
}
468+
size_t total = 0;
469+
size_t count = length / sizeof(*processes);
470+
for (size_t i = 0; i < count; i++) {
471+
if (processes[i].kp_eproc.e_pgid != process->pgid) {
472+
continue;
473+
}
474+
struct rusage_info_v2 usage;
475+
memset(&usage, 0, sizeof(usage));
476+
if (proc_pid_rusage(processes[i].kp_proc.p_pid, RUSAGE_INFO_V2, (rusage_info_t *)&usage) ==
477+
0) {
478+
size_t resident = (size_t)usage.ri_resident_size;
479+
total = resident > SIZE_MAX - total ? SIZE_MAX : total + resident;
480+
}
481+
}
482+
free(processes);
483+
*bytes_out = total;
484+
return true;
485+
#else
486+
DIR *directory = opendir("/proc");
487+
if (!directory) {
488+
return false;
489+
}
490+
long page_size = sysconf(_SC_PAGESIZE);
491+
size_t total = 0;
492+
struct dirent *entry = NULL;
493+
while (page_size > 0 && (entry = readdir(directory)) != NULL) {
494+
char *end = NULL;
495+
long pid = strtol(entry->d_name, &end, 10);
496+
if (pid <= 0 || !end || *end != '\0') {
497+
continue;
498+
}
499+
char path[128];
500+
(void)snprintf(path, sizeof(path), "/proc/%ld/stat", pid);
501+
FILE *stat_file = fopen(path, "r");
502+
char stat_line[4096];
503+
if (!stat_file || !fgets(stat_line, sizeof(stat_line), stat_file)) {
504+
if (stat_file) {
505+
(void)fclose(stat_file);
506+
}
507+
continue;
508+
}
509+
(void)fclose(stat_file);
510+
char *command_end = strrchr(stat_line, ')');
511+
char state = 0;
512+
long parent = 0;
513+
long group = 0;
514+
if (!command_end || sscanf(command_end + 1, " %c %ld %ld", &state, &parent, &group) != 3 ||
515+
group != process->pgid) {
516+
continue;
517+
}
518+
(void)snprintf(path, sizeof(path), "/proc/%ld/statm", pid);
519+
FILE *memory_file = fopen(path, "r");
520+
unsigned long pages = 0;
521+
unsigned long resident_pages = 0;
522+
bool read = memory_file && fscanf(memory_file, "%lu %lu", &pages, &resident_pages) == 2;
523+
if (memory_file) {
524+
(void)fclose(memory_file);
525+
}
526+
if (read) {
527+
size_t resident = resident_pages > SIZE_MAX / (size_t)page_size
528+
? SIZE_MAX
529+
: (size_t)resident_pages * (size_t)page_size;
530+
total = resident > SIZE_MAX - total ? SIZE_MAX : total + resident;
531+
}
532+
}
533+
(void)closedir(directory);
534+
*bytes_out = total;
535+
return true;
536+
#endif
537+
}
538+
539+
bool cbm_subprocess_observed_tree_rss(cbm_subprocess_t *process, size_t *bytes_out) {
540+
if (!bytes_out) {
541+
return false;
542+
}
543+
*bytes_out = 0;
544+
if (!process) {
545+
return false;
546+
}
547+
size_t sample = 0;
548+
bool sampled = cbm_subprocess_tree_rss_sample(process, &sample);
549+
if (sampled && sample > process->observed_tree_rss) {
550+
process->observed_tree_rss = sample;
551+
}
552+
*bytes_out = process->observed_tree_rss;
553+
return sampled || process->observed_tree_rss > 0;
554+
}
555+
439556
static void cbm_subprocess_result_init(cbm_proc_result_t *result) {
440557
result->outcome = CBM_PROC_SPAWN_FAILED;
441558
result->exit_code = -1;

src/foundation/subprocess.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,12 @@ cbm_proc_poll_t cbm_subprocess_poll(cbm_subprocess_t *process, cbm_proc_result_t
125125
* performs signal delivery/escalation. */
126126
bool cbm_subprocess_request_cancel(cbm_subprocess_t *process);
127127

128+
/* Return the largest observed resident byte total for the complete contained
129+
* process tree. The sample is refreshed on each call while the tree is live;
130+
* after exit the cached high-water mark remains available until destroy.
131+
* false means the platform sample was unavailable and leaves *bytes_out zero. */
132+
bool cbm_subprocess_observed_tree_rss(cbm_subprocess_t *process, size_t *bytes_out);
133+
128134
/* Release a terminal handle. This never waits or implicitly cancels; passing a
129135
* still-running handle violates the API contract. NULL is a no-op. */
130136
void cbm_subprocess_destroy(cbm_subprocess_t *process);

src/mcp/index_supervisor.c

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -779,6 +779,10 @@ bool cbm_index_worker_request_cancel(cbm_index_worker_handle_t *handle) {
779779
cbm_subprocess_request_cancel(handle->process);
780780
}
781781

782+
bool cbm_index_worker_observed_tree_rss(cbm_index_worker_handle_t *handle, size_t *bytes_out) {
783+
return handle && cbm_subprocess_observed_tree_rss(handle->process, bytes_out);
784+
}
785+
782786
const char *cbm_index_worker_response_path(const cbm_index_worker_handle_t *handle) {
783787
return handle ? handle->response_path : NULL;
784788
}

src/mcp/index_supervisor.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,7 @@ cbm_index_worker_poll_t cbm_index_worker_poll(cbm_index_worker_handle_t *handle,
179179
* bounded containment failure is explicitly surfaced in the result). The
180180
* owner must stop concurrent cancellation producers before destroy. */
181181
bool cbm_index_worker_request_cancel(cbm_index_worker_handle_t *handle);
182+
bool cbm_index_worker_observed_tree_rss(cbm_index_worker_handle_t *handle, size_t *bytes_out);
182183

183184
/* Borrowed diagnostic paths, stable until destroy. Every start uses securely
184185
* created unique files, so concurrent jobs in one daemon cannot collide. The

0 commit comments

Comments
 (0)