From 3fdf875716b294ffecbb792a34c64b43ab4c3aeb Mon Sep 17 00:00:00 2001
From: Willy Tarreau <w@1wt.eu>
Date: Wed, 8 Jan 2025 09:04:09 +0100
Subject: [PATCH] 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.
---
 dev/ncpu/ncpu.c | 87 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 87 insertions(+)

diff --git a/dev/ncpu/ncpu.c b/dev/ncpu/ncpu.c
index 62f9fde79..cdc91eb33 100644
--- a/dev/ncpu/ncpu.c
+++ b/dev/ncpu/ncpu.c
@@ -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);
+}