/******************************************************************************/
/*                                                                            */
/* Copyright (c) 2010 FUJITSU LIMITED                                         */
/*                                                                            */
/* This program is free software;  you can redistribute it and/or modify      */
/* it under the terms of the GNU General Public License as published by       */
/* the Free Software Foundation; either version 2 of the License, or          */
/* (at your option) any later version.                                        */
/*                                                                            */
/* This program is distributed in the hope that it will be useful,            */
/* but WITHOUT ANY WARRANTY;  without even the implied warranty of            */
/* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See                  */
/* the GNU General Public License for more details.                           */
/*                                                                            */
/* You should have received a copy of the GNU General Public License          */
/* along with this program;  if not, write to the Free Software               */
/* Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA    */
/*                                                                            */
/* Author: Liu Bo <liubo2009@cn.fujitsu.com>                                  */
/*                                                                            */
/******************************************************************************/


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

#define BUFFERSIZE 8192

#ifndef O_DIRECT
#define O_DIRECT	00040000	/* direct disk access hint */
#endif

static int rw_test(char *type, int fd, int nullfd, int zerofd, char *size);


/* Direct-io Test 
 *
 * Return Value:
 * 	0: succeed
 * 	1: fail
 *
 * Parse command line arguments 
 * flag: 
 * 	O_DIRECT, O_DIRECT | O_SYNC
 * rwtype:
 * 	read, readv, write, writev
 * test file name:
 * 	testrw
 * test file size:
 * 	0 | 1K | 4K | 4M | 4G | 4G+1	
 */
int main(int argc, char *argv[])
{
	char *rwtype;
	char *testfile;
	char *size;
	int fd, nullfd, zerofd;
	int flag = O_RDWR;
	int ret;

	if (--argc != 4) {
		fprintf(stderr, "4 arguments are needed!\n");
		exit(1);
	}

	if (!strcmp(argv[1], "O_DIRECT"))
		flag |= O_DIRECT;
	else if (!strcmp(argv[1], "O_DIRECT | O_SYNC") ||
				!strcmp(argv[1], "O_DIRECT|O_SYNC"))
		flag |= O_DIRECT | O_SYNC;

	rwtype = argv[2];
	testfile = argv[3];
	size = argv[4];

	if ((fd = open(testfile, flag)) < 0) {
		fprintf(stderr, "Cannot open the test file! \
					: %s\n\n", strerror(errno));
		exit(1);
	}
	if ((nullfd = open("/dev/null", O_WRONLY)) < 0) {
		fprintf(stderr, "Cannot open the test file! \
					: %s\n\n", strerror(errno));
		exit(1);
	}
	if ((zerofd = open("/dev/zero", O_RDONLY)) < 0) {
		fprintf(stderr, "Cannot open the test file! \
					: %s\n\n", strerror(errno));
		exit(1);
	}

	ret = rw_test(rwtype, fd, nullfd, zerofd, size);

	close(fd);
	close(nullfd);
	close(zerofd);

	if (ret != 0)
		return 1;
	return 0;
}

/* 
 echo "=========================" >> $RESULTFILE
 * rw_test is used to test read/write/readv/writev.
 *
 * Arguments:
 * 	  type: read/write/readv/writev
 * 	    fd: read/write file descripter
 * 	nullfd: write file descripter binded to /dev/null
 *	zerofd: read file descripter binded to /dev/zero
 *	  size: 0 | 1K | 4K | 4M | 4G | 4G+1
 *
 * Return value:
 * 	 0: succeed
 * 	 1: write error
 * 	-1: read error
 */
static int rw_test(char *type, int fd, int nullfd, int zerofd, char *size)
{
	ssize_t n;
	int iovcnt;
	long limit = 0;

	if(!strcmp(size, "1K")) {
		limit = 1024;
	} else if (!strcmp(size, "4K")) {
		limit = 1024 * 4;
	} else if (!strcmp(size, "4M")) {
		limit = 1024 * 1024 * 4;
	} else if (!strcmp(size, "4G")) {
		limit = 1024 * 1024 * 4;
		limit *= 1024;
	} else if (!strcmp(size, "4G+1")) {
		limit = 1024 * 1024 * 4;
		limit *= 1024;
		limit += 1;
	}

	if (!strcmp(type, "read")) {
		char *buf;

		buf = valloc(BUFFERSIZE);
		if (!buf) {
			perror("valloc error");
			exit(1);
		}
		memset(buf, 0, BUFFERSIZE);

		while ((n = read(fd, buf, BUFFERSIZE)) > 0) {
			if (write(nullfd, buf, n) != n) {
				fprintf(stderr, "write error! :\
						 %s\n", strerror(errno));
				return 1;
			}
		}

		if (n < 0) {
			fprintf(stderr, "read error : %s\n\n", \
							strerror(errno));
			return -1;
		}

	} else if (!strcmp(type, "write")) {
		char *buf;
		long count = 0;

		buf = valloc(BUFFERSIZE);
		if (!buf) {
			perror("valloc error");
			exit(1);
		}
		memset(buf, 0, BUFFERSIZE);

		while ((n = read(zerofd, buf, BUFFERSIZE)) > 0) {
			if (write(fd, buf, n) != n) {
				fprintf(stderr, "write error! :\
						 %s\n", strerror(errno));
				return 1;
			}
			count += n;
			if (count >= limit)
				break;
		}

		if (count < limit) {
			fprintf(stderr, "read error : %s\n\n", \
							strerror(errno));
			return -1;
		}

	} else if (!strcmp(type, "readv")) {
		iovcnt = 2;
		char *buf[iovcnt];
		struct iovec iov[iovcnt];
		int i;

		for (i = 0; i < iovcnt; i++) {
			buf[i] = valloc(BUFFERSIZE);
			if (!buf[i]) {
				perror("valloc error");
				exit(1);
			}

			memset(buf[i], 0, BUFFERSIZE);

			iov[i].iov_base = buf[i];
			iov[i].iov_len = BUFFERSIZE;
		}

		while ((n = readv(fd, iov, iovcnt)) > 0) {
			if (writev(nullfd, iov, iovcnt) != 
					(iovcnt * BUFFERSIZE)) {
				fprintf(stderr, "writev error! :\
						 %s\n", strerror(errno));
				return 1;
			}
		}

		if (n < 0) {
			fprintf(stderr, "readv error : %s\n\n", \
							strerror(errno));
			return -1;
		}

	} else if (!strcmp(type, "writev")) {
		iovcnt = 2;
		char *buf[iovcnt];
		struct iovec iov[iovcnt];
		long count = 0;
		int i;

		for (i = 0; i < iovcnt; i++) {
			buf[i] = valloc(BUFFERSIZE);
			if (!buf[i]) {
				perror("valloc error");
				exit(1);
			}

			memset(buf[i], 0, BUFFERSIZE);

			iov[i].iov_base = buf[i];
			iov[i].iov_len = BUFFERSIZE;
		}

		while ((n = readv(zerofd, iov, iovcnt)) > 0) {
			if (writev(fd, iov, iovcnt) != 
					(iovcnt * BUFFERSIZE)) {
				fprintf(stderr, "writev error! :\
						 %s\n", strerror(errno));
				return 1;
			}
			count += n;
			if (count >= limit)
				break;
		}

		if (count < limit) {
			fprintf(stderr, "readv error : %s\n\n", \
							strerror(errno));
			return -1;
		}

	}

	return 0;
}