/*
 * Copyright (C) 2000, 2001  Nominum, Inc.
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND INTERNET SOFTWARE CONSORTIUM
 * DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL
 * INTERNET SOFTWARE CONSORTIUM BE LIABLE FOR ANY SPECIAL, DIRECT,
 * INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING
 * FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
 * NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
 * WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

/*
 * Copyright (C) 2004 - 2006 Nominum, Inc.
 *
 * Permission to use, copy, modify, and distribute this software and its
 * documentation for any purpose with or without fee is hereby granted,
 * provided that the above copyright notice and this permission notice
 * appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND NOMINUM DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL NOMINUM BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
 * OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

/***
 ***	DNS Performance Testing Tool
 ***
 ***	Version $Id: dnsperf.c 73079 2006-08-15 23:02:37Z bwelling $
 ***/

#include "common.h"

/*
 * Configuration defaults
 */

#define DEF_MAX_QUERIES_OUTSTANDING	20
#define DEF_QUERY_TIMEOUT		5		/* in seconds */


/*
 * Other constants / definitions
 */

#define RESPONSE_BLOCKING_WAIT_TIME_US	100000		/* microseconds */
#define THROTTLE_DELAY			25000		/* microseconds */
#define UPDATE_TIME_FREQUENCY		100

/* Categories into which responses get categorized in the histogram. */

#define CATEGORY_SUCCESS 0
#define CATEGORY_FAILURE 1
#define N_CATEGORIES 2

struct histogram_bucket {
	unsigned int categories[2];
};

static char category_markers[N_CATEGORIES] = { '#', '-' };

/*
 * Configuration options (global)
 */

static isc_uint32_t max_queries_outstanding = DEF_MAX_QUERIES_OUTSTANDING;
static double query_timeout = DEF_QUERY_TIMEOUT;
static isc_uint32_t max_qps_limit;			/* init 0 */

static isc_boolean_t run_only_once = ISC_FALSE;
static isc_boolean_t use_timelimit = ISC_FALSE;
static unsigned int run_timelimit;			/* init 0 */

static isc_boolean_t countrcodes = ISC_FALSE;
static int rcodecounts[16] = {0};

static isc_uint32_t histogram_nbuckets = 0;
static isc_uint32_t histogram_time = 1;			/* in seconds */
static struct histogram_bucket *histogram_buckets;
static struct histogram_bucket histogram_overflow;
static double histogram_bps;				/* buckets per second */

static double latency_sum = 0;
static double latency_sum_squares = 0;
static double latency_min = 5000000;
static double latency_max = 0;

/*
 * Other global stuff
 */

static isc_boolean_t setup_phase = ISC_TRUE;

static unsigned int runs_through_file;			/* init 0 */

static unsigned int num_queries_timed_out;		/* init 0 */
static unsigned int num_queries_interrupted;		/* init 0 */
static int num_responses_since_update_time;		/* init 0 */

static isc_boolean_t interrupted = ISC_FALSE;

static isc_boolean_t want_stop = ISC_FALSE;

/*
 * Handle a SIGINT.
 */
static void
handle_interrupt(int sig)
{
	(void)sig;
	interrupted = ISC_TRUE;
}

/*
 * Install interrupt handler.  If this fails, we don't care.
 */
static void
install_interrupt_handler(void)
{
	struct sigaction sa;

	memset(&sa, 0, sizeof sa);
	sa.sa_handler = handle_interrupt;

	if (sigfillset(&sa.sa_mask) != 0)
		return;
	if (sigaction(SIGINT, &sa, NULL) < 0)
		return;
}

/*
 * show_startup_info:
 *   Show name/version
 */
static void
show_startup_info(void) {
	printf("\n"
"DNS Performance Testing Tool\n\n"
"Nominum Version " VERSION "\n\n");
}

/* Print program-specific usage */

void
show_usage_prog(void) {
	fprintf(stderr,
"               [-q num_queries] [-t timeout] [-Q max_qps] [-1] [-l limit]\n"
"               [-c] [-H histogram_buckets] [-T histogram_seconds] [-u]\n");
}

/* Print program-specific option descriptions */

void
show_options_prog(void) {
	fprintf(stderr,
"  -q specifies the maximum number of queries outstanding (default: %d)\n"
"  -t specifies the timeout for query completion in seconds (default: %d)\n"
"  -Q limit the number of queries per second (no default)\n"
"  -1 run through input only once (default: multiple iff limit given)\n"
"  -l specifies how a limit for how long to run tests in seconds (no default)\n"
"  -c print the number of packets with each rcode\n"
"  -H print a response latency histogram with this many buckets\n"
"  -T include latencies up to this many seconds in histogram (use with -H)\n"
"  -u send dynamic updates instead of queries\n"
"  -v verbose: report the RCODE of each response on stdout\n",
		DEF_MAX_QUERIES_OUTSTANDING,
		DEF_QUERY_TIMEOUT);
}

