/*
   Copyright (c) 2008-2012 Red Hat, Inc. <http://www.redhat.com>
   This file is part of GlusterFS.

   This file is licensed to you under your choice of the GNU Lesser
   General Public License, version 3 or any later version (LGPLv3 or
   later), or the GNU General Public License, version 2 (GPLv2), in all
   cases as published by the Free Software Foundation.
*/
#include <stdio.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <unistd.h>
#include <pthread.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <argp.h>

#define TWO_POWER(power) (2UL << (power))

#define RDD_INTEGER_VALUE ((TWO_POWER((sizeof(int) * 8))) - 1)

#ifndef UNIX_PATH_MAX
#define UNIX_PATH_MAX 108
#endif

#define UNIT_KB 1024ULL
#define UNIT_MB UNIT_KB * 1024ULL
#define UNIT_GB UNIT_MB * 1024ULL
#define UNIT_TB UNIT_GB * 1024ULL
#define UNIT_PB UNIT_TB * 1024ULL

#define UNIT_KB_STRING "KB"
#define UNIT_MB_STRING "MB"
#define UNIT_GB_STRING "GB"
#define UNIT_TB_STRING "TB"
#define UNIT_PB_STRING "PB"

struct rdd_file {
    char path[UNIX_PATH_MAX];
    struct stat st;
    int fd;
};

struct rdd_config {
    long iters;
    long max_ops_per_seq;
    size_t max_bs;
    size_t min_bs;
    int thread_count;
    pthread_t *threads;
    pthread_barrier_t barrier;
    pthread_mutex_t lock;
    struct rdd_file in_file;
    struct rdd_file out_file;
    ssize_t file_size;
};
static struct rdd_config rdd_config;

enum rdd_keys {
    RDD_MIN_BS_KEY = 1,
    RDD_MAX_BS_KEY,
};

static error_t
rdd_parse_opts(int key, char *arg, struct argp_state *_state)
{
    switch (key) {
        case 'o': {
            int len = 0;
            len = strlen(arg);
            if (len > UNIX_PATH_MAX) {
                fprintf(stderr, "output file name too long (%s)\n", arg);
                return -1;
            }

            strncpy(rdd_config.out_file.path, arg, len);
        } break;

        case 'i': {
            int len = 0;
            len = strlen(arg);
            if (len > UNIX_PATH_MAX) {
                fprintf(stderr, "input file name too long (%s)\n", arg);
                return -1;
            }

            strncpy(rdd_config.in_file.path, arg, len);
            rdd_config.in_file.path[len] = '\0';
        } break;

        case 'f': {
            char *tmp = NULL;
            unsigned long long fs = 0;
            if (string2bytesize(arg, &fs) == -1) {
                fprintf(stderr,
                        "invalid argument for file size "
                        "(%s)\n",
                        arg);
                return -1;
            }

            rdd_config.file_size = fs;
        } break;

        case RDD_MIN_BS_KEY: {
            char *tmp = NULL;
            long bs = 0;
            bs = strtol(arg, &tmp, 10);
            if ((bs == LONG_MAX) || (bs == LONG_MIN) || (tmp && *tmp)) {
                fprintf(stderr,
                        "invalid argument for minimum block"
                        "size (%s)\n",
                        arg);
                return -1;
            }

            rdd_config.min_bs = bs;
        } break;

        case RDD_MAX_BS_KEY: {
            char *tmp = NULL;
            long bs = 0;
            bs = strtol(arg, &tmp, 10);
            if ((bs == LONG_MAX) || (bs == LONG_MIN) || (tmp && *tmp)) {
                fprintf(stderr,
                        "invalid argument for maximum block"
                        "size (%s)\n",
                        arg);
                return -1;
            }

            rdd_config.max_bs = bs;
        } break;

        case 'r': {
            char *tmp = NULL;
            long iters = 0;
            iters = strtol(arg, &tmp, 10);
            if ((iters == LONG_MAX) || (iters == LONG_MIN) || (tmp && *tmp)) {
                fprintf(stderr,
                        "invalid argument for iterations"
                        "(%s)\n",
                        arg);
                return -1;
            }

            rdd_config.iters = iters;
        } break;

        case 'm': {
            char *tmp = NULL;
            long max_ops = 0;
            max_ops = strtol(arg, &tmp, 10);
            if ((max_ops == LONG_MAX) || (max_ops == LONG_MIN) ||
                (tmp && *tmp)) {
                fprintf(stderr,
                        "invalid argument for max-ops"
                        "(%s)\n",
                        arg);
                return -1;
            }

            rdd_config.max_ops_per_seq = max_ops;
        } break;

        case 't': {
            char *tmp = NULL;
            long threads = 0;
            threads = strtol(arg, &tmp, 10);
            if ((threads == LONG_MAX) || (threads == LONG_MIN) ||
                (tmp && *tmp)) {
                fprintf(stderr,
                        "invalid argument for thread count"
                        "(%s)\n",
                        arg);
                return -1;
            }

            rdd_config.thread_count = threads;
        } break;

        case ARGP_KEY_NO_ARGS:
            break;
        case ARGP_KEY_ARG:
            break;
        case ARGP_KEY_END:
            if (_state->argc == 1) {
                argp_usage(_state);
            }
    }

    return 0;
}

