#include <sys/types.h>
#include <sys/event.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/sysctl.h>
#include <sys/time.h>
#include <sys/usched.h>
#include <sys/wait.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <err.h>
#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
static void mainloop(const struct sockaddr_in *, int, int, long, u_long *,
int, int);
static int bindcpu;
static int cpucnt;
static void
usage(const char *cmd)
{
fprintf(stderr, "%s -4 inet4 [-4 inet4_1] -p port "
"[-i n_instance] [-c conn_max] [-l duration] [-u] [-B]\n", cmd);
exit(1);
}
int
main(int argc, char *argv[])
{
struct sockaddr_in *in, *tmp;
int opt, ninst, nconn, i, in_max, in_cnt, do_udp;
long dur;
u_long *result, sum;
u_short port;
size_t prm_len;
prm_len = sizeof(cpucnt);
if (sysctlbyname("hw.ncpu", &cpucnt, &prm_len, NULL, 0) != 0)
err(2, "sysctl hw.ncpu failed");
ninst = cpucnt;
nconn = 8;
dur = 10;
port = 0;
do_udp = 0;
in_max = 8;
in = calloc(in_max, sizeof(struct sockaddr_in));
if (in == NULL)
err(1, "calloc failed");
in_cnt = 0;
while ((opt = getopt(argc, argv, "4:Bc:i:l:p:u")) != -1) {
switch (opt) {
case '4':
if (in_cnt >= in_max) {
struct sockaddr_in *old_in = in;
int old_in_max = in_max;
in_max <<= 1;
in = calloc(in_max, sizeof(struct sockaddr_in));
if (in == NULL)
err(1, "calloc failed");
memcpy(in, old_in,
old_in_max * sizeof(struct sockaddr_in));
free(old_in);
}
tmp = &in[in_cnt];
if (inet_pton(AF_INET, optarg, &tmp->sin_addr) <= 0) {
fprintf(stderr, "invalid inet address %s\n",
optarg);
usage(argv[0]);
}
++in_cnt;
break;
case 'B':
bindcpu = 1;
break;
case 'c':
nconn = atoi(optarg);
break;
case 'i':
ninst = atoi(optarg);
break;
case 'l':
dur = strtol(optarg, NULL, 10);
break;
case 'p':
port = htons(atoi(optarg));
break;
case 'u':
do_udp = 1;
break;
default:
usage(argv[0]);
}
}
if (ninst < 1 || dur < 1 || nconn < 1 || port == 0 || in_cnt == 0)
usage(argv[0]);
for (i = 0; i < in_cnt; ++i) {
tmp = &in[i];
tmp->sin_family = AF_INET;
tmp->sin_port = port;
}
result = mmap(NULL, ninst * sizeof(u_long), PROT_READ | PROT_WRITE,
MAP_ANON | MAP_SHARED, -1, 0);
if (result == MAP_FAILED)
err(1, "mmap failed");
memset(result, 0, ninst * sizeof(u_long));
for (i = 0; i < ninst; ++i) {
pid_t pid;
pid = fork();
if (pid == 0) {
mainloop(in, in_cnt, nconn, dur, &result[i], do_udp, i);
exit(0);
} else if (pid < 0) {
err(1, "fork failed");
}
}
for (i = 0; i < ninst; ++i) {
pid_t pid;
pid = waitpid(-1, NULL, 0);
if (pid < 0)
err(1, "waitpid failed");
}
sum = 0;
for (i = 0; i < ninst; ++i)
sum += result[i];
printf("%lu\n", sum);
exit(0);
}
static void
udp_send(const struct sockaddr_in *in)
{
uint8_t d[18];
int s;
s = socket(AF_INET, SOCK_DGRAM, 0);
if (s < 0) {
warn("socket DGRAM failed");
return;
}
if (connect(s, (const struct sockaddr *)in, sizeof(*in)) < 0) {
warn("connect DGRAM failed");
goto done;
}
write(s, d, sizeof(d));
done:
close(s);
}
static void
mainloop(const struct sockaddr_in *in, int in_cnt, int nconn_max,
long dur, u_long *res, int do_udp, int inst)
{
struct timespec start, end;
struct kevent *evt_change0, *evt;
int kq, nchange = 0, nconn = 0, nevt_max;
u_long count = 0;
double time_us;
u_int in_idx = 0;
#ifndef SOCK_NONBLOCK
int nblock = 1;
#endif
if (bindcpu) {
int cpu = inst % cpucnt;
usched_set(getpid(), USCHED_SET_CPU, &cpu, sizeof(cpu));
}
kq = kqueue();
if (kq < 0)
err(1, "kqueue failed");
nevt_max = nconn_max + 1;
evt_change0 = malloc(nevt_max * sizeof(struct kevent));
if (evt_change0 == NULL)
err(1, "malloc evt_change failed");
evt = malloc(nevt_max * sizeof(struct kevent));
if (evt == NULL)
err(1, "malloc evt failed");
EV_SET(&evt_change0[0], 0, EVFILT_TIMER, EV_ADD | EV_ONESHOT, 0,
dur * 1000L, NULL);
nchange = 1;
clock_gettime(CLOCK_MONOTONIC_PRECISE, &start);
for (;;) {
struct kevent *evt_change = NULL;
int n, i, done = 0;
while (nconn < nconn_max) {
const struct sockaddr_in *tmp;
int s;
tmp = &in[in_idx % in_cnt];
++in_idx;
if (do_udp)
udp_send(tmp);
#ifndef SOCK_NONBLOCK
s = socket(AF_INET, SOCK_STREAM, 0);
if (s < 0)
err(1, "socket failed");
if (ioctl(s, FIONBIO, &nblock, sizeof(nblock)) < 0)
err(1, "ioctl failed");
#else
s = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
if (s < 0)
err(1, "socket failed");
#endif
n = connect(s, (const struct sockaddr *)tmp,
sizeof(*tmp));
if (n == 0) {
++count;
close(s);
continue;
} else {
int error = errno;
if (error != EINPROGRESS)
errc(1, error, "connect failed");
}
++nconn;
if (nchange >= nevt_max) {
fprintf(stderr, "invalid nchange %d, max %d\n",
nchange, nevt_max);
abort();
}
EV_SET(&evt_change0[nchange], s, EVFILT_WRITE, EV_ADD,
0, 0, NULL);
++nchange;
}
if (nchange)
evt_change = evt_change0;
n = kevent(kq, evt_change, nchange, evt, nevt_max, NULL);
if (n < 0)
err(1, "kevent failed");
nchange = 0;
for (i = 0; i < n; ++i) {
struct kevent *e = &evt[i];
if (e->filter == EVFILT_TIMER) {
done = 1;
continue;
}
if ((e->flags & EV_EOF) && e->fflags) {
} else {
++count;
}
close(e->ident);
--nconn;
}
if (done)
break;
}
clock_gettime(CLOCK_MONOTONIC_PRECISE, &end);
timespecsub(&end, &start, &end);
time_us = ((double)end.tv_sec * 1000000.0) +
((double)end.tv_nsec / 1000.0);
*res = ((double)count * 1000000.0) / time_us;
}
