/* iorate.c - measure rates of sequential IO, showing incremental 
   bandwidth written by Mark Hahn (hahn@mcmaster.ca) 2003-2008
   the main point of this code is to illustrate the danger of 
   running naive bandwidth tests on files that are small relative
   to the memory/disk bandwidth ratio of your system.  that is,
   on any system, the incremental bandwidth will start out huge,
   since IO is purely to the page cache.  once you exceed that size,
   bandwidth will be dominated by the real disk performance.  but 
   using the average of these two modes is a mistake, even if 
   you use very large files.
*/
#define _LARGEFILE64_SOURCE 1
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/time.h>
#include <sys/fcntl.h>
#include <sys/stat.h>
#include <stdarg.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/resource.h>

#ifdef O_LARGEFILE
#define LF O_LARGEFILE
#elif defined(_O_LARGEFILE)
#define LF _O_LARGEFILE
#else
#define LF 0
#endif

#ifndef O_DIRECT
#define O_DIRECT 040000
#endif

typedef unsigned long long u64;

u64 bytes = 0, bytesLast = 0;
double timeStart = 0, timeLast = 0;
FILE *fp_results;
/* whether time is printed as elapsed from start of test or absolute. */
int absolute_time = 0;

/* default reporting interval is every 2 seconds;
   in 2004, an entry-level desktop disk will sustain around 50 MB/s,
   so the default bytes interval is 100 MB.  whichever comes first. */
u64 byteInterval = 100<<20;
double timeInterval = 2;

static inline double tvToSeconds(const struct timeval *tv) {
    return tv->tv_sec + 1e-6 * tv->tv_usec;
}

double gtod() {
    struct timeval tv;
    gettimeofday(&tv,0);
    return tvToSeconds(&tv);
}

void dumpstats(int force) {
    u64 db = bytes - bytesLast;
    double now = gtod();
    double dt;
    static int first = 1;
    struct rusage ru;
    double u,s,e;
    char elapsed[100];

    if (absolute_time) {
	snprintf(elapsed,sizeof(elapsed),"%.6f",now);
    } else {
	snprintf(elapsed,sizeof(elapsed),"%7.3f",now - timeStart);
    }

    getrusage(RUSAGE_SELF, &ru);
    u = tvToSeconds(&ru.ru_utime);
    s = tvToSeconds(&ru.ru_stime);

    if (timeLast == 0) timeStart = timeLast = now;

    dt = now - timeLast;

    if (!force && db < byteInterval && dt < timeInterval) return;

    e = now - timeStart;

    if (first) {
	fprintf(fp_results,"#%7s %7s %7s %7s\n", 
	       "secs",
	       "MB",
	       "MB/sec",
	       "MB/sec");
	first = 0;
    }

    fprintf(fp_results,"%s %7.3f %7.3f %7.3f\n", 
		elapsed,
		1e-6 * bytes, 
		1e-6 * db / dt,
		1e-6 * bytes / e);
    timeLast = now;
    bytesLast = bytes;
    if (force) {
	fprintf(fp_results,
		"%7.3f user + %7.3f system = %7.3f elapsed (%7.3f %)\n",
		u,s,e,100*(u+s)/e);
    }
}

void usage() {
    fprintf(stderr,"iorate [-r/w filename] [-d] [-c chunksz][-b byteivl][-t ivl][-l szlim] [-r in] [-w out]\n");
    fprintf(stderr,"-r in or -w out select which file is read or written ('-' for stdin/out)\n");
    fprintf(stderr,"-c chunksz - size of chunks written (KB);\n");
    fprintf(stderr,"-t timeinterval - collect rate each timeinterval seconds;\n");
    fprintf(stderr,"-b byteinterval - collect rate each byteinterval MB;\n");
    fprintf(stderr,"-l limit - total output size limit (GB);\n");
    fprintf(stderr,"-d use O_DIRECT\n");
    fprintf(stderr,"defaults are: '-c 8 -b 20 -t 10 -l 10'\n");
    exit(1);
}

void fatal(char *format, ...) {
    va_list ap;
    va_start(ap,format);

    vfprintf(stderr,format,ap);
    fprintf(stderr,": errno=%d (%s)\n",errno,strerror(errno));
    va_end(ap);
    dumpstats(1);
    exit(1);
}

