/*
 * case: stress file-backed THP
 */
#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif

#include <assert.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <sched.h>
#include <time.h>
#include <string.h>
#include <fcntl.h>
#include <signal.h> /* for signal */
#include <sys/mman.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <errno.h>

#define PATH_MAX	1024
#define BUFF_MAX	1024
#define TIME_DFL	180	/* seconds */

void signal_handler(int signo)
{
	/* Restore env */
	system("echo never > /sys/kernel/mm/transparent_hugepage/enabled");
	system("echo 10000 > /sys/kernel/mm/transparent_hugepage/khugepaged/scan_sleep_millisecs");

	printf("\nrestore env:\n");
	printf("	echo never > /sys/kernel/mm/transparent_hugepage/enabled\n");
	printf("	echo 10000 > /sys/kernel/mm/transparent_hugepage/khugepaged/scan_sleep_millisecs\n");
	exit(-1);
}

/* in KB */
#define text_size	(14UL << 10)

#define PROCMAP_SZ	8
struct procmap {
	uint64_t vm_start;
	uint64_t vm_end;
	uint64_t pgoff;
	uint32_t maj;
	uint32_t min;
	uint32_t ino;
#define PROT_SZ		5
	char     prot[PROT_SZ];
	char     fname[PATH_MAX];
};

unsigned long sleep_secs = 0;

/*
 * Routines of procmap, i.e., /proc/pid/(s)maps
 */
static int get_memory_map(pid_t pid, struct procmap *procmap,
		const char *fname)
{
	char path[PATH_MAX];
	char line[BUFF_MAX];
	FILE *fp = NULL;
	char *end = NULL;
	char *pos, *sp = NULL, *in[PROCMAP_SZ];
	char dlm[] = "-   :   ";
	uint64_t counter;
	int i;

	snprintf(path, PATH_MAX, "/proc/%u/maps", pid);

	fp = fopen(path, "r");
	if (fp == NULL) {
		printf("fopen: %s: %s\n", path, strerror(errno));
		return -1;
	}

	if (procmap == NULL || fname == NULL) {
		perror("fail: procmap or fname is NULL");
		goto failed;
	}

	while (fgets(line, BUFF_MAX, fp)) {
		/* Split line into fields */
		pos = line;
		for (i = 0; i < PROCMAP_SZ; i++) {
			in[i] = strtok_r(pos, &dlm[i], &sp);
			if (in[i] == NULL)
				break;
			pos = NULL;
		}

		/* Check this line is procmap item header */
		if (i != PROCMAP_SZ)
			continue;

		memcpy(procmap->prot, in[2], PROT_SZ);
		memcpy(procmap->fname, in[7], PATH_MAX);

		/* Find the target entry */
		if (strcmp(procmap->prot, "r-xp") ||
				!strstr(procmap->fname, fname))
			continue;

		/* Convert/Copy each field as needed */
		errno = 0;
		procmap->vm_start = strtoull(in[0], &end, 16);
		if ((in[0] == '\0') || (end == NULL) || (*end != '\0') ||
				(errno != 0))
			goto failed;

		procmap->vm_end = strtoull(in[1], &end, 16);
		if ((in[1] == '\0') || (end == NULL) || (*end != '\0') ||
				(errno != 0))
			goto failed;

		procmap->pgoff = strtoull(in[3], &end, 16);
		if ((in[3] == '\0') || (end == NULL) || (*end != '\0') ||
				(errno != 0))
			goto failed;

		procmap->ino = strtoul(in[6], &end, 16);
		if ((in[6] == '\0') || (end == NULL) || (*end != '\0') ||
				(errno != 0))
			goto failed;
	}

	if (fp)
		fclose(fp);
	return 0;

failed:
	if (fp)
		fclose(fp);
	printf("fail: exit\n");

	return -1;
}

#define NR_CPU 32
uint64_t gettimeofday_sec(void);
inline uint64_t gettimeofday_sec(void)
{
	struct timeval tv;

	gettimeofday(&tv, NULL);
	return tv.tv_sec;
}

