/*
 * testcode/perf.c - debug program to estimate name server performance.
 *
 * Copyright (c) 2008, NLnet Labs. All rights reserved.
 *
 * This software is open source.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 
 * Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 * 
 * Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 * 
 * Neither the name of the NLNET LABS nor the names of its contributors may
 * be used to endorse or promote products derived from this software without
 * specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
 * TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/**
 * \file
 *
 * This program estimates DNS name server performance.
 */

#include "config.h"
#ifdef HAVE_GETOPT_H
#include <getopt.h>
#endif
#include <signal.h>
#include "util/log.h"
#include "util/locks.h"
#include "util/net_help.h"
#include "util/data/msgencode.h"
#include "util/data/msgreply.h"
#include "util/data/msgparse.h"
#include "sldns/sbuffer.h"
#include "sldns/wire2str.h"
#include "sldns/str2wire.h"
#include <sys/time.h>

/** usage information for perf */
static void usage(char* nm) 
{
	printf("usage: %s [options] server\n", nm);
	printf("server: ip address of server, IP4 or IP6.\n");
	printf("	If not on port %d add @port.\n", UNBOUND_DNS_PORT);
	printf("-d sec	duration of test in whole seconds (0: wait for ^C)\n");
	printf("-a str	query to ask, interpreted as a line from qfile\n");
	printf("-f fnm	query list to read from file\n");
	printf("	every line has format: qname qclass qtype [+-]{E}\n");
	printf("	where + means RD set, E means EDNS enabled\n");
	printf("-q 	quiet mode, print only final qps\n");
	exit(1);
}

struct perfinfo;
struct perfio;

/** Global info for perf */
struct perfinfo { 
	/** need to exit */
	volatile int exit;
	/** all purpose buffer (for UDP send and receive) */
	sldns_buffer* buf;

	/** destination */
	struct sockaddr_storage dest;
	/** length of dest socket addr */
	socklen_t destlen;

	/** when did this time slice start */
	struct timeval since;
	/** number of queries received in that time */
	size_t numrecv;
	/** number of queries sent out in that time */
	size_t numsent;

	/** duration of test in seconds */
	int duration;
	/** quiet mode? */
	int quiet;

	/** when did the total test start */
	struct timeval start;
	/** total number recvd */
	size_t total_recv;
	/** total number sent */
	size_t total_sent;
	/** numbers by rcode */
	size_t by_rcode[32];
	
	/** number of I/O ports */
	size_t io_num;
	/** I/O ports array */
	struct perfio* io;
	/** max fd value in io ports */
	int maxfd;
	/** readset */
	fd_set rset;

	/** size of querylist */
	size_t qlist_size;
	/** allocated size of qlist array */
	size_t qlist_capacity;
	/** list of query packets (data) */
	uint8_t** qlist_data;
	/** list of query packets (length of a packet) */
	size_t* qlist_len;
	/** index into querylist, for walking the list */
	size_t qlist_idx;
};

/** I/O port for perf */
struct perfio {
	/** id number */
	size_t id;
	/** file descriptor of socket */
	int fd;
	/** timeout value */
	struct timeval timeout;
	/** ptr back to perfinfo */
	struct perfinfo* info;
};

/** number of msec between starting io ports */
#define START_IO_INTERVAL 10
/** number of msec timeout on io ports */
#define IO_TIMEOUT 10

/** signal handler global info */
static struct perfinfo* sig_info;

/** signal handler for user quit */
static RETSIGTYPE perf_sigh(int sig)
{
	log_assert(sig_info);
	if(!sig_info->quiet)
		printf("exit on signal %d\n", sig);
	sig_info->exit = 1;
}

/** timeval compare, t1 < t2 */
static int
perf_tv_smaller(struct timeval* t1, struct timeval* t2) 
{
#ifndef S_SPLINT_S
	if(t1->tv_sec < t2->tv_sec)
		return 1;
	if(t1->tv_sec == t2->tv_sec &&
		t1->tv_usec < t2->tv_usec)
		return 1;
#endif
	return 0;
}