int
string2bytesize(const char *str, unsigned long long *n)
{
    unsigned long long value = 0ULL;
    char *tail = NULL;
    int old_errno = 0;
    const char *s = NULL;

    if (str == NULL || n == NULL) {
        errno = EINVAL;
        return -1;
    }

    for (s = str; *s != '\0'; s++) {
        if (isspace(*s)) {
            continue;
        }
        if (*s == '-') {
            return -1;
        }
        break;
    }

    old_errno = errno;
    errno = 0;
    value = strtoull(str, &tail, 10);

    if (errno == ERANGE || errno == EINVAL) {
        return -1;
    }

    if (errno == 0) {
        errno = old_errno;
    }

    if (tail[0] != '\0') {
        if (strcasecmp(tail, UNIT_KB_STRING) == 0) {
            value *= UNIT_KB;
        } else if (strcasecmp(tail, UNIT_MB_STRING) == 0) {
            value *= UNIT_MB;
        } else if (strcasecmp(tail, UNIT_GB_STRING) == 0) {
            value *= UNIT_GB;
        } else if (strcasecmp(tail, UNIT_TB_STRING) == 0) {
            value *= UNIT_TB;
        } else if (strcasecmp(tail, UNIT_PB_STRING) == 0) {
            value *= UNIT_PB;
        }

        else {
            return -1;
        }
    }

    *n = value;

    return 0;
}

static struct argp_option rdd_options[] = {
    {"if", 'i', "INPUT_FILE", 0, "input-file"},
    {"of", 'o', "OUTPUT_FILE", 0, "output-file"},
    {"threads", 't', "COUNT", 0, "number of threads to spawn (defaults to 2)"},
    {"min-bs", RDD_MIN_BS_KEY, "MIN_BLOCK_SIZE", 0,
     "Minimum block size in bytes (defaults to 1024)"},
    {"max-bs", RDD_MAX_BS_KEY, "MAX_BLOCK_SIZE", 0,
     "Maximum block size in bytes (defaults to 4096)"},
    {"iters", 'r', "ITERS", 0,
     "Number of read-write sequences (defaults to 1000000)"},
    {"max-ops", 'm', "MAXOPS", 0,
     "maximum number of read-writes to be performed in a sequence (defaults to "
     "1)"},
    {"file-size", 'f', "FILESIZE", 0,
     "the size of the file which will be created and upon it I/O will be done"
     " (defaults to 100MB"},
    {0, 0, 0, 0, 0}};

static struct argp argp = {
    rdd_options, rdd_parse_opts, "",
    "random dd - tool to do a sequence of random block-sized continuous"
    "read writes starting at a random offset"};

static void
rdd_default_config(void)
{
    char *tmp_path = "rdd.in";

    rdd_config.thread_count = 2;
    rdd_config.iters = 1000000;
    rdd_config.max_bs = 4096;
    rdd_config.min_bs = 1024;
    rdd_config.in_file.fd = rdd_config.out_file.fd = -1;
    rdd_config.max_ops_per_seq = 1;
    strncpy(rdd_config.in_file.path, tmp_path, strlen(tmp_path));
    rdd_config.file_size = 104857600;

    return;
}

