heatmap_disk_console.py

heatmap_disk_console.py

import argparse
from subprocess import Popen, PIPE
import json
import sys
from operator import itemgetter
import time
import string
from signal import signal, SIGINT

COLOURS = [16, 17, 18, 19, 20, 21, 57, 56, 55, 54, 53, 52, 88, 124, 125, 160,
           161, 196, 197, 198, 199, 162, 163, 90, 91, 92, 93, 96, 97, 98, 99,
           132, 133, 134, 135, 140, 141, 177, 176, 175, 174, 173, 172, 215,
           214, 226, 227, 228, 229, 230, 231]

COLOUR_CODE = "\033[48;5;{}m"
COLOUR_DEFAULT = "\033[0m"


def handler(signal_received, frame):
    # Handle any cleanup here
    print('exiting')
    exit(0)


def get_iostat_disk_stats(old=False):
    if old:
        cmd = "iostat -xz |egrep 'sd|nvm' |awk '{print $1, $14}'"
    else:
        cmd = 'iostat -xz -o JSON'
    iostat = Popen(cmd, stdout=PIPE, shell=True)
    data = iostat.communicate()[0]
    try:
        if old:
            tmp_data = []
            for line in data.splitlines():
                line = line.strip()
                if " " not in line:
                    continue
                dev, util = line.split()
                tmp_data.append((dev, float(util)))
            data = tmp_data
        else:
            data = json.loads(data)
    except ValueError:
        sys.exit(1)
    if old:
        disks = [dict(disk_device=dev, util=util) for dev, util in data]
        result = {'host': 'old host', 'disks': sorted(disks,
                                                      key=itemgetter("disk_device"))}
    else:
        host = data['sysstat']['hosts'][0]
        result = {'host': host['nodename'],
                  'disks': sorted(host['statistics'][0]['disk'],
                                  key=itemgetter("disk_device"))}
    return result


def main(args):
    signal(SIGINT, handler)
    parser = argparse.ArgumentParser(
        description='Display ShardRange graphs')
    parser.add_argument('-i', '--interval', default=60, type=int,
                        help="Interval between checks")
    parser.add_argument('-k', '--key', default=False, action="store_true",
                        help="print heat map key")
    parser.add_argument('-o', '--old', default=False, action="store_true",
                        help="Enable old mode, when iostat doesn't support json")
    args = parser.parse_args(args)

    if args.key:
        print("Key:")
        for i in range(100):
            print("  [{} {}] {}%".format(COLOUR_CODE.format(COLOURS[i // 2]),
                                         COLOUR_DEFAULT, i))
        print('')

    print_title = False
    while True:
        data = get_iostat_disk_stats(args.old)
        if not print_title:
            print_title = True
            title = string.ascii_letters[:len(data['disks'])]
            for i, disk in enumerate(data['disks']):
                print(" {} = {}".format(title[i], disk['disk_device']))
            print('')
            print("|{}|".format(title))
            print("-" * (len(title) + 2))
        out = "|"
        for disk in data['disks']:
            util = disk['util']
            colour = COLOUR_CODE.format(COLOURS[int(util // 2)])
            out += "{} ".format(colour)
        out += COLOUR_DEFAULT + "|"
        print(out)
        time.sleep(args.interval)


if __name__ == '__main__':
    exit(main(sys.argv[1:]))