1 // SPDX-License-Identifier: GPL-2.0+
2 
3 /*
4  * Context switch microbenchmark.
5  *
6  * Copyright 2018, Anton Blanchard, IBM Corp.
7  */
8 
9 #define _GNU_SOURCE
10 #include <assert.h>
11 #include <errno.h>
12 #include <getopt.h>
13 #include <limits.h>
14 #include <linux/futex.h>
15 #include <pthread.h>
16 #include <sched.h>
17 #include <signal.h>
18 #include <stdio.h>
19 #include <stdlib.h>
20 #include <string.h>
21 #include <sys/shm.h>
22 #include <sys/syscall.h>
23 #include <sys/time.h>
24 #include <sys/types.h>
25 #include <sys/wait.h>
26 #include <unistd.h>
27 
28 static unsigned int timeout = 30;
29 
30 static void set_cpu(int cpu)
31 {
32 	cpu_set_t cpuset;
33 
34 	if (cpu == -1)
35 		return;
36 
37 	CPU_ZERO(&cpuset);
38 	CPU_SET(cpu, &cpuset);
39 
40 	if (sched_setaffinity(0, sizeof(cpuset), &cpuset)) {
41 		perror("sched_setaffinity");
42 		exit(1);
43 	}
44 }
45 
46 static void start_process_on(void *(*fn)(void *), void *arg, int cpu)
47 {
48 	int pid;
49 
50 	pid = fork();
51 	if (pid == -1) {
52 		perror("fork");
53 		exit(1);
54 	}
55 
56 	if (pid)
57 		return;
58 
59 	set_cpu(cpu);
60 
61 	fn(arg);
62 
63 	exit(0);
64 }
65 
66 static int cpu;
67 static int do_fork = 0;
68 static int do_vfork = 0;
69 static int do_exec = 0;
70 static char *exec_file;
71 static int exec_target = 0;
72 static unsigned long iterations;
73 static unsigned long iterations_prev;
74 
75 static void run_exec(void)
76 {
77 	char *const argv[] = { "./exec_target", NULL };
78 
79 	if (execve("./exec_target", argv, NULL) == -1) {
80 		perror("execve");
81 		exit(1);
82 	}
83 }
84 
85 static void bench_fork(void)
86 {
87 	while (1) {
88 		pid_t pid = fork();
89 		if (pid == -1) {
90 			perror("fork");
91 			exit(1);
92 		}
93 		if (pid == 0) {
94 			if (do_exec)
95 				run_exec();
96 			_exit(0);
97 		}
98 		pid = waitpid(pid, NULL, 0);
99 		if (pid == -1) {
100 			perror("waitpid");
101 			exit(1);
102 		}
103 		iterations++;
104 	}
105 }
106 
107 static void bench_vfork(void)
108 {
109 	while (1) {
110 		pid_t pid = vfork();
111 		if (pid == -1) {
112 			perror("fork");
113 			exit(1);
114 		}
115 		if (pid == 0) {
116 			if (do_exec)
117 				run_exec();
118 			_exit(0);
119 		}
120 		pid = waitpid(pid, NULL, 0);
121 		if (pid == -1) {
122 			perror("waitpid");
123 			exit(1);
124 		}
125 		iterations++;
126 	}
127 }
128 
129 static void *null_fn(void *arg)
130 {
131 	pthread_exit(NULL);
132 }
133 
134 static void bench_thread(void)
135 {
136 	pthread_t tid;
137 	cpu_set_t cpuset;
138 	pthread_attr_t attr;
139 	int rc;
140 
141 	rc = pthread_attr_init(&attr);
142 	if (rc) {
143 		errno = rc;
144 		perror("pthread_attr_init");
145 		exit(1);
146 	}
147 
148 	if (cpu != -1) {
149 		CPU_ZERO(&cpuset);
150 		CPU_SET(cpu, &cpuset);
151 
152 		rc = pthread_attr_setaffinity_np(&attr, sizeof(cpu_set_t), &cpuset);
153 		if (rc) {
154 			errno = rc;
155 			perror("pthread_attr_setaffinity_np");
156 			exit(1);
157 		}
158 	}
159 
160 	while (1) {
161 		rc = pthread_create(&tid, &attr, null_fn, NULL);
162 		if (rc) {
163 			errno = rc;
164 			perror("pthread_create");
165 			exit(1);
166 		}
167 		rc = pthread_join(tid, NULL);
168 		if (rc) {
169 			errno = rc;
170 			perror("pthread_join");
171 			exit(1);
172 		}
173 		iterations++;
174 	}
175 }
176 
177 static void sigalrm_handler(int junk)
178 {
179 	unsigned long i = iterations;
180 
181 	printf("%ld\n", i - iterations_prev);
182 	iterations_prev = i;
183 
184 	if (--timeout == 0)
185 		kill(0, SIGUSR1);
186 
187 	alarm(1);
188 }
189 
190 static void sigusr1_handler(int junk)
191 {
192 	exit(0);
193 }
194 
195 static void *bench_proc(void *arg)
196 {
197 	signal(SIGALRM, sigalrm_handler);
198 	alarm(1);
199 
200 	if (do_fork)
201 		bench_fork();
202 	else if (do_vfork)
203 		bench_vfork();
204 	else
205 		bench_thread();
206 
207 	return NULL;
208 }
209 
210 static struct option options[] = {
211 	{ "fork", no_argument, &do_fork, 1 },
212 	{ "vfork", no_argument, &do_vfork, 1 },
213 	{ "exec", no_argument, &do_exec, 1 },
214 	{ "timeout", required_argument, 0, 's' },
215 	{ "exec-target", no_argument, &exec_target, 1 },
216 	{ NULL },
217 };
218 
219 static void usage(void)
220 {
221 	fprintf(stderr, "Usage: fork <options> CPU\n\n");
222 	fprintf(stderr, "\t\t--fork\tUse fork() (default threads)\n");
223 	fprintf(stderr, "\t\t--vfork\tUse vfork() (default threads)\n");
224 	fprintf(stderr, "\t\t--exec\tAlso exec() (default no exec)\n");
225 	fprintf(stderr, "\t\t--timeout=X\tDuration in seconds to run (default 30)\n");
226 	fprintf(stderr, "\t\t--exec-target\tInternal option for exec workload\n");
227 }
228 
229 int main(int argc, char *argv[])
230 {
231 	signed char c;
232 
233 	while (1) {
234 		int option_index = 0;
235 
236 		c = getopt_long(argc, argv, "", options, &option_index);
237 
238 		if (c == -1)
239 			break;
240 
241 		switch (c) {
242 		case 0:
243 			if (options[option_index].flag != 0)
244 				break;
245 
246 			usage();
247 			exit(1);
248 			break;
249 
250 		case 's':
251 			timeout = atoi(optarg);
252 			break;
253 
254 		default:
255 			usage();
256 			exit(1);
257 		}
258 	}
259 
260 	if (do_fork && do_vfork) {
261 		usage();
262 		exit(1);
263 	}
264 	if (do_exec && !do_fork && !do_vfork) {
265 		usage();
266 		exit(1);
267 	}
268 
269 	if (do_exec) {
270 		char *dirname = strdup(argv[0]);
271 		int i;
272 		i = strlen(dirname) - 1;
273 		while (i) {
274 			if (dirname[i] == '/') {
275 				dirname[i] = '\0';
276 				if (chdir(dirname) == -1) {
277 					perror("chdir");
278 					exit(1);
279 				}
280 				break;
281 			}
282 			i--;
283 		}
284 	}
285 
286 	if (exec_target) {
287 		exit(0);
288 	}
289 
290 	if (((argc - optind) != 1)) {
291 		cpu = -1;
292 	} else {
293 		cpu = atoi(argv[optind++]);
294 	}
295 
296 	if (do_exec)
297 		exec_file = argv[0];
298 
299 	set_cpu(cpu);
300 
301 	printf("Using ");
302 	if (do_fork)
303 		printf("fork");
304 	else if (do_vfork)
305 		printf("vfork");
306 	else
307 		printf("clone");
308 
309 	if (do_exec)
310 		printf(" + exec");
311 
312 	printf(" on cpu %d\n", cpu);
313 
314 	/* Create a new process group so we can signal everyone for exit */
315 	setpgid(getpid(), getpid());
316 
317 	signal(SIGUSR1, sigusr1_handler);
318 
319 	start_process_on(bench_proc, NULL, cpu);
320 
321 	while (1)
322 		sleep(3600);
323 
324 	return 0;
325 }
326