summaryrefslogtreecommitdiffstats
path: root/src/setcpuaffinity.c
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--src/setcpuaffinity.c155
1 files changed, 155 insertions, 0 deletions
diff --git a/src/setcpuaffinity.c b/src/setcpuaffinity.c
new file mode 100644
index 0000000..77b1621
--- /dev/null
+++ b/src/setcpuaffinity.c
@@ -0,0 +1,155 @@
+/* ==========================================================================
+ * setcpuaffinity.c - Linux/BSD setcpuaffinity.
+ * --------------------------------------------------------------------------
+ * Copyright (C) 2020 zhenwei pi
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the
+ * "Software"), to deal in the Software without restriction, including
+ * without limitation the rights to use, copy, modify, merge, publish,
+ * distribute, sublicense, and/or sell copies of the Software, and to permit
+ * persons to whom the Software is furnished to do so, subject to the
+ * following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included
+ * in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+ * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
+ * NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+ * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+ * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+ * USE OR OTHER DEALINGS IN THE SOFTWARE.
+ * ==========================================================================
+ */
+#ifndef _GNU_SOURCE
+#define _GNU_SOURCE
+#endif
+#include <stdlib.h>
+#include <string.h>
+#include <ctype.h>
+#ifdef __linux__
+#include <sched.h>
+#endif
+#ifdef __FreeBSD__
+#include <sys/param.h>
+#include <sys/cpuset.h>
+#endif
+#ifdef __DragonFly__
+#include <pthread.h>
+#include <pthread_np.h>
+#endif
+#ifdef __NetBSD__
+#include <pthread.h>
+#include <sched.h>
+#endif
+#include "config.h"
+
+#ifdef USE_SETCPUAFFINITY
+static const char *next_token(const char *q, int sep) {
+ if (q)
+ q = strchr(q, sep);
+ if (q)
+ q++;
+
+ return q;
+}
+
+static int next_num(const char *str, char **end, int *result) {
+ if (!str || *str == '\0' || !isdigit(*str))
+ return -1;
+
+ *result = strtoul(str, end, 10);
+ if (str == *end)
+ return -1;
+
+ return 0;
+}
+
+/* set current thread cpu affinity to cpu list, this function works like
+ * taskset command (actually cpulist parsing logic reference to util-linux).
+ * example of this function: "0,2,3", "0,2-3", "0-20:2". */
+void setcpuaffinity(const char *cpulist) {
+ const char *p, *q;
+ char *end = NULL;
+#ifdef __linux__
+ cpu_set_t cpuset;
+#endif
+#if defined (__FreeBSD__) || defined(__DragonFly__)
+ cpuset_t cpuset;
+#endif
+#ifdef __NetBSD__
+ cpuset_t *cpuset;
+#endif
+
+ if (!cpulist)
+ return;
+
+#ifndef __NetBSD__
+ CPU_ZERO(&cpuset);
+#else
+ cpuset = cpuset_create();
+#endif
+
+ q = cpulist;
+ while (p = q, q = next_token(q, ','), p) {
+ int a, b, s;
+ const char *c1, *c2;
+
+ if (next_num(p, &end, &a) != 0)
+ return;
+
+ b = a;
+ s = 1;
+ p = end;
+
+ c1 = next_token(p, '-');
+ c2 = next_token(p, ',');
+
+ if (c1 != NULL && (c2 == NULL || c1 < c2)) {
+ if (next_num(c1, &end, &b) != 0)
+ return;
+
+ c1 = end && *end ? next_token(end, ':') : NULL;
+ if (c1 != NULL && (c2 == NULL || c1 < c2)) {
+ if (next_num(c1, &end, &s) != 0)
+ return;
+
+ if (s == 0)
+ return;
+ }
+ }
+
+ if ((a > b))
+ return;
+
+ while (a <= b) {
+#ifndef __NetBSD__
+ CPU_SET(a, &cpuset);
+#else
+ cpuset_set(a, cpuset);
+#endif
+ a += s;
+ }
+ }
+
+ if (end && *end)
+ return;
+
+#ifdef __linux__
+ sched_setaffinity(0, sizeof(cpuset), &cpuset);
+#endif
+#ifdef __FreeBSD__
+ cpuset_setaffinity(CPU_LEVEL_WHICH, CPU_WHICH_TID, -1, sizeof(cpuset), &cpuset);
+#endif
+#ifdef __DragonFly__
+ pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
+#endif
+#ifdef __NetBSD__
+ pthread_setaffinity_np(pthread_self(), cpuset_size(cpuset), cpuset);
+ cpuset_destroy(cpuset);
+#endif
+}
+
+#endif /* USE_SETCPUAFFINITY */