#!/usr/bin/python
#  Copyright (C) 2007 Yannick Gingras <ygingras@ygingras.net>

#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.

#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.


""" silly file system benchmark """

import md5
import sys
import os
from time import time
from random import sample, shuffle
from pprint import pprint
import beaker.util
from subprocess import call

NB_SEEK  = 2**14
SIZE = 2**10
DATA = chr(0)*SIZE

FS_TYPES = ["mkfs.ext3 -O ^dir_index",
            "mkfs.ext3 -O dir_index",
            "mkfs.reiserfs -f",
            "mkfs.xfs -f",
            ]

NB_FILES = [2**16, 2**18, 2**20]

def hexdigest(i):
    return md5.new(str(i)).hexdigest()

SAMPLES = dict([(nb, map(hexdigest, sample(range(nb), NB_SEEK)))
                for nb in NB_FILES])
for lst in SAMPLES.values():
    shuffle(lst)

def hashpath(dir, hash):
    return beaker.util.encoded_path(dir,
                                    [hash],
                                    digest=False,
                                    extension="")

def fill(dir, nb, mkpath):
    for i in xrange(nb):
        h = hexdigest(i)
        path = mkpath(dir, h)
        open(path, "w").write(DATA)


def seek(dir, nb, mkpath):
    for h in SAMPLES[nb]:
        path = mkpath(dir, h)
        open(path)


def seek_read(dir, nb, mkpath):
    for h in SAMPLES[nb]:
        path = mkpath(dir, h)
        assert open(path).read() == DATA


def delete(dir, nb, mkpath):
    for h in SAMPLES[nb]:
        path = mkpath(dir, h)
        os.unlink(path)


def create(dir, nb, mkpath):
    for h in SAMPLES[nb]:
        path = mkpath(dir, h)
        open(path, "w").write(DATA)


def cleanup(dir, nb, mkpath):
    for i in xrange(nb):
        h = hexdigest(i)
        path = mkpath(dir, h)
        os.unlink(path)

def fail(msg):
    print msg
    sys.exit(1)

def bench(dev, dir):
    stats = {}
    for nb in NB_FILES:
        for fs in FS_TYPES:
            for mkpath in [hashpath, os.path.join]:
                if call(fs.split(" ")+[dev]):
                   fail("can't make fs: %s" % fs)
                for test in [fill, seek, seek_read, delete, create]:
                    
                    call(["mount", dev, dir])
                    beg = time()
                    test(dir, nb, mkpath)
                    key = (nb,
                           fs,
                           mkpath.__name__,
                           test.__name__)
                    stats[key] = "%0.2f" % (time()-beg)
                    print "**", key, stats[key]

                    # delete the FS cache
                    call(["umount", dir])

    pprint(stats)

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print "USAGE: %s DEV MNT_POINT" % sys.argv[0]
        sys.exit(1)

    bench(sys.argv[1], sys.argv[2])