void thread_read(int cpu, char *args)
{
	int fd;
	char *dso_path = args;
	char buf[0x800000];
	struct procmap maps;
	pid_t pid = getpid();

	cpu_set_t mask;
	CPU_ZERO(&mask);
	CPU_SET(cpu, &mask);
	if (sched_setaffinity(0, sizeof(cpu_set_t), &mask) == -1) {
		printf("warning: can not set CPU affinity\n");
	}

	printf("read %s\n", dso_path);
	fd = open(dso_path, O_RDONLY);
	/* The start addr must be alignment with 2M */
	void *p = mmap((void *)0x40000dc00000UL, 0x800000, PROT_READ | PROT_EXEC,
			MAP_PRIVATE, fd, 0);
	if (p == MAP_FAILED) {
		perror("mmap");
		goto out;
	}

	/* get the mapping address (ONLY r-xp) of the DSO */
	get_memory_map(pid, &maps, dso_path);
	printf("pid: %d\n", pid);
	printf("text vm_start: 0x%lx\n", maps.vm_start);
	printf("text vm_end: 0x%lx\n", maps.vm_end);
	madvise((void *)maps.vm_start, maps.vm_end - maps.vm_start, MADV_HUGEPAGE);
	lseek(fd, 0, SEEK_SET);
	for(;;) {
		memcpy(buf, p, 0x800000 - 1);
		sleep(1);
	}

	sleep(100);

out:
	/* Restore env */
	system("echo never > /sys/kernel/mm/transparent_hugepage/enabled");
	system("echo 10000 > /sys/kernel/mm/transparent_hugepage/khugepaged/scan_sleep_millisecs");

	printf("read exit %s\n", dso_path);
}

void thread_write(int cpu, char *args)
{
	void *p = NULL;
	char buf[32];
	uint64_t sec = 1;
	uint64_t count = 0;
	char *dso_path = args;

	cpu_set_t mask;
	CPU_ZERO(&mask);
	CPU_SET(cpu, &mask);
	if (sched_setaffinity(0, sizeof(cpu_set_t), &mask) == -1) {
		printf("warning: can not set CPU affinity\n");
	}

	sleep(3);
	printf("write %s\n", dso_path);
	for (;;) {
		sec = gettimeofday_sec();
		while ((sec % 10) >= 3) {
			sec = gettimeofday_sec();
		}

		int fd = open(dso_path, O_RDWR);
		p = mmap(NULL, 0x800000, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
		if (p == MAP_FAILED) {
			perror("mmap");
			goto out; /* fail */
		}

		lseek(fd, 0x1600, SEEK_SET);
		for(long i=1; i <= 2; i++){
			memcpy(p + 0x10, buf, 16);
		}

		munmap(p, 0x800000);
		close(fd);

		sleep(2);
		count++;
		if (count >= sleep_secs) {
			printf("test finish: %ld\n", count);
			break;
		}
	} /* end for */

out:
	printf("write exit %s\n", dso_path);
}

/*
 * usage:
 *		stress_madvise_dso <test time> <libtest.so>
 */
int main(int argc, char *argv[])
{
	struct timeval start, end;
	char dso_path[80];
	int ret = 0;
	pid_t pid;

	if (argc > 2) {
		sleep_secs = strtoul(argv[1], NULL, 10);
		realpath(argv[2], dso_path);
	}
	else {
		printf("usage error:\n"\
				"	stress_madvise_dso <test time> <libtest.so>\n"\
				"	e.g. stress_madvise_dso 10000 libtest.so\n");
		exit(-1);
	}

	/* Set env */
	system("ulimit -s unlimited");
	system("echo madvise > /sys/kernel/mm/transparent_hugepage/enabled");
	system("echo 1000 > /sys/kernel/mm/transparent_hugepage/khugepaged/scan_sleep_millisecs");

	gettimeofday(&start, NULL);

	/*
	 * fork 32 task to read and write the same DSO:
	 *		task 0: read dso;
	 *		task 1 - 31: write dso;
	 */
	for (int i = 0; i < NR_CPU; ++i) {
		pid = fork();
		if (pid == 0) {
			if (i == 0)
				thread_read(i, dso_path);
			else
				thread_write(i, dso_path);
			break; /* forbid child fork */
		}
		else {
			/* parent */
		}
	}

	if (pid != 0) {
		signal(SIGINT, signal_handler);
		signal(SIGSEGV, signal_handler);
		signal(SIGABRT, signal_handler);
		/* wait */
		while (1) {
			int status;

			pid_t done = wait(&status);
			if (done == -1) {
				if (errno == ECHILD)
					break; /* No more child processes */
			} else {
				if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
					printf("Pid:%d failed\n", done);
					goto out;
				}
			}
		}
	}

out:
	if (ret == 0) {
		gettimeofday(&end, NULL);
		printf("time to collapse file thp: %ld ms\n",
				1000 * (end.tv_sec - start.tv_sec) +
				(end.tv_usec - start.tv_usec) / 1000);
	}

	return ret;
}