From 20cc969e5db3467bddaa35f48ba051a7ae3b3653 Mon Sep 17 00:00:00 2001 From: Brad King Date: Wed, 6 Sep 2017 15:01:50 -0400 Subject: [PATCH] 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 Reviewed-By: Santiago Gimeno --- docs/src/process.rst | 19 ++++++++ include/uv.h | 13 ++++++ src/unix/process.c | 45 ++++++++++++++++++ src/win/process.c | 56 ++++++++++++++++++++++ test/run-tests.c | 49 +++++++++++++++++++ test/task.h | 7 +++ test/test-list.h | 4 ++ test/test-spawn.c | 109 +++++++++++++++++++++++++++++++++++++++++++ 8 files changed, 302 insertions(+) diff --git a/docs/src/process.rst b/docs/src/process.rst index 283a3c24..1ff74d5b 100644 --- a/docs/src/process.rst +++ b/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 diff --git a/include/uv.h b/include/uv.h index b7e4798c..7ca6ace7 100644 --- a/include/uv.h +++ b/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; /* diff --git a/src/unix/process.c b/src/unix/process.c index 7f39a66d..999afc81 100644 --- a/src/unix/process.c +++ b/src/unix/process.c @@ -32,6 +32,7 @@ #include #include #include +#include #if defined(__APPLE__) && !TARGET_OS_IPHONE # include @@ -44,6 +45,14 @@ extern char **environ; # include #endif +#if defined(__linux__) +# define uv__cpu_set_t cpu_set_t +#elif defined(__FreeBSD__) +# include +# include +# include +# 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/src/win/process.c b/src/win/process.c index 026e2e09..89a923c0 100644 --- a/src/win/process.c +++ b/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. */ diff --git a/test/run-tests.c b/test/run-tests.c index da4ac82e..68aff458 100644 --- a/test/run-tests.c +++ b/test/run-tests.c @@ -27,6 +27,13 @@ # include #else # include +# include +#endif + +#if defined(__FreeBSD__) +# include +# include +# include #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); } diff --git a/test/task.h b/test/task.h index f99f3cbb..b04b06e5 100644 --- a/test/task.h +++ b/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_ */ diff --git a/test/test-list.h b/test/test-list.h index 4406f550..ab015193 100644 --- a/test/test-list.h +++ b/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) diff --git a/test/test-spawn.c b/test/test-spawn.c index 38b38bb0..0842a42e 100644 --- a/test/test-spawn.c +++ b/test/test-spawn.c @@ -37,6 +37,12 @@ #else # include # include +# include +# if defined(__FreeBSD__) +# include +# include +# include +# 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