/*
 * Simple AIO cancellation test. Submits random reads to the specified file and
 * attempts to cancel as many as possible; upon completion it reports the number
 * of IOs that were reported cancelled.
 *
 * Note that currently (6/26/2013), the block layer can't cancel IOs after
 * they've been submitted to the device - it only checks if the IO has been
 * cancelled when pulling it off the request queue, which means in practice for
 * IO to be cancelled your queue depth has to be higher than the device's
 * internal queue depth (for SATA disks with TCQ, I think that's 32?).
 *
 * Compile it with
 *  gcc -o aio-cancel aio-cancel.c -laio
 */

#define _GNU_SOURCE

#include <linux/fs.h>
#include <unistd.h>
#include <stdio.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/param.h>
#include <fcntl.h>
#include <errno.h>
#include <stdlib.h>

#include <libaio.h>

static unsigned long canceled = 0, completed = 0;

static off_t getblocks(int fd)
{
	off_t ret;
	struct stat stat;

	if (fstat(fd, &stat)) {
		perror("stat error\n");
		exit(EXIT_FAILURE);
	}

	ret = stat.st_size / 512;
	if (S_ISBLK(stat.st_mode))
		if (ioctl(fd, BLKGETSIZE, &ret)) {
			perror("ioctl error");
			exit(EXIT_FAILURE);
		}

	return ret;
}

static void die(const char *msg, int err)
{
	fprintf(stderr, "%s error: %s\n", msg, strerror(err));
	exit(EXIT_FAILURE);
}

static void usage(void)
{
	printf("Usage: aio-cancel <test file> <iodepth> <iterations> <iosize>\n"
	       "Simple AIO cancellation test\n");
	exit(EXIT_FAILURE);
}

static unsigned reap_events(io_context_t io_ctx, unsigned long iosize)
{
	/* Wait on/reap some completions */
	struct io_event events[16], *ev;
	int nr_events = io_getevents(io_ctx, 1, 16, events, NULL);

	if (nr_events < 0)
		die("io_getevents", -nr_events);

	for (ev = events; ev < events + nr_events; ev++) {
		if ((ev->res != iosize &&
		     ev->res != -ECANCELED) ||
		    ev->res2 != 0)
			printf("something wrong with event: res = %ld res2 = %ld\n",
			       ev->res, ev->res2);

		if (ev->res == -ECANCELED)
			canceled++;
		else
			completed++;

		free(ev->obj);
	}

	return nr_events;
}

int main(int argc, char const *argv[])
{
	io_context_t io_ctx;
	struct iocb *iocb;
	struct io_event cancelled;
	void *buffer;
	unsigned long nr_ios, iodepth, iosize, in_flight = 0;
	off_t blocks;
	int ret, fd;

	if (argc != 5)
		usage();

	fd = open(argv[1], O_RDONLY|O_DIRECT);
	if (fd < 0) {
		perror("open");
		exit(1);
	}

	nr_ios = strtoul(argv[2], NULL, 10);
	if (!nr_ios) {
		printf("Can't parse number of IOs\n");
		usage();
	}

	iodepth = strtoul(argv[3], NULL, 10);
	if (!iodepth) {
		printf("Can't parse iodepth\n");
		usage();
	}

	iosize = strtoul(argv[4], NULL, 10);
	if (!iosize) {
		printf("Can't parse iosize\n");
		usage();
	}

	ret = posix_memalign(&buffer, 4096, iosize);
	if (ret)
		die("posix_memalign", ret);

	blocks = getblocks(fd);
	if (blocks < iosize / 512) {
		printf("File too small\n");
		exit(EXIT_FAILURE);
	}

	blocks -= iosize / 512; /* avoid reading past the end of the file */

	printf("Doing %lu random reads to %s with iodepth %lu, iosize %lu\n",
	       nr_ios, argv[1], iodepth, iosize);

	memset(&io_ctx, 0, sizeof(io_ctx));
	ret = io_queue_init(iodepth, &io_ctx);
	if (ret)
		die("io_queue_init", -ret);

	while (nr_ios) {
		while (in_flight >= iodepth)
			in_flight -= reap_events(io_ctx, iosize);

		iocb = malloc(sizeof(struct iocb));
		io_prep_pread(iocb, fd, buffer, iosize,
			      (random() % (blocks + 1)) * 512);

		ret = io_submit(io_ctx, 1, &iocb);
		if (ret != 1)
			die("io_submit", -ret);

		ret = io_cancel(io_ctx, iocb, &cancelled);
		if (ret && ret != -EINPROGRESS)
			printf("io_cancel error: %s\n", strerror(-ret));

		in_flight++;
		nr_ios--;
		//printf("submitted to %llu: nr_ios %lu in_flight %lu iodepth %lu\n",
		//       iocb->u.c.offset, nr_ios, in_flight, iodepth);
	}

	while (in_flight)
		in_flight -= reap_events(io_ctx, iosize);

	close(fd);

	printf("canceled %lu completed %lu\n", canceled, completed);

	return 0;
}