1/*-
2 * Copyright (c) 2008-2009 Robert N. M. Watson
3 * Copyright (c) 2010 Juniper Networks, Inc.
4 * All rights reserved.
5 *
6 * This software was developed by Robert N. M. Watson under contract
7 * to Juniper Networks, Inc.
8 *
9 * Redistribution and use in source and binary forms, with or without
10 * modification, are permitted provided that the following conditions
11 * are met:
12 * 1. Redistributions of source code must retain the above copyright
13 *    notice, this list of conditions and the following disclaimer.
14 * 2. Redistributions in binary form must reproduce the above copyright
15 *    notice, this list of conditions and the following disclaimer in the
16 *    documentation and/or other materials provided with the distribution.
17 *
18 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
19 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
20 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
21 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
22 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
23 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
24 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
25 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
26 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
27 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
28 * SUCH DAMAGE.
29 */
30
31#include <sys/types.h>
32#include <sys/event.h>
33#include <sys/resource.h>
34#include <sys/sched.h>
35#include <sys/socket.h>
36#include <sys/sysctl.h>
37#include <sys/time.h>
38#include <sys/uio.h>
39#include <sys/wait.h>
40
41#include <netinet/in.h>
42#include <netinet/tcp.h>
43
44#include <err.h>
45#include <errno.h>
46#include <fcntl.h>
47#include <inttypes.h>
48#include <signal.h>
49#include <stdio.h>
50#include <stdlib.h>
51#include <string.h>
52#include <unistd.h>
53
54#include "tcpp.h"
55
56#define	min(x, y)	(x < y ? x : y)
57
58
59/*
60 * Gist of each client worker: build up to mflag connections at a time, and
61 * pump data in to them somewhat fairly until tflag connections have been
62 * completed.
63 */
64#define	CONNECTION_MAGIC	0x87a3f56e
65struct connection {
66	uint32_t	conn_magic;		/* Just magic. */
67	int		conn_fd;
68	struct tcpp_header	conn_header;	/* Header buffer. */
69	u_int		conn_header_sent;	/* Header bytes sent. */
70	u_int64_t	conn_data_sent;		/* Data bytes sent.*/
71};
72
73static u_char			 buffer[256 * 1024];	/* Buffer to send. */
74static pid_t			*pid_list;
75static int			 kq;
76static int			 started;	/* Number started so far. */
77static int			 finished;	/* Number finished so far. */
78static int			 counter;	/* IP number offset. */
79static uint64_t			 payload_len;
80
81static struct connection *
82tcpp_client_newconn(void)
83{
84	struct sockaddr_in sin;
85	struct connection *conn;
86	struct kevent kev;
87	int fd, i;
88
89	/*
90	 * Spread load over available IPs, rotating through them as we go.  No
91	 * attempt to localize IPs to particular workers.
92	 */
93	sin = localipbase;
94	sin.sin_addr.s_addr = htonl(ntohl(localipbase.sin_addr.s_addr) +
95	    (counter++ % Mflag));
96
97	fd = socket(PF_INET, SOCK_STREAM, 0);
98	if (fd < 0)
99		err(-1, "socket");
100
101	if (fcntl(fd, F_SETFL, O_NONBLOCK) < 0)
102		err(-1, "fcntl");
103
104	i = 1;
105	if (setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &i, sizeof(i)) < 0)
106		err(-1, "setsockopt");
107	i = 1;
108	if (setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &i, sizeof(i)) < 0)
109		err(-1, "setsockopt");
110#if 0
111	i = 1;
112	if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &i, sizeof(i)) < 0)
113		err(-1, "setsockopt");
114#endif
115
116	if (lflag) {
117		if (bind(fd, (struct sockaddr *)&sin, sizeof(sin)) < 0)
118			err(-1, "bind");
119	}
120
121	if (connect(fd, (struct sockaddr *)&remoteip, sizeof(remoteip)) < 0 &&
122	    errno != EINPROGRESS)
123		err(-1, "connect");
124
125	conn = malloc(sizeof(*conn));
126	if (conn == NULL)
127		return (NULL);
128	bzero(conn, sizeof(*conn));
129	conn->conn_magic = CONNECTION_MAGIC;
130	conn->conn_fd = fd;
131	conn->conn_header.th_magic = TCPP_MAGIC;
132	conn->conn_header.th_len = payload_len;
133	tcpp_header_encode(&conn->conn_header);
134
135	EV_SET(&kev, fd, EVFILT_WRITE, EV_ADD, 0, 0, conn);
136	if (kevent(kq, &kev, 1, NULL, 0, NULL) < 0)
137		err(-1, "newconn kevent");
138
139	started++;
140	return (conn);
141}
142
143static void
144tcpp_client_closeconn(struct connection *conn)
145{
146
147	close(conn->conn_fd);
148	bzero(conn, sizeof(*conn));
149	free(conn);
150	finished++;
151}
152
153static void
154tcpp_client_handleconn(struct kevent *kev)
155{
156	struct connection *conn;
157	struct iovec iov[2];
158	ssize_t len, header_left;
159
160	conn = kev->udata;
161	if (conn->conn_magic != CONNECTION_MAGIC)
162		errx(-1, "tcpp_client_handleconn: magic");
163
164	if (conn->conn_header_sent < sizeof(conn->conn_header)) {
165		header_left = sizeof(conn->conn_header) -
166		    conn->conn_header_sent;
167		iov[0].iov_base = ((u_char *)&conn->conn_header) +
168		    conn->conn_header_sent;
169		iov[0].iov_len = header_left;
170		iov[1].iov_base = buffer;
171		iov[1].iov_len = min(sizeof(buffer), payload_len);
172		len = writev(conn->conn_fd, iov, 2);
173		if (len < 0) {
174			tcpp_client_closeconn(conn);
175			err(-1, "tcpp_client_handleconn: header write");
176		}
177		if (len == 0) {
178			tcpp_client_closeconn(conn);
179			errx(-1, "tcpp_client_handleconn: header write "
180			    "premature EOF");
181		}
182		if (len > header_left) {
183			conn->conn_data_sent += (len - header_left);
184			conn->conn_header_sent += header_left;
185		} else
186			conn->conn_header_sent += len;
187	} else {
188		len = write(conn->conn_fd, buffer, min(sizeof(buffer),
189		    payload_len - conn->conn_data_sent));
190		if (len < 0) {
191			tcpp_client_closeconn(conn);
192			err(-1, "tcpp_client_handleconn: data write");
193		}
194		if (len == 0) {
195			tcpp_client_closeconn(conn);
196			errx(-1, "tcpp_client_handleconn: data write: "
197			    "premature EOF");
198		}
199		conn->conn_data_sent += len;
200	}
201	if (conn->conn_data_sent >= payload_len) {
202		/*
203		 * All is well.
204		 */
205		tcpp_client_closeconn(conn);
206	}
207}
208
209static void
210tcpp_client_worker(int workernum)
211{
212	struct kevent *kev_array;
213	int i, numevents, kev_bytes;
214#if defined(CPU_SETSIZE) && 0
215	cpu_set_t mask;
216	int ncpus;
217	size_t len;
218
219	if (Pflag) {
220		len = sizeof(ncpus);
221		if (sysctlbyname(SYSCTLNAME_CPUS, &ncpus, &len, NULL, 0) < 0)
222			err(-1, "sysctlbyname: %s", SYSCTLNAME_CPUS);
223		if (len != sizeof(ncpus))
224			errx(-1, "sysctlbyname: %s: len %jd", SYSCTLNAME_CPUS,
225			    (intmax_t)len);
226
227		CPU_ZERO(&mask);
228		CPU_SET(workernum % ncpus, &mask);
229		if (sched_setaffinity(0, CPU_SETSIZE, &mask) < 0)
230			err(-1, "sched_setaffinity");
231	}
232#endif
233	setproctitle("tcpp_client %d", workernum);
234
235	/*
236	 * Add the worker number to the remote port.
237	 */
238	remoteip.sin_port = htons(rflag + workernum);
239
240	kev_bytes = sizeof(*kev_array) * mflag;
241	kev_array = malloc(kev_bytes);
242	if (kev_array == NULL)
243		err(-1, "malloc");
244	bzero(kev_array, kev_bytes);
245
246	kq = kqueue();
247	if (kq < 0)
248		err(-1, "kqueue");
249
250	while (finished < tflag) {
251		while ((started - finished < mflag) && (started < tflag))
252			(void)tcpp_client_newconn();
253		numevents = kevent(kq, NULL, 0, kev_array, mflag, NULL);
254		if (numevents < 0)
255			err(-1, "kevent");
256		if (numevents > mflag)
257			errx(-1, "kevent: %d", numevents);
258		for (i = 0; i < numevents; i++)
259			tcpp_client_handleconn(&kev_array[i]);
260	}
261	/* printf("Worker %d done - %d finished\n", workernum, finished); */
262}
263
264void
265tcpp_client(void)
266{
267	struct timespec ts_start, ts_finish;
268	long cp_time_start[CPUSTATES], cp_time_finish[CPUSTATES];
269	long ticks;
270	size_t size;
271	pid_t pid;
272	int i, failed, status;
273
274	if (bflag < sizeof(struct tcpp_header))
275		errx(-1, "Can't use -b less than %zu\n",
276		   sizeof(struct tcpp_header));
277	payload_len = bflag - sizeof(struct tcpp_header);
278
279	pid_list = malloc(sizeof(*pid_list) * pflag);
280	if (pid_list == NULL)
281		err(-1, "malloc pid_list");
282	bzero(pid_list, sizeof(*pid_list) * pflag);
283
284	/*
285	 * Start workers.
286	 */
287	size = sizeof(cp_time_start);
288	if (sysctlbyname(SYSCTLNAME_CPTIME, &cp_time_start, &size, NULL, 0)
289	    < 0)
290		err(-1, "sysctlbyname: %s", SYSCTLNAME_CPTIME);
291	if (clock_gettime(CLOCK_REALTIME, &ts_start) < 0)
292		err(-1, "clock_gettime");
293	for (i = 0; i < pflag; i++) {
294		pid = fork();
295		if (pid < 0) {
296			warn("fork");
297			for (i = 0; i < pflag; i++) {
298				if (pid_list[i] != 0)
299					(void)kill(pid_list[i], SIGKILL);
300			}
301			exit(-1);
302		}
303		if (pid == 0) {
304			tcpp_client_worker(i);
305			exit(0);
306		}
307		pid_list[i] = pid;
308	}
309
310	/*
311	 * GC workers.
312	 */
313	failed = 0;
314	for (i = 0; i < pflag; i++) {
315		if (pid_list[i] != 0) {
316			while (waitpid(pid_list[i], &status, 0) != pid_list[i]);
317			if (WEXITSTATUS(status) != 0)
318				failed = 1;
319		}
320	}
321	if (clock_gettime(CLOCK_REALTIME, &ts_finish) < 0)
322		err(-1, "clock_gettime");
323	size = sizeof(cp_time_finish);
324	if (sysctlbyname(SYSCTLNAME_CPTIME, &cp_time_finish, &size, NULL, 0)
325	    < 0)
326		err(-1, "sysctlbyname: %s", SYSCTLNAME_CPTIME);
327	timespecsub(&ts_finish, &ts_start, &ts_finish);
328
329	if (failed)
330		errx(-1, "Too many errors");
331
332	if (hflag)
333		printf("bytes,seconds,conn/s,Gb/s,user%%,nice%%,sys%%,"
334		    "intr%%,idle%%\n");
335
336	/*
337	 * Configuration parameters.
338	 */
339	printf("%jd,", bflag * tflag * pflag);
340	printf("%jd.%09jd,", (intmax_t)ts_finish.tv_sec,
341	    (intmax_t)(ts_finish.tv_nsec));
342
343	/*
344	 * Effective transmit rates.
345	 */
346	printf("%f,", (double)(pflag * tflag)/
347	    (ts_finish.tv_sec + ts_finish.tv_nsec * 1e-9));
348	printf("%f,", (double)(bflag * tflag * pflag * 8) /
349	    (ts_finish.tv_sec + ts_finish.tv_nsec * 1e-9) * 1e-9);
350
351	/*
352	 * CPU time (est).
353	 */
354	ticks = 0;
355	for (i = 0; i < CPUSTATES; i++) {
356		cp_time_finish[i] -= cp_time_start[i];
357		ticks += cp_time_finish[i];
358	}
359	printf("%0.02f,", (float)(100 * cp_time_finish[CP_USER]) / ticks);
360	printf("%0.02f,", (float)(100 * cp_time_finish[CP_NICE]) / ticks);
361	printf("%0.02f,", (float)(100 * cp_time_finish[CP_SYS]) / ticks);
362	printf("%0.02f,", (float)(100 * cp_time_finish[CP_INTR]) / ticks);
363	printf("%0.02f", (float)(100 * cp_time_finish[CP_IDLE]) / ticks);
364	printf("\n");
365}
366