360 lines · c
1// SPDX-License-Identifier: GPL-2.0-only2/*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 it10 * under the terms of version 2.1 of the GNU Lesser General Public License as11 * published by the Free Software Foundation.12 *13 * This library is distributed in the hope that it will be useful, but WITHOUT14 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or15 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License16 * for more details.17 *18 * You should have received a copy of the GNU Lesser General Public License19 * along with this library; if not, see <http://www.gnu.org/licenses>.20 */21 22#define _GNU_SOURCE23#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) == 036#include <sys/syscall.h>37static pid_t gettid(void)38{39 return syscall(SYS_gettid);40}41#endif42 43#ifndef PR_SCHED_CORE44#define PR_SCHED_CORE 6245#define PR_SCHED_CORE_GET 046#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 450#endif51 52#define MAX_PROCESSES 12853#define MAX_THREADS 12854 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) < 0338 && errno == EINVAL);339 340 validate(_prctl(PR_SCHED_CORE, PR_SCHED_CORE_SHARE_TO, 0, PIDTYPE_PGID, 1) < 0341 && 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 else353 sleep(delay);354 355 for (pidx = 0; pidx < num_processes; ++pidx)356 kill(procs[pidx].cpid, 15);357 358 return res;359}360