1 | // SPDX-License-Identifier: GPL-2.0-only |
2 | /* |
3 | * Use the core scheduling prctl() to test core scheduling cookies control. |
4 | * |
5 | * Copyright (c) 2021 Oracle and/or its affiliates. |
6 | * Author: Chris Hyser <chris.hyser@oracle.com> |
7 | * |
8 | * |
9 | * This library is free software; you can redistribute it and/or modify it |
10 | * under the terms of version 2.1 of the GNU Lesser General Public License as |
11 | * published by the Free Software Foundation. |
12 | * |
13 | * This library is distributed in the hope that it will be useful, but WITHOUT |
14 | * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or |
15 | * FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License |
16 | * for more details. |
17 | * |
18 | * You should have received a copy of the GNU Lesser General Public License |
19 | * along with this library; if not, see <http://www.gnu.org/licenses>. |
20 | */ |
21 | |
22 | #define _GNU_SOURCE |
23 | #include <sys/eventfd.h> |
24 | #include <sys/wait.h> |
25 | #include <sys/types.h> |
26 | #include <sched.h> |
27 | #include <sys/prctl.h> |
28 | #include <unistd.h> |
29 | #include <time.h> |
30 | #include <errno.h> |
31 | #include <stdio.h> |
32 | #include <stdlib.h> |
33 | #include <string.h> |
34 | |
35 | #if __GLIBC_PREREQ(2, 30) == 0 |
36 | #include <sys/syscall.h> |
37 | static pid_t gettid(void) |
38 | { |
39 | return syscall(SYS_gettid); |
40 | } |
41 | #endif |
42 | |
43 | #ifndef PR_SCHED_CORE |
44 | #define PR_SCHED_CORE 62 |
45 | # define PR_SCHED_CORE_GET 0 |
46 | # define PR_SCHED_CORE_CREATE 1 /* create unique core_sched cookie */ |
47 | # define PR_SCHED_CORE_SHARE_TO 2 /* push core_sched cookie to pid */ |
48 | # define PR_SCHED_CORE_SHARE_FROM 3 /* pull core_sched cookie to pid */ |
49 | # define PR_SCHED_CORE_MAX 4 |
50 | #endif |
51 | |
52 | #define MAX_PROCESSES 128 |
53 | #define MAX_THREADS 128 |
54 | |
55 | static const char USAGE[] = "cs_prctl_test [options]\n" |
56 | " options:\n" |
57 | " -P : number of processes to create.\n" |
58 | " -T : number of threads per process to create.\n" |
59 | " -d : delay time to keep tasks alive.\n" |
60 | " -k : keep tasks alive until keypress.\n" ; |
61 | |
62 | enum pid_type {PIDTYPE_PID = 0, PIDTYPE_TGID, PIDTYPE_PGID}; |
63 | |
64 | const int THREAD_CLONE_FLAGS = CLONE_THREAD | CLONE_SIGHAND | CLONE_FS | CLONE_VM | CLONE_FILES; |
65 | |
66 | struct child_args { |
67 | int num_threads; |
68 | int pfd[2]; |
69 | int cpid; |
70 | int thr_tids[MAX_THREADS]; |
71 | }; |
72 | |
73 | static struct child_args procs[MAX_PROCESSES]; |
74 | static int num_processes = 2; |
75 | static int need_cleanup; |
76 | |
77 | static int _prctl(int option, unsigned long arg2, unsigned long arg3, unsigned long arg4, |
78 | unsigned long arg5) |
79 | { |
80 | int res; |
81 | |
82 | res = prctl(option, arg2, arg3, arg4, arg5); |
83 | printf("%d = prctl(%d, %ld, %ld, %ld, %lx)\n" , res, option, (long)arg2, (long)arg3, |
84 | (long)arg4, arg5); |
85 | return res; |
86 | } |
87 | |
88 | #define STACK_SIZE (1024 * 1024) |
89 | |
90 | #define handle_error(msg) __handle_error(__FILE__, __LINE__, msg) |
91 | static void __handle_error(char *fn, int ln, char *msg) |
92 | { |
93 | int pidx; |
94 | printf("(%s:%d) - " , fn, ln); |
95 | perror(msg); |
96 | if (need_cleanup) { |
97 | for (pidx = 0; pidx < num_processes; ++pidx) |
98 | kill(procs[pidx].cpid, 15); |
99 | need_cleanup = 0; |
100 | } |
101 | exit(EXIT_FAILURE); |
102 | } |
103 | |
104 | static void handle_usage(int rc, char *msg) |
105 | { |
106 | puts(USAGE); |
107 | puts(msg); |
108 | putchar('\n'); |
109 | exit(rc); |
110 | } |
111 | |
112 | static unsigned long get_cs_cookie(int pid) |
113 | { |
114 | unsigned long long cookie; |
115 | int ret; |
116 | |
117 | ret = prctl(PR_SCHED_CORE, PR_SCHED_CORE_GET, pid, PIDTYPE_PID, |
118 | (unsigned long)&cookie); |
119 | if (ret) { |
120 | printf("Not a core sched system\n" ); |
121 | return -1UL; |
122 | } |
123 | |
124 | return cookie; |
125 | } |
126 | |
127 | static int child_func_thread(void __attribute__((unused))*arg) |
128 | { |
129 | while (1) |
130 | usleep(20000); |
131 | return 0; |
132 | } |
133 | |
134 | static void create_threads(int num_threads, int thr_tids[]) |
135 | { |
136 | void *child_stack; |
137 | pid_t tid; |
138 | int i; |
139 | |
140 | for (i = 0; i < num_threads; ++i) { |
141 | child_stack = malloc(STACK_SIZE); |
142 | if (!child_stack) |
143 | handle_error("child stack allocate" ); |
144 | |
145 | tid = clone(child_func_thread, child_stack + STACK_SIZE, THREAD_CLONE_FLAGS, NULL); |
146 | if (tid == -1) |
147 | handle_error("clone thread" ); |
148 | thr_tids[i] = tid; |
149 | } |
150 | } |
151 | |
152 | static int child_func_process(void *arg) |
153 | { |
154 | struct child_args *ca = (struct child_args *)arg; |
155 | int ret; |
156 | |
157 | close(ca->pfd[0]); |
158 | |
159 | create_threads(num_threads: ca->num_threads, thr_tids: ca->thr_tids); |
160 | |
161 | ret = write(ca->pfd[1], &ca->thr_tids, sizeof(int) * ca->num_threads); |
162 | if (ret == -1) |
163 | printf("write failed on pfd[%d] - error (%s)\n" , |
164 | ca->pfd[1], strerror(errno)); |
165 | |
166 | close(ca->pfd[1]); |
167 | |
168 | while (1) |
169 | usleep(20000); |
170 | return 0; |
171 | } |
172 | |
173 | static unsigned char child_func_process_stack[STACK_SIZE]; |
174 | |
175 | void create_processes(int num_processes, int num_threads, struct child_args proc[]) |
176 | { |
177 | pid_t cpid; |
178 | int i, ret; |
179 | |
180 | for (i = 0; i < num_processes; ++i) { |
181 | proc[i].num_threads = num_threads; |
182 | |
183 | if (pipe(proc[i].pfd) == -1) |
184 | handle_error("pipe() failed" ); |
185 | |
186 | cpid = clone(child_func_process, child_func_process_stack + STACK_SIZE, |
187 | SIGCHLD, &proc[i]); |
188 | proc[i].cpid = cpid; |
189 | close(proc[i].pfd[1]); |
190 | } |
191 | |
192 | for (i = 0; i < num_processes; ++i) { |
193 | ret = read(proc[i].pfd[0], &proc[i].thr_tids, sizeof(int) * proc[i].num_threads); |
194 | if (ret == -1) |
195 | printf("read failed on proc[%d].pfd[0] error (%s)\n" , |
196 | i, strerror(errno)); |
197 | close(proc[i].pfd[0]); |
198 | } |
199 | } |
200 | |
201 | void disp_processes(int num_processes, struct child_args proc[]) |
202 | { |
203 | int i, j; |
204 | |
205 | printf("tid=%d, / tgid=%d / pgid=%d: %lx\n" , gettid(), getpid(), getpgid(0), |
206 | get_cs_cookie(pid: getpid())); |
207 | |
208 | for (i = 0; i < num_processes; ++i) { |
209 | printf(" tid=%d, / tgid=%d / pgid=%d: %lx\n" , proc[i].cpid, proc[i].cpid, |
210 | getpgid(proc[i].cpid), get_cs_cookie(pid: proc[i].cpid)); |
211 | for (j = 0; j < proc[i].num_threads; ++j) { |
212 | printf(" tid=%d, / tgid=%d / pgid=%d: %lx\n" , proc[i].thr_tids[j], |
213 | proc[i].cpid, getpgid(0), get_cs_cookie(pid: proc[i].thr_tids[j])); |
214 | } |
215 | } |
216 | puts("\n" ); |
217 | } |
218 | |
219 | static int errors; |
220 | |
221 | #define validate(v) _validate(__LINE__, v, #v) |
222 | void _validate(int line, int val, char *msg) |
223 | { |
224 | if (!val) { |
225 | ++errors; |
226 | printf("(%d) FAILED: %s\n" , line, msg); |
227 | } else { |
228 | printf("(%d) PASSED: %s\n" , line, msg); |
229 | } |
230 | } |
231 | |
232 | int main(int argc, char *argv[]) |
233 | { |
234 | int keypress = 0; |
235 | int num_threads = 3; |
236 | int delay = 0; |
237 | int res = 0; |
238 | int pidx; |
239 | int pid; |
240 | int opt; |
241 | |
242 | while ((opt = getopt(argc, argv, ":hkT:P:d:" )) != -1) { |
243 | switch (opt) { |
244 | case 'P': |
245 | num_processes = (int)strtol(optarg, NULL, 10); |
246 | break; |
247 | case 'T': |
248 | num_threads = (int)strtoul(optarg, NULL, 10); |
249 | break; |
250 | case 'd': |
251 | delay = (int)strtol(optarg, NULL, 10); |
252 | break; |
253 | case 'k': |
254 | keypress = 1; |
255 | break; |
256 | case 'h': |
257 | printf(USAGE); |
258 | exit(EXIT_SUCCESS); |
259 | default: |
260 | handle_usage(rc: 20, msg: "unknown option" ); |
261 | } |
262 | } |
263 | |
264 | if (num_processes < 1 || num_processes > MAX_PROCESSES) |
265 | handle_usage(rc: 1, msg: "Bad processes value" ); |
266 | |
267 | if (num_threads < 1 || num_threads > MAX_THREADS) |
268 | handle_usage(rc: 2, msg: "Bad thread value" ); |
269 | |
270 | if (keypress) |
271 | delay = -1; |
272 | |
273 | srand(time(NULL)); |
274 | |
275 | /* put into separate process group */ |
276 | if (setpgid(0, 0) != 0) |
277 | handle_error("process group" ); |
278 | |
279 | printf("\n## Create a thread/process/process group hierarchy\n" ); |
280 | create_processes(num_processes, num_threads, proc: procs); |
281 | need_cleanup = 1; |
282 | disp_processes(num_processes, proc: procs); |
283 | validate(get_cs_cookie(0) == 0); |
284 | |
285 | printf("\n## Set a cookie on entire process group\n" ); |
286 | if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_CREATE, arg3: 0, arg4: PIDTYPE_PGID, arg5: 0) < 0) |
287 | handle_error("core_sched create failed -- PGID" ); |
288 | disp_processes(num_processes, proc: procs); |
289 | |
290 | validate(get_cs_cookie(0) != 0); |
291 | |
292 | /* get a random process pid */ |
293 | pidx = rand() % num_processes; |
294 | pid = procs[pidx].cpid; |
295 | |
296 | validate(get_cs_cookie(0) == get_cs_cookie(pid)); |
297 | validate(get_cs_cookie(0) == get_cs_cookie(procs[pidx].thr_tids[0])); |
298 | |
299 | printf("\n## Set a new cookie on entire process/TGID [%d]\n" , pid); |
300 | if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_CREATE, arg3: pid, arg4: PIDTYPE_TGID, arg5: 0) < 0) |
301 | handle_error("core_sched create failed -- TGID" ); |
302 | disp_processes(num_processes, proc: procs); |
303 | |
304 | validate(get_cs_cookie(0) != get_cs_cookie(pid)); |
305 | validate(get_cs_cookie(pid) != 0); |
306 | validate(get_cs_cookie(pid) == get_cs_cookie(procs[pidx].thr_tids[0])); |
307 | |
308 | printf("\n## Copy the cookie of current/PGID[%d], to pid [%d] as PIDTYPE_PID\n" , |
309 | getpid(), pid); |
310 | if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_SHARE_TO, arg3: pid, arg4: PIDTYPE_PID, arg5: 0) < 0) |
311 | handle_error("core_sched share to itself failed -- PID" ); |
312 | disp_processes(num_processes, proc: procs); |
313 | |
314 | validate(get_cs_cookie(0) == get_cs_cookie(pid)); |
315 | validate(get_cs_cookie(pid) != 0); |
316 | validate(get_cs_cookie(pid) != get_cs_cookie(procs[pidx].thr_tids[0])); |
317 | |
318 | printf("\n## Copy cookie from a thread [%d] to current/PGID [%d] as PIDTYPE_PID\n" , |
319 | procs[pidx].thr_tids[0], getpid()); |
320 | if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_SHARE_FROM, arg3: procs[pidx].thr_tids[0], |
321 | arg4: PIDTYPE_PID, arg5: 0) < 0) |
322 | handle_error("core_sched share from thread failed -- PID" ); |
323 | disp_processes(num_processes, proc: procs); |
324 | |
325 | validate(get_cs_cookie(0) == get_cs_cookie(procs[pidx].thr_tids[0])); |
326 | validate(get_cs_cookie(pid) != get_cs_cookie(procs[pidx].thr_tids[0])); |
327 | |
328 | printf("\n## Copy cookie from current [%d] to current as pidtype PGID\n" , getpid()); |
329 | if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_SHARE_TO, arg3: 0, arg4: PIDTYPE_PGID, arg5: 0) < 0) |
330 | handle_error("core_sched share to self failed -- PGID" ); |
331 | disp_processes(num_processes, proc: procs); |
332 | |
333 | validate(get_cs_cookie(0) == get_cs_cookie(pid)); |
334 | validate(get_cs_cookie(pid) != 0); |
335 | validate(get_cs_cookie(pid) == get_cs_cookie(procs[pidx].thr_tids[0])); |
336 | |
337 | validate(_prctl(PR_SCHED_CORE, PR_SCHED_CORE_MAX, 0, PIDTYPE_PGID, 0) < 0 |
338 | && errno == EINVAL); |
339 | |
340 | validate(_prctl(PR_SCHED_CORE, PR_SCHED_CORE_SHARE_TO, 0, PIDTYPE_PGID, 1) < 0 |
341 | && errno == EINVAL); |
342 | |
343 | if (errors) { |
344 | printf("TESTS FAILED. errors: %d\n" , errors); |
345 | res = 10; |
346 | } else { |
347 | printf("SUCCESS !!!\n" ); |
348 | } |
349 | |
350 | if (keypress) |
351 | getchar(); |
352 | else |
353 | sleep(delay); |
354 | |
355 | for (pidx = 0; pidx < num_processes; ++pidx) |
356 | kill(procs[pidx].cpid, 15); |
357 | |
358 | return res; |
359 | } |
360 | |