libuv: unix,win: add uv_spawn option to set child CPU affinity mask

Implement it on Linux, FreeBSD, and Windows for now, and fail with
UV_ENOTSUP on other platforms.

Backported from upstream libuv PR 1527, scheduled for inclusion
in libuv 2.0.

3 files changed, 114 insertions, 0 deletions

diff --git a/Utilities/cmlibuv/include/uv.h b/Utilities/cmlibuv/include/uv.h
index a2ba487391..875e30a861 100644
--- a/Utilities/cmlibuv/include/uv.h
+++ b/Utilities/cmlibuv/include/uv.h
@@ -925,6 +925,19 @@ typedef struct uv_process_options_s {
    */
   uv_uid_t uid;
   uv_gid_t gid;
+  /*
+    Libuv can set the child process' CPU affinity mask.  This happens when
+    `cpumask` is non-NULL.  It must point to an array of char values
+    of length `cpumask_size`, whose value must be at least that returned by
+    uv_cpumask_size().  Each byte in the mask can be either zero (false)
+    or non-zero (true) to indicate whether the corresponding processor at
+    that index is included.
+
+    If enabled on an unsupported platform, uv_spawn() will fail with
+    UV_ENOTSUP.
+   */
+  char* cpumask;
+  size_t cpumask_size;
 } uv_process_options_t;
 
 /*
diff --git a/Utilities/cmlibuv/src/unix/process.c b/Utilities/cmlibuv/src/unix/process.c
index 9842710d0e..d7b6ad3109 100644
--- a/Utilities/cmlibuv/src/unix/process.c
+++ b/Utilities/cmlibuv/src/unix/process.c
@@ -32,6 +32,7 @@
 #include <unistd.h>
 #include <fcntl.h>
 #include <poll.h>
+#include <sched.h>
 
 #if defined(__APPLE__) && !TARGET_OS_IPHONE
 # include <crt_externs.h>
@@ -44,6 +45,14 @@ extern char **environ;
 # include <grp.h>
 #endif
 
+#if defined(__linux__)
+# define uv__cpu_set_t cpu_set_t
+#elif defined(__FreeBSD__)
+# include <sys/param.h>
+# include <sys/cpuset.h>
+# include <pthread_np.h>
+# define uv__cpu_set_t cpuset_t
+#endif
 
 static void uv__chld(uv_signal_t* handle, int signum) {
   uv_process_t* process;
@@ -285,6 +294,12 @@ static void uv__process_child_init(const uv_process_options_t* options,
   int err;
   int fd;
   int n;
+#if defined(__linux__) || defined(__FreeBSD__)
+  int r;
+  int i;
+  int cpumask_size;
+  uv__cpu_set_t cpuset;
+#endif
 
   if (options->flags & UV_PROCESS_DETACHED)
     setsid();
@@ -375,6 +390,26 @@ static void uv__process_child_init(const uv_process_options_t* options,
     _exit(127);
   }
 
+#if defined(__linux__) || defined(__FreeBSD__)
+  if (options->cpumask != NULL) {
+    cpumask_size = uv_cpumask_size();
+    assert(options->cpumask_size >= (size_t)cpumask_size);
+
+    CPU_ZERO(&cpuset);
+    for (i = 0; i < cpumask_size; ++i) {
+      if (options->cpumask[i]) {
+        CPU_SET(i, &cpuset);
+      }
+    }
+
+    r = -pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
+    if (r != 0) {
+      uv__write_int(error_fd, r);
+      _exit(127);
+    }
+  }
+#endif
+
   if (options->env != NULL) {
     environ = options->env;
   }
@@ -429,6 +464,16 @@ int uv_spawn(uv_loop_t* loop,
   int i;
   int status;
 
+  if (options->cpumask != NULL) {
+#if defined(__linux__) || defined(__FreeBSD__)
+    if (options->cpumask_size < (size_t)uv_cpumask_size()) {
+      return UV_EINVAL;
+    }
+#else
+    return UV_ENOTSUP;
+#endif
+  }
+
   assert(options->file != NULL);
   assert(!(options->flags & ~(UV_PROCESS_DETACHED |
                               UV_PROCESS_SETGID |
diff --git a/Utilities/cmlibuv/src/win/process.c b/Utilities/cmlibuv/src/win/process.c
index cc06d9e22a..f5f05af147 100644
--- a/Utilities/cmlibuv/src/win/process.c
+++ b/Utilities/cmlibuv/src/win/process.c
@@ -954,6 +954,12 @@ int uv_spawn(uv_loop_t* loop,
     return UV_EINVAL;
   }
 
+  if (options->cpumask != NULL) {
+    if (options->cpumask_size < (size_t)uv_cpumask_size()) {
+      return UV_EINVAL;
+    }
+  }
+
   assert(options->file != NULL);
   assert(!(options->flags & ~(UV_PROCESS_DETACHED |
                               UV_PROCESS_SETGID |
@@ -1084,6 +1090,12 @@ int uv_spawn(uv_loop_t* loop,
     process_flags |= DETACHED_PROCESS | CREATE_NEW_PROCESS_GROUP;
   }
 
+  if (options->cpumask != NULL) {
+    /* Create the child in a suspended state so we have a chance to set
+       its process affinity before it runs.  */
+    process_flags |= CREATE_SUSPENDED;
+  }
+
   if (!CreateProcessW(application_path,
                      arguments,
                      NULL,
@@ -1099,6 +1111,50 @@ int uv_spawn(uv_loop_t* loop,
     goto done;
   }
 
+  if (options->cpumask != NULL) {
+    /* The child is currently suspended.  Set its process affinity
+       or terminate it if we can't.  */
+    int i;
+    int cpumasksize;
+    DWORD_PTR sysmask;
+    DWORD_PTR oldmask;
+    DWORD_PTR newmask;
+
+    cpumasksize = uv_cpumask_size();
+
+    if (!GetProcessAffinityMask(info.hProcess, &oldmask, &sysmask)) {
+      err = GetLastError();
+      TerminateProcess(info.hProcess, 1);
+      goto done;
+    }
+
+    newmask = 0;
+    for (i = 0; i < cpumasksize; i++) {
+      if (options->cpumask[i]) {
+        if (oldmask & (((DWORD_PTR)1) << i)) {
+          newmask |= ((DWORD_PTR)1) << i;
+        } else {
+          err = UV_EINVAL;
+          TerminateProcess(info.hProcess, 1);
+          goto done;
+        }
+      }
+    }
+
+    if (!SetProcessAffinityMask(info.hProcess, newmask)) {
+      err = GetLastError();
+      TerminateProcess(info.hProcess, 1);
+      goto done;
+    }
+
+    /* The process affinity of the child is set.  Let it run.  */
+    if (ResumeThread(info.hThread) == ((DWORD)-1)) {
+      err = GetLastError();
+      TerminateProcess(info.hProcess, 1);
+      goto done;
+    }
+  }
+
   /* Spawn succeeded */
   /* Beyond this point, failure is reported asynchronously. */