gpMgmt/bin/lib/multidd

#!/usr/bin/env python3
"""
multidd - run multiple dd concurrently to determine disk throughput

Usage: multidd {-i infile} {-o outfile} [-B blocksz] [-S filesz]

	-i infile	where infile specifies the path of input file
	-o outfile	where outfile specifies the path of output file
	-B blocksz	where blocksz is the size of each I/O block
			[default=8KB]
	-S filesz	where filesz is the size of the file to write
			[default=2X system RAM]

	Multiple -i and -o arguments may be specified but they must pair up.

	e.g. multidd -i /dev/zero -o /dbfast1/ddtest -B 8kb -S 2gb
"""

import getopt
import math
import os
import sys


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


def run(cmd):
    f = None
    ok = False
    out = None
    try:
        f = os.popen(cmd)
        out = f.read()
        ok = not f.close()
    except:
        f.close()
        ok = False
    return ok, out


def getPlatform():
    if sys.platform.find('linux') >= 0:
        return 'linux'
    if sys.platform.find('darwin') >= 0:
        return 'darwin'
    return '?'


def getMemory():
    if getPlatform() == 'linux':
        ok, out = run("sh -c 'cat /proc/meminfo | grep MemTotal'")
        if not ok:
            return '?'
        mem_string_list = out.strip().split(' ')
        val = int(mem_string_list[len(mem_string_list) - 2])
        factor = mem_string_list[len(mem_string_list) - 1]
        if factor == 'kB':
            return val * 1024
        return '?'

    if getPlatform() == 'darwin':
        ok, out = run("sysctl hw.physmem")
        if not ok:
            return '?'
        mem_string_list = out.strip().split(' ')
        val = int(mem_string_list[1])
        return val

    return '?'


def parseMemorySize(line):
    factor = 1
    try:
        line = line.strip().upper()
        if line.endswith('B'):
            line = line[:-1]
        if line.endswith('G'):
            factor = 1024 * 1024 * 1024
        elif line.endswith('M'):
            factor = 1024 * 1024
        elif line.endswith('K'):
            factor = 1024
        if factor > 1:
            line = line[:-1]
        return int(line) * factor
    except:
        return 0


def parseCommandLine():
    global opt
    try:
        (options, args) = getopt.getopt(sys.argv[1:], '?i:o:B:S:')
    except:
        e = sys.exc_info()
        usage('Error: %s %s' % (e[0], e[1]))

    for (switch, val) in options:
        if switch == '-?':
            usage(0)
        elif switch[1] in 'io':
            opt[switch].append(val)
        elif switch[1] in 'BS':
            opt[switch] = parseMemorySize(val)

    if opt['-S'] == '?':
        sys.exit('Error: unable to obtain system RAM size')
    if opt['-S'] < 0:
        usage('Error: invalid -S filesz parameter')
    if opt['-B'] <= 0:
        usage('Error: invalid -B blocksz parameter')
    if opt['-B'] > 1024 * 1024:
        usage('Error: maximum 1MB for -B parameter')
    if len(opt['-i']) != len(opt['-o']):
        usage('Error: -i and -o parameters must pair up')
    if len(opt['-i']) == 0:
        usage('Error: missing -i and -o parameters')


opt = {'-i': [], '-o': [], '-B': parseMemorySize('8KB'), '-S': 0}
parseCommandLine()
if opt['-S'] == 0:
    opt['-S'] = getMemory() * 2
if opt['-S'] < opt['-B']:
    opt['-S'] = opt['-B']
cmd = []
pfile = []
blocksz = opt['-B']
cnt = int(math.ceil(opt['-S'] / blocksz))
totalBytes = 0
for i in range(len(opt['-i'])):
    ifile = opt['-i'][i]
    ofile = opt['-o'][i]
    cmd.append('dd if=%s of=%s count=%d bs=%d' % (ifile, ofile, cnt, blocksz))
    totalBytes += cnt * blocksz

for c in cmd:
    print(c)
    pfile.append(os.popen(c))

for f in pfile:
    ok = False
    try:
        print(f.read())
        ok = not f.close()
        f = None
    except:
        if f:
            f.close()
        ok = False

    if not ok:
        sys.exit(1)

os.system('sync')
print('multidd total bytes ', totalBytes)