#define _LARGEFILE64_SOURCE

#include <stdio.h>
#include <fcntl.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/timeb.h>
#include <sys/types.h>
#include <linux/fs.h>

#define BLOCKSIZE   512
#define NUMBLOCKS   2048 /* 1MB */
#define READSIZE    (NUMBLOCKS * BLOCKSIZE)
#define READTIME    5 /* sec */
#define SEEKTIME    5 /* sec */


double timediff(struct timeb *before, struct timeb *after)
{
  return (double)(after->time - before->time)
    + (double)(after->millitm - before->millitm) / 1000.0;
}


off64_t determine_filesize(int fd, long long blocksize)
{
  off64_t endoffset;
  unsigned long size_in_blocks;

  if (0 == ioctl(fd, BLKGETSIZE, &size_in_blocks))
    return (long long)size_in_blocks * blocksize;
  else { // try seeking
    lseek64(fd, 0, SEEK_SET);
    endoffset = lseek64(fd, 0, SEEK_END);
    if (endoffset == 0)
      fprintf(stderr, "Cannot determine size of device: %s\n", 
	      strerror(errno));
    return endoffset;
  }
}


void readperf(int fd, char *scratch, int readsize)
{
  int readcount = 0;
  struct timeb before, after;
  double td;
  double megabytes;

  ftime(&before);
  do {
    size_t readbytes = read(fd, scratch, readsize);
    if (readbytes != readsize) {
      if (errno)
	fprintf(stderr, "Error reading data: %s\n",
		strerror(errno));
      else
	fprintf(stderr, "%ld bytes read instead of %ld.\n",
		readbytes, readsize);
      return;
    }
    readcount++;
    ftime(&after);
  } while (after.time - before.time < READTIME);

  megabytes = (double)(readcount * readsize) / 1048576.0;
  td        = timediff(&before, &after);
  printf("%6.2f MB/sec     ", megabytes / td);
  fflush(stdout);
}


void seekperf(int fd, char *scratch, long long blocksize)
{
  int seekcount=0;
  struct timeb before, after;
  double td_ms;
  off64_t endoffset;
  long long numblocks;

  endoffset = determine_filesize(fd, blocksize);
  if (!endoffset) {
    fprintf(stderr, "Cannot determine size of device/file.\n");
    return;
  }
  numblocks = endoffset / blocksize;

  ftime(&before);
  srandom(before.time + before.millitm); /* avoid cache */
  do {
    long long blocknr = (long long)(random()) % numblocks;
    lseek64(fd, blocknr * blocksize, SEEK_SET);
    read(fd, scratch, blocksize);
    seekcount++;
    ftime(&after);
  } while (after.time - before.time < SEEKTIME);

  td_ms     = timediff(&before, &after) * 1000.0;
  printf("%5.1f ms/seek\n", td_ms / seekcount);
}


void perform_test(char *devname, char *scratch)
{
  int fd;
  printf("%-17s ", devname);
  fflush(stdout);
  
  /* Open raw device for reading */
  fd = open(devname, O_RDONLY);
  if (fd < 0) {
    printf("Cannot open for reading: %s\n", strerror(errno));
    return;
  }
  readperf(fd, scratch, READSIZE);
  seekperf(fd, scratch, BLOCKSIZE);
  close(fd);
}


main(int argc, char **argv)
{
  char *scratch;
  int dev;
  
  if (argc < 2) {
    fprintf(stderr, "Usage: %s DEVICE1 DEVICE2 ...\n", argv[0]);
    exit(1);
  }

  /* allocate block-aligned memory for reading data into */
  if (posix_memalign(&scratch, BLOCKSIZE, READSIZE)) {
    fprintf(stderr, "Cannot allocate aligned memory: %s\n",
	    strerror(errno));
    exit(3);
  }

  printf("IO performance test version 1.0.0 "
	 "-- Copyright Mathias Kettner\n"
	 "This software is licensed under GPL "
	 "- use at your own risk.\n"
	 "\n"
	 "DEVICE              CONTINOUS        RANDOM SEEK\n"
	 "---------------------------------------------------\n");
  for (dev=1; dev<argc; dev++)
    perform_test(argv[dev], scratch);

  free(scratch);
}