#include <sys/types.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <signal.h>
#include <errno.h>

#ifdef __i386__
#define rdtscl(low) \
     __asm__ __volatile__ ("rdtsc" : "=A" (low))
#else
#define rdtscl(val) \
do {						\
	struct timespec now;			\
	clock_gettime(CLOCK_REALTIME, &now);	\
	val  = (u_int64_t)now.tv_sec;		\
	val *= 1000000000LL;			\
	val += (u_int64_t)now.tv_nsec;		\
} while(0)
#endif

int afinet = AF_UNSPEC;

void
usage(char *argv0)
{
	fprintf(stderr, "%s [-n nsockets]\n", argv0);
	exit(1);
}

void
display(char *msg, int i, u_int64_t a, u_int64_t b)
{
	printf("%s\t%d\t%lld\n", msg, i, b - a);
}

void
socket_ben(int n, int do_bind)
{
        struct addrinfo hints, *ai;
	int i, *sockets, gaierr;
	u_int64_t a, b;

        memset(&hints, 0, sizeof(hints));
        hints.ai_family = afinet;
        hints.ai_flags = AI_PASSIVE;
        hints.ai_socktype = SOCK_STREAM;
        if ((gaierr = getaddrinfo(NULL, "0", &hints, &ai)) != 0) {
		fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(gaierr));
		exit(1);
	}
	sockets = malloc(sizeof(int) * n);
	if (sockets == NULL) {
		fprintf(stderr, "malloc: %s\n", strerror(errno));
		exit(1);
	}

	for (i = 0; i < n; i++) {
		rdtscl(a);
                if ((sockets[i] = socket(ai->ai_family, ai->ai_socktype,
		    ai->ai_protocol)) < 0) {
			fprintf(stderr, "socket %d: %s\n", i, strerror(errno));
			n = i;
			break;
		}
		rdtscl(b);
		if (do_bind == 0) {
			display("socket", i, a, b);
			continue;
		}
		rdtscl(a);
		if (bind(sockets[i], ai->ai_addr, ai->ai_addrlen)) {
			fprintf(stderr, "bind %d: %s\n", i, strerror(errno));
			exit(2);
		}
		rdtscl(b);
		display("bind", i, a, b);
	}
	free(ai);
	for (i = 0; i < n; i++)
		close(sockets[i]);
	free(sockets);
}

void
fork_ben(int n)
{
	int i, quit = 0, pfd[2];
	char buf[1];
	pid_t *proc;
	u_int64_t a, b;

	proc = malloc(sizeof(pid_t) * n);
	if (proc == NULL) {
		fprintf(stderr, "malloc: %s\n", strerror(errno));
		exit(1);
	}
	if (pipe(pfd) < 0) {
		fprintf(stderr, "pipe: %s\n", strerror(errno));
		exit(1);
	}
	for (i = 0; i < n; i++) {
		rdtscl(a);
		proc[i] = fork();
		switch(proc[i]) {
		case -1:
			n = i;
			quit = 1;
			break;
		case 0:
			/* child */
			write(pfd[1], "+", 1);
			close(pfd[1]);
			sleep(5*60);
			_exit(0);
			break;
		default:
			break;
		}
		if (quit || read(pfd[0], buf, 1) != 1)
			break;	
		rdtscl(b);
		display("fork", i, a, b);
	}

	for (i = 0; i < n; i++) {
		if (proc[i] > 1) {
			rdtscl(a);
			kill(proc[i], SIGTERM);
			rdtscl(b);
			display("kill", i, a, b);
		}
	}
}

int
main(int argc, char **argv)
{
	struct rlimit rl;
	int n, ch, b_flag, s_flag, f_flag;
        struct sockaddr_in sin;

	b_flag = s_flag = f_flag = 0;
	n = 100;

	while ((ch = getopt(argc, argv, "46bfn:s")) != -1) {
		switch (ch) {
		case '4':
			afinet = AF_INET;
			break;
		case '6':
			afinet = AF_INET6;
			break;
		case 'b':
			b_flag = 1;
			break;
		case 'f':
			f_flag = 1;
			break;
		case 'n':
			n = atoi(optarg);
			break;
		case 's':
			s_flag = 1;
			break;
		case '?':
			default:
			usage(argv[0]);
		}
	}
	argc -= optind;
	argv += optind;

	rl.rlim_cur = RLIM_INFINITY;
	rl.rlim_max = RLIM_INFINITY;
	setrlimit(RLIMIT_NOFILE, &rl);
	setrlimit(RLIMIT_NPROC, &rl);

	if (s_flag)
		socket_ben(n, 0);
	else if (b_flag)
		socket_ben(n, 1);
	else if (f_flag)
		fork_ben(n);
	else
		fprintf(stderr, "no benchmark\n");
	exit(0);
}