/* $Id: uio.c,v 1.4 2012/05/20 20:44:58 kostik Exp kostik $ */

#include <sys/types.h>
#include <sys/mman.h>
#include <sys/uio.h>
#include <err.h>
#include <fcntl.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

static int max_iovcnt = 256;
static int iterations = 1024;
static off_t file_size = 4 * 1024 * 1024;
static const char *file_name = "/tmp/1";

static int fd;
static char *src, *filemap;

extern char *__progname;

static void
usage(void)
{

	fprintf(stderr,
"Usage: %s [-f file] [-i max_iovcnt] [-j iterations] [-s file_size]\n",
	    __progname);
	exit(2);
}

static bool
do_write(void)
{
	struct iovec *iov;
	int i, iovcnt;
	off_t base, offset, res, total;

	iovcnt = arc4random_uniform(max_iovcnt) + 1;
	iov = calloc(iovcnt, sizeof(struct iovec));
	if (iov == NULL)
		return (false);
	do {
		offset = arc4random_uniform(file_size);
	} while (offset == 0);
	for (total = 0, base = offset, i = 0; i < iovcnt; i++) {
		iov[i].iov_base = src + base;
		iov[i].iov_len = arc4random_uniform(file_size - base);
		total += iov[i].iov_len;
		base += iov[i].iov_len;
	}

	res = pwritev(fd, iov, iovcnt, offset);
	if (res != total) {
		fprintf(stderr, "pwritev res %jd total %jd\n", res, total);
		return (false);
	}
	free(iov);
	return (true);
}

static bool
verify(void)
{
	off_t i;

	for (i = 0; i < file_size; i++) {
		if (((int *)filemap)[i] != ((int *)src)[i]) {
			fprintf(stderr, "i %jd filemap %x src %x\n", i,
			    ((int *)filemap)[i], ((int *)src)[i]);
			return (false);
		}
	}
	return (true);
}

static void
init_src(void)
{
	off_t i;

	src = mmap(NULL, file_size * sizeof(int), PROT_READ | PROT_WRITE,
	    MAP_ANON, -1, 0);
	if (src == MAP_FAILED)
		err(1, "mmap anon");
	for (i = 0; i < file_size; i++)
		((int *)src)[i] = i;
}

static void
init_file(void)
{
	off_t i;

	if (ftruncate(fd, file_size * sizeof(int)) == -1)
		err(1, "ftruncate");
	filemap = mmap(NULL, file_size * sizeof(int), PROT_READ | PROT_WRITE,
	    MAP_SHARED, fd, 0);
	if (filemap == MAP_FAILED)
		err(1, "mmap");
	for (i = 0; i < file_size; i++)
		((int *)filemap)[i] = i;
}

int
main(int argc, char *argv[])
{
	char *endp;
	int ch, i, ret;

	while ((ch = getopt(argc, argv, "f:i:j:s:")) != -1) {
		switch (ch) {
		case 'f':
			file_name = optarg;
			break;

		case 'i':
			max_iovcnt = strtol(optarg, &endp, 10);
			if (*endp != '\0')
				usage();
			break;

		case 'j':
			iterations = strtol(optarg, &endp, 10);
			if (*endp != '\0')
				usage();
			break;

		case 's':
			file_size = strtol(optarg, &endp, 10);
			if (*endp != '\0')
				usage();
			break;

		case '?':
		default:
			usage();
			break;
		}
	}
	argc -= optind;
	argv += optind;
	if (argc != 0)
		usage();

	init_src();
	fd = open(file_name, O_CREAT | O_RDWR, 0666);
	if (fd == -1)
		err(1, "open %s", file_name);
	init_file();
	ret = 0;
	for (i = 0; i < iterations; i++) {
		printf(".");
		fflush(stdout);
		if (!do_write()) {
			ret = 1;
			break;
		}
		if (!verify()) {
			ret = 1;
			break;
		}
	}
	printf("\n");

	return (ret);
}