#define _GNU_SOURCE
#include <stdio.h>
#include <fcntl.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>
#include <getopt.h>
#include <sys/wait.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/time.h>

#define COUNT 5000
#define SIZE 16384
#define MAX_PROCS 1024

char shm_name[64];
char *wbuf;
int call_sync, use_falloc;
int count = COUNT;
int size = SIZE;

void prepare(char *base, int num)
{
	char dir[128];

	sprintf(dir, "%s/dir-%d", base, num);
	if (mkdir(dir, 0700) < 0 && errno != EEXIST) {
		perror("mkdir");
		exit(1);
	}
}

void teardown(char *base, int num)
{
	char dir[128];

	sprintf(dir, "%s/dir-%d", base, num);
	rmdir(dir);
}

void run_test(char *base, int num)
{
	char name[128];
	int shmfd;
	int i;
	int *shm;
	int fd;

	sprintf(name, "%s/dir-%d/file", base, num);
	shmfd = shm_open(shm_name, O_RDWR, 0);
	if (shmfd < 0) {
		perror("shm_open");
		exit(1);
	}
	shm = mmap(NULL, sizeof(int) * MAX_PROCS, PROT_READ | PROT_WRITE,
		   MAP_SHARED, shmfd, 0);
	if (shm == MAP_FAILED) {
		perror("mmap");
		exit(1);
	}
	shm[num + 1] = 1;
	while (shm[0] == 0)
		usleep(1);

	for (i = 0; i < count; i++) {
		fd = open(name, O_CREAT | O_WRONLY, 0644);
		if (fd < 0) {
			perror("open");
			exit(1);
		}
		if (use_falloc) {
			if (fallocate(fd, 0, 0, size) < 0) {
				perror("fallocate");
				exit(1);
			}
		} else {
			if (write(fd, wbuf, size) != size) {
				perror("write");
				exit(1);
			}
		}
		if (call_sync)
			fsync(fd);
		close(fd);
		unlink(name);
	}
}

pid_t start_flusher(char *base, int num)
{
	int fd;
	char name[128];
	pid_t pid;

	pid = fork();
	if (pid < 0) {
		perror("fork");
		return -1;
	}
	if (pid > 0)
		return pid;
	sprintf(name, "%s/sync-file-%d", base, num);
	fd = open(name, O_CREAT | O_TRUNC | O_WRONLY, 0644);
	if (fd < 0) {
		perror("open");
		exit(1);
	}
	while (1) {
		if (pwrite(fd, wbuf, size, 0) != size) {
			perror("pwrite");
			exit(1);
		}
		fsync(fd);
	}
}

int main(int argc, char **argv)
{
	int procs, sync_procs, i, j;
	pid_t pids[MAX_PROCS];
	pid_t sync_pids[MAX_PROCS];
	int shmfd;
	int *shm;
	struct timeval start, end;
	long long ms;
	int opt;
	char *path;

	while ((opt = getopt(argc, argv, "sc:f:a")) != -1) {
		switch (opt) {
		case 's':
			call_sync = 1;
			break;
		case 'a':
			use_falloc = 1;
			break;
		case 'c':
			count = atoi(optarg);
			break;
		case 'f':
			size = atoi(optarg);
			break;
		default:
usage:
			fprintf(stderr, "Usage: stress-unlink [-s] [-c <count>] [-f <filesize>] <processes> <flushers> <dir>\n");
			return 1;
		}
	}
	if (argc - optind != 3)
		goto usage;

	procs = strtol(argv[optind++], NULL, 10);
	if (procs > MAX_PROCS) {
		fprintf(stderr, "Too many processes!\n");
		return 1;
	}
	sync_procs = strtol(argv[optind++], NULL, 10);
	if (sync_procs > MAX_PROCS) {
		fprintf(stderr, "Too many flusher processes!\n");
		return 1;
	}
	path = argv[optind];

	wbuf = malloc(size);
	memset(wbuf, 0xcc, size);
	sprintf(shm_name, "/stress-unlink-%u", getpid());
	shmfd = shm_open(shm_name, O_CREAT | O_RDWR, 0600);
	if (shmfd < 0) {
		perror("shm_open");
		return 1;
	}
	if (ftruncate(shmfd, sizeof(int) * MAX_PROCS) < 0) {
		perror("ftruncate shm");
		return 1;
	}
	shm = mmap(NULL, sizeof(int) * MAX_PROCS, PROT_READ | PROT_WRITE,
		   MAP_SHARED, shmfd, 0);
	if (shm == MAP_FAILED) {
		perror("mmap");
		return 1;
	}
	
	shm[0] = 0;
	for (i = 0; i < procs; i++) {
		shm[i+1] = 0;
		prepare(path, i);
	}

	for (i = 0; i < procs; i++) {
		pids[i] = fork();
		if (pids[i] < 0) {
			perror("fork");
			for (j = 0; j < i; j++)
				kill(pids[j], SIGKILL);
			return 1;
		}
		if (pids[i] == 0) {
			run_test(path, i);
			exit(0);
		}
	}

	for (i = 0; i < sync_procs; i++) {
		sync_pids[i] = start_flusher(path, i);
		if (sync_pids[i] < 0) {
			for (j = 0; j < procs; j++)
				kill(pids[j], SIGKILL);
			for (j = 0; j < i; i++)
				kill(sync_pids[i], SIGKILL);
			exit(1);
		}
	}

	do {
		for (i = 0; i < procs && shm[i + 1]; i++);
	} while (i != procs);
	gettimeofday(&start, NULL);
	shm[0] = 1;
	fprintf(stderr, "Processes started.\n");

	for (i = 0; i < procs; i++)
		waitpid(pids[i], NULL, 0);
	gettimeofday(&end, NULL);
	for (i = 0; i < procs; i++)
		teardown(path, i);
	shm_unlink(shm_name);
	for (i = 0; i < sync_procs; i++)
		kill(sync_pids[i], SIGKILL);
	for (i = 0; i < sync_procs; i++)
		waitpid(sync_pids[i], NULL, 0);
	ms = (((long long)(end.tv_sec - start.tv_sec) * 1000000) +
		(end.tv_usec - start.tv_usec)) / 1000;
	printf("%lld.%03lld\n", ms/1000, ms%1000);
	
	return 0;
}