diff options
Diffstat (limited to 'src/setcpuaffinity.c')
| -rw-r--r-- | src/setcpuaffinity.c | 155 |
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 */ |