fork-exit.c revision 1.2
1/*	$OpenBSD: fork-exit.c,v 1.2 2021/04/29 13:39:22 bluhm Exp $	*/
2
3/*
4 * Copyright (c) 2021 Alexander Bluhm <bluhm@openbsd.org>
5 *
6 * Permission to use, copy, modify, and distribute this software for any
7 * purpose with or without fee is hereby granted, provided that the above
8 * copyright notice and this permission notice appear in all copies.
9 *
10 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
11 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
12 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
13 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
14 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
15 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
16 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
17 */
18
19#include <sys/select.h>
20#include <sys/wait.h>
21
22#include <err.h>
23#include <errno.h>
24#include <fcntl.h>
25#include <limits.h>
26#include <pthread.h>
27#include <signal.h>
28#include <stdio.h>
29#include <stdlib.h>
30#include <unistd.h>
31
32int execute = 0;
33int daemonize = 0;
34int procs = 1;
35int threads = 0;
36int timeout = 30;
37
38pthread_barrier_t thread_barrier;
39char timeoutstr[sizeof("-2147483647")];
40
41static void __dead
42usage(void)
43{
44	fprintf(stderr, "fork-exit [-ed] [-p procs] [-t threads] [-T timeout]\n"
45	    "    -e          child execs sleep(1), default call sleep(3)\n"
46	    "    -d          daemonize, use if already process group leader\n"
47	    "    -p procs    number of processes to fork, default 1\n"
48	    "    -t threads   number of threads to create, default 0\n"
49	    "    -T timeout  parent and children will exit, default 30 sec\n");
50	exit(2);
51}
52
53static void * __dead
54run_thread(void *arg)
55{
56	int error;
57
58	error = pthread_barrier_wait(&thread_barrier);
59	if (error && error != PTHREAD_BARRIER_SERIAL_THREAD)
60		errc(1, error, "pthread_barrier_wait");
61
62	if (sleep(timeout) != 0)
63		err(1, "sleep %d", timeout);
64
65	/* should not happen */
66	_exit(0);
67}
68
69static void
70create_threads(void)
71{
72	pthread_t *thrs;
73	int i, error;
74
75	error = pthread_barrier_init(&thread_barrier, NULL, threads + 1);
76	if (error)
77		errc(1, errno, "pthread_barrier_init");
78
79	thrs = reallocarray(NULL, threads, sizeof(pthread_t));
80	if (thrs == NULL)
81		err(1, "thrs");
82
83	for (i = 0; i < threads; i++) {
84		error = pthread_create(&thrs[i], NULL, run_thread, NULL);
85		if (error)
86			errc(1, error, "pthread_create");
87	}
88
89	error = pthread_barrier_wait(&thread_barrier);
90	if (error && error != PTHREAD_BARRIER_SERIAL_THREAD)
91		errc(1, error, "pthread_barrier_wait");
92
93	/* return to close child's pipe and sleep */
94}
95
96static void __dead
97exec_sleep(void)
98{
99	execl("/bin/sleep", "sleep", timeoutstr, NULL);
100	err(1, "exec sleep");
101}
102
103static void __dead
104run_child(int fd)
105{
106	/* close pipe to parent and sleep until killed */
107	if (execute) {
108		if (fcntl(fd, F_SETFD, FD_CLOEXEC))
109			err(1, "fcntl FD_CLOEXEC");
110		exec_sleep();
111	} else {
112		if (threads)
113			create_threads();
114		if (close(fd) == -1)
115			err(1, "close child");
116		if (sleep(timeout) != 0)
117			err(1, "sleep %d", timeout);
118	}
119
120	/* should not happen */
121	_exit(0);
122}
123
124static void
125sigexit(int sig)
126{
127	int i, status;
128	pid_t pid;
129
130	/* all children must terminate in time */
131	if ((int)alarm(timeout) == -1)
132		err(1, "alarm");
133
134	for (i = 0; i < procs; i++) {
135		pid = wait(&status);
136		if (pid == -1)
137			err(1, "wait");
138		if (!WIFSIGNALED(status))
139			errx(1, "child %d not killed", pid);
140		if(WTERMSIG(status) != SIGTERM)
141			errx(1, "child %d signal %d", pid, WTERMSIG(status));
142	}
143	exit(0);
144}
145
146int
147main(int argc, char *argv[])
148{
149	const char *errstr;
150	int ch, i, fdmax, fdlen, *rfds, waiting;
151	fd_set *fdset;
152	pid_t pgrp;
153	struct timeval tv;
154
155	while ((ch = getopt(argc, argv, "edp:T:t:")) != -1) {
156	switch (ch) {
157		case 'e':
158			execute = 1;
159			break;
160		case 'd':
161			daemonize = 1;
162			break;
163		case 'p':
164			procs = strtonum(optarg, 0, INT_MAX, &errstr);
165			if (errstr != NULL)
166				errx(1, "number of procs is %s: %s", errstr,
167				    optarg);
168			break;
169		case 't':
170			threads = strtonum(optarg, 0, INT_MAX, &errstr);
171			if (errstr != NULL)
172				errx(1, "number of threads is %s: %s", errstr,
173				    optarg);
174			break;
175		case 'T':
176			timeout = strtonum(optarg, 0, INT_MAX, &errstr);
177			if (errstr != NULL)
178				errx(1, "timeout is %s: %s", errstr, optarg);
179			break;
180		default:
181			usage();
182		}
183	}
184	if (execute) {
185		int ret;
186
187		if (threads > 0)
188			errx(1, "execute sleep cannot be used with threads");
189
190		ret = snprintf(timeoutstr, sizeof(timeoutstr), "%d", timeout);
191		if (ret < 0 || (size_t)ret >= sizeof(timeoutstr))
192			err(1, "snprintf");
193	}
194
195	/* become process group leader */
196	if (daemonize) {
197		/* get rid of process leadership */
198		switch (fork()) {
199		case -1:
200			err(1, "fork parent");
201		case 0:
202			break;
203		default:
204			/* parent leaves orphan behind to do the work */
205			_exit(0);
206		}
207	}
208	pgrp = setsid();
209	if (pgrp == -1) {
210		if (!daemonize)
211			warnx("try -d to become process group leader");
212		err(1, "setsid");
213	}
214
215	/* create pipes to keep in contact with children */
216	rfds = reallocarray(NULL, procs, sizeof(int));
217	if (rfds == NULL)
218		err(1, "rfds");
219	fdmax = 0;
220
221	/* fork child processes and pass writing end of pipe */
222	for (i = 0; i < procs; i++) {
223		int pipefds[2], error;
224
225		if (pipe(pipefds) == -1)
226			err(1, "pipe");
227		if (fdmax < pipefds[0])
228			fdmax = pipefds[0];
229		rfds[i] = pipefds[0];
230
231		switch (fork()) {
232		case -1:
233			/* resource temporarily unavailable may happen */
234			error = errno;
235			/* reap children, but not parent */
236			signal(SIGTERM, SIG_IGN);
237			kill(-pgrp, SIGTERM);
238			errc(1, error, "fork child");
239		case 0:
240			/* child closes reading end, read is for the parent */
241			if (close(pipefds[0]) == -1)
242				err(1, "close read");
243			run_child(pipefds[1]);
244			/* cannot happen */
245			_exit(0);
246		default:
247			/* parent closes writing end, write is for the child */
248			if (close(pipefds[1]) == -1)
249				err(1, "close write");
250			break;
251		}
252	}
253
254	/* create select mask with all reading ends of child pipes */
255	fdlen = howmany(fdmax + 1, NFDBITS);
256	fdset = calloc(fdlen, sizeof(fd_mask));
257	if (fdset == NULL)
258		err(1, "fdset");
259	waiting = 0;
260	for (i = 0; i < procs; i++) {
261		FD_SET(rfds[i], fdset);
262		waiting = 1;
263	}
264
265	/* wait until all child processes are waiting */
266	while (waiting) {
267		tv.tv_sec = timeout;
268		tv.tv_usec = 0;
269		errno = ETIMEDOUT;
270		if (select(fdmax + 1, fdset, NULL, NULL, &tv) <= 0)
271			err(1, "select");
272
273		waiting = 0;
274		/* remove fd of children that closed their end  */
275		for (i = 0; i < procs; i++) {
276			if (rfds[i] >= 0) {
277				if (FD_ISSET(rfds[i], fdset)) {
278					if (close(rfds[i]) == -1)
279						err(1, "close parent");
280					FD_CLR(rfds[i], fdset);
281					rfds[i] = -1;
282				} else {
283					FD_SET(rfds[i], fdset);
284					waiting = 1;
285				}
286			}
287		}
288	}
289
290	/* kill all children simultaneously, parent exits in signal handler */
291	if (signal(SIGTERM, sigexit) == SIG_ERR)
292		err(1, "signal SIGTERM");
293	if (kill(-pgrp, SIGTERM) == -1)
294		err(1, "kill %d", -pgrp);
295
296	errx(1, "alive");
297}
298