#!/usr/bin/env python
#
# Copyright (c) 2008-2009 Benjamin Schweizer and others. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#   * Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#   * 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.
#   * Neither the name of Benjamin Schweizer nor the names of its
#     contributors may be used to endorse or promote products derived
#     from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
# TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#
# Abstract
# ~~~~~~~~
# Benchmark disk IOs
#
# Authors
# ~~~~~~~
# Benjamin Schweizer, http://benjamin-schweizer.de/contact
# Uwe Menges
#
# Changes
# ~~~~~~~
# 2009-09-16, uwe: changed formatting, fixed last block bug
# 2008-10-16, benjamin: initial release
#
# Todo
# ~~~~
# - we'll see
#

import sys
import random
import time

def usage():
    print """Copyright (c) 2008-2009 Benjamin Schweizer and others. All rights reserved.

usage:

    iostat <device> [time]

    device  := some block device
    time    := time in seconds

example:

    iostat /dev/sda

"""


def greek(size):
    """Return a string representing the greek/metric suffix of a size"""
    # Copyright (c) 1999 Martin Pohl, copied from
    # http://mail.python.org/pipermail/python-list/1999-December/018519.html
    _abbrevs = [
        (1<<50L, 'Pi'),
        (1<<40L, 'Ti'), 
        (1<<30L, 'Gi'), 
        (1<<20L, 'Mi'), 
        (1<<10L, 'Ki'),
        (1     , '  ')
    ]

    for factor, suffix in _abbrevs:
        if size >= factor:
            break

    return "%3.d %s" % (int(size/factor), suffix)


def iotest(fh, eof, blocksize=512, t=10):
    """io test"""

    io_num = 0
    start_ts = time.time()
    while time.time() < start_ts+t:
        io_num += 1
        pos = random.randint(0, eof - blocksize)
        fh.seek(pos)
        blockdata = fh.read(blocksize)
    end_ts = time.time()

    total_ts = end_ts - start_ts

    io_s = io_num/total_ts
    by_s = int(blocksize*io_num/total_ts)
    print " %sB blocks: %6.1f IOs/s, %sB/s" % (greek(blocksize), io_s, greek(by_s))

    return io_num/total_ts

if __name__ == '__main__':
    if len(sys.argv) < 2:
        usage()
        raise SystemExit

    dev = sys.argv[1]
    t = 10
    if len(sys.argv) == 3:
        t = int(sys.argv[2])

    blocksize = 512
    try:
        fh = open(dev, 'r')
        fh.seek(0,2)
        eof = fh.tell()

        print("%s, %sB:" % (dev, greek(eof)))

        iops = 2
        while iops > 1:
            iops = iotest(fh, eof, blocksize, t)
            blocksize *= 2
    except IOError, (err_no, err_str):
        raise SystemExit(err_str)
    except KeyboardInterrupt:
        print "caught ctrl-c, bye."

# eof.