/*
 * Do sequential/random read/write benchmarks.
 *
 * Copyright (C) 2007 Avishay Traeger, Kumar Thangavelu, Erez Zadok
 * Copyright (C) 2007 Stony Brook University (SUNY)
 *
 */

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>

#define READ	0
#define WRITE	1

#define SEQ	0
#define RAND	1

/* Options */
char *filename = NULL;
unsigned int chunksize;
unsigned int numchunks;
unsigned int rwmode = READ;
unsigned int rndmode = SEQ;

/* Misc globals */
off_t *pattern;

void usage(char *prog)
{
	fprintf(stderr, "usage: %s -f <filename> -c <chunk size B> -n <num chunks> [-w] [-r]\n", prog);
	fprintf(stderr, "       -w: write (default is read)\n");
	fprintf(stderr, "       -r: random (default is sequential)\n");
	exit(1);
}

void read_args(int argc, char **argv)
{
	int c;
	int set_c = 0, set_n = 0;

	while((c = getopt(argc, argv, "f:c:n:wrh")) != -1) {
		switch(c) {
			case 'f':
				if (!optarg) usage(argv[0]);
				filename = strdup(optarg);
				break;
			case 'c':
				if (!optarg) usage(argv[0]);
				chunksize = atoi(optarg);
				set_c = 1;
				break;
			case 'n':
				if (!optarg) usage(argv[0]);
				numchunks = atoi(optarg);
				set_n = 1;
				break;
			case 'w':
				rwmode = WRITE;
				break;
			case 'r':
				rndmode = RAND;
				break;
			case 'h':
				usage(argv[0]);
				break;
			case '?':
				fprintf(stderr, "ERROR: Unknown option %c\n", c);
				usage(argv[0]);
				break;
		}
	}

	if (!set_c || !set_n || !filename) {
		fprintf(stderr, "ERROR: Must set -f, -n, and -c options\n");
		usage(argv[0]);
	}
}

void make_pattern() {
	int i;
	int j;
	off_t tmp;

	pattern = (off_t *)malloc(sizeof(off_t) * numchunks);
	if (!pattern) {
		fprintf(stderr, "ERROR: Failed to allocate memory\n");
		exit(1);
	}

	for (i = 0 ; i < numchunks ; i++)
		pattern[i] = i;

	if (rndmode == RAND) {
		srandom(1);
		for (i = 0 ; i < numchunks - 1; i++) {
			j = i + (int)(((double)(numchunks -i) * rand()) / (RAND_MAX + 1.0));
			tmp = pattern[i];
			pattern[i] = pattern[j];
			pattern[j] = tmp;
		}
	}
}

void rw()
{
	int i;
	off_t off;
	char *buf;
	int fd = -1;
	ssize_t ret;
	int flags;
	mode_t mode = S_IRWXU | S_IRWXG | S_IRWXO;
	ssize_t (*rwop) ();

	if (rwmode == WRITE) {
		flags = O_WRONLY | O_CREAT;
		rwop = write;
	} else {
		flags = O_RDONLY;
		rwop = read;
	}

	buf = (char *)malloc(chunksize);
	if (!buf) {
		fprintf(stderr, "ERROR: Failed to allocate memory.\n");
		exit(1);
	}

	fd = open(filename, flags, mode);
	if (fd == -1) {
		perror("open");
		exit(1);
	}
	for (i = 0 ; i < numchunks ; i++) {
		off = lseek(fd, (pattern[i] * chunksize), SEEK_SET);
		if (off == -1) {
			perror("lseek");
			exit(1);
		}
		ret = rwop(fd, buf, chunksize);
		if (ret != chunksize) {
			perror("rw");
			exit(1);
		}
	}

	if (rwmode == WRITE) {
		fsync(fd);
	}

	close(fd);
	free(buf);
	free(filename);
	free(pattern);
}

int main (int argc, char **argv)
{
	read_args(argc, argv);
	make_pattern();
	rw();
	return 0;
}