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>
 36static 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
 54static 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
 61enum pid_type {PIDTYPE_PID = 0, PIDTYPE_TGID, PIDTYPE_PGID};
 62
 63const int THREAD_CLONE_FLAGS = CLONE_THREAD | CLONE_SIGHAND | CLONE_FS | CLONE_VM | CLONE_FILES;
 64
 65struct child_args {
 66	int num_threads;
 67	int pfd[2];
 68	int cpid;
 69	int thr_tids[MAX_THREADS];
 70};
 71
 72static struct child_args procs[MAX_PROCESSES];
 73static int num_processes = 2;
 74static int need_cleanup = 0;
 75
 76static 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)
 90static 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
103static void handle_usage(int rc, char *msg)
104{
105	puts(USAGE);
106	puts(msg);
107	putchar('\n');
108	exit(rc);
109}
110
111static 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
126static int child_func_thread(void __attribute__((unused))*arg)
127{
128	while (1)
129		usleep(20000);
130	return 0;
131}
132
133static 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
151static 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
167static unsigned char child_func_process_stack[STACK_SIZE];
168
169void 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
192void 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
210static int errors;
211
212#define validate(v) _validate(__LINE__, v, #v)
213void _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
223int 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}