1 | /* ========================================================================== |
2 | * setcpuaffinity.c - Linux/BSD setcpuaffinity. |
3 | * -------------------------------------------------------------------------- |
4 | * Copyright (C) 2020 zhenwei pi |
5 | * |
6 | * Permission is hereby granted, free of charge, to any person obtaining a |
7 | * copy of this software and associated documentation files (the |
8 | * "Software"), to deal in the Software without restriction, including |
9 | * without limitation the rights to use, copy, modify, merge, publish, |
10 | * distribute, sublicense, and/or sell copies of the Software, and to permit |
11 | * persons to whom the Software is furnished to do so, subject to the |
12 | * following conditions: |
13 | * |
14 | * The above copyright notice and this permission notice shall be included |
15 | * in all copies or substantial portions of the Software. |
16 | * |
17 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS |
18 | * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF |
19 | * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN |
20 | * NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, |
21 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR |
22 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE |
23 | * USE OR OTHER DEALINGS IN THE SOFTWARE. |
24 | * ========================================================================== |
25 | */ |
26 | #ifndef _GNU_SOURCE |
27 | #define _GNU_SOURCE |
28 | #endif |
29 | #include <stdlib.h> |
30 | #include <string.h> |
31 | #include <ctype.h> |
32 | #ifdef __linux__ |
33 | #include <sched.h> |
34 | #endif |
35 | #ifdef __FreeBSD__ |
36 | #include <sys/param.h> |
37 | #include <sys/cpuset.h> |
38 | #endif |
39 | #ifdef __DragonFly__ |
40 | #include <pthread.h> |
41 | #include <pthread_np.h> |
42 | #endif |
43 | #ifdef __NetBSD__ |
44 | #include <pthread.h> |
45 | #include <sched.h> |
46 | #endif |
47 | #include "config.h" |
48 | |
49 | #ifdef USE_SETCPUAFFINITY |
50 | static const char *next_token(const char *q, int sep) { |
51 | if (q) |
52 | q = strchr(q, sep); |
53 | if (q) |
54 | q++; |
55 | |
56 | return q; |
57 | } |
58 | |
59 | static int next_num(const char *str, char **end, int *result) { |
60 | if (!str || *str == '\0' || !isdigit(*str)) |
61 | return -1; |
62 | |
63 | *result = strtoul(str, end, 10); |
64 | if (str == *end) |
65 | return -1; |
66 | |
67 | return 0; |
68 | } |
69 | |
70 | /* set current thread cpu affinity to cpu list, this function works like |
71 | * taskset command (actually cpulist parsing logic reference to util-linux). |
72 | * example of this function: "0,2,3", "0,2-3", "0-20:2". */ |
73 | void setcpuaffinity(const char *cpulist) { |
74 | const char *p, *q; |
75 | char *end = NULL; |
76 | #ifdef __linux__ |
77 | cpu_set_t cpuset; |
78 | #endif |
79 | #if defined (__FreeBSD__) || defined(__DragonFly__) |
80 | cpuset_t cpuset; |
81 | #endif |
82 | #ifdef __NetBSD__ |
83 | cpuset_t *cpuset; |
84 | #endif |
85 | |
86 | if (!cpulist) |
87 | return; |
88 | |
89 | #ifndef __NetBSD__ |
90 | CPU_ZERO(&cpuset); |
91 | #else |
92 | cpuset = cpuset_create(); |
93 | #endif |
94 | |
95 | q = cpulist; |
96 | while (p = q, q = next_token(q, ','), p) { |
97 | int a, b, s; |
98 | const char *c1, *c2; |
99 | |
100 | if (next_num(p, &end, &a) != 0) |
101 | return; |
102 | |
103 | b = a; |
104 | s = 1; |
105 | p = end; |
106 | |
107 | c1 = next_token(p, '-'); |
108 | c2 = next_token(p, ','); |
109 | |
110 | if (c1 != NULL && (c2 == NULL || c1 < c2)) { |
111 | if (next_num(c1, &end, &b) != 0) |
112 | return; |
113 | |
114 | c1 = end && *end ? next_token(end, ':') : NULL; |
115 | if (c1 != NULL && (c2 == NULL || c1 < c2)) { |
116 | if (next_num(c1, &end, &s) != 0) |
117 | return; |
118 | |
119 | if (s == 0) |
120 | return; |
121 | } |
122 | } |
123 | |
124 | if ((a > b)) |
125 | return; |
126 | |
127 | while (a <= b) { |
128 | #ifndef __NetBSD__ |
129 | CPU_SET(a, &cpuset); |
130 | #else |
131 | cpuset_set(a, cpuset); |
132 | #endif |
133 | a += s; |
134 | } |
135 | } |
136 | |
137 | if (end && *end) |
138 | return; |
139 | |
140 | #ifdef __linux__ |
141 | sched_setaffinity(0, sizeof(cpuset), &cpuset); |
142 | #endif |
143 | #ifdef __FreeBSD__ |
144 | cpuset_setaffinity(CPU_LEVEL_WHICH, CPU_WHICH_TID, -1, sizeof(cpuset), &cpuset); |
145 | #endif |
146 | #ifdef __DragonFly__ |
147 | pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset); |
148 | #endif |
149 | #ifdef __NetBSD__ |
150 | pthread_setaffinity_np(pthread_self(), cpuset_size(cpuset), cpuset); |
151 | cpuset_destroy(cpuset); |
152 | #endif |
153 | } |
154 | |
155 | #endif /* USE_SETCPUAFFINITY */ |
156 | |