benchmarker.c

#include <arpa/inet.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <time.h>
#include <unistd.h>
#include <sys/resource.h>

#define MILLION 1000000
#define BILLION 1000000000

long tv_sub(struct timeval * end, struct timeval * start) {
	return (end->tv_sec - start->tv_sec) * MILLION +
		(end->tv_usec - start->tv_usec);
}

int main(int argc, char **argv) {
	if(argc < 3) {
		fprintf(stderr, "usage: %s <tcp|unix> <client|server> message_size\n", argv[0]);
		return 1;
	}

	int domain = AF_INET;
	if(strcmp(argv[1], "unix") == 0) {
		domain = AF_UNIX;
	}

	int fd = socket(domain, SOCK_STREAM, 0);

	int addr_size = domain == AF_UNIX ? 
		sizeof(struct sockaddr_un) : sizeof(struct sockaddr_in);
	struct sockaddr * addr = alloca(addr_size);
	memset(addr, 0, addr_size);
	addr->sa_family = domain;

	if(domain == AF_UNIX) {
		strncpy(((struct sockaddr_un *)addr)->sun_path, "socket", 7);
	} else {
		struct sockaddr_in * tcp_addr = (struct sockaddr_in *)addr;
		tcp_addr->sin_addr.s_addr = inet_addr("127.0.0.1");
		tcp_addr->sin_port = htons(8082);
	}

	char * buf = malloc(1048576);
	ssize_t (*op)(int, void *, size_t, int);
	int connfd = fd;
	if(strcmp(argv[2], "server") == 0) {
		if(domain == AF_UNIX) {
			unlink("socket");
		}
		if(bind(fd, addr, addr_size) == -1) {
			perror("bind");
			return 1;
		}
		if(listen(fd, 1) == -1) {
			perror("listen");
			return 1;
		}
	
		connfd = accept(fd, NULL, NULL);
		op = recv;
	} else {
		if(connect(fd, addr, addr_size) == -1) {
			perror("connect");
			return 1;
		}

		op = send;
	}

	int msg_size = atoi(argv[3]);

	struct timespec start;
	struct timespec end;
	struct rusage start_usage;
	struct rusage end_usage;
	clock_gettime(CLOCK_MONOTONIC, &start);
	getrusage(RUSAGE_SELF, &start_usage);
	long total = 50000L * msg_size;
	long remaining = total;
	int sndbuf[1] = { msg_size };
	setsockopt(connfd, SOL_SOCKET, SO_SNDBUF, sndbuf, sizeof(sndbuf));
	while(remaining > 0) {
		int did = op(connfd, buf, remaining > msg_size ? msg_size : remaining, 0);
		if(did == 0) {
			perror("send/recv");
			break;
		}
		remaining -= did;
	}
	getrusage(RUSAGE_SELF, &end_usage);
	clock_gettime(CLOCK_MONOTONIC, &end);

	long nsecs = (end.tv_sec - start.tv_sec) * BILLION + 
		(end.tv_nsec - start.tv_nsec);

	double utime = tv_sub(&end_usage.ru_utime, &start_usage.ru_utime) / (double)MILLION;
	double stime = tv_sub(&end_usage.ru_stime, &start_usage.ru_stime) / (double)MILLION;

	printf(
		"%lu bytes in %f secs (user=%f, system=%f). %f GB/s\n", 
		total,
		(double)nsecs / BILLION, utime, stime,
		(double)total / nsecs
	);
}