static char
rdd_valid_config(void)
{
    char ret = 1;
    int fd = -1;

    fd = open(rdd_config.in_file.path, O_RDONLY);
    if (fd == -1 && (errno != ENOENT)) {
        fprintf(stderr, "open: (%s)", strerror(errno));
        ret = 0;
        goto out;
    }
    close(fd);

    if (rdd_config.min_bs > rdd_config.max_bs) {
        fprintf(stderr,
                "minimum blocksize %ld is greater than the "
                "maximum blocksize %ld",
                rdd_config.min_bs, rdd_config.max_bs);
        ret = 0;
        goto out;
    }

    if (strlen(rdd_config.out_file.path) == 0) {
        sprintf(rdd_config.out_file.path, "%s.rddout", rdd_config.in_file.path);
    }

out:
    return ret;
}

static void *
rdd_read_write(void *arg)
{
    int i = 0, ret = 0;
    size_t bs = 0;
    off_t offset = 0;
    long rand = 0;
    long max_ops = 0;
    char *buf = NULL;

    buf = calloc(1, rdd_config.max_bs);
    if (!buf) {
        fprintf(stderr, "calloc failed (%s)\n", strerror(errno));
        ret = -1;
        goto out;
    }

    for (i = 0; i < rdd_config.iters; i++) {
        pthread_mutex_lock(&rdd_config.lock);
        {
            int bytes = 0;
            rand = random();

            if (rdd_config.min_bs == rdd_config.max_bs) {
                bs = rdd_config.max_bs;
            } else {
                bs = rdd_config.min_bs +
                     (rand % (rdd_config.max_bs - rdd_config.min_bs));
            }

            offset = rand % rdd_config.in_file.st.st_size;
            max_ops = rand % rdd_config.max_ops_per_seq;
            if (!max_ops) {
                max_ops++;
            }

            ret = lseek(rdd_config.in_file.fd, offset, SEEK_SET);
            if (ret != offset) {
                fprintf(stderr, "lseek failed (%s)\n", strerror(errno));
                ret = -1;
                goto unlock;
            }

            ret = lseek(rdd_config.out_file.fd, offset, SEEK_SET);
            if (ret != offset) {
                fprintf(stderr, "lseek failed (%s)\n", strerror(errno));
                ret = -1;
                goto unlock;
            }

            while (max_ops--) {
                bytes = read(rdd_config.in_file.fd, buf, bs);
                if (!bytes) {
                    break;
                }

                if (bytes == -1) {
                    fprintf(stderr, "read failed (%s)\n", strerror(errno));
                    ret = -1;
                    goto unlock;
                }

                if (write(rdd_config.out_file.fd, buf, bytes) != bytes) {
                    fprintf(stderr, "write failed (%s)\n", strerror(errno));
                    ret = -1;
                    goto unlock;
                }
            }
        }
    unlock:
        pthread_mutex_unlock(&rdd_config.lock);
        if (ret == -1) {
            goto out;
        }
        ret = 0;
    }
out:
    free(buf);
    pthread_barrier_wait(&rdd_config.barrier);

    return NULL;
}

static void
cleanup(void)
{
    close(rdd_config.in_file.fd);
    close(rdd_config.out_file.fd);
    rdd_config.in_file.fd = rdd_config.out_file.fd = -1;
}

static int
check_and_create(void)
{
    int ret = -1;
    char buf[4096] = {
        0,
    };
    struct stat stbuf = {
        0,
    };
    int fd[2] = {
        -1,
    };
    size_t total_size = -1;

    total_size = rdd_config.file_size;

    ret = stat(rdd_config.in_file.path, &stbuf);
    if (ret == -1 && (errno != ENOENT))
        goto out;

    fd[1] = open(rdd_config.in_file.path, O_CREAT | O_WRONLY | O_TRUNC);
    if (fd[1] == -1)
        goto out;

    fd[0] = open("/dev/urandom", O_RDONLY);
    if (fd[0] == -1)
        goto out;

    while (total_size > 0) {
        if (total_size >= 4096) {
            ret = read(fd[0], buf, 4096);
            if (ret == -1)
                goto out;
            ret = write(fd[1], buf, 4096);
            if (ret == -1)
                goto out;
            total_size = total_size - 4096;
        } else {
            ret = read(fd[0], buf, total_size);
            if (ret == -1)
                goto out;
            ret = write(fd[1], buf, total_size);
            if (ret == -1)
                goto out;
            total_size = total_size - total_size;
        }
    }

    ret = 0;

out:
    if (fd[0] > 0)
        close(fd[0]);
    if (fd[1] > 0)
        close(fd[1]);
    return ret;
}

