/*
 * Copyright (c) 2009 David Teigland
 * All Rights Reserved.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License V2
 * as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it would 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 the Free Software Foundation,
 * Inc.,  51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <errno.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/file.h>
#include <fcntl.h>
#include <unistd.h>
#include <time.h>

#define die(fmt, args...) \
{ \
	fprintf(stderr, "%s: ", prog_name); \
	fprintf(stderr, fmt, ##args); \
	exit(EXIT_FAILURE); \
}

#define MAX_FILES 1000
int fda[MAX_FILES];
unsigned int inodea[MAX_FILES];

int do_create;
int use_flock;
int use_plock;
int filecount;
int iterations;
int our_pid;
char *prog_name;
int blocking;
int delay = 10; /* milliseconds */
int quiet;

void read_flock(int fd)
{
	int error;
	int mode = LOCK_SH;

	if (!blocking)
		mode |= LOCK_NB;

	error = flock(fd, mode);
	if (error < 0) {
		printf("%d read lock failed %d errno %d\n",
			our_pid, error, errno);
	}
}

void write_flock(int fd)
{
	int error;
	int mode = LOCK_EX;

	if (!blocking)
		mode |= LOCK_NB;

	error = flock(fd, mode);
	if (error < 0) {
		printf("%d write lock failed %d errno %d\n",
			our_pid, error, errno);
	}
}

void unlock_flock(int fd)
{
	int error;

	error = flock(fd, LOCK_UN);
	if (error < 0) {
		printf("%d unlock failed %d errno %d\n",
			our_pid, error, errno);
	}
}

int do_plock(int fd, int offset, int len, int type)
{
	struct flock lock;
	int action, error;

	lock.l_type = type;
	lock.l_start = offset;
	lock.l_whence = SEEK_SET;
	lock.l_len = len;

	if (blocking)
		action = F_SETLKW;
	else
		action = F_SETLK;

	error = fcntl(fd, action, &lock);

	if (error < 0 && errno != EAGAIN) {
		fprintf(stderr, "fcntl type %d action %d errno %d\n",
			type, action, errno);
		exit(-1);
	}

	return error;
}

int rand_int(int a, int b)
{
	return a + (int) (((float)(b - a + 1)) * random() / (RAND_MAX+1.0)); 
}

void rand_range(int *offset, int *len)
{
	switch (rand_int(0, 2)) {
	case 0:
		*offset = 0;
		*len = 10;
		return;
	case 1:
		*offset = 0;
		*len = 5;
		return;
	case 2:
		*offset = 5;
		*len = 5;
		return;
	}
}

static uint64_t dt_usec(struct timeval *start, struct timeval *stop)
{
	uint64_t dt;

	dt = stop->tv_sec - start->tv_sec;
	dt *= 1000000;
	dt += stop->tv_usec - start->tv_usec;
	return dt;
}

void loop_plock(void)
{
	struct timeval t1, t2;
	int i, fdi, fd, offset, len, rv;

	for (i = 0;; i++) {

		fdi = rand_int(0, filecount-1);
		fd = fda[fdi];

		rand_range(&offset, &len);

		switch (rand_int(0, 2)) {
		case 0:
			gettimeofday(&t1, NULL);
			rv = do_plock(fd, offset, len, F_UNLCK);
			gettimeofday(&t2, NULL);

			if (!quiet)
				printf("%06d file%04d ino %04x U %02d-%02d pid %d err %02d sec %.6f\n",
					i, fdi, inodea[fdi], offset, offset+len-1, our_pid,
					rv ? errno : 0, 1.e-6 * dt_usec(&t1, &t2));
			break;

		case 1:
			gettimeofday(&t1, NULL);
			rv = do_plock(fd, offset, len, F_RDLCK);
			gettimeofday(&t2, NULL);

			if (!quiet)
				printf("%06d file%04d ino %04x R %02d-%02d pid %d err %02d sec %.6f\n",
					i, fdi, inodea[fdi], offset, offset+len-1, our_pid,
					rv ? errno : 0, 1.e-6 * dt_usec(&t1, &t2));

			break;
		case 2:
			gettimeofday(&t1, NULL);
			rv = do_plock(fd, offset, len, F_WRLCK);
			gettimeofday(&t2, NULL);

			if (!quiet)
				printf("%06d file%04d ino %04x W %02d-%02d pid %d err %02d sec %.6f\n",
					i, fdi, inodea[fdi], offset, offset+len-1, our_pid,
					rv ? errno : 0, 1.e-6 * dt_usec(&t1, &t2));
			break;
		}

		if (iterations && i == iterations)
			break;

		if (delay)
			usleep(delay * 1000);
	}
}

void loop_flock(void)
{
	return;
}

int open_all_files(void)
{
	int i, fd;
	struct stat st;
	char path[64];

	for (i = 0; i < filecount; i++) {
		memset(path, 0, sizeof(path));
		sprintf(path, "file%04u", i);

		fd = open(path, O_RDWR | O_CREAT, 0644);
		if (fd < 0)
			die("can't open file %s\n", strerror(errno));
		fda[i] = fd;

		memset(&st, 0, sizeof(st));
		fstat(fd, &st);
		inodea[i] = (unsigned int)st.st_ino;
	}
	return 0;
}

void close_all_files(void)
{
	int i;

	for (i = 0; i < filecount; i++)
		close(fda[i]);
}

void print_usage(void)
{
	fprintf(stderr, "lock_load\n");
	fprintf(stderr, " -i <n>   iterations (default 0, no limit)\n");
	fprintf(stderr, " -n <n>   number of files (default 100)\n");
	fprintf(stderr, " -c       only create files\n");
	fprintf(stderr, " -f       use flock\n");
	fprintf(stderr, " -p       use plock (default)\n");
	fprintf(stderr, " -d <ms>  millisecond delay between each op (default 10)\n");
	fprintf(stderr, " -b       use blocking locks (default non-blocking)\n");
	fprintf(stderr, "          (warning: can easily deadlock)\n");
	fprintf(stderr, " -q       quiet output\n");
	exit(EXIT_FAILURE);
}

int main(int argc, char *argv[])
{
	int cont = 1;
	int optchar;

	do_create = 0;
	use_flock = 0;
	use_plock = 1;
	filecount = 100;
	iterations = 0;
	blocking = 0;
	delay = 10;

	while (cont) {
		optchar = getopt(argc, argv, "cfpn:i:d:bhq");

		switch (optchar) {
		case 'c':
			do_create = 1;
			break;

		case 'f':
			use_flock = 1;
                        break;

		case 'p':
			use_plock = 1;
                        break;

		case 'n':
			filecount = atoi(optarg);
			break;

		case 'i':
			iterations = atoi(optarg);
			break;

		case 'd':
			delay = atoi(optarg);
			break;

		case 'b':
			blocking = 1;
			break;

		case 'q':
			quiet = 1;
			break;

		case EOF:
			cont = 0;
			break;

		case 'h':
		default:
			print_usage();
			break;
		};
	}

	srandom(time(NULL));
	prog_name = argv[0];
	our_pid = getpid();

	if (filecount > MAX_FILES) {
		fprintf(stderr, "max files is %d\n", MAX_FILES);
		return -1;
	}

	if (do_create) {
		open_all_files();
		close_all_files();
		return 0;
	}

	open_all_files();

	if (use_flock)
		loop_flock();
	else if (use_plock)
		loop_plock();
	else
		fprintf(stderr, "use -f or -p for flock or plock\n");

	close_all_files();

	return 0;
}