/*
 * Copyright (C) 2007, Ingo Molnar
 *
 * fd-scale-bench.c: check the cost of opening a file descriptor
 *
 * Compile with: gcc -Wall -O2 -o fd-scale-bench fd-scale-bench.c -lrt
 */
#include <stdio.h>
#include <stdarg.h>
#include <stdlib.h>
#include <signal.h>
#include <sys/wait.h>
#include <linux/unistd.h>
#include <unistd.h>
#include <string.h>
#include <pwd.h>
#include <grp.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/io.h>
#include <regex.h>
#include <fcntl.h>
#include <time.h>
#include <sys/mman.h>
#include <dlfcn.h>
#include <popt.h>
#include <sys/socket.h>
#include <ctype.h>
#include <assert.h>
#include <sched.h>
#include <time.h>

#if DEBUG
# define Printf(x...) printf(x)
#else
# define Printf(x...) do { } while (0)
#endif

#define BUG_ON(c) assert(!(c))

typedef unsigned long long cycles_t;
typedef unsigned long long usecs_t;

#define rdtod(val)					\
({							\
	struct timeval tv;				\
							\
	gettimeofday(&tv, NULL);			\
	tv.tv_sec * 1000000ULL + tv.tv_usec;		\
})

static int num_fds, max_fds = 1000, flush_cache;

#define LOOPS 5

static int one_more_fd(void)
{
	int fd;

	fd = open("/dev/null", O_RDONLY);
	if (fd < 0) {
		perror("open");
		exit(-1);
	}

	return fd;
}

#define MAX_CACHE_SIZE (8*1024*1024)

static char cache_array[MAX_CACHE_SIZE];

static void flush(void)
{
	int i;

	if (!flush_cache)
		return;

	for (i = 0; i < MAX_CACHE_SIZE; i++)
		cache_array[i]++;
}

static inline void test_fd(void)
{
	usecs_t t0, t1, delta, best_delta = 1000000, worst_delta = 0;
	int i, fd1 = 0, fd2 = 0, middle_fd, prev_num_fds, more_fds;

	prev_num_fds = num_fds;
	num_fds += num_fds/4 + 1;
	if (num_fds > max_fds)
		num_fds = max_fds;
	more_fds = num_fds - prev_num_fds;

	/*
	 * Allocate more fds:
	 */
	for (i = 0; i < more_fds; i++)
		fd1 = one_more_fd();

	/*
	 * We free an fd in the middle:
	 */
	middle_fd = fd1/2 + 1;
	close(middle_fd);

	for (i = 0; i < LOOPS; i++) {
		flush();

		t0 = rdtod();
		fd1 = one_more_fd();
		fd2 = one_more_fd();
		close(fd1);
		close(fd2);
		t1 = rdtod();
//		printf("%d %d\n", fd1, fd2);
		delta = t1 - t0;
		if (delta < best_delta)
			best_delta = delta;
		if (delta > worst_delta)
			worst_delta = delta;
	}

	/*
	 * fill in the hole:
	 */
	one_more_fd();

	printf("num_fds: %d, best cost: %.2f us, worst cost: %.2f us\n",
		num_fds,
		(double)best_delta, (double)worst_delta);

	if (num_fds == max_fds)
		exit(0);
}

/*
 * Open/close max # of fds, to check whether we can allocate
 * them and to also prime the kernel's fd bitmap:
 */
static void prime_fd_array(void)
{
	int i, first_fd, last_fd = 0;

	first_fd = one_more_fd();

	for (i = 1; i < max_fds; i++) {
		last_fd = one_more_fd();
		if (last_fd < 0) {
			perror("open");
			exit(-1);
		}
	}
	for (i = first_fd; i <= last_fd; i++)
		close(i);
}

int main(int argc, char **argv)
{
	struct rlimit rlim;

	if (argc > 3) {
usage:
		fprintf(stderr,
			"usage: fd-scale-bench [<max_files> [<flush_cache>]]\n");
		exit(-1);
	}
	getrlimit(RLIMIT_NOFILE, &rlim);

	if (argc >= 2) {
		max_fds = atol(argv[1]);
		if (max_fds <= 0)
			goto usage;
		rlim.rlim_cur = max_fds + 10;
		rlim.rlim_max = max_fds + 10;

		setrlimit(RLIMIT_NOFILE, &rlim);
		getrlimit(RLIMIT_NOFILE, &rlim);
	} else {
		max_fds = rlim.rlim_cur - 10;
	}
	if (max_fds > rlim.rlim_cur)
		max_fds = rlim.rlim_cur - 10;

	if (argc == 3)
		flush_cache = atol(argv[2]);

	printf("checking the %s performance of open()-ing %d fds.\n",
		flush_cache ? "cache-cold" : "cache-hot", max_fds);

	prime_fd_array();

	while (1) {
		test_fd();
	}

	return 0;
}