/* cyrusdb_bench.c: cyrusdb benchmark tool.
 *                  [Based on the db_bench tool in LevelDB]
 *
 * Copyright (c) 1994-2019 Carnegie Mellon University.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 *
 * 3. The name "Carnegie Mellon University" must not be used to
 *    endorse or promote products derived from this software without
 *    prior written permission. For permission or any legal
 *    details, please contact
 *      Carnegie Mellon University
 *      Center for Technology Transfer and Enterprise Creation
 *      4615 Forbes Avenue
 *      Suite 302
 *      Pittsburgh, PA  15213
 *      (412) 268-7393, fax: (412) 268-7395
 *      innovation@andrew.cmu.edu
 *
 * 4. Redistributions of any form whatsoever must retain the following
 *    acknowledgment:
 *    "This product includes software developed by Computing Services
 *     at Carnegie Mellon University (http://www.cmu.edu/computing/)."
 *
 * CARNEGIE MELLON UNIVERSITY DISCLAIMS ALL WARRANTIES WITH REGARD TO
 * THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
 * AND FITNESS, IN NO EVENT SHALL CARNEGIE MELLON UNIVERSITY BE LIABLE
 * FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN
 * AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING
 * OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <assert.h>
#include <errno.h>
#include <getopt.h>
#include <ftw.h>
#include <inttypes.h>
#include <libgen.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>

#include "cyrusdb.h"
#include "strarray.h"
#include "util.h"
#include "xmalloc.h"

/* Globals */
static char *DBNAME;
static char *BACKEND;
static const char *BENCHMARKS;
static int NUMRECS = 1000;
static int new_db = 0;          /* set to 1 if we created a new db */
static size_t VALLEN = 0;

#define ALLBENCHMARKS "writeseq,writeseqtxn,writerandom,writerandomtxn,write100k"

enum {
        BATCHED,
        NOTBATCHED,
        SEQUENTIAL,
        RANDOM,
};

static struct option long_options[] = {
        {"benchmarks", required_argument, NULL, 'b'},
        {"db", required_argument, NULL, 'd'},
        {"backend", required_argument, NULL, 't'},
        {"numrecs", optional_argument, NULL, 'n'},
        {"help", no_argument, NULL, 'h'},
        {NULL, 0, NULL, 0}
};

EXPORTED void fatal(const char *message, int code)
{
  static int recurse_code = 0;

  if (recurse_code) {
    exit(code);
  }

  recurse_code = code;
  fprintf(stderr, "fatal error: %s\n", message);
  exit(code);
}

static char *create_tmp_dir_name(void)
{
    static const char charset[] =
        "abcdefghijklmnopqrstuvwxyz"
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "0123456789";
    static const int num_chars = 62;
    uint64_t value;
    struct timeval tv;
    const char *tmpdir;
    char path[PATH_MAX];
    char *dirname_template;
    int i, ret = 1;
    struct stat sb;

    gettimeofday(&tv, NULL);
    value = ((size_t)(tv.tv_usec << 16)) ^ tv.tv_sec ^ getpid();

    tmpdir = getenv("TMPDIR");
    if (!tmpdir)
        tmpdir = "/tmp";

    snprintf(path, sizeof(path), "%s/zsbench-XXXXXX", tmpdir);

    dirname_template = &path[strlen(path) - 6];

    /* TMP_MAX: The minimum number of unique filenames generated by
     * tmpnam() */
    for (i = 0; i < TMP_MAX; ++i) {
        uint64_t v = value;
        int j;

        ret = 1;

        /* Fill in the random bits. */
        for (j = 0; j < 6; j++) {
            dirname_template[j] = charset[v % num_chars];
            v /= num_chars;
        }

        if (stat(path, &sb) == -1) {
            if (errno == ENOENT) {
                /* FOUND an unused path! */
                ret = 0;
                break;
            }
        }

        value += 9999;
    }

    if (!ret)
        return xstrdup(path);
    else
        return NULL;
}

static char *generate_random_string(char *str, size_t length)
{
    static const char charset[] =
        "abcdefghijklmnopqrstuvwxyz"
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "0123456789"
        "!@#$%^&*()-=_+|{}[];<>,./?:";

    if (length) {
        --length;

        for (size_t n = 0; n < length; n++) {
            int pos = rand() % (int) (sizeof charset - 1);
            str[n] = charset[pos];
        }

        str[length] = '\0';
    }

    return str;
}

static char *random_string(size_t length)
{
    char *s = xmalloc(length + 1);

    generate_random_string(s, length);

    return s;
}

#define LINUX
#if defined(LINUX) || defined(DARWIN) || defined(BSD)
static int recursive_rm_cb(const char *path,
                           const struct stat *sb __attribute__((__unused__)),
                           int typeflag __attribute__((__unused__)),
                           struct FTW *ftwbuf __attribute__((__unused__)))
{
    int ret = remove(path);

