#include <sys/types.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/sysctl.h>
#include <sys/time.h>
#include <sys/wait.h>

#include <arpa/inet.h>
#include <netinet/in.h>

#include <err.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

static void     mainloop(const struct sockaddr_in *, long, u_long *);

static int      global_stopped;

static void
usage(const char *cmd)
{
        fprintf(stderr, "%s -4 inet4 -p port "
            "[-i n_instance] [-l duration]\n", cmd);
        exit(1);
}

int
main(int argc, char *argv[])
{
        struct sockaddr_in in;
        int opt, ninst, i;
        long dur;
        u_long *result, sum;
        size_t prm_len;

        prm_len = sizeof(ninst);
        if (sysctlbyname("hw.ncpu", &ninst, &prm_len, NULL, 0) != 0)
                err(2, "sysctl hw.ncpu failed");

        memset(&in, 0, sizeof(in));
        in.sin_family = AF_INET;

        dur = 10;

        while ((opt = getopt(argc, argv, "4:p:i:l:")) != -1) {
                switch (opt) {
                case '4':
                        if (inet_pton(AF_INET, optarg, &in.sin_addr) <= 0) {
                                fprintf(stderr, "invalid inet address %s\n",
                                    optarg);
                                usage(argv[0]);
                        }
                        break;

                case 'p':
                        in.sin_port = htons(atoi(optarg));
                        break;

                case 'i':
                        ninst = atoi(optarg);
                        break;

                case 'l':
                        dur = strtol(optarg, NULL, 10);
                        break;

                default:
                        usage(argv[0]);
                }
        }

        if (ninst < 1 || dur < 1 ||
            in.sin_port == 0 || in.sin_addr.s_addr == INADDR_ANY)
                usage(argv[0]);

        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, dur, &result[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("%.2f\n", (double)sum / (double)dur);

        exit(0);
}

static void
signal_handler(int signum __unused)
{
        global_stopped = 1;
}

static void
mainloop(const struct sockaddr_in *in, long dur, u_long *res)
{
        struct itimerval it;
        u_long count = 0;

        if (signal(SIGALRM, signal_handler) == SIG_ERR)
                err(1, "signal failed");

        it.it_interval.tv_sec = 0;
        it.it_interval.tv_usec = 0;
        it.it_value.tv_sec = dur;
        it.it_value.tv_usec = 0;

        if (setitimer(ITIMER_REAL, &it, NULL) < 0)
                err(1, "setitimer failed");

        while (!global_stopped) {
                int s;

                s = socket(AF_INET, SOCK_STREAM, 0);
                if (s < 0)
                        continue;

                if (connect(s,
                    (const struct sockaddr *)in, sizeof(*in)) == 0)
                        ++count;
                close(s);
        }
        *res = count;
}