/* Handle program-specific options */

isc_result_t
handle_prog_opt(int c) {
	isc_result_t result;

	switch (c) {
	case 'q':
		result = isc_parse_uint32(&max_queries_outstanding,
					  optarg, 10);
		if (result != ISC_R_SUCCESS) {
			fprintf(stderr, "Invalid number of "
				"outstanding queries: %s\n",
				optarg);
			return (ISC_R_FAILURE);
		}
		break;
	case 't':
		result = parse_udouble(&query_timeout, optarg);
		if (result != ISC_R_SUCCESS) {
			fprintf(stderr, "Invalid query timeout "
				"(positive floating point value "
				"expected): %s\n", optarg);
			return (ISC_R_FAILURE);
		}
		break;
	case 'Q':
		result = isc_parse_uint32(&max_qps_limit, optarg, 10);
		if (result != ISC_R_SUCCESS) {
			fprintf(stderr, "Invalid max queries per "
				"second: %s\n", optarg);
			return (ISC_R_FAILURE);
		}
		break;
	case '1':
		run_only_once = ISC_TRUE;
		break;
	case 'l':
		result = isc_parse_uint32(&run_timelimit, optarg, 10);
		if (result != ISC_R_SUCCESS) {
			fprintf(stderr, "Invalid time limit: %s\n",
				optarg);
			return (ISC_R_FAILURE);
		}
		use_timelimit = ISC_TRUE;
		break;
	case 'c':
		countrcodes = ISC_TRUE;
		break;
	case 'H':
		result = isc_parse_uint32(&histogram_nbuckets,
					  optarg, 10);
		if (result != ISC_R_SUCCESS) {
			fprintf(stderr, "Invalid number of histogram "
				"buckets: %s\n",
				optarg);
			return (ISC_R_FAILURE);
		}
		break;
	case 'T':
		result = isc_parse_uint32(&histogram_time,
					  optarg, 10);
		if (result != ISC_R_SUCCESS) {
			fprintf(stderr, "Invalid histogram time: %s\n",
				optarg);
			return (ISC_R_FAILURE);
		}
		break;
	case 'u':
		update_mode = ISC_TRUE;
		break;
	case 'v':
		verbose = ISC_TRUE;
		break;
	default:
		return (ISC_R_NOTFOUND);
	}


	return (ISC_R_SUCCESS);
}

static void
init_histogram(void) {
	if (histogram_nbuckets == 0)
		return;
	histogram_buckets = calloc(histogram_nbuckets,
				   sizeof(struct histogram_bucket));
	if (histogram_buckets == NULL) {
		fprintf(stderr, "Out of memory\n");
		exit(1);
	}
	histogram_bps = (double) histogram_nbuckets / histogram_time;
}

/* Print a single histogram bucket. */

static void
print_bucket(struct histogram_bucket *bucket, int maxval, const char *op,
	     double time)
{
	int hashes;
	int j, k;

	printf("%s%8.3fs", op, time);

	for (k = 0; k < N_CATEGORIES; k++) {
		printf("%6d", bucket->categories[k]);
	}
	printf(" |");

	for (k = 0; k < N_CATEGORIES; k++) {
		hashes = (int) floor(60.0 * bucket->categories[k] / maxval);
		for (j = 0; j < hashes; j++)
			putchar(category_markers[k]);
	}
	printf("\n");
}

static void
print_histogram(void) {
	int total = 0;
	int maxval = 1;
	unsigned int i, j;
	int bucket_total;
	int count;

	if (histogram_nbuckets == 0)
		return;

	for (i = 0; i < histogram_nbuckets; i++) {
		bucket_total = 0;
		for (j = 0; j < N_CATEGORIES; j++) {
			count = histogram_buckets[i].categories[j];
			total += count;
			bucket_total += count;
		}
		if (bucket_total > maxval)
			maxval = bucket_total;
	}

	bucket_total = 0;
	for (j = 0; j < N_CATEGORIES; j++) {
		count = histogram_overflow.categories[j];
		total += count;
		bucket_total += count;
	}
	if (bucket_total > maxval)
		maxval = bucket_total;

	printf("Latency: Min: %f s; Max: %f s; Avg: %f s; StdDev: %f\n",
	       latency_min, latency_max, latency_sum / total,
	       sqrt((latency_sum_squares - ((latency_sum * latency_sum)/total))
		    /(total-1)));

	printf("\nResponse latency distribution (total %d responses):\n\n",
	       total);

	printf("Latency   Success  Fail |\n");

	for (i = 0; i < histogram_nbuckets; i++) {
		print_bucket(&histogram_buckets[i], maxval, "< ",
			     (i + 1) / histogram_bps);

	}
	print_bucket(&histogram_overflow, maxval, ">=",
		     histogram_nbuckets / histogram_bps);

	printf("\nLegend:\n\n");
	printf("##### = success responses (RCODE was NOERROR or NXDOMAIN)\n");
	printf("----- = failure responses (any other RCODE)\n");
}

