1 #include <sys/types.h>
2 #include <sys/event.h>
3 #include <sys/ioctl.h>
4 #include <sys/socket.h>
5 #include <sys/sysctl.h>
6 #include <sys/time.h>
7 #include <sys/usched.h>
8 
9 #include <arpa/inet.h>
10 #include <netinet/in.h>
11 
12 #include <err.h>
13 #include <stdio.h>
14 #include <stdlib.h>
15 #include <string.h>
16 #include <unistd.h>
17 
18 #define EVENT_MAX	128
19 
20 static void	mainloop(int, const struct sockaddr_in *, int);
21 
22 static void
23 usage(const char *cmd)
24 {
25 	fprintf(stderr, "%s -p port [-i n_instance] [-r] [-B]\n", cmd);
26 	exit(1);
27 }
28 
29 static int
30 create_socket(const struct sockaddr_in *in, int reuseport)
31 {
32 	int serv_s, on;
33 
34 	serv_s = socket(AF_INET, SOCK_STREAM, 0);
35 	if (serv_s < 0)
36 		err(1, "socket failed");
37 
38 	on = 1;
39 	if (!reuseport) {
40 		if (setsockopt(serv_s, SOL_SOCKET, SO_REUSEADDR,
41 		    &on, sizeof(on)) < 0)
42 			err(1, "setsockopt(REUSEADDR) failed");
43 	} else {
44 		if (setsockopt(serv_s, SOL_SOCKET, SO_REUSEPORT,
45 		    &on, sizeof(on)) < 0)
46 			err(1, "setsockopt(REUSEPORT) failed");
47 	}
48 
49 	on = 1;
50 	if (ioctl(serv_s, FIONBIO, &on, sizeof(on)) < 0)
51 		err(1, "ioctl(FIONBIO) failed");
52 
53 	if (bind(serv_s, (const struct sockaddr *)in, sizeof(*in)) < 0)
54 		err(1, "bind failed");
55 
56 	if (listen(serv_s, -1) < 0)
57 		err(1, "listen failed");
58 	return serv_s;
59 }
60 
61 int
62 main(int argc, char *argv[])
63 {
64 	struct sockaddr_in in;
65 	int opt, ninst, serv_s, i, reuseport, bindcpu;
66 	size_t prm_len;
67 
68 	prm_len = sizeof(ninst);
69 	if (sysctlbyname("hw.ncpu", &ninst, &prm_len, NULL, 0) != 0)
70 		err(2, "sysctl hw.ncpu failed");
71 
72 	memset(&in, 0, sizeof(in));
73 	in.sin_family = AF_INET;
74 	in.sin_addr.s_addr = INADDR_ANY;
75 
76 	bindcpu = 0;
77 	reuseport = 0;
78 
79 	while ((opt = getopt(argc, argv, "Bp:i:r")) != -1) {
80 		switch (opt) {
81 		case 'p':
82 			in.sin_port = htons(atoi(optarg));
83 			break;
84 
85 		case 'i':
86 			ninst = atoi(optarg);
87 			break;
88 
89 		case 'r':
90 			reuseport = 1;
91 			break;
92 
93 		case 'B':
94 			bindcpu = 1;
95 			break;
96 
97 		default:
98 			usage(argv[0]);
99 		}
100 	}
101 
102 	if (ninst < 1 || in.sin_port == 0)
103 		usage(argv[0]);
104 	if (!reuseport)
105 		bindcpu = 0;
106 
107 	serv_s = -1;
108 	if (!reuseport)
109 		serv_s = create_socket(&in, 0);
110 
111 	for (i = 1; i < ninst; ++i) {
112 		pid_t pid;
113 
114 		pid = fork();
115 		if (pid == 0) {
116 			mainloop(serv_s, &in, bindcpu);
117 			exit(0);
118 		} else if (pid < 0) {
119 			err(1, "fork failed");
120 		}
121 	}
122 
123 	mainloop(serv_s, &in, bindcpu);
124 	exit(0);
125 }
126 
127 static void
128 mainloop(int serv_s, const struct sockaddr_in *in, int bindcpu)
129 {
130 	struct kevent change_evt0[EVENT_MAX];
131 	int kq, nchange;
132 
133 	if (serv_s < 0)
134 		serv_s = create_socket(in, 1);
135 
136 	if (bindcpu) {
137 		socklen_t cpu_len;
138 		int cpu;
139 
140 		cpu_len = sizeof(cpu);
141 		if (getsockopt(serv_s, SOL_SOCKET, SO_CPUHINT,
142 		    &cpu, &cpu_len) < 0)
143 			err(1, "getsockopt(CPUHINT) failed");
144 		if (cpu >= 0)
145 			usched_set(getpid(), USCHED_SET_CPU, &cpu, sizeof(cpu));
146 	}
147 
148 	kq = kqueue();
149 	if (kq < 0)
150 		err(1, "kqueue failed");
151 
152 	EV_SET(&change_evt0[0], serv_s, EVFILT_READ, EV_ADD, 0, 0, NULL);
153 	nchange = 1;
154 
155 	for (;;) {
156 		const struct kevent *change_evt = NULL;
157 		struct kevent evt[EVENT_MAX];
158 		int i, nevt;
159 
160 		if (nchange > 0)
161 			change_evt = change_evt0;
162 
163 		nevt = kevent(kq, change_evt, nchange, evt, EVENT_MAX, NULL);
164 		if (nevt < 0)
165 			err(1, "kevent failed");
166 		nchange = 0;
167 
168 		for (i = 0; i < nevt; ++i) {
169 			if (evt[i].ident == (u_int)serv_s) {
170 				while (nchange < EVENT_MAX) {
171 					int s;
172 
173 					s = accept(serv_s, NULL, NULL);
174 					if (s < 0)
175 						break;
176 					EV_SET(&change_evt0[nchange], s,
177 					    EVFILT_READ, EV_ADD, 0, 0, NULL);
178 					++nchange;
179 				}
180 			} else {
181 				close(evt[i].ident);
182 			}
183 		}
184 	}
185 	/* NEVER REACHED */
186 }
187