1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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>
38static 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
48# define PR_SCHED_CORE_SHARE_TO 2
49# define PR_SCHED_CORE_SHARE_FROM 3
50# define PR_SCHED_CORE_MAX 4
51#endif
52
53#define MAX_PROCESSES 128
54#define MAX_THREADS 128
55
56static 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
63enum pid_type {PIDTYPE_PID = 0, PIDTYPE_TGID, PIDTYPE_PGID};
64
65const int THREAD_CLONE_FLAGS = CLONE_THREAD | CLONE_SIGHAND | CLONE_FS | CLONE_VM | CLONE_FILES;
66
67static 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)
81static 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
88static void handle_usage(int rc, char *msg)
89{
90 puts(USAGE);
91 puts(msg);
92 putchar('\n');
93 exit(rc);
94}
95
96static 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
111struct child_args {
112 int num_threads;
113 int pfd[2];
114 int cpid;
115 int thr_tids[MAX_THREADS];
116};
117
118static int child_func_thread(void __attribute__((unused))*arg)
119{
120 while (1)
121 usleep(20000);
122 return 0;
123}
124
125static 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
143static 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
159static unsigned char child_func_process_stack[STACK_SIZE];
160
161void 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
184void 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
202static int errors;
203
204#define validate(v) _validate(__LINE__, v, #v)
205void _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
215int 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
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
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