/*
 * timelimit_reached:
 *   Have we reached the time limit (if any)?
 *
 *   Returns ISC_FALSE if there is no time limit or if we have not reached it
 *   Returns ISC_TRUE otherwise
 */
static isc_boolean_t
timelimit_reached(void) {
	if (! use_timelimit || setup_phase)
		return (ISC_FALSE);

	return ISC_TF(difftv(&time_now, &time_of_program_start) >
		      (double)run_timelimit);
}

/*
 * keep_sending:
 *   Should we keep sending queries or stop here?
 *
 *   Return ISC_TRUE if we should keep on sending queries
 *   Return ISC_FALSE if we should stop
 *
 *   Side effects:
 *   Rewinds the input and clears reached_end_input if we have reached the
 *   end of the input, but we are meant to run through it multiple times
 *   and have not hit the time limit yet (if any is set).
 */
static isc_boolean_t
keep_sending(isc_boolean_t *reached_end_input) {
	int at_timelimit;

	if (want_stop)
		return (ISC_FALSE);

	at_timelimit = timelimit_reached();

	if (!(*reached_end_input) && !at_timelimit)
		return (ISC_TRUE);
	else if (*reached_end_input && !run_only_once && !at_timelimit) {
		rewind(datafile);
		*reached_end_input = ISC_FALSE;
		runs_through_file++;
		return (ISC_TRUE);
	} else {
		if (*reached_end_input)
			runs_through_file++;
		want_stop = ISC_TRUE;
		return (ISC_FALSE);
	}
}

/*
 * register_response:
 *   Register receipt of a query
 *
 *   Removes the record for the given query id in status[] if any exists.
 *
 *   Returns ISC_TRUE if the response was for an outstanding query
 *   Returns ISC_FALSE otherwise.
 */
isc_boolean_t
register_response(unsigned int id, unsigned int rcode) {
	struct query_status *s = &status[id];

	if (s->list != &outstanding_list) {
		fprintf(stderr, "Warning: Received a response with an "
		        "unexpected (maybe timed out) id: %u\n", id);
		return (ISC_FALSE);
	}

	ISC_LIST_UNLINK(outstanding_list, s, link);
	ISC_LIST_APPEND(instanding_list, s, link);
	s->list = &instanding_list;

	num_queries_outstanding--;
	if (s->desc != NULL) {
		printf("> %s %s\n", rcode_strings[rcode], s->desc);
		free(s->desc);
		s->desc = NULL;
	}

	if (countrcodes)
		rcodecounts[rcode]++;

	if (histogram_nbuckets) {
		double latency;
		int buckno;
		struct histogram_bucket *bucket;
		int category;

		latency = difftv(&time_now, &s->sent_timestamp);
		latency_sum += latency;
		latency_sum_squares += (latency * latency);
		if (latency < latency_min)
			latency_min = latency;
		if (latency > latency_max)
			latency_max = latency;

		buckno = (int) floor(latency * histogram_bps);
		if (buckno < 0)
			buckno = 0;
		if ((unsigned int) buckno < histogram_nbuckets)
			bucket = &histogram_buckets[buckno];
		else
			bucket = &histogram_overflow;

		if (rcode == dns_rcode_noerror || rcode == dns_rcode_nxdomain)
			category = CATEGORY_SUCCESS;
		else
			category = CATEGORY_FAILURE;

		bucket->categories[category]++;
	}
	return (ISC_TRUE);
}


/*
 * Wait for a response from the server.
 *
 *   Returns whether we have a response to read.
 */
