From 9cc4bcfdeaa4d15e39aad36cb30af2f8b430154a Mon Sep 17 00:00:00 2001
From: Ben Noordhuis <info@bnoordhuis.nl>
Date: Fri, 4 Mar 2022 22:35:14 +0100
Subject: [PATCH] unix,win: add uv_available_parallelism() (#3499)

Replacement for the usage pattern where people use uv_cpu_info() as an
imperfect heuristic for determining the amount of parallelism that is
available to their programs.

Fixes #3493.
---
 docs/src/misc.rst           | 19 +++++++++++++++++++
 include/uv.h                |  1 +
 src/unix/core.c             | 33 +++++++++++++++++++++++++++++++++
 src/win/util.c              | 17 +++++++++++++++++
 test/test-platform-output.c |  5 +++++
 5 files changed, 75 insertions(+)

diff --git a/docs/src/misc.rst b/docs/src/misc.rst
index 38fbb56081e..bae44814f19 100644
--- a/docs/src/misc.rst
+++ b/docs/src/misc.rst
@@ -334,11 +334,30 @@ API
 
     .. versionadded:: 1.16.0
 
+.. c:function:: unsigned int uv_available_parallelism(void)
+
+    Returns an estimate of the default amount of parallelism a program should
+    use. Always returns a non-zero value.
+
+    On Linux, inspects the calling thread's CPU affinity mask to determine if
+    it has been pinned to specific CPUs.
+
+    On Windows, the available parallelism may be underreported on systems with
+    more than 64 logical CPUs.
+
+    On other platforms, reports the number of CPUs that the operating system
+    considers to be online.
+
+    .. versionadded:: 1.44.0
+
 .. c:function:: int uv_cpu_info(uv_cpu_info_t** cpu_infos, int* count)
 
     Gets information about the CPUs on the system. The `cpu_infos` array will
     have `count` elements and needs to be freed with :c:func:`uv_free_cpu_info`.
 
+    Use :c:func:`uv_available_parallelism` if you need to know how many CPUs
+    are available for threads or child processes.
+
 .. c:function:: void uv_free_cpu_info(uv_cpu_info_t* cpu_infos, int count)
 
     Frees the `cpu_infos` array previously allocated with :c:func:`uv_cpu_info`.
diff --git a/include/uv.h b/include/uv.h
index 9a418b995cd..ee1c94ccd38 100644
--- a/include/uv.h
+++ b/include/uv.h
@@ -1242,6 +1242,7 @@ UV_EXTERN uv_pid_t uv_os_getppid(void);
 UV_EXTERN int uv_os_getpriority(uv_pid_t pid, int* priority);
 UV_EXTERN int uv_os_setpriority(uv_pid_t pid, int priority);
 
+UV_EXTERN unsigned int uv_available_parallelism(void);
 UV_EXTERN int uv_cpu_info(uv_cpu_info_t** cpu_infos, int* count);
 UV_EXTERN void uv_free_cpu_info(uv_cpu_info_t* cpu_infos, int count);
 
diff --git a/src/unix/core.c b/src/unix/core.c
index 0bb334bd952..c0448b31e6c 100644
--- a/src/unix/core.c
+++ b/src/unix/core.c
@@ -84,6 +84,7 @@ extern char** environ;
 #endif
 
 #if defined(__linux__)
+# include <sched.h>
 # include <sys/syscall.h>
 # define uv__accept4 accept4
 #endif
@@ -1648,3 +1649,35 @@ int uv__search_path(const char* prog, char* buf, size_t* buflen) {
   /* Out of tokens (path entries), and no match found */
   return UV_EINVAL;
 }
+
+
+unsigned int uv_available_parallelism(void) {
+#ifdef __linux__
+  cpu_set_t set;
+  long rc;
+
+  memset(&set, 0, sizeof(set));
+
+  /* sysconf(_SC_NPROCESSORS_ONLN) in musl calls sched_getaffinity() but in
+   * glibc it's... complicated... so for consistency try sched_getaffinity()
+   * before falling back to sysconf(_SC_NPROCESSORS_ONLN).
+   */
+  if (0 == sched_getaffinity(0, sizeof(set), &set))
+    rc = CPU_COUNT(&set);
+  else
+    rc = sysconf(_SC_NPROCESSORS_ONLN);
+
+  if (rc < 1)
+    rc = 1;
+
+  return (unsigned) rc;
+#else  /* __linux__ */
+  long rc;
+
+  rc = sysconf(_SC_NPROCESSORS_ONLN);
+  if (rc < 1)
+    rc = 1;
+
+  return (unsigned) rc;
+#endif  /* __linux__ */
+}
diff --git a/src/win/util.c b/src/win/util.c
index ee70ebc2083..99432053cc3 100644
--- a/src/win/util.c
+++ b/src/win/util.c
@@ -536,6 +536,23 @@ int uv_uptime(double* uptime) {
 }
 
 
+unsigned int uv_available_parallelism(void) {
+  SYSTEM_INFO info;
+  unsigned rc;
+
+  /* TODO(bnoordhuis) Use GetLogicalProcessorInformationEx() to support systems
+   * with > 64 CPUs? See https://github.com/libuv/libuv/pull/3458
+   */
+  GetSystemInfo(&info);
+
+  rc = info.dwNumberOfProcessors;
+  if (rc < 1)
+    rc = 1;
+
+  return rc;
+}
+
+
 int uv_cpu_info(uv_cpu_info_t** cpu_infos_ptr, int* cpu_count_ptr) {
   uv_cpu_info_t* cpu_infos;
   SYSTEM_PROCESSOR_PERFORMANCE_INFORMATION* sppi;
diff --git a/test/test-platform-output.c b/test/test-platform-output.c
index 341c7ae54ed..5827dca1cb2 100644
--- a/test/test-platform-output.c
+++ b/test/test-platform-output.c
@@ -41,6 +41,7 @@ TEST_IMPL(platform_output) {
   uv_interface_address_t* interfaces;
   uv_passwd_t pwd;
   uv_utsname_t uname;
+  unsigned par;
   int count;
   int i;
   int err;
@@ -88,6 +89,10 @@ TEST_IMPL(platform_output) {
   printf("  maximum resident set size: %llu\n",
          (unsigned long long) rusage.ru_maxrss);
 
+  par = uv_available_parallelism();
+  ASSERT_GE(par, 1);
+  printf("uv_available_parallelism: %u\n", par);
+
   err = uv_cpu_info(&cpus, &count);
 #if defined(__CYGWIN__) || defined(__MSYS__)
   ASSERT(err == UV_ENOSYS);