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