mdbench.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
#
# Copyright (C) 2013-2020 Deutsches Elektronen-Synchroton,
# Member of the Helmholtz Association, (DESY), HAMBURG, GERMANY
#
# 2020 added median calculation johannes.veh@fau.de
#
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2 or any later version.
#

'''
Simple filesystem metadata operations benchmark

Usage: mdbench [options] <PATH>

  where options are:
    -f, --files <N>  : number of generated files per directory
    -d, --dirs  <N>  : number of generated directories to generate
    -s, --size  <N>  : size of generated files in B/K/M/G
    -n, --no-clean   : do not delete created files and directories
    --no-container   : do not create the 'mdbench.<name>.<pid>' directory
    -h, --help       : help message

  and PATH points to the directory where tests should run. Current directory
  is used if not specified.

If number of generated directories is 0 then files are created within
the container directory ("mdbench.<name>.<pid>") or in PATH if
--no-container is specified.

The file size can be specified in human friendly format, e.g.: 1K, 256M. 4G.
'''

from __future__ import division
import sys
import os
import socket
import getopt
import string
from datetime import datetime
import math
from heapq import heappush, heappop

DIR = 'dir.'
FILE = 'file.'

gen_dir = lambda base, gen : '%s/%s%d' % (base, DIR, gen)
gen_file = lambda base, gen : '%s/%s%d' % (base, FILE, gen)

B = 1
K = 1024
M = K*K
G = K*M

DATA_SIZES = {'b':B, 'k': K, 'm': M, 'g': G}


class MovingAvg:
	"""
	Calculate moving average.
	"""

	def __init__(self):
		self._avg = 0.0
		self._count = 0
		self._sigma = 0.0
		self._total = 0
		self.lowers, self.highers = [],[]
	def avg(self):
		"""
		Return the current value of average.
		"""
		return self._avg

	def std(self):
		"""
		Returns the the current value of standard deviation.
		"""
		return math.sqrt(self._sigma - self._avg**2)

	def update(self, v):
		"""
		Update average with new value.
		"""
		self._avg = (self._avg * self._count + v) / (self._count + 1)
		self._sigma = (self._sigma * self._count + v**2) / (self._count + 1)
		self._count += 1
		self._total += v
		if not self.highers or v > self.highers[0]:
            		heappush(self.highers, v)
		else:
            		heappush(self.lowers, -v)  # for lowers we need a max heap
		self.rebalance()

	def rebalance(self):
		if len(self.lowers) - len(self.highers) > 1:
			heappush(self.highers, -heappop(self.lowers))
		elif len(self.highers) - len(self.lowers) > 1:
			heappush(self.lowers, -heappop(self.highers))

	def median(self):
		if len(self.lowers) == len(self.highers):
			return (-self.lowers[0] + self.highers[0])/2
		elif len(self.lowers) > len(self.highers):
			return -self.lowers[0]
		else:
			return self.highers[0]

	def sum(self):
		"""
		Return the sum of all values.
		"""
		return self._total

	def count(self):
		"""
		Return the total number up updates.
		"""
		return self._count


dir_creates = MovingAvg()
file_creates = MovingAvg()
file_stats = MovingAvg()
dir_stats = MovingAvg()
dir_removes = MovingAvg()
file_removes = MovingAvg()

def get_size(s):

	last_symbol = s[-1:].lower()
	if last_symbol in string.digits:
		return int(s)

	if last_symbol not in DATA_SIZES:
		raise Exception('Invalid format: %s' % s)

	return int(s[:-1])*DATA_SIZES[last_symbol]

def make_dirs(root, count):
	for i in range(count):
		mkdir( gen_dir(root, i) )

def make_files(root, dir_count, file_count, size = 0):

	data = bytearray(size)

	for j in range(file_count):
		if dir_count > 0:
			for i in range(dir_count):
				mkfile(gen_file( gen_dir(root, i), j ), data, 1024)
		else:
			mkfile(gen_file(root, j), data, 1024)

