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