    if (ret)
        perror(path);

    return ret;
}

int recursive_rm(const char *path)
{
    return nftw(path, recursive_rm_cb, 64, FTW_DEPTH | FTW_PHYS);
}
#else
int recursive_rm(const char *path __attribute__((__unused__)))
{
    return 0;
}
#endif


static void cleanup_db_dir(void)
{
        recursive_rm(DBNAME);
}
static uint64_t get_time_now(void)
{
    struct timeval tv;
    gettimeofday(&tv, NULL);

    return tv.tv_sec * 1000000 + tv.tv_usec;
}


static void usage(const char *progname)
{
    printf("Usage: %s [OPTION]... [DB]...\n", progname);

    printf("  -b, --benchmarks     comma separated list of benchmarks to run\n");
    printf("                       [will run all the benchmarks by default]");
    printf("                       Available benchmarks:\n");
    printf("                       * writeseq       - write values in sequential key order\n");
    printf("                       * writeseqtxn    - write values in sequential key order in separate transactions\n");
    printf("                       * writerandom    - write values in random key order\n");
    printf("                       * writerandomtxn - write values in random key order in separate transactions\n");
    printf("                       * write100k      - write values 100K long in random key order\n");
    printf("\n");
    printf("  -d, --db             the db to run the benchmarks on\n");
    printf("                       (if not provided, will create a new db)\n");
    printf("  -t, --backend        type of the db backend to run benchmarks on\n");
    printf("                       Available Cyrus DB's: twoskip, zeroskip\n");
    printf("  -n, --numrecs        number of records to write[default: 1000]\n");
    printf("  -h, --help           display this help and exit\n");
}


static void print_environment(void)
{
    fprintf(stderr, "Cyrus DB:       %s\n", DBNAME);

    time_t now = time(NULL);
    fprintf(stderr, "Date:           %s", ctime(&now));

    FILE *cpuinfo = fopen("/proc/cpuinfo", "r");
    if (cpuinfo != NULL) {
        char line[1000];
        int num_cpus = 0;
        struct buf cpu_type = BUF_INITIALIZER;
        struct buf cache_size = BUF_INITIALIZER;

        while (fgets(line, sizeof(line), cpuinfo) != NULL) {
            const char *sep = strchr(line, ':');
            struct buf key = BUF_INITIALIZER;
            struct buf val = BUF_INITIALIZER;

            if (sep == NULL)
                continue;

            buf_setmap(&key, line, (sep - 1 - line));
            buf_trim(&key);

            buf_setcstr(&val, (sep + 1));
            buf_trim(&val);

            if (strcmp("model name", key.s) == 0) {
                ++num_cpus;
                buf_truncate(&cpu_type, 0);
                buf_setcstr(&cpu_type, val.s);
            } else if (strcmp("cache size", key.s) == 0) {
                buf_truncate(&cache_size, 0);
                buf_setcstr(&cache_size, val.s);
            }

            buf_free(&key);
            buf_free(&val);
        }

        fclose(cpuinfo);

        const char *ctype = buf_cstring(&cpu_type);
        const char *csize = buf_cstring(&cache_size);
        fprintf(stderr, "CPU:            %d * %s",
                num_cpus, ctype);
        fprintf(stderr, "CPUCache:       %s", csize);

        buf_free(&cpu_type);
        buf_free(&cache_size);
    }
}

static void print_header(void)
{
    print_environment();
    fprintf(stdout, "------------------------------------------------\n");
}

static size_t do_write(int txnmode, int insmode)
{
    int i;
    int ret;
    struct db *db = NULL;
    size_t bytes = 0;
    struct txn *txn = NULL;

    /* Open DB */
    ret = cyrusdb_open(BACKEND, DBNAME, new_db ? CYRUSDB_CREATE : 0, &db);
    assert(ret == CYRUSDB_OK);
    assert(db != NULL);

    for (i = 0; i < NUMRECS; i++) {
        char key[100];
        size_t keylen, vallen;
        char *val;
        int k;

        k = (insmode == SEQUENTIAL) ? i : ((time(NULL) *  i) % NUMRECS);

        snprintf(key, sizeof(key), "%016d", k);
        keylen = strlen(key);
        vallen = VALLEN ? VALLEN : keylen * 2;
        val = random_string(vallen);

        ret = cyrusdb_store(db, key, keylen, val, vallen,
                            (txnmode == BATCHED) ? &txn : NULL);
        assert(ret == CYRUSDB_OK);

        if (txnmode == BATCHED) {
                assert(txn != NULL);
        }

        bytes += (keylen + vallen);

        if (txnmode == BATCHED) {
            ret = cyrusdb_commit(db, txn);
            assert(ret == CYRUSDB_OK);
            txn = NULL;
        }
    }

    if (txnmode == NOTBATCHED) {
        ret = cyrusdb_commit(db, txn);
        assert(ret == CYRUSDB_OK);
    }

    /* Close DB */
    ret = cyrusdb_close(db);
    assert(ret == CYRUSDB_OK);

    return bytes;
}

