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 <stdio.h>
31 #include <stdlib.h>
32 #include <string.h>
33
34 #if __GLIBC_PREREQ(2, 30) == 0
35 #include <sys/syscall.h>
gettid(void)36 static pid_t gettid(void)
37 {
38 return syscall(SYS_gettid);
39 }
40 #endif
41
42 #ifndef PR_SCHED_CORE
43 #define PR_SCHED_CORE 62
44 # define PR_SCHED_CORE_GET 0
45 # define PR_SCHED_CORE_CREATE 1 /* create unique core_sched cookie */
46 # define PR_SCHED_CORE_SHARE_TO 2 /* push core_sched cookie to pid */
47 # define PR_SCHED_CORE_SHARE_FROM 3 /* pull core_sched cookie to pid */
48 # define PR_SCHED_CORE_MAX 4
49 #endif
50
51 #define MAX_PROCESSES 128
52 #define MAX_THREADS 128
53
54 static const char USAGE[] = "cs_prctl_test [options]\n"
55 " options:\n"
56 " -P : number of processes to create.\n"
57 " -T : number of threads per process to create.\n"
58 " -d : delay time to keep tasks alive.\n"
59 " -k : keep tasks alive until keypress.\n";
60
61 enum pid_type {PIDTYPE_PID = 0, PIDTYPE_TGID, PIDTYPE_PGID};
62
63 const int THREAD_CLONE_FLAGS = CLONE_THREAD | CLONE_SIGHAND | CLONE_FS | CLONE_VM | CLONE_FILES;
64
_prctl(int option,unsigned long arg2,unsigned long arg3,unsigned long arg4,unsigned long arg5)65 static int _prctl(int option, unsigned long arg2, unsigned long arg3, unsigned long arg4,
66 unsigned long arg5)
67 {
68 int res;
69
70 res = prctl(option, arg2, arg3, arg4, arg5);
71 printf("%d = prctl(%d, %ld, %ld, %ld, %lx)\n", res, option, (long)arg2, (long)arg3,
72 (long)arg4, arg5);
73 return res;
74 }
75
76 #define STACK_SIZE (1024 * 1024)
77
78 #define handle_error(msg) __handle_error(__FILE__, __LINE__, msg)
__handle_error(char * fn,int ln,char * msg)79 static void __handle_error(char *fn, int ln, char *msg)
80 {
81 printf("(%s:%d) - ", fn, ln);
82 perror(msg);
83 exit(EXIT_FAILURE);
84 }
85
handle_usage(int rc,char * msg)86 static void handle_usage(int rc, char *msg)
87 {
88 puts(USAGE);
89 puts(msg);
90 putchar('\n');
91 exit(rc);
92 }
93
get_cs_cookie(int pid)94 static unsigned long get_cs_cookie(int pid)
95 {
96 unsigned long long cookie;
97 int ret;
98
99 ret = prctl(PR_SCHED_CORE, PR_SCHED_CORE_GET, pid, PIDTYPE_PID,
100 (unsigned long)&cookie);
101 if (ret) {
102 printf("Not a core sched system\n");
103 return -1UL;
104 }
105
106 return cookie;
107 }
108
109 struct child_args {
110 int num_threads;
111 int pfd[2];
112 int cpid;
113 int thr_tids[MAX_THREADS];
114 };
115
child_func_thread(void * arg)116 static int child_func_thread(void __attribute__((unused))*arg)
117 {
118 while (1)
119 usleep(20000);
120 return 0;
121 }
122
create_threads(int num_threads,int thr_tids[])123 static void create_threads(int num_threads, int thr_tids[])
124 {
125 void *child_stack;
126 pid_t tid;
127 int i;
128
129 for (i = 0; i < num_threads; ++i) {
130 child_stack = malloc(STACK_SIZE);
131 if (!child_stack)
132 handle_error("child stack allocate");
133
134 tid = clone(child_func_thread, child_stack + STACK_SIZE, THREAD_CLONE_FLAGS, NULL);
135 if (tid == -1)
136 handle_error("clone thread");
137 thr_tids[i] = tid;
138 }
139 }
140
child_func_process(void * arg)141 static int child_func_process(void *arg)
142 {
143 struct child_args *ca = (struct child_args *)arg;
144
145 close(ca->pfd[0]);
146
147 create_threads(ca->num_threads, ca->thr_tids);
148
149 write(ca->pfd[1], &ca->thr_tids, sizeof(int) * ca->num_threads);
150 close(ca->pfd[1]);
151
152 while (1)
153 usleep(20000);
154 return 0;
155 }
156
157 static unsigned char child_func_process_stack[STACK_SIZE];
158
create_processes(int num_processes,int num_threads,struct child_args proc[])159 void create_processes(int num_processes, int num_threads, struct child_args proc[])
160 {
161 pid_t cpid;
162 int i;
163
164 for (i = 0; i < num_processes; ++i) {
165 proc[i].num_threads = num_threads;
166
167 if (pipe(proc[i].pfd) == -1)
168 handle_error("pipe() failed");
169
170 cpid = clone(child_func_process, child_func_process_stack + STACK_SIZE,
171 SIGCHLD, &proc[i]);
172 proc[i].cpid = cpid;
173 close(proc[i].pfd[1]);
174 }
175
176 for (i = 0; i < num_processes; ++i) {
177 read(proc[i].pfd[0], &proc[i].thr_tids, sizeof(int) * proc[i].num_threads);
178 close(proc[i].pfd[0]);
179 }
180 }
181
disp_processes(int num_processes,struct child_args proc[])182 void disp_processes(int num_processes, struct child_args proc[])
183 {
184 int i, j;
185
186 printf("tid=%d, / tgid=%d / pgid=%d: %lx\n", gettid(), getpid(), getpgid(0),
187 get_cs_cookie(getpid()));
188
189 for (i = 0; i < num_processes; ++i) {
190 printf(" tid=%d, / tgid=%d / pgid=%d: %lx\n", proc[i].cpid, proc[i].cpid,
191 getpgid(proc[i].cpid), get_cs_cookie(proc[i].cpid));
192 for (j = 0; j < proc[i].num_threads; ++j) {
193 printf(" tid=%d, / tgid=%d / pgid=%d: %lx\n", proc[i].thr_tids[j],
194 proc[i].cpid, getpgid(0), get_cs_cookie(proc[i].thr_tids[j]));
195 }
196 }
197 puts("\n");
198 }
199
200 static int errors;
201
202 #define validate(v) _validate(__LINE__, v, #v)
_validate(int line,int val,char * msg)203 void _validate(int line, int val, char *msg)
204 {
205 if (!val) {
206 ++errors;
207 printf("(%d) FAILED: %s\n", line, msg);
208 } else {
209 printf("(%d) PASSED: %s\n", line, msg);
210 }
211 }
212
main(int argc,char * argv[])213 int main(int argc, char *argv[])
214 {
215 struct child_args procs[MAX_PROCESSES];
216
217 int keypress = 0;
218 int num_processes = 2;
219 int num_threads = 3;
220 int delay = 0;
221 int res = 0;
222 int pidx;
223 int pid;
224 int opt;
225
226 while ((opt = getopt(argc, argv, ":hkT:P:d:")) != -1) {
227 switch (opt) {
228 case 'P':
229 num_processes = (int)strtol(optarg, NULL, 10);
230 break;
231 case 'T':
232 num_threads = (int)strtoul(optarg, NULL, 10);
233 break;
234 case 'd':
235 delay = (int)strtol(optarg, NULL, 10);
236 break;
237 case 'k':
238 keypress = 1;
239 break;
240 case 'h':
241 printf(USAGE);
242 exit(EXIT_SUCCESS);
243 default:
244 handle_usage(20, "unknown option");
245 }
246 }
247
248 if (num_processes < 1 || num_processes > MAX_PROCESSES)
249 handle_usage(1, "Bad processes value");
250
251 if (num_threads < 1 || num_threads > MAX_THREADS)
252 handle_usage(2, "Bad thread value");
253
254 if (keypress)
255 delay = -1;
256
257 srand(time(NULL));
258
259 /* put into separate process group */
260 if (setpgid(0, 0) != 0)
261 handle_error("process group");
262
263 printf("\n## Create a thread/process/process group hiearchy\n");
264 create_processes(num_processes, num_threads, procs);
265 disp_processes(num_processes, procs);
266 validate(get_cs_cookie(0) == 0);
267
268 printf("\n## Set a cookie on entire process group\n");
269 if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_CREATE, 0, PIDTYPE_PGID, 0) < 0)
270 handle_error("core_sched create failed -- PGID");
271 disp_processes(num_processes, procs);
272
273 validate(get_cs_cookie(0) != 0);
274
275 /* get a random process pid */
276 pidx = rand() % num_processes;
277 pid = procs[pidx].cpid;
278
279 validate(get_cs_cookie(0) == get_cs_cookie(pid));
280 validate(get_cs_cookie(0) == get_cs_cookie(procs[pidx].thr_tids[0]));
281
282 printf("\n## Set a new cookie on entire process/TGID [%d]\n", pid);
283 if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_CREATE, pid, PIDTYPE_TGID, 0) < 0)
284 handle_error("core_sched create failed -- TGID");
285 disp_processes(num_processes, procs);
286
287 validate(get_cs_cookie(0) != get_cs_cookie(pid));
288 validate(get_cs_cookie(pid) != 0);
289 validate(get_cs_cookie(pid) == get_cs_cookie(procs[pidx].thr_tids[0]));
290
291 printf("\n## Copy the cookie of current/PGID[%d], to pid [%d] as PIDTYPE_PID\n",
292 getpid(), pid);
293 if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_SHARE_TO, pid, PIDTYPE_PID, 0) < 0)
294 handle_error("core_sched share to itself failed -- PID");
295 disp_processes(num_processes, procs);
296
297 validate(get_cs_cookie(0) == get_cs_cookie(pid));
298 validate(get_cs_cookie(pid) != 0);
299 validate(get_cs_cookie(pid) != get_cs_cookie(procs[pidx].thr_tids[0]));
300
301 printf("\n## Copy cookie from a thread [%d] to current/PGID [%d] as PIDTYPE_PID\n",
302 procs[pidx].thr_tids[0], getpid());
303 if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_SHARE_FROM, procs[pidx].thr_tids[0],
304 PIDTYPE_PID, 0) < 0)
305 handle_error("core_sched share from thread failed -- PID");
306 disp_processes(num_processes, procs);
307
308 validate(get_cs_cookie(0) == get_cs_cookie(procs[pidx].thr_tids[0]));
309 validate(get_cs_cookie(pid) != get_cs_cookie(procs[pidx].thr_tids[0]));
310
311 printf("\n## Copy cookie from current [%d] to current as pidtype PGID\n", getpid());
312 if (_prctl(PR_SCHED_CORE, PR_SCHED_CORE_SHARE_TO, 0, PIDTYPE_PGID, 0) < 0)
313 handle_error("core_sched share to self failed -- PGID");
314 disp_processes(num_processes, procs);
315
316 validate(get_cs_cookie(0) == get_cs_cookie(pid));
317 validate(get_cs_cookie(pid) != 0);
318 validate(get_cs_cookie(pid) == get_cs_cookie(procs[pidx].thr_tids[0]));
319
320 if (errors) {
321 printf("TESTS FAILED. errors: %d\n", errors);
322 res = 10;
323 } else {
324 printf("SUCCESS !!!\n");
325 }
326
327 if (keypress)
328 getchar();
329 else
330 sleep(delay);
331
332 for (pidx = 0; pidx < num_processes; ++pidx)
333 kill(procs[pidx].cpid, 15);
334
335 return res;
336 }
337