1 files changed, 68 insertions, 34 deletions

diff --git a/bin/pwait/pwait.c b/bin/pwait/pwait.c
index a78c0bb84ca8..59bf0eb93ced 100644
--- a/bin/pwait/pwait.c
+++ b/bin/pwait/pwait.c
@@ -33,7 +33,9 @@
 
 #include <sys/types.h>
 #include <sys/event.h>
+#include <sys/sysctl.h>
 #include <sys/time.h>
+#include <sys/tree.h>
 #include <sys/wait.h>
 
 #include <err.h>
@@ -46,10 +48,25 @@
 #include <sysexits.h>
 #include <unistd.h>
 
+struct pid {
+	RB_ENTRY(pid) entry;
+	pid_t pid;
+};
+
+static int
+pidcmp(const struct pid *a, const struct pid *b)
+{
+	return (a->pid > b->pid ? 1 : a->pid < b->pid ? -1 : 0);
+}
+
+RB_HEAD(pidtree, pid);
+static struct pidtree pids = RB_INITIALIZER(&pids);
+RB_GENERATE_STATIC(pidtree, pid, entry, pidcmp);
+
 static void
 usage(void)
 {
-	fprintf(stderr, "usage: pwait [-t timeout] [-ov] pid ...\n");
+	fprintf(stderr, "usage: pwait [-t timeout] [-opv] pid ...\n");
 	exit(EX_USAGE);
 }
 
@@ -61,21 +78,28 @@ main(int argc, char *argv[])
 {
 	struct itimerval itv;
 	struct kevent *e;
+	struct pid k, *p;
 	char *end, *s;
 	double timeout;
+	size_t sz;
 	long pid;
-	int i, kq, n, nleft, opt, status;
-	bool oflag, tflag, verbose;
+	pid_t mypid;
+	int i, kq, n, ndone, nleft, opt, pid_max, ret, status;
+	bool oflag, pflag, tflag, verbose;
 
 	oflag = false;
+	pflag = false;
 	tflag = false;
 	verbose = false;
 	memset(&itv, 0, sizeof(itv));
 
-	while ((opt = getopt(argc, argv, "ot:v")) != -1) {
+	while ((opt = getopt(argc, argv, "opt:v")) != -1) {
 		switch (opt) {
 		case 'o':
-			oflag = 1;
+			oflag = true;
+			break;
+		case 'p':
+			pflag = true;
 			break;
 		case 't':
 			tflag = true;
@@ -127,16 +151,18 @@ main(int argc, char *argv[])
 		usage();
 	}
 
-	kq = kqueue();
-	if (kq == -1) {
+	if ((kq = kqueue()) < 0)
 		err(EX_OSERR, "kqueue");
-	}
 
-	e = malloc((argc + tflag) * sizeof(struct kevent));
-	if (e == NULL) {
+	sz = sizeof(pid_max);
+	if (sysctlbyname("kern.pid_max", &pid_max, &sz, NULL, 0) != 0) {
+		pid_max = 99999;
+	}
+	if ((e = malloc((argc + tflag) * sizeof(*e))) == NULL) {
 		err(EX_OSERR, "malloc");
 	}
-	nleft = 0;
+	ndone = nleft = 0;
+	mypid = getpid();
 	for (n = 0; n < argc; n++) {
 		s = argv[n];
 		/* Undocumented Solaris compat */
@@ -145,35 +171,37 @@ main(int argc, char *argv[])
 		}
 		errno = 0;
 		pid = strtol(s, &end, 10);
-		if (pid < 0 || *end != '\0' || errno != 0) {
+		if (pid < 0 || pid > pid_max || *end != '\0' || errno != 0) {
 			warnx("%s: bad process id", s);
 			continue;
 		}
-		if (pid == getpid()) {
+		if (pid == mypid) {
 			warnx("%s: skipping my own pid", s);
 			continue;
 		}
-		for (i = 0; i < nleft; i++) {
-			if (e[i].ident == (uintptr_t)pid) {
-				break;
-			}
+		if ((p = malloc(sizeof(*p))) == NULL) {
+			err(EX_OSERR, NULL);
 		}
-		if (i < nleft) {
+		p->pid = pid;
+		if (RB_INSERT(pidtree, &pids, p) != NULL) {
 			/* Duplicate. */
+			free(p);
 			continue;
 		}
 		EV_SET(e + nleft, pid, EVFILT_PROC, EV_ADD, NOTE_EXIT, 0, NULL);
 		if (kevent(kq, e + nleft, 1, NULL, 0, NULL) == -1) {
+			if (errno != ESRCH)
+				err(EX_OSERR, "kevent()");
 			warn("%ld", pid);
-			if (oflag) {
-				exit(EX_OK);
-			}
+			RB_REMOVE(pidtree, &pids, p);
+			free(p);
+			ndone++;
 		} else {
 			nleft++;
 		}
 	}
 
-	if (nleft > 0 && tflag) {
+	if ((ndone == 0 || !oflag) && nleft > 0 && tflag) {
 		/*
 		 * Explicitly detect SIGALRM so that an exit status of 124
 		 * can be returned rather than 142.
@@ -188,7 +216,8 @@ main(int argc, char *argv[])
 			err(EX_OSERR, "setitimer");
 		}
 	}
-	while (nleft > 0) {
+	ret = EX_OK;
+	while ((ndone == 0 || !oflag) && ret == EX_OK && nleft > 0) {
 		n = kevent(kq, NULL, 0, e, nleft + tflag, NULL);
 		if (n == -1) {
 			err(EX_OSERR, "kevent");
@@ -198,29 +227,34 @@ main(int argc, char *argv[])
 				if (verbose) {
 					printf("timeout\n");
 				}
-				exit(124);
+				ret = 124;
 			}
+			pid = e[i].ident;
 			if (verbose) {
 				status = e[i].data;
 				if (WIFEXITED(status)) {
 					printf("%ld: exited with status %d.\n",
-					    (long)e[i].ident,
-					    WEXITSTATUS(status));
+					    pid, WEXITSTATUS(status));
 				} else if (WIFSIGNALED(status)) {
 					printf("%ld: killed by signal %d.\n",
-					    (long)e[i].ident,
-					    WTERMSIG(status));
+					    pid, WTERMSIG(status));
 				} else {
-					printf("%ld: terminated.\n",
-					    (long)e[i].ident);
+					printf("%ld: terminated.\n", pid);
 				}
 			}
-			if (oflag) {
-				exit(EX_OK);
+			k.pid = pid;
+			if ((p = RB_FIND(pidtree, &pids, &k)) != NULL) {
+				RB_REMOVE(pidtree, &pids, p);
+				free(p);
+				ndone++;
 			}
 			--nleft;
 		}
 	}
-
-	exit(EX_OK);
+	if (pflag) {
+		RB_FOREACH(p, pidtree, &pids) {
+			printf("%d\n", p->pid);
+		}
+	}
+	exit(ret);
 }