process: add CPU affinity mask option to uv_spawn

It allows setting the child process' CPU affinity mask.
Implement it on Linux, FreeBSD, and Windows for now, and fail with
UV_ENOTSUP on other platforms.

Fixes: https://github.com/libuv/libuv/issues/1389
PR-URL: https://github.com/libuv/libuv/pull/1527
Reviewed-By: Ben Noordhuis <info@bnoordhuis.nl>
Reviewed-By: Santiago Gimeno <santiago.gimeno@gmail.com>

docs/src/process.rst

@@ -32,6 +32,8 @@ Data types
             uv_stdio_container_t* stdio;
             uv_uid_t uid;
             uv_gid_t gid;
+            char* cpumask;
+            size_t cpumask_size;
         } uv_process_options_t;
 
 .. c:type:: void (*uv_exit_cb)(uv_process_t*, int64_t exit_status, int term_signal)
@@ -172,6 +174,23 @@ Public members
         This is not supported on Windows, :c:func:`uv_spawn` will fail and set the error
         to ``UV_ENOTSUP``.
 
+.. c:member:: uv_process_options_t.cpumask
+.. c:member:: uv_process_options_t.cpumask_size
+
+    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
+    :c:func:`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.
+
+    .. note::
+
+        If enabled on an unsupported platform, :c:func:`uv_spawn` will fail
+        with ``UV_ENOTSUP``.
+
+    .. versionadded:: 2.0.0
+
 .. c:member:: uv_stdio_container_t.flags
 
     Flags specifying how the stdio container should be passed to the child. See

include/uv.h

@@ -974,6 +974,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;
 
 /*

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 |

src/win/process.c

@@ -949,6 +949,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. */
 

test/run-tests.c

@@ -27,6 +27,13 @@
 # include <io.h>
 #else
 # include <unistd.h>
+# include <sched.h>
+#endif
+
+#if defined(__FreeBSD__)
+# include <sys/param.h>
+# include <sys/cpuset.h>
+# include <pthread_np.h>
 #endif
 
 #include "uv.h"
@@ -200,5 +207,47 @@ static int maybe_run_test(int argc, char **argv) {
   }
 #endif  /* !_WIN32 */
 
+#if !defined(NO_CPU_AFFINITY)
+  if (strcmp(argv[1], "spawn_helper_affinity") == 0) {
+    int i;
+    int r;
+    int cpu;
+    int cpumask_size;
+#ifdef _WIN32
+    DWORD_PTR procmask;
+    DWORD_PTR sysmask;
+#elif defined(__linux__)
+    cpu_set_t cpuset;
+#else
+    cpuset_t cpuset;
+#endif
+
+    cpumask_size = uv_cpumask_size();
+    ASSERT(cpumask_size > 0);
+
+    cpu = atoi(argv[2]);
+    ASSERT(cpu >= 0);
+    ASSERT(cpu < cpumask_size);
+
+    /* verify the mask has the cpu we expect */
+#ifdef _WIN32
+    r = GetProcessAffinityMask(GetCurrentProcess(), &procmask, &sysmask);
+    ASSERT(r != 0);
+    for (i = 0; i < cpumask_size; ++i) {
+      ASSERT(((procmask & (((DWORD_PTR)1) << i)) != 0) == (i == cpu));
+    }
+#else
+    CPU_ZERO(&cpuset);
+    r = pthread_getaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
+    ASSERT(r == 0);
+    for (i = 0; i < cpumask_size; ++i) {
+      ASSERT(CPU_ISSET(i, &cpuset) == (i == cpu));
+    }
+#endif
+
+    return 1;
+  }
+#endif
+
   return run_test(argv[1], 0, 1);
 }

test/task.h

@@ -215,4 +215,11 @@ UNUSED static int can_ipv6(void) {
   "Cygwin runtime hangs on listen+connect in same process."
 #endif
 
+#if !defined(__linux__) && \
+    !defined(__FreeBSD__) && \
+    !defined(_WIN32)
+# define NO_CPU_AFFINITY \
+  "affinity not supported on this platform."
+#endif
+
 #endif /* TASK_H_ */

test/test-list.h

@@ -264,6 +264,8 @@ TEST_DECLARE   (spawn_and_ping)
 TEST_DECLARE   (spawn_preserve_env)
 TEST_DECLARE   (spawn_setuid_fails)
 TEST_DECLARE   (spawn_setgid_fails)
+TEST_DECLARE   (spawn_affinity)
+TEST_DECLARE   (spawn_affinity_invalid_mask)
 TEST_DECLARE   (spawn_stdout_to_file)
 TEST_DECLARE   (spawn_stdout_and_stderr_to_file)
 TEST_DECLARE   (spawn_stdout_and_stderr_to_file2)
@@ -769,6 +771,8 @@ TASK_LIST_START
   TEST_ENTRY  (spawn_preserve_env)
   TEST_ENTRY  (spawn_setuid_fails)
   TEST_ENTRY  (spawn_setgid_fails)
+  TEST_ENTRY  (spawn_affinity)
+  TEST_ENTRY  (spawn_affinity_invalid_mask)
   TEST_ENTRY  (spawn_stdout_to_file)
   TEST_ENTRY  (spawn_stdout_and_stderr_to_file)
   TEST_ENTRY  (spawn_stdout_and_stderr_to_file2)

test/test-spawn.c

@@ -37,6 +37,12 @@
 #else
 # include <unistd.h>
 # include <sys/wait.h>
+# include <sched.h>
+# if defined(__FreeBSD__)
+#  include <sys/param.h>
+#  include <sys/cpuset.h>
+#  include <pthread_np.h>
+# endif
 #endif
 
 
@@ -1443,6 +1449,109 @@ TEST_IMPL(spawn_setgid_fails) {
 }
 #endif
 
+TEST_IMPL(spawn_affinity) {
+#if defined(NO_CPU_AFFINITY)
+  RETURN_SKIP(NO_CPU_AFFINITY);
+#else
+  int i;
+  int r;
+  int cpu;
+  char cpustr[11];
+  char* newmask;
+  int cpumask_size;
+#if defined(_WIN32)
+  DWORD_PTR procmask;
+  DWORD_PTR sysmask;
+#elif defined(__linux__)
+  cpu_set_t cpuset;
+#else
+  cpuset_t cpuset;
+#endif
+
+  cpumask_size = uv_cpumask_size();
+  ASSERT(cpumask_size > 0);
+
+  /* find a cpu we can use */
+  cpu = cpumask_size;
+#ifdef _WIN32
+  r = GetProcessAffinityMask(GetCurrentProcess(), &procmask, &sysmask);
+  ASSERT(r != 0);
+  for (i = 0; i < cpumask_size; ++i) {
+    if (procmask & (((DWORD_PTR)1) << i)) {
+      cpu = i;
+      break;
+    }
+  }
+#else
+  CPU_ZERO(&cpuset);
+  r = pthread_getaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
+  ASSERT(r == 0);
+  for (i = 0; i < cpumask_size; ++i) {
+    if (CPU_ISSET(i, &cpuset)) {
+      cpu = i;
+      break;
+    }
+  }
+#endif
+  ASSERT(cpu < cpumask_size);
+  snprintf(cpustr, sizeof(cpustr), "%d", cpu);
+
+  init_process_options("spawn_helper_affinity", exit_cb);
+
+  /* mask the child to just one cpu */
+  newmask = (char*)calloc(cpumask_size, 1);
+  ASSERT(newmask != NULL);
+  newmask[cpu] = 1;
+  options.cpumask_size = (size_t)cpumask_size;
+  options.cpumask = newmask;
+
+  /* tell the child which one it should get */
+  options.args[2] = cpustr;
+  options.args[3] = "dummy"; /* need 4 args for test/run-tests.c dispatch */
+
+  r = uv_spawn(uv_default_loop(), &process, &options);
+  ASSERT(r == 0);
+
+  r = uv_run(uv_default_loop(), UV_RUN_DEFAULT);
+  ASSERT(r == 0);
+
+  ASSERT(exit_cb_called == 1);
+  ASSERT(close_cb_called == 1);
+
+  free(newmask);
+
+  MAKE_VALGRIND_HAPPY();
+  return 0;
+#endif
+}
+
+TEST_IMPL(spawn_affinity_invalid_mask) {
+#if defined(NO_CPU_AFFINITY)
+  RETURN_SKIP(NO_CPU_AFFINITY);
+#else
+  int r;
+  char newmask[1];
+  int cpumask_size;
+
+  cpumask_size = uv_cpumask_size();
+  ASSERT(cpumask_size > 0);
+
+  init_process_options("", exit_cb);
+
+  /* provide a mask that is too small */
+  newmask[0] = 0;
+  options.cpumask_size = 0;
+  options.cpumask = newmask;
+
+  r = uv_spawn(uv_default_loop(), &process, &options);
+  ASSERT(r == UV_EINVAL);
+
+  ASSERT(exit_cb_called == 0);
+
+  MAKE_VALGRIND_HAPPY();
+  return 0;
+#endif
+}
 
 #ifdef _WIN32