/* allocate a buffer using mmap to ensure it's page-aligned.  
   O_DIRECT *could* be more strict than that, but probably isn't */
void *myalloc(unsigned size) {
    unsigned s = (size + 4095) & ~4095U;
    void *p = mmap(0, 
                   s, 
                   PROT_READ|PROT_WRITE, 
                   MAP_ANONYMOUS|MAP_PRIVATE, 
                   -1, 0);
    if (p == MAP_FAILED)
        return 0;
    return p;
}

u64 interpret_size(char *string) {
    u64 v;
    if (strcmp(string,"-1")==0) {
	return (u64)-1;
    }
    char suffix;
    sscanf(string,"%lu%c",&v,&suffix);
//    printf("v(%lu) c(%c)\n",v,suffix);

    switch(suffix) {
    case 'g': case 'G': return v*1024*1024*1024;
    case 'm': case 'M': return v*1024*1024;
    case 'k': case 'K': return v*1024;
    default: return v;
    }
}

void t(char *s) {
    printf("%s -> %Lu\n",s,interpret_size(s));
}


int main(int argc, char *argv[]) {
   unsigned size = 8*1024;
    char *buffer;
    char *fnameI = 0;
    char *fnameO = 0;
    char *fnameL = 0;
    int fdI = 0;
    int fdO = 1;
    int doRead = 0;
    int doWrite = 0;
    int doDirect = 0;
    u64 limit_bytes = 1LU << 20;

    int letter;
    while ((letter = getopt(argc,argv,"ar:w:l:b:c:x:t:d")) != -1) {
	switch(letter) {
	case 'a':
	    absolute_time = 1;
	    break;
	case 'r':
	    fnameI = optarg;
	    doRead = 1;
	    break;
	case 'w':
	    fnameO = optarg;
	    doWrite = 1;
	    break;
	case 'l':
	    fnameL = optarg;
	    break;

	case 'x':
	    limit_bytes = interpret_size(optarg);
	    break;

	case 'b':
	    byteInterval = interpret_size(optarg);
	    break;

	case 'c':
	    size = interpret_size(optarg);
	    break;

	case 't':
	    printf("optarg -t %s\n",optarg);
	    timeInterval = atof(optarg);
	    break;

	case 'd':
	    doDirect = 1;
	    break;
	default:
	    usage();
	}
    }
    
    if (argc != optind) usage();
    if (!byteInterval) byteInterval = size;

    printf("size(%lu) limit(%lu) byteinterval(%lu)\n",
	   size,limit_bytes,byteInterval);

    if (doRead && fnameI && strcmp(fnameI,"-")) {
	int flags = O_RDONLY | LF;
	if (doDirect) flags |= O_DIRECT;

	fdI = open(fnameI, flags);
	if (fdI == -1) fatal("open(read) failed");
    }
    if (doWrite && fnameO && strcmp(fnameO,"-")) {
	int flags = O_RDWR | O_CREAT | LF;
	if (doDirect) flags |= O_DIRECT;

	fdO = open(fnameO, flags, 0600);
	if (fdO == -1) fatal("open(write) failed");
    }

    fp_results = stderr;
    if (fnameL) {
	fp_results = fopen(fnameL,"w+");
	if (!fp_results) fatal("fopen(%s) failed",fnameL);
    }

    fprintf(fp_results,"chunk %luK, byteInterval %uM, timeInterval %f, limit %uM\n",
	    size>>10,
	    (unsigned)(byteInterval>>20),
	    timeInterval,
	    (unsigned)(limit_bytes >>20));
    setvbuf(fp_results, 0, _IOLBF, 0);

    buffer = myalloc(size);
    memset(buffer,'m',size);

    timeStart = timeLast = gtod();

    bytes = 0;

    while (bytes < limit_bytes) {
	int c = size;

    	dumpstats(0);

	if (doRead) {
	    c = read(fdI,buffer,c);	
	    if (c == -1) fatal("read failed");
	}

	if (doWrite) {
	    c = write(fdO,buffer,c);
	    if (c == -1) fatal("write failed");
	}

	bytes += c;

	/* short read/write means EOF. */
	if (c == 0) break;
    }

    dumpstats(1);

    return 0;
}