static isc_boolean_t
wait_for_response(isc_boolean_t throttling) {
	fd_set read_fds;
	struct timeval tv;
	int retval;

	FD_ZERO(&read_fds);
	FD_SET(query_socket, &read_fds);

	tv.tv_sec = 0;
	if (throttling)
		tv.tv_usec = THROTTLE_DELAY;
	else
		tv.tv_usec = RESPONSE_BLOCKING_WAIT_TIME_US;

	retval = select(query_socket + 1, &read_fds, NULL, NULL, &tv);
	if (retval < 0) {
		if (errno == EINTR)
			return (ISC_FALSE);
		fprintf(stderr, "select() failed\n");
		exit(1);
	}

	update_current_time();

	return ISC_TF(FD_ISSET(query_socket, &read_fds));
}

/*
 * need_throttle:
 *   Determine if we've exceeded the maximum queries per second limit.
 */
static isc_boolean_t
need_throttle(void) {
	double runtime;

	runtime = difftv(&time_now, &time_of_program_start);
	return ISC_TF((runtime > 0.1 &&
		       num_queries_sent > max_qps_limit * runtime));
}

/*
 * retire_old_queries:
 *   Go through the list of open queries and remove any queries which are
 *   older than the timeout, decrementing the number of queries
 *   outstanding for each one removed.
 */
static void
retire_old_queries(void) {
	struct query_status *s;

	while (ISC_TRUE) {
		s = ISC_LIST_HEAD(outstanding_list);
		if (s == NULL ||
		    (difftv(&time_now, &s->sent_timestamp) < query_timeout))
			break;
		ISC_LIST_UNLINK(outstanding_list, s, link);
		ISC_LIST_APPEND(instanding_list, s, link);
		s->list = &instanding_list;

		num_queries_outstanding--;
		num_queries_timed_out++;

		if (s->desc != NULL) {
			printf("> T %s\n", s->desc);
			free(s->desc);
			s->desc = NULL;
		} else {
			printf("[Timeout] %s timed out: msg id %u\n",
			       update_mode ? "Update" : "Query",
			       (unsigned int)(s - status));
		}
	}
}

/*
 * cancel_queries:
 *   Go through the list of open queries and remove any queries which are
 *   outstanding, decrementing the number of queries outstanding for each
 *   one removed.
 */
static void
cancel_queries(void) {
	struct query_status *s;

	while (ISC_TRUE) {
		s = ISC_LIST_HEAD(outstanding_list);
		if (s == NULL)
			break;
		ISC_LIST_UNLINK(outstanding_list, s, link);
		ISC_LIST_APPEND(instanding_list, s, link);
		s->list = &instanding_list;

		num_queries_outstanding--;
		num_queries_interrupted++;

		if (s->desc != NULL) {
			printf("> I %s\n", s->desc);
			free(s->desc);
			s->desc = NULL;
		} else {
			printf("[Interrupt] %s interrupted: msg id %u\n",
			       update_mode ? "Update" : "Query",
			       (unsigned int)(s - status));
		}
	}
}

/*
 * print_statistics:
 *   Print out statistics based on the results of the test
 */
static void
print_statistics(void) {
	unsigned int num_queries_completed;
	double per_lost, per_completed;
	double run_time, queries_per_sec;
	time_t start_secs, end_secs;
	isc_uint64_t avg_request_size, avg_response_size;
	const char *units_cap, *units;
	unsigned int i;

	units = update_mode ? "updates" : "queries";
	units_cap = update_mode ? "Updates" : "Queries";

	num_queries_completed = num_queries_sent - num_queries_timed_out -
			        num_queries_interrupted;

	if (num_queries_completed == 0) {
		per_lost = 0.0;
		per_completed = 0.0;
		avg_response_size = 0;
	} else {
		per_lost = 100.0 * (double)num_queries_timed_out
			   / (double)num_queries_sent;
		per_completed = 100.0 - per_lost;
		avg_response_size = total_response_size /
				    num_queries_completed;
	}

	run_time = difftv(&time_of_end_of_run, &time_of_program_start);
	
	if (num_queries_sent == 0) {
		queries_per_sec = 0.0;
		avg_request_size = 0;
	} else {
		queries_per_sec = (double)num_queries_completed / run_time;
		avg_request_size = total_request_size / num_queries_sent;
	}

	printf("\n");

	printf("Statistics:\n\n");

	printf("  Parse input file:     %s\n",
	       (run_only_once ? "once" : "multiple times"));
	if (use_timelimit)
		printf("  Run time limit:       %u seconds\n", run_timelimit);
	if (!run_only_once)
		printf("  Ran through file:     %u times\n",
		       runs_through_file);
	else {
		const char *reason;
		if (interrupted)
			reason = "interruption";
		else if (runs_through_file == 0)
			reason = "reaching time limit";
		else
			reason = "reaching end of file";
		printf("  Ended due to:         %s\n", reason);
	}

	printf("\n");

	printf("  %s sent:         %u %s\n",
	       units_cap, num_queries_sent, units);
	printf("  %s completed:    %u %s\n",
	       units_cap, num_queries_completed, units);
	printf("  %s lost:         %u %s\n",
	       units_cap, num_queries_timed_out, units);
	if (num_queries_interrupted > 0)
		printf("  %s interrupted:  %u %s\n",
		       units_cap, num_queries_interrupted, units);
	printf("\n");

	if (countrcodes) {
		for (i = 0; i < 16; i++) {
			if (rcodecounts[i] == 0)
				continue;
			printf("  Returned %8s:    %u %s\n",
			       rcode_strings[i], rcodecounts[i], units);
		}
		printf("\n");
	}

	printf("  Avg request size:     %u bytes\n",
	       (unsigned int)avg_request_size);
	printf("  Avg response size:    %u bytes\n",
	       (unsigned int)avg_response_size);
	printf("\n");

	printf("  Percentage completed: %6.2lf%%\n", per_completed);
	printf("  Percentage lost:      %6.2lf%%\n", per_lost);
	printf("\n");

	start_secs = time_of_program_start.tv_sec;
	end_secs = time_of_end_of_run.tv_sec;
	printf("  Started at:           %s", ctime(&start_secs));
	printf("  Finished at:          %s", ctime(&end_secs));
	printf("  Ran for:              %.6lf seconds\n", run_time);

	printf("\n");

	printf("  %s per second:   %.6lf %cps\n",
	       units_cap, queries_per_sec, units[0]);

	printf("\n");
}

