/* Copyright Ingo Molnar (c) 2006 */

#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <pthread.h>
#include <unistd.h>
#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <errno.h>
#include <stdlib.h>
#include <assert.h>
#include <sys/times.h>
#include <sys/wait.h>
#include <sys/ioctl.h>
#include <linux/unistd.h>
#include <unistd.h>
#include <sys/mman.h>

#define _GNU_SOURCE
#include <unistd.h>
#include <sys/mman.h>

#define DEBUG 0

#if DEBUG
# define dprintf(x...) printf(x)
#else
# define dprintf(x...)
#endif

#define PAGE_SIZE 4096UL

#define MIN_WINDOW_SIZE (800 * 1024UL * 1024UL)

static unsigned long window_size;
static void *window;

enum {
	OP_MMAP,
	OP_MUNMAP,
	OP_MREMAP,
	NR_OPS
};

static void print_maps(void)
{
	char buf[8192];

        int fd, n;
        if ((fd = open("/proc/self/maps", O_RDONLY)) < 0) {
                perror("open");
                exit(1);
        }
        if ((n = read(fd, &buf[0], sizeof(buf))) < 0) {
                perror("read");
                exit(2);
        }
        close(fd);
        if (n > 0)
                write(1, &buf[0], n);
}

unsigned long get_rand_size_idx(unsigned long start_idx)
{
	unsigned long rnd1 = random(), rnd2 = random();
	unsigned long max_idx, size_idx;
	unsigned long factor, range;

	max_idx = window_size / PAGE_SIZE;
	/*
	 * limit the size of remapping to 1-256 pages, but still leave
	 * a (small) chance for larger remaps:
	 */
	size_idx = rnd1 & 0xf;
	if (!size_idx) {
		size_idx = rnd1 & 0xff;
		if (!size_idx) {
			size_idx = rnd1 & 0xfff;
			if (!size_idx) {
				factor = sqrt(rnd2 % max_idx) * 100 /
							sqrt(max_idx) + 1;
				range = (max_idx - start_idx) / factor + 1;
				size_idx = rnd2 % range + 1;
			} else
				size_idx /= 32;
		} else
			size_idx /= 16;
	} else
		size_idx /= 2;
	if (!size_idx)
		size_idx = 1;

	if (size_idx > max_idx - start_idx)
		size_idx = max_idx - start_idx;

	return size_idx;
}

static void
do_munmap_op(unsigned long start_idx, unsigned long size_idx)
{
	void *start_addr = window + start_idx * PAGE_SIZE;
	unsigned long size = size_idx * PAGE_SIZE;
	int ret;

	ret = munmap(start_addr, size);
	assert(ret == 0);
	dprintf("munmap: %p [%ld]\n", start_addr, size);
}

static void
do_mmap_op(unsigned long start_idx, unsigned long size_idx, unsigned long rnd)
{
	void *start_addr = window + start_idx * PAGE_SIZE;
	unsigned long size = size_idx * PAGE_SIZE;
	int prot_rnd = (rnd/256) % 3, prot = 0, map_fixed = 0,
	    touch_it_read, touch_it_write;
	void *addr;

	if (!prot_rnd)
		prot = PROT_NONE;
	/*
	 * Both read and write mapped has a 50% chance to be set:
	 */
	else {
		if (prot_rnd & 1)
			prot |= PROT_READ;
		if (prot_rnd & 2)
			prot |= PROT_WRITE;
	}
	/*
	 * map it fixed, with a 50% chance:
	 */
	if (rnd & 8192)
		map_fixed = MAP_FIXED;
	/*
	 * Touch the mapping with a 25% chance:
	 */
	touch_it_read = 0;
	if (((rnd >> 10) & 3) == 3)
		touch_it_read = 1;

	touch_it_write = 0;
	if (((rnd >> 6) & 3) == 3)
		touch_it_write = 1;

	addr = mmap(start_addr, size, prot,
		    MAP_ANONYMOUS|MAP_PRIVATE|map_fixed, -1, 0);
	dprintf("mmap: %d, %d, %p - %p => %p\n",
		prot, map_fixed ? 1 : 0, start_addr, start_addr+size, addr);
	assert(!map_fixed || (addr == start_addr));
	if (addr != (void *)-1) {
		volatile char data = 0xf;

		if ((prot & PROT_READ) && touch_it_read)
			data = *(char *)addr;
		if ((prot & PROT_WRITE) && touch_it_write)
			*(char *)addr = data;
	}
	/*
	 * If it's a non-fixed mmap then unmap it, if it's outside
	 * of the window (otherwise we'd slowly leak memory):
	 */
	if (addr < window || addr >= window + window_size) {
		int ret = munmap(addr, size);
		assert(ret == 0);
		dprintf("mmap-munmap: %p [%ld]\n", addr, size);
	}
}