/** timeval add, t1 += t2 */
static void
perf_tv_add(struct timeval* t1, struct timeval* t2) 
{
#ifndef S_SPLINT_S
	t1->tv_sec += t2->tv_sec;
	t1->tv_usec += t2->tv_usec;
	while(t1->tv_usec >= 1000000) {
		t1->tv_usec -= 1000000;
		t1->tv_sec++;
	}
#endif
}

/** timeval subtract, t1 -= t2 */
static void
perf_tv_subtract(struct timeval* t1, struct timeval* t2) 
{
#ifndef S_SPLINT_S
	t1->tv_sec -= t2->tv_sec;
	if(t1->tv_usec >= t2->tv_usec) {
		t1->tv_usec -= t2->tv_usec;
	} else {
		t1->tv_sec--;
		t1->tv_usec = 1000000-(t2->tv_usec-t1->tv_usec);
	}
#endif
}


/** setup perf test environment */
static void
perfsetup(struct perfinfo* info)
{
	size_t i;
	if(gettimeofday(&info->start, NULL) < 0)
		fatal_exit("gettimeofday: %s", strerror(errno));
	sig_info = info;
	if( signal(SIGINT, perf_sigh) == SIG_ERR || 
#ifdef SIGQUIT
		signal(SIGQUIT, perf_sigh) == SIG_ERR ||
#endif
#ifdef SIGHUP
		signal(SIGHUP, perf_sigh) == SIG_ERR ||
#endif
#ifdef SIGBREAK
		signal(SIGBREAK, perf_sigh) == SIG_ERR ||
#endif
		signal(SIGTERM, perf_sigh) == SIG_ERR)
		fatal_exit("could not bind to signal");
	info->io = (struct perfio*)calloc(sizeof(struct perfio), info->io_num);
	if(!info->io) fatal_exit("out of memory");
#ifndef S_SPLINT_S
	FD_ZERO(&info->rset);
#endif
	info->since = info->start;
	for(i=0; i<info->io_num; i++) {
		info->io[i].id = i;
		info->io[i].info = info;
		info->io[i].fd = socket(
			addr_is_ip6(&info->dest, info->destlen)?
			AF_INET6:AF_INET, SOCK_DGRAM, 0);
		if(info->io[i].fd == -1) {
			fatal_exit("socket: %s", sock_strerror(errno));
		}
		if(info->io[i].fd > info->maxfd)
			info->maxfd = info->io[i].fd;
#ifndef S_SPLINT_S
		FD_SET(FD_SET_T info->io[i].fd, &info->rset);
		info->io[i].timeout.tv_usec = ((START_IO_INTERVAL*i)%1000)
						*1000;
		info->io[i].timeout.tv_sec = (START_IO_INTERVAL*i)/1000;
		perf_tv_add(&info->io[i].timeout, &info->since);
#endif
	}
}

/** cleanup perf test environment */
static void
perffree(struct perfinfo* info)
{
	size_t i;
	if(!info) return;
	if(info->io) {
		for(i=0; i<info->io_num; i++) {
			sock_close(info->io[i].fd);
		}
		free(info->io);
	}
	for(i=0; i<info->qlist_size; i++)
		free(info->qlist_data[i]);
	free(info->qlist_data);
	free(info->qlist_len);
}

/** send new query for io */
static void
perfsend(struct perfinfo* info, size_t n, struct timeval* now)
{
	ssize_t r;
	r = sendto(info->io[n].fd, (void*)info->qlist_data[info->qlist_idx],
		info->qlist_len[info->qlist_idx], 0,
		(struct sockaddr*)&info->dest, info->destlen);
	/*log_hex("send", info->qlist_data[info->qlist_idx],
		info->qlist_len[info->qlist_idx]);*/
	if(r == -1) {
		log_err("sendto: %s", sock_strerror(errno));
	} else if(r != (ssize_t)info->qlist_len[info->qlist_idx]) {
		log_err("partial sendto");
	}
	info->qlist_idx = (info->qlist_idx+1) % info->qlist_size;
	info->numsent++;

	info->io[n].timeout.tv_sec = IO_TIMEOUT/1000;
	info->io[n].timeout.tv_usec = (IO_TIMEOUT%1000)*1000;
	perf_tv_add(&info->io[n].timeout, now);
}

