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>
37static 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
55static 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
62enum pid_type {PIDTYPE_PID = 0, PIDTYPE_TGID, PIDTYPE_PGID};
63
64const int THREAD_CLONE_FLAGS = CLONE_THREAD | CLONE_SIGHAND | CLONE_FS | CLONE_VM | CLONE_FILES;
65
66struct child_args {
67	int num_threads;
68	int pfd[2];
69	int cpid;
70	int thr_tids[MAX_THREADS];
71};
72
73static struct child_args procs[MAX_PROCESSES];
74static int num_processes = 2;
75static int need_cleanup;
76
77static 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)
91static 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
104static void handle_usage(int rc, char *msg)
105{
106	puts(USAGE);
107	puts(msg);
108	putchar('\n');
109	exit(rc);
110}
111
112static 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
127static int child_func_thread(void __attribute__((unused))*arg)
128{
129	while (1)
130		usleep(20000);
131	return 0;
132}
133
134static 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
152static 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
173static unsigned char child_func_process_stack[STACK_SIZE];
174
175void 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
201void 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
219static int errors;
220
221#define validate(v) _validate(__LINE__, v, #v)
222void _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
232int 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