tools/testing/selftests/sched/cs_prctl_test.c

   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 <[email protected]>
   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;
  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 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 }