/** got reply for io */
static void
perfreply(struct perfinfo* info, size_t n, struct timeval* now)
{
	ssize_t r;
	r = recv(info->io[n].fd, (void*)sldns_buffer_begin(info->buf),
		sldns_buffer_capacity(info->buf), 0);
	if(r == -1) {
		log_err("recv: %s", sock_strerror(errno));
	} else {
		info->by_rcode[LDNS_RCODE_WIRE(sldns_buffer_begin(
			info->buf))]++;
		info->numrecv++;
	}
	/*sldns_buffer_set_limit(info->buf, r);
	log_buf(0, "reply", info->buf);*/
	perfsend(info, n, now);
}

/** got timeout for io */
static void
perftimeout(struct perfinfo* info, size_t n, struct timeval* now)
{
	/* may not be a dropped packet, this is also used to start
	 * up the sending IOs */
	perfsend(info, n, now);
}

/** print nice stats about qps */
static void
stat_printout(struct perfinfo* info, struct timeval* now, 
	struct timeval* elapsed)
{
	/* calculate qps */
	double dt, qps = 0;
#ifndef S_SPLINT_S
	dt = (double)(elapsed->tv_sec*1000000 + elapsed->tv_usec) / 1000000;
#endif
	if(dt > 0.001)
		qps = (double)(info->numrecv) / dt;
	if(!info->quiet)
		printf("qps: %g\n", qps);
	/* setup next slice */
	info->since = *now;
	info->total_sent += info->numsent;
	info->total_recv += info->numrecv;
	info->numrecv = 0;
	info->numsent = 0;
}

/** wait for new events for performance test */
static void
perfselect(struct perfinfo* info)
{
	fd_set rset = info->rset;
	struct timeval timeout, now;
	int num;
	size_t i;
	if(gettimeofday(&now, NULL) < 0)
		fatal_exit("gettimeofday: %s", strerror(errno));
	/* time to exit? */
	if(info->duration > 0) {
		timeout = now;
		perf_tv_subtract(&timeout, &info->start);
		if((int)timeout.tv_sec >= info->duration) {
			info->exit = 1;
			return;
		}
	}
	/* time for stats printout? */
	timeout = now;
	perf_tv_subtract(&timeout, &info->since);
	if(timeout.tv_sec > 0) {
		stat_printout(info, &now, &timeout);
	}
	/* see what is closest port to timeout; or if there is a timeout */
	timeout = info->io[0].timeout;
	for(i=0; i<info->io_num; i++) {
		if(perf_tv_smaller(&info->io[i].timeout, &now)) {
			perftimeout(info, i, &now);
			return;
		}
		if(perf_tv_smaller(&info->io[i].timeout, &timeout)) {
			timeout = info->io[i].timeout;
		}
	}
	perf_tv_subtract(&timeout, &now);
	
	num = select(info->maxfd+1, &rset, NULL, NULL, &timeout);
	if(num == -1) {
		if(errno == EAGAIN || errno == EINTR)
			return;
		log_err("select: %s", strerror(errno));
	}

	/* handle new events */
	for(i=0; num && i<info->io_num; i++) {
		if(FD_ISSET(info->io[i].fd, &rset)) {
			perfreply(info, i, &now);
			num--;
		}
	}
}