static void
do_mremap_op(unsigned long start_idx, unsigned long size_idx, unsigned long rnd)
{
	unsigned long old_size = size_idx * PAGE_SIZE, new_size, may_move;
	void *start_addr = window + start_idx * PAGE_SIZE;
	void *addr;

	/*
	 * Resize or just plain move?
	 */
	may_move = 0;
	if (rnd & 1024) {
		new_size = get_rand_size_idx(start_idx);
		if (rnd & 2048)
			may_move = 1;
	} else
		new_size = old_size;

	addr = (void *)mremap(start_addr, old_size, new_size, may_move);
	dprintf("mremap: [%ld, %ld] => %p.\n", new_size, may_move, addr);
}


static void do_random_mem_op(void)
{
	unsigned long rnd1 = random(), rnd2 = random();
	unsigned long max_idx, start_idx, end_idx, size_idx;
	int op;

	max_idx = window_size / PAGE_SIZE;
	start_idx = rnd1 % (max_idx-1);
	size_idx = get_rand_size_idx(start_idx);

	end_idx = start_idx + size_idx;

	op = rnd2 % NR_OPS;

	dprintf("op: %d, %6lu - %6lu [%lu, %.1f%%] [%p - %p]\n",
		op, start_idx, end_idx, size_idx,
		100.0*(double)size_idx/(double)max_idx,
		window + start_idx*PAGE_SIZE, window + end_idx*PAGE_SIZE);

	assert(end_idx <= max_idx);

	switch (op) {
		case OP_MMAP:
			do_mmap_op(start_idx, size_idx, rnd2);
			break;

		case OP_MUNMAP:
			do_munmap_op(start_idx, size_idx);
			break;

		case OP_MREMAP:
			do_mremap_op(start_idx, size_idx, rnd2);
			break;

		default:
			assert(0);
	}
}

static void *threadfn(void *arg)
{
	int thread = (int)(long)arg;
	int i;

	printf("thread %d started.\n", thread); fflush(stdout);

	for (i = 0; i < 1000000; i++) {
		do_random_mem_op();
	}
//	print_maps();

	return NULL;
}

/*
 * We may get segfaults, if one thread unmaps an area that we
 * are touching. We ignore these segfaults:
 */
static void segfault_handler(int sig, siginfo_t *info, void *uc)
{
}

int main (int argc, char **argv)
{
	unsigned int nr_threads = 0, i, ret;
	struct sigaction act;
	pthread_t *t;
	char random_state[32];

	initstate(0, random_state, sizeof(random_state));

	sigfillset(&act.sa_mask);
	act.sa_sigaction = segfault_handler;
	sigaction(SIGSEGV, &act, NULL);
	sigaction(SIGBUS, &act, NULL);

	if (argc != 1 && argc != 2 && argc != 3) {
usage:
		printf("usage: loop_print2 [<nr threads>] [<test window size [MB]>]\n");
		exit(-1);
	}
	/*
	 * The default is to use # of CPUs threads:
	 */
	if (argc == 1) {
		nr_threads = system("exit `grep -w processor /proc/cpuinfo  | wc -l`");
		nr_threads = WEXITSTATUS(nr_threads);
	}
	if (argc >= 2) {
		nr_threads = atol(argv[1]);
		if (!nr_threads)
			goto usage;
	}
	if (argc >= 3)
		window_size = atol(argv[2]) * 1024UL * 1024UL;
	if (window_size < MIN_WINDOW_SIZE)
		window_size = MIN_WINDOW_SIZE;

	/*
	 * Create the mmap() playground:
	 */
	window = mmap(0, window_size, PROT_READ|PROT_WRITE,
			MAP_ANONYMOUS|MAP_PRIVATE, -1, 0);
	if ((long)window == -1) {
		printf("test window mmap failed!\n");
		return -1;
	}

	/*
	 * Create and start all threads:
	 */
	t = calloc(nr_threads, sizeof(*t));

	for (i = 1; i < nr_threads; i++) {
		ret = pthread_create (t+i, NULL, threadfn, (void *)(long)i);
		if (ret)
			exit(-1);
	}
	threadfn((void *)0);
	for (i = 1; i < nr_threads; i++) {
		ret = pthread_join (t[i], NULL);
		if (ret)
			exit(-1);
	}

	return 0;
}