static int
rdd_spawn_threads(void)
{
    int i = 0, ret = -1, fd = -1;
    char buf[4096];

    ret = check_and_create();
    if (ret == -1)
        goto out;

    fd = open(rdd_config.in_file.path, O_RDONLY);
    if (fd < 0) {
        fprintf(stderr, "cannot open %s (%s)\n", rdd_config.in_file.path,
                strerror(errno));
        ret = -1;
        goto out;
    }
    ret = fstat(fd, &rdd_config.in_file.st);
    if (ret != 0) {
        close(fd);
        fprintf(stderr, "cannot stat %s (%s)\n", rdd_config.in_file.path,
                strerror(errno));
        ret = -1;
        goto out;
    }
    rdd_config.in_file.fd = fd;

    fd = open(rdd_config.out_file.path, O_WRONLY | O_CREAT | O_TRUNC,
              S_IRWXU | S_IROTH);
    if (fd < 0) {
        close(rdd_config.in_file.fd);
        rdd_config.in_file.fd = -1;
        fprintf(stderr, "cannot open %s (%s)\n", rdd_config.out_file.path,
                strerror(errno));
        ret = -1;
        goto out;
    }
    rdd_config.out_file.fd = fd;

    while ((ret = read(rdd_config.in_file.fd, buf, 4096)) > 0) {
        if (write(rdd_config.out_file.fd, buf, ret) != ret) {
            fprintf(stderr, "write failed (%s)\n", strerror(errno));
            cleanup();
            ret = -1;
            goto out;
        }
    }

    rdd_config.threads = calloc(rdd_config.thread_count, sizeof(pthread_t));
    if (rdd_config.threads == NULL) {
        fprintf(stderr, "calloc() failed (%s)\n", strerror(errno));

        ret = -1;
        cleanup();
        goto out;
    }

    ret = pthread_barrier_init(&rdd_config.barrier, NULL,
                               rdd_config.thread_count + 1);
    if (ret != 0) {
        fprintf(stderr, "pthread_barrier_init() failed (%s)\n", strerror(ret));

        free(rdd_config.threads);
        cleanup();
        ret = -1;
        goto out;
    }

    ret = pthread_mutex_init(&rdd_config.lock, NULL);
    if (ret != 0) {
        fprintf(stderr, "pthread_mutex_init() failed (%s)\n", strerror(ret));

        free(rdd_config.threads);
        pthread_barrier_destroy(&rdd_config.barrier);
        cleanup();
        ret = -1;
        goto out;
    }

    for (i = 0; i < rdd_config.thread_count; i++) {
        ret = pthread_create(&rdd_config.threads[i], NULL, rdd_read_write,
                             NULL);
        if (ret != 0) {
            fprintf(stderr, "pthread_create failed (%s)\n", strerror(errno));
            exit(1);
        }
    }

out:
    return ret;
}

static void
rdd_wait_for_completion(void)
{
    pthread_barrier_wait(&rdd_config.barrier);
}

int
main(int argc, char *argv[])
{
    int ret = -1;

    rdd_default_config();

    ret = argp_parse(&argp, argc, argv, 0, 0, NULL);
    if (ret != 0) {
        ret = -1;
        fprintf(stderr, "%s: argp_parse() failed\n", argv[0]);
        goto err;
    }

    if (!rdd_valid_config()) {
        ret = -1;
        fprintf(stderr, "%s: configuration validation failed\n", argv[0]);
        goto err;
    }

    ret = rdd_spawn_threads();
    if (ret != 0) {
        fprintf(stderr, "%s: spawning threads failed\n", argv[0]);
        goto err;
    }

    rdd_wait_for_completion();

err:
    return ret;
}