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>
gettid(void)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
_prctl(int option,unsigned long arg2,unsigned long arg3,unsigned long arg4,unsigned long arg5)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)
__handle_error(char * fn,int ln,char * 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
handle_usage(int rc,char * msg)104 static void handle_usage(int rc, char *msg)
105 {
106 puts(USAGE);
107 puts(msg);
108 putchar('\n');
109 exit(rc);
110 }
111
get_cs_cookie(int pid)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
child_func_thread(void * arg)127 static int child_func_thread(void __attribute__((unused))*arg)
128 {
129 while (1)
130 usleep(20000);
131 return 0;
132 }
133
create_threads(int num_threads,int thr_tids[])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
child_func_process(void * arg)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(ca->num_threads, 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
create_processes(int num_processes,int num_threads,struct child_args proc[])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
disp_processes(int num_processes,struct child_args proc[])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(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(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(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)
_validate(int line,int val,char * msg)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
main(int argc,char * argv[])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(20, "unknown option");
261 }
262 }
263
264 if (num_processes < 1 || num_processes > MAX_PROCESSES)
265 handle_usage(1, "Bad processes value");
266
267 if (num_threads < 1 || num_threads > MAX_THREADS)
268 handle_usage(2, "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, procs);
281 need_cleanup = 1;
282 disp_processes(num_processes, 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, 0, PIDTYPE_PGID, 0) < 0)
287 handle_error("core_sched create failed -- PGID");
288 disp_processes(num_processes, 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, pid, PIDTYPE_TGID, 0) < 0)
301 handle_error("core_sched create failed -- TGID");
302 disp_processes(num_processes, 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, pid, PIDTYPE_PID, 0) < 0)
311 handle_error("core_sched share to itself failed -- PID");
312 disp_processes(num_processes, 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, procs[pidx].thr_tids[0],
321 PIDTYPE_PID, 0) < 0)
322 handle_error("core_sched share from thread failed -- PID");
323 disp_processes(num_processes, 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, 0, PIDTYPE_PGID, 0) < 0)
330 handle_error("core_sched share to self failed -- PGID");
331 disp_processes(num_processes, 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