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