/** show end stats */
static void
perfendstats(struct perfinfo* info)
{
	double dt, qps;
	struct timeval timeout, now;
	int i, lost; 
	if(gettimeofday(&now, NULL) < 0)
		fatal_exit("gettimeofday: %s", strerror(errno));
	timeout = now;
	perf_tv_subtract(&timeout, &info->since);
	stat_printout(info, &now, &timeout);
	
	timeout = now;
	perf_tv_subtract(&timeout, &info->start);
	dt = (double)(timeout.tv_sec*1000000 + timeout.tv_usec) / 1000000.0;
	qps = (double)(info->total_recv) / dt;
	lost = (int)(info->total_sent - info->total_recv) - (int)info->io_num;
	if(!info->quiet) {
		printf("overall time: 	%g sec\n", 
			(double)timeout.tv_sec + 
			(double)timeout.tv_usec/1000000.);
		if(lost > 0) 
			printf("Packets lost: 	%d\n", (int)lost);
	
		for(i=0; i<(int)(sizeof(info->by_rcode)/sizeof(size_t)); i++)
		{
			if(info->by_rcode[i] > 0) {
				char rc[16];
				sldns_wire2str_rcode_buf(i, rc, sizeof(rc));
				printf("%d(%5s): 	%u replies\n",
					i, rc, (unsigned)info->by_rcode[i]);
			}
		}
	}
	printf("average qps: 	%g\n", qps);
}

/** perform the performance test */
static void
perfmain(struct perfinfo* info)
{
	perfsetup(info);
	while(!info->exit) {
		perfselect(info);
	}
	perfendstats(info);
	perffree(info);
}

/** parse a query line to a packet into buffer */
static int
qlist_parse_line(sldns_buffer* buf, char* p)
{
	char nm[1024], cl[1024], tp[1024], fl[1024];
	int r; 
	int rec = 1, edns = 0;
	struct query_info qinfo;
	nm[0] = 0; cl[0] = 0; tp[0] = 0; fl[0] = 0;
	r = sscanf(p, " %1023s %1023s %1023s %1023s", nm, cl, tp, fl);
	if(r != 3 && r != 4)
		return 0;
	/*printf("nm='%s', cl='%s', tp='%s', fl='%s'\n", nm, cl, tp, fl);*/
	if(strcmp(tp, "IN") == 0 || strcmp(tp, "CH") == 0) {
		qinfo.qtype = sldns_get_rr_type_by_name(cl);
		qinfo.qclass = sldns_get_rr_class_by_name(tp);
	} else {
		qinfo.qtype = sldns_get_rr_type_by_name(tp);
		qinfo.qclass = sldns_get_rr_class_by_name(cl);
	}
	if(fl[0] == '+') rec = 1;
	else if(fl[0] == '-') rec = 0;
	else if(fl[0] == 'E') edns = 1;
	if((fl[0] == '+' || fl[0] == '-') && fl[1] == 'E')
		edns = 1;
	qinfo.qname = sldns_str2wire_dname(nm, &qinfo.qname_len);
	if(!qinfo.qname)
		return 0;
	qinfo.local_alias = NULL;
	qinfo_query_encode(buf, &qinfo);
	sldns_buffer_write_u16_at(buf, 0, 0); /* zero ID */
	if(rec) LDNS_RD_SET(sldns_buffer_begin(buf));
	if(edns) {
		struct edns_data ed;
		memset(&ed, 0, sizeof(ed));
		ed.edns_present = 1;
		ed.udp_size = EDNS_ADVERTISED_SIZE;
		/* Set DO bit in all EDNS datagrams ... */
		ed.bits = EDNS_DO;
		attach_edns_record(buf, &ed);
	}
	free(qinfo.qname);
	return 1;
}

/** grow query list capacity */
static void
qlist_grow_capacity(struct perfinfo* info)
{
	size_t newcap = (size_t)((info->qlist_capacity==0)?16:
		info->qlist_capacity*2);
	uint8_t** d = (uint8_t**)calloc(sizeof(uint8_t*), newcap);
	size_t* l = (size_t*)calloc(sizeof(size_t), newcap);
	if(!d || !l) fatal_exit("out of memory");
	if(info->qlist_data && info->qlist_capacity)
		memcpy(d, info->qlist_data, sizeof(uint8_t*)*
			info->qlist_capacity);
	if(info->qlist_len && info->qlist_capacity)
		memcpy(l, info->qlist_len, sizeof(size_t)*
			info->qlist_capacity);
	free(info->qlist_data);
	free(info->qlist_len);
	info->qlist_data = d;
	info->qlist_len = l;
	info->qlist_capacity = newcap;
}

