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.

diff --git a/dev/ncpu/ncpu.c b/dev/ncpu/ncpu.c
index 62f9fde79aaa7efc53a807d6ee063eb932af196d..cdc91eb33eb8fd85b59cea01a012ff660f46da86 100644 (file)
--- 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);
+}