static int parse_options(int argc, char **argv, const struct option *options)
{
    int option;
    int option_index;

    while ((option = getopt_long(argc, argv, "d:b:t:h?",
                                 long_options, &option_index)) != -1) {
        switch (option) {
            case 'b':
                BENCHMARKS = optarg;
                break;
            case 'd':
                DBNAME = optarg;
                break;
            case 't':
                BACKEND = optarg;
                break;
            case 'n':
                NUMRECS = atoi(optarg);
                break;
            case 'h':
                GCC_FALLTHROUGH
            case '?':
                usage(basename(argv[0]));
                exit(option == 'h');
        }
    }

    return 0;
}

static int run_benchmarks(void)
{
    strarray_t *benchmarks = NULL;
    int i;
    uint64_t start, finish;
    size_t bytes;

    print_header();

    benchmarks = strarray_split(BENCHMARKS, ",", 0);

    for (i = 0; i < benchmarks->count; i++) {
        const char *benchname = strarray_nth(benchmarks, i);
        if (strcmp(benchname, "writeseq") == 0) {
            start = get_time_now();
            bytes = do_write(NOTBATCHED, SEQUENTIAL);
            finish = get_time_now();

            fprintf(stderr, "writeseq        : %zu bytes written in %" PRIu64 " μs.\n",
                    bytes, (finish - start));
        } else if (strcmp(benchname, "writeseqtxn") == 0) {
            start = get_time_now();
            bytes = do_write(BATCHED, SEQUENTIAL);
            finish = get_time_now();

            fprintf(stderr, "writeseqtxn     : %zu bytes written in %" PRIu64 " μs.\n",
                    bytes, (finish - start));
        } else if (strcmp(benchname, "writerandom") == 0) {
            start = get_time_now();
            bytes = do_write(NOTBATCHED, RANDOM);
            finish = get_time_now();

            fprintf(stderr, "writerandom     : %zu bytes written in %" PRIu64 " μs.\n",
                    bytes, (finish - start));
        } else if (strcmp(benchname, "writerandomtxn") == 0) {
            start = get_time_now();
            bytes = do_write(BATCHED, RANDOM);
            finish = get_time_now();

            fprintf(stderr, "writerandomtxn  : %zu bytes written in %" PRIu64 " μs.\n",
                    bytes, (finish - start));
        } else if (strcmp(benchname, "write100k") == 0) {
            VALLEN = 100 * 1000;
            start = get_time_now();
            bytes = do_write(NOTBATCHED, RANDOM);
            finish = get_time_now();

            fprintf(stderr, "write100k       : %zu bytes written in %" PRIu64 " μs.\n",
                    bytes, (finish - start));
            VALLEN = 0;
        } else {
            fprintf(stderr, "Unknown benchmark '%s'\n", benchname);
        }

        if (new_db) {
            cleanup_db_dir();
        }
    }

    strarray_free(benchmarks);

    return 0;
}

int main(int argc, char *argv[])
{
    int ret = EXIT_SUCCESS;
    int seed = 1103515245;

    /* Random Seed */
    srand(time(NULL) * seed);

    ret = parse_options(argc, argv, long_options);

    if (BACKEND == NULL) {
        fprintf(stderr, "Require a backend to run benchmarks.\n");
        usage(basename(argv[0]));
        ret = EXIT_FAILURE;
        goto done;
    }

    if (strncmp(BACKEND, "twoskip", strlen("twoskip")) == 0 ||
        strncmp(BACKEND, "zeroskip", strlen("zeroskip")) == 0) {
        fprintf(stderr, "Running benchmarks for `%s` backend\n", BACKEND);
    } else {
        fprintf(stderr, "%s is not a valid CyrusDB backend. ", BACKEND);
        fprintf(stderr, "Choose between `twoskip` or `zeroskip`.\n");
        ret = EXIT_FAILURE;
        goto done;
    }

    if (BENCHMARKS == NULL) {
        fprintf(stderr, "Running all benchmarks\n");
        BENCHMARKS = ALLBENCHMARKS;
    }

    if (DBNAME == NULL) {
        new_db = 1;
        DBNAME = create_tmp_dir_name();
        assert(DBNAME != NULL);
        printf("Creating a new DB: %s\n", DBNAME);
    } else {
        printf("Using existing DB: %s\n", DBNAME);
    }

    ret = run_benchmarks();

    if (new_db) {
        free(DBNAME);
        DBNAME = NULL;
    }

 done:
    exit(ret);
}