/** setup query list in info */
static void
qlist_add_line(struct perfinfo* info, char* line, int no)
{
	if(!qlist_parse_line(info->buf, line)) {
		printf("error parsing query %d: %s\n", no, line);
		exit(1);
	}
	sldns_buffer_write_u16_at(info->buf, 0, (uint16_t)info->qlist_size); 
	if(info->qlist_size + 1 > info->qlist_capacity) {
		qlist_grow_capacity(info);
	}
	info->qlist_len[info->qlist_size] = sldns_buffer_limit(info->buf);
	info->qlist_data[info->qlist_size] = memdup(
		sldns_buffer_begin(info->buf), sldns_buffer_limit(info->buf));
	if(!info->qlist_data[info->qlist_size])
		fatal_exit("out of memory");
	info->qlist_size ++;
}

/** setup query list in info */
static void
qlist_read_file(struct perfinfo* info, char* fname)
{
	char buf[1024];
	char *p;
	FILE* in = fopen(fname, "r");
	int lineno = 0;
	if(!in) {
		perror(fname);
		exit(1);
	}
	while(fgets(buf, (int)sizeof(buf), in)) {
		lineno++;
		buf[sizeof(buf)-1] = 0;
		p = buf;
		while(*p == ' ' || *p == '\t')
			p++;
		if(p[0] == 0 || p[0] == '\n' || p[0] == ';' || p[0] == '#')
			continue;
		qlist_add_line(info, p, lineno);
	}
	printf("Read %s, got %u queries\n", fname, (unsigned)info->qlist_size);
	fclose(in);
}

/** getopt global, in case header files fail to declare it. */
extern int optind;
/** getopt global, in case header files fail to declare it. */
extern char* optarg;

/** main program for perf */
int main(int argc, char* argv[]) 
{
	char* nm = argv[0];
	int c;
	struct perfinfo info;
#ifdef USE_WINSOCK
	int r;
	WSADATA wsa_data;
#endif

	/* defaults */
	memset(&info, 0, sizeof(info));
	info.io_num = 16;

	log_init(NULL, 0, NULL);
	log_ident_set("perf");
	checklock_start();
#ifdef USE_WINSOCK
	if((r = WSAStartup(MAKEWORD(2,2), &wsa_data)) != 0)
		fatal_exit("WSAStartup failed: %s", wsa_strerror(r));
#endif

	info.buf = sldns_buffer_new(65553);
	if(!info.buf) fatal_exit("out of memory");

	/* parse the options */
	while( (c=getopt(argc, argv, "d:ha:f:q")) != -1) {
		switch(c) {
		case 'q':
			info.quiet = 1;
			break;
		case 'd':
			if(atoi(optarg)==0 && strcmp(optarg, "0")!=0) {
				printf("-d not a number %s", optarg);
				exit(1);
			}
			info.duration = atoi(optarg);
			break;
		case 'a':
			qlist_add_line(&info, optarg, 0);
			break;
		case 'f':
			qlist_read_file(&info, optarg);
			break;
		case '?':
		case 'h':
		default:
			usage(nm);
		}
	}
	argc -= optind;
	argv += optind;

	if(argc != 1) {
		printf("error: pass server IP address on commandline.\n");
		usage(nm);
	}
	if(!extstrtoaddr(argv[0], &info.dest, &info.destlen)) {
		printf("Could not parse ip: %s\n", argv[0]);
		exit(1);
	}
	if(info.qlist_size == 0) {
		printf("No queries to make, use -f or -a.\n");
		exit(1);
	}
	
	/* do the performance test */
	perfmain(&info);

	sldns_buffer_free(info.buf);
#ifdef USE_WINSOCK
	WSACleanup();
#endif
	checklock_stop();
	return 0;
}