/*
  Copyright (C) 2007 Arnaldo Carvalho de Melo <acme@redhat.com>

  This program is free software; you can redistribute it and/or modify it
  under the terms of version 2 of the GNU General Public License as
  published by the Free Software Foundation.

  Build it with:

  gcc tcp_nagle_client.c -o tcp_nagle_client -lrt
*/

#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <ctype.h>
#include <errno.h>
#include <string.h>
#include <netdb.h>
#include <time.h>
#include <unistd.h>
#include <poll.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>

#define NSEC_PER_SEC 1000000000L

static long timespec_delta(const struct timespec *large,
                           const struct timespec *small)
{
        time_t secs = large->tv_sec - small->tv_sec;
        long nsecs = large->tv_nsec - small->tv_nsec;

        if (nsecs < 0) {
                secs--;
                nsecs += NSEC_PER_SEC;
        }
        return secs * NSEC_PER_SEC + nsecs;
}

#define NR_DATA_ENTRIES 15
#define SIZE_DATA_ENTRY 2
#define SIZE_RESPONSE 2

static char data[NR_DATA_ENTRIES][SIZE_DATA_ENTRY];

static int value;

static void tcp_nodelay(int fd)
{
	value = 1;
	if (setsockopt(fd, SOL_TCP, TCP_NODELAY, &value, sizeof(value)) != 0)
		perror("setsockopt(TCP_NODELAY):");
}

static void __tcp_cork(int fd, int cork)
{
	value = cork;
	if (setsockopt(fd, SOL_TCP, TCP_CORK, &value, sizeof(value)) != 0)
		perror("setsockopt(TCP_CORK):");
}

static void tcp_cork(int fd)
{
	__tcp_cork(fd, 1);
}

static void tcp_uncork(int fd)
{
	__tcp_cork(fd, 0);
}

static void nread(int fd, char *bf, int len)
{
	do {
		int n = read(fd, bf, len);

		if (n < 0) {
			perror("client read:");
			exit(1);
		}

		len -= n;
		bf += n;
	} while (len != 0);
}

static void exchange_packet(int fd, int cork)
{
	int i;
	char response[SIZE_RESPONSE];

	if (cork)
		tcp_cork(fd);

	for (i = 0; i < NR_DATA_ENTRIES; ++i) {
		if (write(fd, data[i], SIZE_DATA_ENTRY) != SIZE_DATA_ENTRY) {
			fprintf(stderr, "client: write failed!\n");
			exit(1);
		}
	}

	if (cork)
		tcp_uncork(fd);
#if 0
	if (0) {
		fd_set rfds;
		struct timeval tv = {
			.tv_sec = 0,
			.tv_usec = 5,
		};

		FD_ZERO(&rfds);
		FD_SET(fd, &rfds);
		select(fd + 1, &rfds, NULL, NULL, &tv);
	} else {
		struct pollfd fds[1];

		fds[0].fd = fd;
		fds[0].events = POLLIN;
		poll(fds, 1, 1);
	}
#endif
	nread(fd, response, SIZE_RESPONSE);
}

int main(int argc, char *argv[])
{
	struct timespec start, finish;
	float delta;
	float rate;
	struct addrinfo *host;
	struct addrinfo hints = {
		.ai_family   = AF_INET,
		.ai_socktype = SOCK_STREAM,
		.ai_protocol = IPPROTO_TCP,
	};
	char *hostname, *port;
	int nr_logical_packets, fd, i, rc, cork = 0, no_delay = 0;

	if (argc < 4) {
		fprintf(stderr, "usage: %s <host> <port> <nr_logical_packets> "
				"[no_delay|cork]\n", argv[0]);
		return 1;
	}

	if (argc > 4) {
		cork = strcmp(argv[4], "cork") == 0;
		no_delay = strcmp(argv[4], "no_delay") == 0;
	}

	hostname = argv[1];
	port = argv[2];
	nr_logical_packets = atoi(argv[3]);

	rc = getaddrinfo(hostname, port, &hints, &host);
	if (rc != 0) {
		fprintf(stderr, "error using getaddrinfo: %s\n", gai_strerror(rc));
		goto out;
	}

	fd = socket(host->ai_family, host->ai_socktype, host->ai_protocol);
	if (fd < 0) {
		perror("socket: ");
		goto out_freeaddrinfo;
	}

	if (connect(fd, host->ai_addr, host->ai_addrlen) < 0) {
		perror("connect: ");
		goto out_close;
	}

	for (i = 0; i < NR_DATA_ENTRIES; ++i)
		memset(data[i], 'A' + i, SIZE_DATA_ENTRY);

	clock_gettime(CLOCK_MONOTONIC, &start);

	if (no_delay)
		tcp_nodelay(fd);

	i = nr_logical_packets;
	while (i--)
		exchange_packet(fd, cork);

	clock_gettime(CLOCK_MONOTONIC, &finish);
	delta = timespec_delta(&finish, &start) / 1000000.0;

	rate = (nr_logical_packets * NR_DATA_ENTRIES * SIZE_DATA_ENTRY) / delta;

	printf("%d packets of %d bytes sent in %f ms: %f bytes/ms %s\n",
	       nr_logical_packets, NR_DATA_ENTRIES * SIZE_DATA_ENTRY, delta, rate,
	       cork ? "using TCP_CORK" : no_delay ? "using TCP_NODELAY" : "");

	rc = 0;
out_close:
	close(fd);
out_freeaddrinfo:
	freeaddrinfo(host);
out:
	return rc;
}