static void
maybe_update_current_time() {
	/*
	 * We don't want to call gettimeofday() every time a packet is sent
	 * or received, as it's expensive, so we keep a cached copy of the
	 * time when we last checked.
	 *
	 * We check the current time after blocking, and also keep a count
	 * of the number of packets we've read since we last blocked so that
	 * if we never block, we can still update the cached time periodically.
	 */
	num_responses_since_update_time++;
	if (num_responses_since_update_time > UPDATE_TIME_FREQUENCY) {
		num_responses_since_update_time = 0;
		update_current_time();
	}
}

/*
 * Main program.
 */
int
main(int argc, char **argv) {
	isc_boolean_t got_eof = ISC_FALSE, keepsending;
	char line[MAX_INPUT_LEN + 1];
	isc_result_t result;

	show_startup_info();

	if (setup(argc, argv, "q:t:Q:1l:cH:T:uv") != ISC_R_SUCCESS)
		exit(1);

	if (!run_only_once && !use_timelimit)
		run_only_once = ISC_TRUE;

	init_histogram();

	install_interrupt_handler();

	printf("[Status] Processing input data\n");

	update_current_time();
	time_of_program_start = time_now;

	while (!interrupted) {
		int burst;
		isc_boolean_t throttle;

		keepsending = keep_sending(&got_eof);
		if (! keepsending && num_queries_outstanding == 0)
			break;

		burst = 0;
		throttle = ISC_FALSE;
		while (keepsending && !interrupted &&
		       num_queries_outstanding < max_queries_outstanding)
		{
			if (max_qps_limit > 0) {
				throttle = need_throttle();
				if (throttle)
					break;
			}
			if (get_input_line(line, sizeof(line)) !=
					   ISC_R_SUCCESS)
			{
				got_eof = ISC_TRUE;
				if (setup_phase) {
					fprintf(stderr, "Input file "
						"contains no queries.\n");
					exit(1);
				}
				break;
			}
			result = process_line(line);
			if (result == ISC_R_SUCCESS && setup_phase) {
				setup_phase = ISC_FALSE;
				printf("[Status] Sending %s (to %s)\n",
				       update_mode ? "updates" : "queries", server_addr_str);
			}
			
			/*
			 * This implements TCP-like "slow start".
			 */
			if (++burst == 2)
				break;
		}

		if (num_queries_outstanding || throttle) {
			for (;;) {
				if (try_process_response(query_socket)) {
					maybe_update_current_time();
					break;
				} else {
					num_responses_since_update_time = 0;
				}

				/*
				 * There was nothing on the socket.  We have
				 * free time.  Use it to retire old queries,
				 * and then drop into select().
				 */
				retire_old_queries();

				if (!wait_for_response(throttle))
					break;
			}
		}
	}

	update_current_time();
	time_of_end_of_run = time_now;

	if (interrupted)
		cancel_queries();

	printf("[Status] Testing complete\n");

	maybe_print_args(argc, argv);

	print_statistics();
	print_histogram();

	cleanup();

	return (0);
}