/* Stress test for Xen Store: multiple people hammering transactions */
#include "xs.h"
#include "utils.h"
#include <stdlib.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>

#define NUM_HANDLES 2
#define DIR_FANOUT 3
#define DIR_DEPTH 3

/* How often to print progress */
static int print;

/* Layout looks like /<num>/<num>/count. */
static void work(unsigned int cycles, unsigned int childnum)
{
	unsigned int i;
	struct xs_handle *handles[NUM_HANDLES];
	char id;

	if (childnum < 10)
		id = '0' + childnum;
	else
		id = 'A' + childnum - 10;

	for (i = 0; i < NUM_HANDLES; i++) {
		handles[i] = xs_daemon_open();
		if (!handles[i])
			barf_perror("Opening handle %i", i);
	}

	srandom(childnum);
	for (i = 0; i < cycles; i++) {
		unsigned int j, len;
		char file[100] = "";
		char *contents, tmp[100];
		struct xs_handle *h = handles[random() % NUM_HANDLES];

		for (j = 0; j < DIR_DEPTH; j++)
			sprintf(file + strlen(file), "/%li",
				random()%DIR_FANOUT);

		if (!xs_transaction_start(h))
			barf_perror("%i: starting transaction %i",
				    childnum, i);

		sprintf(file + strlen(file), "/count");
		contents = xs_read(h, file, &len);
		if (!contents)
			barf_perror("%i: can't read %s iter %i",
				    childnum, file, i);
		sprintf(tmp, "%i", atoi(contents) + 1);
		if (!xs_write(h, file, tmp, strlen(tmp)+1))
			barf_perror("%i: can't write %s iter %i",
				    childnum, file, i);

		/* Abandon 1 in 10 */
		if (random() % 10 == 0) {
			if (!xs_transaction_end(h, true))
				barf_perror("%i: can't abort transact",
					    childnum);
			i--;
		} else {
			if (!xs_transaction_end(h, false)) {
				if (errno == EAGAIN) {
					write(STDOUT_FILENO, "!", 1);
					i--;
				} else
					barf_perror("%i: can't commit trans",
						    childnum);
			} else {
				/* Offset when we print . so kids don't all
				 * print at once. */
				if ((i + print/(childnum+1)) % print == 0)
					write(STDOUT_FILENO, &id, 1);
			}
		}
	}
}

static void create_dirs(struct xs_handle *h, const char *base, int togo)
{
	unsigned int i;
	char filename[100];

	if (togo == 0) {
		sprintf(filename, "%s/count", base);
		if (!xs_write(h, filename, "0", 1))
			barf_perror("Writing to %s", filename);
		return;
	}

	for (i = 0; i < DIR_FANOUT; i++) {
		sprintf(filename, "%s/%i", base, i);
		if (!xs_mkdir(h, filename))
			barf_perror("xs_mkdir %s", filename);
		create_dirs(h, filename, togo-1);
	}
}

static unsigned int add_count(struct xs_handle *h, const char *base, int togo)
{
	unsigned int i, count;
	char filename[100];

	if (togo == 0) {
		char *answer;
		unsigned int len;

		sprintf(filename, "%s/count", base);
		answer = xs_read(h, filename, &len);
		if (!answer)
			barf_perror("Reading %s", filename);
		count = atoi(answer);
		free(answer);
		return count;
	}

	count = 0;
	for (i = 0; i < DIR_FANOUT; i++) {
		sprintf(filename, "%s/%i", base, i);
		count += add_count(h, filename, togo-1);
	}
	return count;
}

static void setup(void)
{
	struct xs_handle *h;

	/* Do setup. */
	h = xs_daemon_open();
	if (!h)
		barf_perror("Contacting daemon");
	create_dirs(h, "", DIR_DEPTH);
	xs_daemon_close(h);
}

static unsigned int tally_counts(void)
{
	struct xs_handle *h;
	unsigned int ret;
	
	h = xs_daemon_open();
	if (!h)
		barf_perror("Contacting daemon");

	ret = add_count(h, "", DIR_DEPTH);
	xs_daemon_close(h);
	return ret;
}	

int main(int argc, char *argv[])
{
	unsigned int i;
	bool failed = false;
	int kids[10];

	if (argc != 2)
		barf("Usage: xs_stress <iterations>");

	printf("Setting up directories...\n");
	setup();

	print = atoi(argv[1]) / 76;
	if (!print)
		print = 1;

	printf("Running %i children...\n", ARRAY_SIZE(kids));
	for (i = 0; i < ARRAY_SIZE(kids); i++) {
		kids[i] = fork();
		if (kids[i] == -1)
			barf_perror("fork");
		if (kids[i] == 0) {
			work(atoi(argv[1]) / ARRAY_SIZE(kids), i);
			exit(0);
		}
	}

	for (i = 0; i < ARRAY_SIZE(kids); i++) {
		int status;
		if (waitpid(kids[i], &status, 0) == -1)
			barf_perror("waitpid");
		if (!WIFEXITED(status))
			barf("Kid %i died via signal %i\n",
			     i, WTERMSIG(status));
		if (WEXITSTATUS(status) != 0) {
			printf("Child %i exited %i\n", i, WEXITSTATUS(status));
			failed = true;
		}
	}
	if (failed)
		exit(1);

	printf("\nCounting results...\n");
	i = tally_counts();
	if (i != (unsigned)atoi(argv[1]))
		barf("Total counts %i not %s", i, argv[1]);
	printf("Success!\n");
	exit(0);
}