/**
 * Copyright 2022 Joshua Bakita
 * This program clocks how long it takes to page-fault in a 1GiB buffer, making
 * efforts to exclude userspace overheads.
 *
 * More precisely, this program clocks:
 *   mmap(big_buffer);
 *   sequentially_walk(big_buffer);
 * where `big_buffer` is randomly filled. It then clocks another:
 *   sequentially_walk(big_buffer);
 * subtracts that time from the first one, and outputs the result. /dev/nvme0n1
 * is preinitialized with random non-zero bytes.
 */
#define _GNU_SOURCE

#include <assert.h>
#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <time.h>
#include <unistd.h>

#define GiB 1024l*1024l*1024l
#define s2ns(s) ((s)*1000l*1000l*1000l)
#define ns2us(ns) ((ns)/1000l)
#define PAGED_FILE "/dev/nvme0n1"
#define CLEAR_PAGECACHE_DENTRIES_INODES "3"
int max(int x, int y) {return x > y ? x : y;}

int seq_walk(char* mem, int len, char to_find) {
	int num_42 = 0;
	// Stride of 4096 bytes (one 4k page)
	for (int i = 4096; i < len; i += 4096)
		if (mem[i] == to_find)
			num_42++;
	return num_42;
}

// Original function from copy_only.cu
void fill_rand(char* buf, uint64_t buf_len) {
        uint64_t i = 0;
        for (; i < buf_len; i++)
                buf[i] = max((rand() & 0xff), 1);
}

uint64_t count_zero(char* buf, uint64_t buf_len) {
        uint64_t i = 0;
        uint64_t num_zeros = 0;
        for (; i < buf_len; i++)
                num_zeros += (!buf[i]);
        return num_zeros;
}

// Subtract first parameter from second parameter. Return as nanoseconds.
long time_diff_ns(struct timespec start, struct timespec stop) {
	return (s2ns(stop.tv_sec) + stop.tv_nsec) - (s2ns(start.tv_sec) + start.tv_nsec);
}

int main(int argc, char **argv) {
	struct timespec start, stop, stop2;
	int iters, res;

	if (argc != 2 || argv[1][0] == '-') {
		fprintf(stderr, "Usage: %s <number of iterations>\n", argv[0]);
		return 1;
	}
	iters = atoi(argv[1]);

        // If output is redirected, add comment with source details
        if (!isatty(fileno(stdout)))
                fprintf(stdout, "# Generated by '%s %s'\n", argv[0], argv[1]);

	// Open control device for use to reset between iterations any caches
	// that the demand paging system might use.
	int clear_fd = open("/proc/sys/vm/drop_caches", O_WRONLY);
	if (clear_fd == -1) {
		perror("Unable to open /proc/sys/vm/drop_caches");
		return 1;
	}

	// Fill the area that we'll read from with random data
	// (Direct I/O just used here for convenience.)
	char *mem_in, *mem_out; // In and out of SSD
	int fd = open(PAGED_FILE, O_RDWR | O_DIRECT | O_SYNC);
	if (fd == -1) {
		perror("Unable to open " PAGED_FILE);
		return 1;
	}
	// Aligned malloc(GiB) basicially
	res = posix_memalign((void**)&mem_in, 4096, GiB);
	fill_rand(mem_in, GiB);
	res = write(fd, mem_in, GiB);
	if (res == -1) {
		perror("Unable to write inital 1GiB random buffer to " PAGED_FILE);
		return 1;
	}
	if (res != GiB) {
		fprintf(stderr, "Unable to write the buffer all at once!");
		return 2;
	}
	free(mem_in);
	close(fd);

	// Output table header. One read sample per row.
	printf("in (us)\n");
	// Perform iterations of demand paging in
	for (int i = 0; i < iters; i++) {
		int fd = open(PAGED_FILE, O_RDWR);
		if (fd == -1) {
			perror("Unable to open " PAGED_FILE);
			return 1;
		}
		// Clear page cache
		write(clear_fd, CLEAR_PAGECACHE_DENTRIES_INODES, 1);

		// Begin paging by mapping the file into (unbacked) virtual memory
		clock_gettime(CLOCK_MONOTONIC_RAW, &start);
		mem_out = mmap(NULL, GiB, PROT_READ, MAP_PRIVATE, fd, 0);
		if (mem_out == MAP_FAILED) {
			perror("Unable to mmap " PAGED_FILE);
			return 1;
		}
		// Page fault in all the data via a sequential walk
		res = seq_walk(mem_out, GiB, 42); // Made up work to fool optimizer
		assert(res > 0);
		clock_gettime(CLOCK_MONOTONIC_RAW, &stop);

		// This benchmark attempts to capture data comparable to the
		// GPU or Direct I/O paging benchmarks. As they don't have to
		// perform a sequential walk, we emulate that here by measuring
		// how long a sequential walk takes on its own, and subtract
		// that from `stop - start`.
		res = seq_walk(mem_out, GiB, 7); // Made up work to fool optimizer
		assert(res > 0);
		clock_gettime(CLOCK_MONOTONIC_RAW, &stop2);

		// `stop2 - stop` is how long just a sequential walk takes
		long seq_walk_time = time_diff_ns(stop, stop2);
		// `stop - start` is how long the faulting sequential walk took
		long demand_paging_raw_duration = time_diff_ns(start, stop);
		// `(stop - start) - (stop2 - stop)` is emulated duration
		long demand_paging_duration = demand_paging_raw_duration - seq_walk_time;
		printf("%ld\n", ns2us(demand_paging_duration));

		munmap(mem_out, GiB);
		close(fd);
	}
	return 0;
}