/*
 * wb-latency.c
 *
 * This file may be redistributed under the terms of the GNU Public
 * License, version 2.
 */
#define _FILE_OFFSET_BITS 64
#define _XOPEN_SOURCE 600
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/wait.h>
#include <signal.h>
#include <time.h>
#include <fcntl.h>
#include <string.h>

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

static int page_size = 4096;

static float timeval_subtract(struct timeval *tv1, struct timeval *tv2)
{
	return ((tv1->tv_sec - tv2->tv_sec) +
		((float) (tv1->tv_usec - tv2->tv_usec)) / 1000000);
}

/*
 * the magic offset is where we write our timestamps.
 * The idea is that we write constantly to the magic offset
 * and then pull the power.
 * After the OS comes back, we read the timestamp stored and compare
 * it with the time stamp printed.  Any difference over 1s is time the
 * IO spent stalled in cache.
 */
static loff_t magic_offset(loff_t total)
{
	loff_t cur = total - ((loff_t)64) * 1024;
	cur = cur / page_size;
	cur = cur * page_size;
	return cur;
}

/*
 * this function runs in a loop overwriting two nearby
 * sectors.  The idea is to create something the
 * drive is likely to store in cache and not send down very often.
 *
 * It writes a timestamp to the sector and to stderr.  After
 * crashing, compare the output of wb-latency -c with the last
 * thing printed on stderr.
 */
static void timestamp_io(int fd, char *buf, loff_t total)
{
	loff_t cur = magic_offset(total);
	struct timeval tv;
	struct timeval print_tv;
	int ret;

	cur = cur / page_size;
	cur = cur * page_size;

	printf("starting tester run\n");
	gettimeofday(&print_tv, NULL);
	while(1) {
		gettimeofday(&tv, NULL);
		memcpy(buf, &tv, sizeof(tv));

		if (timeval_subtract(&tv, &print_tv) >= 1) {
			fprintf(stderr, "current time %lu.%lu\n",
				tv.tv_sec, tv.tv_usec);
			gettimeofday(&print_tv, NULL);
		}

		ret = pwrite(fd, buf, page_size, cur);
		if (ret < page_size) {
			fprintf(stderr, "short write ret %d cur %llu\n",
				ret, (unsigned long long)cur);
			exit(1);
		}

		ret = pwrite(fd, buf, page_size, cur + page_size * 2);
		if (ret < page_size) {
			fprintf(stderr, "short write ret %d cur %llu\n",
				ret, (unsigned long long)cur);
			exit(1);
		}

	}
}

/*
 * just print out the timestamp in our magic sector
 */
static void check_timestamp_io(int fd, char *buf, loff_t total)
{
	int ret;
	struct timeval tv;
	loff_t cur = magic_offset(total);

	ret = pread(fd, buf, page_size, cur);
	if (ret < page_size) {
		perror("read");
		exit(1);
	}
	memcpy(&tv, buf, sizeof(tv));
	printf("Found tv %lu.%lu\n", tv.tv_sec, tv.tv_usec);
}

int main(int argc, char **argv)
{
	int	fd;
	struct stat st;
	pid_t pid;
	int ret;
	int i;
	int status;
	loff_t total_size = 128 * 1024 * 1024;
	loff_t hot_size = 26 * 1024 * 1024;
	loff_t cur;
	char *buf;
	char *filename = NULL;
	int check_only = 0;

	ret = posix_memalign((void *)(&buf), page_size, page_size);
	if (ret) {
		perror("memalign\n");
		exit(1);
	}

	memset(buf, 0, page_size);

	if (argc < 2) {
		fprintf(stderr, "usage: wb-latency [-c] file\n");
		exit(1);
	}
	for (i = 1; i < argc; i++) {
		if (strcmp(argv[i], "-c") == 0)
			check_only = 1;
		else
			filename = argv[i];
	}

	fd = open(filename, O_RDWR | O_DIRECT | O_CREAT);
	if (fd < 0) {
		perror("open");
		exit(1);
	}

	ret = fstat(fd, &st);
	if (ret < 0) {
		perror("fstat");
		exit(1);
	}

	check_timestamp_io(fd, buf, total_size);

	if (check_only)
		exit(0);

	/* setup the file if we aren't doing a block device */
	if (!S_ISBLK(st.st_mode) && st.st_size < total_size) {
		printf("setting up file %s\n", filename);
		while(cur < total_size) {
			ret = write(fd, buf, page_size);
			if (ret <= 0) {
				fprintf(stderr, "short write\n");
				exit(1);
			}
			cur += ret;
		}
		printf("done setting up %s\n", filename);
	}

	pid = fork();
	if (pid == 0) {
		timestamp_io(fd, buf, total_size);
		exit(0);
	}
	waitpid(pid, &status, WNOHANG);

	/*
	 * here we run the hot IO.  This is something the drive isn't
	 * going to bypass the cache on, but something the drive will
	 * tend to allow to dominate the cache.
	 */
	printf("starting hot writes run\n");
	cur = 0;
	while(1) {
		pwrite(fd, buf, page_size, cur);
		cur += page_size;
		if (cur > hot_size)
			cur = 0;
	}

	return 0;
}