DEV: ncpu: implement a wrapper mode

The wrapper mode allows to present itself as LD_PRELOAD before loading
haproxy, which is often more convenient since it allows to pass the
number of CPUs in argument. However, this mode is no longer supported by
modern glibcs, so a future patch will come to implement a trick that was
tested to work at least on x86.
This commit is contained in:
Willy Tarreau 2025-01-08 09:04:09 +01:00
parent 25c08562cb
commit 3fdf875716

View File

@ -1,11 +1,16 @@
#define _GNU_SOURCE
#include <limits.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
// gcc -fPIC -shared -O2 -o ncpu{.so,.c}
// NCPU=16 LD_PRELOAD=$PWD/ncpu.so command args...
static char prog_full_path[PATH_MAX];
/* return a cpu_set having the first $NCPU set */
int sched_getaffinity(pid_t pid, size_t cpusetsize, cpu_set_t *mask)
{
@ -30,3 +35,85 @@ int sched_setaffinity(pid_t pid, size_t cpusetsize, const cpu_set_t *mask)
{
return 0;
}
void usage(const char *argv0)
{
fprintf(stderr,
"Usage: %s [-n ncpu] [cmd [args...]]\n"
" Will install itself in LD_PRELOAD before calling <cmd> with args.\n"
" The number of CPUs may also come from variable NCPU or default to %d.\n"
"\n"
"",
argv0, CPU_SETSIZE);
exit(1);
}
/* Called in wrapper mode, no longer supported on recent glibc */
int main(int argc, char **argv)
{
const char *argv0 = argv[0];
char *preload;
int plen;
prog_full_path[0] = 0;
plen = readlink("/proc/self/exe", prog_full_path, sizeof(prog_full_path) - 1);
if (plen != -1)
prog_full_path[plen] = 0;
else
plen = snprintf(prog_full_path, sizeof(prog_full_path), "%s", argv[0]);
while (1) {
argc--;
argv++;
if (argc < 1)
usage(argv0);
if (strcmp(argv[0], "--") == 0) {
argc--;
argv++;
break;
}
else if (strcmp(argv[0], "-n") == 0) {
if (argc < 2)
usage(argv0);
if (setenv("NCPU", argv[1], 1) != 0)
usage(argv0);
argc--;
argv++;
}
else {
/* unknown arg, that's the command */
break;
}
}
/* here the only args left start with the cmd name */
/* now we'll concatenate ourselves at the end of the LD_PRELOAD variable */
preload = getenv("LD_PRELOAD");
if (preload) {
int olen = strlen(preload);
preload = realloc(preload, olen + 1 + plen + 1);
if (!preload) {
perror("realloc");
exit(2);
}
preload[olen] = ' ';
memcpy(preload + olen + 1, prog_full_path, plen);
preload[olen + 1 + plen] = 0;
}
else {
preload = prog_full_path;
}
if (setenv("LD_PRELOAD", preload, 1) < 0) {
perror("setenv");
exit(2);
}
execvp(*argv, argv);
perror("execve");
exit(2);
}