def del_files(root, dir_count, file_count):
	for j in range(file_count):
		if dir_count > 0:
			for i in range(dir_count):
				rmfile(gen_file( gen_dir(root, i), j ))
		else:
			rmfile(gen_file(root, j))

def del_dirs(root, count):
	for i in range(count):
		rmdir( gen_dir(root, i) )

def stat_dirs(root, count):
	for i in range(count):
		statdir( gen_dir(root, i) )

def stat_files(root, dir_count, file_count):
	for j in range(file_count):
		if dir_count > 0:
			for i in range(dir_count):
				statfile(gen_file( gen_dir(root, i), j ))
		else:
			statfile(gen_file(root, j))

def rmfile(f):
	start = datetime.now()
	os.remove(f)
	end = datetime.now()
	file_removes.update(total_millis(end - start))

def rmdir(d):
	start = datetime.now()
	os.rmdir(d)
	end = datetime.now()
	dir_removes.update(total_millis(end - start))

def mkdir(d):
	start = datetime.now()
	os.mkdir(d)
	end = datetime.now()
	dir_creates.update(total_millis(end - start))

def statfile(f):
	start = datetime.now()
	os.stat(f)
	end = datetime.now()
	file_stats.update(total_millis(end - start))

def statdir(f):
	start = datetime.now()
	os.stat(f)
	end = datetime.now()
	dir_stats.update(total_millis(end - start))

def mkfile(fname, data, chunk = 65536, sync = False) :

	off = 0
	remaining = len(data)

	start = datetime.now()
	with open(fname, "wb") as f:

		while remaining > 0:
			wsize = min(chunk, remaining)
			f.write(data[off:off+wsize])
			off += wsize
			remaining -= wsize

		if sync:
			f.flush()
			os.fsync(f.fileno())

	end = datetime.now()
	file_creates.update(total_millis(end - start))

def total_micros(td):
	return (td.microseconds + (td.seconds + td.days * 24 * 3600) * 10**6)

def total_millis(td):
	return total_micros(td)/1000

def report(title, counter):
	print('{:16}: {:6.3f}ms ±{:=6.3f}ms, median {:=6.3f}ms {:6.3f} op/s' \
		.format(title, counter.avg(), counter.std(), counter.median() ,counter.count()/counter.sum()*10**3))

DIR_COUNT = 1000
FILE_COUNT = 10
FILE_SIZE = 0

def usage():
	print(__doc__)
	sys.exit(1)

def main():

	dir_count = DIR_COUNT
	file_count = FILE_COUNT
	file_size = FILE_SIZE
	cleanup = True
	createContainer = True

	try:
		options, remainder = getopt.gnu_getopt(sys.argv[1:], 'f:d:s:nh', \
					 ['files=','dirs=','size=','no-clean','no-container','help'])
	except getopt.GetoptError as err:
		print(str(err))
		usage()

	for opt, arg in options:
		if opt in ('-f', '--files'):
			file_count = int(arg)
		elif opt in ('-d', '--dirs'):
			dir_count = int(arg)
		elif opt in ('-s', '--size'):
			file_size = get_size(arg)
		elif opt in ('-n', '--no-clean'):
			cleanup = False
		elif opt in ('--no-container'):
			createContainer = False
		elif opt in ('-h', '--help'):
			usage()

	if len(remainder) != 1 :
		usage()

	path = remainder[0]

	root = '%s/mdbench.%s.%d' % (path, socket.gethostname(), os.getpid()) if createContainer else path

	if createContainer:
		os.mkdir(root)

	make_dirs(root, dir_count)
	report("dir creates", dir_creates)
	make_files(root, dir_count, file_count , file_size)
	report("file creates", file_creates)
	stat_files(root, dir_count, file_count)
	report("file stats", file_stats)
	stat_dirs(root, dir_count)
	report("dir stats", dir_stats)

	if cleanup:
		del_files(root, dir_count, file_count )
		report("file removes", file_removes)
		del_dirs(root, dir_count )
		report("dir removes", dir_removes)
		if createContainer:
			os.rmdir(root)

if __name__ == '__main__':
	main()