/*
 * TUXLOAD.C
 *
 * (c)Copyright 2012 Antonio Huete Jimenez <tuxillo@quantumachine.net>,
 *    this code is hereby placed in the public domain.
 *
 * As a safety the directory 'tmpfiles/' must exist.  This program will
 * create 500 x 8MB files in tmpfiles/*, memory map the files MAP_SHARED,
 * R+W, make random modifications, and msync() in a loop.
 *
 * The purpose is to stress the VM system.
 *
 */
#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <err.h>
#include <errno.h>

static void randomfill(int fd);
static int mmap01(void *);
static int mmap02(void *);
static int build_files(void);

int opt_verbose;
int opt_testno;
int opt_nofiles;
int opt_mbfile;
int opt_count;
int fliparg;

int *fd;
struct stat *st;
char **pp;

static const struct test {
        char testdesc[128];
        int (*testfn)(void *);
} testlist[] = {
        { "mmap01 - Massive mmap / msync (flushing all pages)", mmap01 },
        { "mmap02 - Massive mmap / msync (flushing only specified pages)", mmap02 },
        { "", NULL }
};

static int
mmap01(void *arg)
{
        int i;
        int *bug = arg;
        long jump;
        size_t size, len;

        if ((build_files()) != 0)
                err(1, "Failed to create the files");

        if (opt_verbose)
                printf("\n");

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

                if (opt_verbose) {
                        fflush(stdout);
                        fprintf(stdout, "\rDoing mmap() + msync() [%d/%d] ", i+1, opt_nofiles);
                }
                size = st[i].st_size;
                pp[i] = mmap(0, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd[i], 0);
                if (pp[i] == MAP_FAILED)
                        err(1, "mmap");

                for (jump = 0; jump < size; jump += 65535) {
                        pp[i][jump] = jump + i;
                }

                if (fliparg)
                        len = MS_SYNC;
                else
                        len = 0;

                if ((msync(pp[i], len, MS_SYNC)) == -1) {
                        printf("fd %d %p\n", fd[i], pp[i]);
                        err(1, "msync");

                }
        }
        printf("\n");

        return 0;
}

static int
mmap02(void *arg)
{
        fliparg = 1;
        mmap01(&fliparg);

        return 0;
}

static void
usage(void)
{
        int i;
        const struct test *tp;

        printf("tuxload: [-v] [-c count] [-m megs_per_file] [-n no_of_files] [-t test_to_run] \n"
            "Available tests: \n");

        for (tp = testlist; tp->testfn != NULL; tp++)
                printf("\t%s\n", tp->testdesc);
        exit(1);
}

int
main(int argc, char *argv[])
{
        int forever;
        char c;

        opt_verbose = opt_testno = 0;
        opt_nofiles = 500;
        opt_mbfile = 8;
        opt_count = 1;
        fliparg = 0;
        forever = 0;

        while ((c = getopt(argc, argv, "n:t:m:c:v")) != -1) {
                switch(c) {
                case 'v':
                        opt_verbose++;
                        break;
                case 'n':
                        opt_nofiles = (int)strtol(optarg, NULL, 0);
                        break;
                case 't':
                        opt_testno = (int)strtol(optarg, NULL, 0);
                        opt_testno--;
                        break;
                case 'm':
                        opt_mbfile = (int)strtol(optarg, NULL, 0);
                        break;
                case 'c':
                        opt_count = (int)strtol(optarg, NULL, 0);
                        if (opt_count == 0)
                                forever++;
                        break;
                default:
                        usage();
                        ;
                }
        }
        argc -= optind;
        argv += optind;

        if (argc != 0)
                usage();

        st = malloc(opt_nofiles * sizeof(*st));
        fd = malloc(opt_nofiles * sizeof(*fd));
        pp = malloc(opt_nofiles * sizeof(*pp));

        while (opt_count-- || forever)
                testlist[opt_testno].testfn(0);

        return 0;
}

static int
build_files(void)
{
        char name[128];
        int i;
        int error;

        for (i = 0, error = 0; i <  opt_nofiles; i++) {
                snprintf(name, 128, "tmpfiles/file%d", i);
                if ((fd[i] = open(name, O_RDWR)) < 1) {
                        if ((fd[i] = open(name, O_RDWR | O_CREAT, 0644)) < 1) {
                                error = errno;
                                break;
                        }
                        randomfill(fd[i]);
                }
                if ((fstat(fd[i], &st[i])) == -1) {
                        error = errno;
                        break;
                }
                if (opt_verbose) {
                        fprintf(stdout, "\rFile creation, random data filled [%d/%d] ", i+1, opt_nofiles);
                        fflush(stdout);
                }
        }

        return error;
}

static void
randomfill(int fd)
{
        char buf[32768];
        long tot;
        int i;

        srandomdev();
        tot = opt_mbfile * 1024L;
        for (i = 0; i < 32768; ++i)
                buf[i] = random();
        for (i = 0; i < tot; i += 32)   /* 8MB by default */
                write(fd, buf, 32768);
        fsync(fd);
        lseek(fd, 0L, 0);
}