/*
 * rate.c:
 * Simple IO rate measurement.
 *
 * Syntax: rate (read|write|writesync) WHAT
 *
 * Perform an IO rate measurement on WHAT (either a file or a block device),
 * performing either reads, writes, or synchronous writes according to its
 * first argument. The output is a series of measurements of the cost of
 * transactions at random offsets in the file or device; each line of output
 * lists the size of the transaction and the number of seconds of wall clock,
 * user CPU and system CPU time elapsed to complete the transaction.
 *
 * To get useful results from this you will probably want to compute a median
 * and other percentile points for each transaction size.
 *
 * Written by Chris Lightfoot <chris@mythic-beasts.com>
 * Donated to the public domain, January 2006
 *
 */

#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#include <sys/mount.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/time.h>

#define MAXBLOCK    (16 * 1024 * 1024)
#define MAXLOG      15

void quit(const char *why) {
    fprintf(stderr, "rate: %s\n", why);
    exit(1);
}

double timediff(const struct timeval *before, const struct timeval *after) {
    return (double)(after->tv_sec - before->tv_sec) + 1e-6 * (after->tv_usec - before->tv_usec);
}

int main(int argc, char *argv[]) {
    enum { op_unknown = 0, op_read, op_write, op_writesync } operation = op_unknown;
    char *what;
    int fd, f = 0;
    unsigned char *buffer, *p;
    struct stat st;
    off_t size;
    int N;

    srand(time(NULL));
    
    if (!(buffer = malloc(MAXBLOCK))) {
        fprintf(stderr, "rate: malloc: %s\n", strerror(errno));
        return 1;
    }

    for (p = buffer; p < buffer + MAXBLOCK; ) {
        *p++ = 0xaa;
        *p++ = 0x55;
    }
    
    if (argc != 3)
        quit("two arguments are operation and file/device for testing");
    if (!strcmp(argv[1], "read"))
        operation = op_read;
    else if (!strcmp(argv[1], "write"))
        operation = op_write;
    else if (!strcmp(argv[1], "writesync"))
        operation = op_writesync;
    else
        quit("operation should be one of: read, write, writesync");

    what = argv[2];
    f = (operation == op_read ? O_RDONLY : O_RDWR);
    if (operation == op_writesync)
        f |= O_SYNC;
    if (-1 == (fd = open(what, f))) {
        fprintf(stderr, "rate: %s: %s\n", what, strerror(errno));
        return 1;
    } else if (-1 == fstat(fd, &st)) {
        fprintf(stderr, "rate: %s: stat: %s\n", what, strerror(errno));
        return 1;
    }

    if (S_ISBLK(st.st_mode)) {
        /* Block special device; get size by ioctl. */
        unsigned long b;
        if (-1 == ioctl(fd, BLKGETSIZE, &b)) {
            fprintf(stderr, "rate: %s: ioctl(BLKGETSIZE): %s\n", what, strerror(errno));
            return 1;
        }
        size = (long long)b * 512;
        fprintf(stderr, "%s: block device; size = %llu bytes\n", what, (unsigned long long)size);
    } else {
        size = st.st_size;
        fprintf(stderr, "%s: regular file; size = %llu bytes\n", what, (unsigned long long)size);
    }

    fprintf(stderr, "%s: %.2f GB\n", what, (double)size / (1024. * 1024. * 1024.));

    for (N = 0; N < 4000; ++N) {
        int j;
        size_t howmuch, done = 0;
        off_t where;
        struct timeval real_1, real_2, user_1, user_2, sys_1, sys_2;
        struct rusage ru;

        if (0 == (N % 100))
            fprintf(stderr, "rate: %s: %d/4000 transactions\n", what, N);
        
        /* Pick a size of transaction. */
        do {
            j = (int)((rand() / (float)RAND_MAX) * MAXLOG);
            howmuch = 512 << j;
        } while (howmuch > size);

        if (howmuch > MAXBLOCK)
            quit("oops");

        /* Pick a place to do it. */
        where = 512 * ((rand() / (float)RAND_MAX) * ((size - howmuch) / 512));

        gettimeofday(&real_1, NULL);
        getrusage(RUSAGE_SELF, &ru);
        user_1 = ru.ru_utime;
        sys_1 = ru.ru_stime;

        if (-1 == lseek(fd, where, SEEK_SET)) {
            fprintf(stderr, "rate: %s: lseek(%u): %s\n", what, (unsigned)where, strerror(errno));
            return 1;
        }

        while (done < howmuch) {
            ssize_t n;
            if (operation == op_read)
                n = read(fd, buffer, howmuch - done);
            else
                n = write(fd, buffer, howmuch - done);
            if (operation == op_writesync)
                fdatasync(fd);
            if (n == -1) {
                fprintf(stderr, "rate: %s: %s: %s\n", what, (operation == op_read ? "read" : "write"), strerror(errno));
                return 1;
            } else if (n == 0) {
                fprintf(stderr, "rate: %s: read: unexpected EOF\n", what);
                return 1;
            }
            done += n;
        }

        getrusage(RUSAGE_SELF, &ru);
        user_2 = ru.ru_utime;
        sys_2 = ru.ru_stime;
        gettimeofday(&real_2, NULL);

        /* Print results. */
        printf("%u %f %f %f\n", howmuch, timediff(&real_1, &real_2), timediff(&user_1, &user_2), timediff(&sys_1, &sys_2));
    }

    return 0;
}