CPU affinity

NAMESPACE

@@ -22,6 +22,7 @@ export(ps_environ)
 export(ps_environ_raw)
 export(ps_exe)
 export(ps_find_tree)
+export(ps_get_cpu_affinity)
 export(ps_get_nice)
 export(ps_gids)
 export(ps_handle)
@@ -45,6 +46,7 @@ export(ps_pids)
 export(ps_ppid)
 export(ps_resume)
 export(ps_send_signal)
+export(ps_set_cpu_affinity)
 export(ps_set_nice)
 export(ps_shared_lib_users)
 export(ps_shared_libs)

NEWS.md

@@ -3,6 +3,9 @@
 
 * ps now compiles on platforms that enable OpenMP (#109).
 
+* New functions `ps_get_cpu_affinity()` and `ps_set_cpu_affinity()` to query
+  and set CPU affinity (#123).
+
 * `ps_memory_info()` now reports memory in bytes instead of pages on Linux (#115)
 
 # ps 1.6.0

R/low-level.R

@@ -1093,3 +1093,66 @@ ps_shared_libs <- function(p = ps_handle()) {
   class(d) <- unique(c("tbl_df", "tbl", class(d)))
   d
 }
+
+#' Query or set CPU affinity
+#'
+#' `ps_get_cpu_affinity()` queries the
+#' [CPU affinity](https://www.linuxjournal.com/article/6799?page=0,0) of
+#' a process. `ps_set_cpu_affinity()` sets the CPU affinity of a process.
+#'
+#' CPU affinity consists in telling the OS to run a process on a limited
+#' set of CPUs only (on Linux cmdline, the `taskset` command is typically
+#' used).
+#'
+#' These functions are only supported on Linux and Windows. They error on macOS.
+#'
+#' @param p Process handle.
+#' @param affinity Integer vector of CPU numbers to restrict a process to.
+#' CPU numbers start with zero, and they have to be smaller than the
+#' number of (logical) CPUs, see [ps_cpu_count()].
+#'
+#' @return `ps_get_cpu_affinity()` returns an integer vector of CPU
+#' numbers, starting with zero.
+#'
+#' `ps_set_cpu_affinity()` returns `NULL`, invisibly.
+#'
+#' @export
+#' @examplesIf ps::ps_is_supported() && ! ps:::is_cran_check() && ! ps::ps_os_type()[["MACOS"]]
+#' # current
+#' orig <- ps_get_cpu_affinity()
+#' orig
+#'
+#' # restrict
+#' ps_set_cpu_affinity(affinity = 0:0)
+#' ps_get_cpu_affinity()
+#'
+#' # restore
+#' ps_set_cpu_affinity(affinity = orig)
+#' ps_get_cpu_affinity()
+
+ps_get_cpu_affinity <- function(p = ps_handle()) {
+  assert_ps_handle(p)
+  type <- ps_os_type()
+  if (!type[["LINUX"]] && !type[["WINDOWS"]]) {
+    stop("`ps_cpu_affinity()` is only supported on Windows and Linux")
+  }
+
+  .Call(psll_get_cpu_aff, p)
+}
+
+#' @export
+#' @rdname ps_get_cpu_affinity
+
+ps_set_cpu_affinity <- function(p = ps_handle(), affinity) {
+  assert_ps_handle(p)
+  type <- ps_os_type()
+  if (!type[["LINUX"]] && !type[["WINDOWS"]]) {
+    stop("`ps_cpu_affinity()` is only supported on Windows and Linux")
+  }
+
+  # check affinity values
+  cnt <- ps_cpu_count()
+  stopifnot(is.integer(affinity), all(affinity < cnt))
+
+  invisible(.Call(psll_set_cpu_aff, p, affinity))
+}

_pkgdown.yml

@@ -30,6 +30,7 @@ reference:
   - ps_handle
   - ps_is_running
   - ps_memory_info
+  - ps_memory_full_info
   - ps_name
   - ps_num_threads
   - ps_pid
@@ -54,6 +55,8 @@ reference:
   - ps_send_signal
   - ps_suspend
   - ps_terminate
+  - ps_get_cpu_affinity
+  - ps_set_cpu_affinity
   - ps_get_nice
   - ps_set_nice
   - ps_windows_nice_values

src/api-linux.c

@@ -15,6 +15,7 @@
 #include <utmp.h>
 #include <mntent.h>
 #include <sys/sysinfo.h>
+#include <sched.h>
 
 #include <Rinternals.h>
 
@@ -1114,6 +1115,100 @@ SEXP ps__memory_maps(SEXP p) {
   return result;
 }
 
+static const int NCPUS_START = sizeof(unsigned long) * CHAR_BIT;
+
+SEXP psll_get_cpu_aff(SEXP p) {
+  ps_handle_t *handle = R_ExternalPtrAddr(p);
+  int cpu, ncpus, count, cpucount_s;
+  pid_t pid;
+  size_t setsize;
+  cpu_set_t *mask = NULL;
+
+  if (!handle) error("Process pointer cleaned up already");
+
+  PS__CHECK_HANDLE(handle);
+
+  ncpus = NCPUS_START;
+  while (1) {
+    setsize = CPU_ALLOC_SIZE(ncpus);
+    mask = CPU_ALLOC(ncpus);
+    if (mask == NULL) {
+      ps__no_memory("");
+      goto error;
+    }
+    if (sched_getaffinity(pid, setsize, mask) == 0) {
+      break;
+    }
+    CPU_FREE(mask);
+    mask = NULL;
+    if (errno != EINVAL) {
+      ps__set_error_from_errno();
+      goto error;
+    }
+    if (ncpus > INT_MAX / 2) {
+      ps__set_error("could not allocate a large enough CPU set");
+      goto error;
+    }
+    ncpus = ncpus * 2;
+  }
+
+  cpucount_s = CPU_COUNT_S(setsize, mask);
+
+  /* This is unsafe, in that in case of an R exception, CPU_FREE will
+     not be called. But it is also very hard to fix without cleancall. */
+  PROTECT_INDEX pidx;
+  SEXP result;
+  PROTECT_WITH_INDEX(result = allocVector(INTSXP, cpucount_s), &pidx);
+  int pp = 0;
+
+  for (cpu = 0, count = cpucount_s; count; cpu++) {
+    if (CPU_ISSET_S(cpu, setsize, mask)) {
+      INTEGER(result)[pp++] = cpu;
+      --count;
+    }
+  }
+  CPU_FREE(mask);
+
+  REPROTECT(result = Rf_lengthgets(result, pp), pidx);
+  UNPROTECT(1);
+  return result;
+
+error:
+  if (mask) CPU_FREE(mask);
+  ps__throw_error();
+  return R_NilValue;
+}
+
+SEXP psll_set_cpu_aff(SEXP p, SEXP affinity) {
+  ps_handle_t *handle = R_ExternalPtrAddr(p);
+  cpu_set_t cpu_set;
+  pid_t pid;
+  int i, seq_len = LENGTH(affinity);
+
+  if (!handle) error("Process pointer cleaned up already");
+
+  pid = handle->pid;
+
+  CPU_ZERO(&cpu_set);
+  for (i = 0; i < seq_len; i++) {
+    int value = INTEGER(affinity)[i];
+    CPU_SET(value, &cpu_set);
+  }
+
+  PS__CHECK_HANDLE(handle);
+
+  if (sched_setaffinity(pid, sizeof(cpu_set), &cpu_set)) {
+    ps__set_error_from_errno();
+    goto error;
+  }
+
+  return R_NilValue;
+
+error:
+  ps__throw_error();
+  return R_NilValue;
+}
+
 SEXP ps__users() {
   struct utmp *ut;
   SEXP result;