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