/*
 * Test performance of allocating a zeroed file via various methods on
 * various file systems
 *
 *   Copyright (C) 2009 Red Hat Inc.
 *   Author(s): Amit Shah
 *
 *    This program is free software; you can redistribute it and/or
 *    modify it under the terms of the GNU General Public License
 *    version 2 as published by the Free Software Foundation
 *
 *  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., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

#include <stdio.h>
#include <fcntl.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <asm/unistd.h>
#include <sys/time.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/mount.h>
#include <sys/mman.h>

#define GB(x) ((unsigned long long)x << 30)

int pre_test_setup(char *source, char *target, char *fstype,
		   unsigned long mntflags, char *mntopts, char *name, int *fd)
{
	int r;

	if (source) {
		r = mount(source, target, fstype, mntflags, mntopts);
		if (r < 0) {
			perror("mount");
			return -1;
		}
	}
	unlink(name);
	*fd = open(name, O_RDWR|O_CREAT, S_IRUSR|S_IWUSR);
	if (*fd < 0) {
		perror("open");
		return -1;
	}

	/* Flush out all other data that might be pending writing to
	 * the disk. We don't want to count the time we wait for the
	 * buffers of other apps to clear. (Ideally, this test would
	 * be run in single user mode with nothing else running on the
	 * system.)
	 */
	sync();
	sync();
	/* Just to be sure - http://lwn.net/Articles/325420/ */
	sync();

	return 0;
}

int post_test_cleanup(char *target, int *fd)
{
	int r;

	r = close(*fd);
	if (r < 0)
		perror("close");

	if (strcmp(target, ".")) {
		r = umount(target);
		if (r < 0)
			perror("unmount");
	}
	return r;
}

void run_test(int (test)(char *, int *, off_t, size_t),
	      char *source, char *target, char *fstype, unsigned long mntflags,
	      char *mntopts, char *filename, off_t len, size_t data)
{
	int r, fd;

	r = pre_test_setup(source, target, fstype, mntflags, mntopts,
			   filename, &fd);
	if (r < 0)
		goto error_exit;

	(test)(target, &fd, len, data);
	return;

error_exit:
	post_test_cleanup(target, &fd);
	return;
}

int do_posix_fallocate(char *target, int *fd, off_t len, size_t data)
{
	int r;
	struct timeval tv1, tv2;

	printf("posix-fallocate run time:\n");
	gettimeofday(&tv1, NULL);
	r = posix_fallocate(*fd, 0, len);
	post_test_cleanup(target, fd);
	gettimeofday(&tv2, NULL);
	if (r < 0) {
		printf("posix_fallocate, error %d\n", r);
		return r;
	}
	printf("\tseconds:microseconds: %llu:%llu\n", tv1.tv_sec, tv1.tv_usec);
	printf("\tseconds:microseconds: %llu:%llu\n", tv2.tv_sec, tv2.tv_usec);
	printf("\t(approx %us)\n", tv2.tv_sec - tv1.tv_sec);

	return 0;
}

int do_fallocate(char *target, int *fd, off_t len, size_t data)
{
	int r;
	struct timeval tv1, tv2;

	printf("fallocate run time:\n");
	gettimeofday(&tv1, NULL);
	r = syscall(__NR_fallocate, *fd, 0, len);
	post_test_cleanup(target, fd);
	gettimeofday(&tv2, NULL);
	if (r < 0) {
		perror("fallocate");
		return r;
	}
	printf("\tseconds:microseconds: %llu:%llu\n", tv1.tv_sec, tv1.tv_usec);
	printf("\tseconds:microseconds: %llu:%llu\n", tv2.tv_sec, tv2.tv_usec);
	printf("\t(approx %us)\n", tv2.tv_sec - tv1.tv_sec);

	return 0;
}

int do_mmap(char *target, int *fd, off_t len, size_t data)
{
	struct timeval tv1, tv2;
	char *addr;

	/* memset has to have the mmap'ed file backed by something on
	 * disk -- bus error otherwise
	 */
	ftruncate(*fd, len);

	addr = mmap(NULL, len, PROT_READ|PROT_WRITE, MAP_SHARED, *fd, 0);
	if (addr == MAP_FAILED) {
		perror("mmap");
		return -2;
	}

	printf("mmap run time:\n");
	gettimeofday(&tv1, NULL);
	memset(addr, 0, len);
	munmap(addr, len);
	post_test_cleanup(target, fd);
	gettimeofday(&tv2, NULL);

	printf("\tseconds:microseconds: %llu:%llu\n", tv1.tv_sec, tv1.tv_usec);
	printf("\tseconds:microseconds: %llu:%llu\n", tv2.tv_sec, tv2.tv_usec);
	printf("\t(approx %us)\n", tv2.tv_sec - tv1.tv_sec);

	return 0;
}

int do_write_chunks(char *target, int *fd, off_t len, size_t data)
{
	struct timeval tv1, tv2;

	unsigned long long remain = len;
	char *zeros;

	zeros = calloc(1, data);

	printf("%llu-sized chunk run time:\n", data);
	gettimeofday(&tv1, NULL);
	while (remain) {
		int bytes = data;
		if (bytes > remain)
			bytes = remain;
		if ((bytes = write(*fd, zeros, bytes)) < 0) {
			perror("write");
			return -2;
		}
		remain -= bytes;
	}
	post_test_cleanup(target, fd);
	gettimeofday(&tv2, NULL);

	printf("\tseconds:microseconds: %llu:%llu\n", tv1.tv_sec, tv1.tv_usec);
	printf("\tseconds:microseconds: %llu:%llu\n", tv2.tv_sec, tv2.tv_usec);
	printf("\t(approx %us)\n", tv2.tv_sec - tv1.tv_sec);
	free(zeros);

	return 0;
}

int main(int argc, char **argv)
{
	char *source, *target, *fstype, *mntopts;
	char filename[200];
	unsigned long mntflags;
	unsigned long long filesize;

	/* FIXME! use getopt */
	if (argc < 2 || argc > 2 && argc < 5) {
		printf("usage: %s <filesize-in-gb> <device> <mountpoint> <fstype> [mntflags] [mntopts]\n",
		       argv[0]);
		return -1;
	}

	filesize = GB(atol(argv[1]));
	if (argc > 2) {
		source = argv[2];
		target = argv[3];
		fstype = argv[4];
	} else {
		source = NULL;
		target = ".";
		fstype = "native";
	}
	if (argc > 5)
		mntflags = atol(argv[5]);
	else
		mntflags = 0;
	if (argc > 6)
		mntopts = argv[6];
	else
		mntopts = NULL;

	sprintf(filename, "%s/%s-pf", target, fstype);
	run_test(do_posix_fallocate, source, target, fstype, mntflags, mntopts,
		 filename, filesize, 0);

#if 0
	sprintf(filename, "%s/%s-fallocate", target, fstype);
	run_test(do_fallocate, source, target, fstype, mntflags, mntopts,
		 filename, filesize, 0);
#endif

	sprintf(filename, "%s/%s-mmap", target, fstype);
	run_test(do_mmap, source, target, fstype, mntflags, mntopts,
		 filename, filesize, 0);

	sprintf(filename, "%s/%s-chunk4k", target, fstype);
	run_test(do_write_chunks, source, target, fstype, mntflags, mntopts,
		 filename, filesize, 4 * 1024);

	sprintf(filename, "%s/%s-chunk8k", target, fstype);
	run_test(do_write_chunks, source, target, fstype, mntflags, mntopts,
		 filename, filesize, 8 * 1024);

	return 0;
}