#!/usr/bin/env python
#
# isi_rbm_ping
#
# by Jon Lasser <jon.lasser@emc.com>
#
# Script to ping all 'alive' nodes in a cluster over the back end and 
# report back the times in a pretty grid. (Pretty is dependent upon the
# number of nodes in the cluster and your terminal width...)

import getopt
import os
import sys
import re
import subprocess
import csv
from threading import Thread
from Queue import Queue, Empty

import isi.config.array as array
import isi.fs.restripe.group as group
import isi.misc as misc

# Some internal variables - not exposing these for configuration but we
# might need to do so in the future, and at any rate they should be easy to
# find and change.
rdoCmd = '/usr/bin/isi_rdo'
pythonCmd = '/usr/bin/python'
allHandlers = [ 'Non-threaded handler', 'Threaded handler' ]
allSizes = [ '1 byte msg', '8k msg' ]
whichHandlers = [ 'Threaded handler' ]

# Actual config options - exposed via CLI
# Keep this list in same order as displayed in usage, below
config = {
    'useArrayID': False,
    'outputCSV': False,
    'pingCount': 100,
    'suppressDown': True,
    'localMode': False,
    'explicitNodeList': False,
    'whichSizes': [ '8k msg' ],
    'printHeader': True,
    'useAverage': True,
    'maxWorkers': 16
    }

def usage(exitCode):
    sys.stderr.write(r'''
Usage: ''' + os.path.basename(sys.argv[0]) + r''' [OPTION]... NODE

Print a table of back-end ping times between nodes, with times reported in
microseconds. By default, print the average ping time from source to target,
measured across 100 pings, excluding nodes marked down within the group.

Options:
  -i, --array-id            Use Array ID instead of LNN
  -C, --csv                 Produce CSV output
  -c, --count <number>      Perform <number> pings per target (default 100)
  -d, --down                Try to ping down (not-in-group) nodes
  -l, --local               Execute only on the local node
  -n, --nodes               Specify nodes to ping on the command line
  -o, --one-byte            Use 1 byte pings (default 8k)
  -s, --suppress-header     Do not print the header line
  -t, --total               Report total (rather than average) ping times
  -w, --workers <number>    Maximum workers to execute in parallel
  -h, --help                Print this message and exit
''')
    sys.exit(exitCode)

########################################################################
########################################################################
########################################################################
########################################################################
#
# It's so simple, so very simple, that even a manager can write it...
#
def main():
   
    # Parse our configuration, update our global variable blob
    global config
    config, args = parse_options(sys.argv[1:], config)

    # Get a list of nodes to ping
    pingList = nodes_to_ping(args)

    # If we're just asked to return results for one node, do that.
    # Otherwise, return results for all nodes.
    if config['localMode']:
        pingResults = line_of_pings(pingList)
    else:
        pingResults = do_remote_pings(pingList)

    # Print those results
    print_results(pingList, pingResults)

    # Exit cleanly
    sys.exit(0)
########################################################################
########################################################################
########################################################################
########################################################################

# Parse the CLI options and return a nice config blob.
# Keep the options in the same order as the usage message, above.
def parse_options(args,config):

    try:
        opts, args = getopt.getopt(args, 'iCc:dlnostw:h',
            [ 'array-id', 'CSV', 'count=', 'down', 'local', 'nodes',
              'one-byte', 'suppress-header', 'total', 'workers=', 'help' ])
        for o, a in opts:
            if o in [ '-i', '--array-id' ]:
                config['useArrayID'] = True
            if o in [ '-C', '--csv']:
                config['outputCSV'] = True
            if o in [ '-c', '--count' ]: 
                try:
                    config['pingCount'] = int(a)
                except ValueError:
                    print >>sys.stderr, "Ping count must be an integer."
                    print >>sys.stderr, "Provided value:", a
                    sys.exit(1)
                if config['pingCount'] < 1:
                    print >>sys.stderr, "Ping count must be positive!"
                    print >>sys.stderr, "Provided value:", a
                    sys.exit(1)
            if o in [ '-d', '--down' ]:
                config['suppressDown'] = False
            if o in [ '-l', '--local' ]:
                config['localMode'] = True
            if o in [ '-n', '--nodes' ]:
                config['explicitNodeList'] = True
            if o in [ '-o', '--one-byte' ]:
                config['whichSizes'] = [ '1 byte msg' ]
            if o in [ '-s', '--suppress-header' ]:
                config['printHeader'] = False
            if o in [ '-t', '--total' ]:
                config['useAverage'] = False
            if o in [ '-w', '--workers' ]:
                try:
                    config['maxWorkers'] = int(a)
                except ValueError:
                    print >>sys.stderr, "Worker count must be an integer."
                    print >>sys.stderr, "Provided value:", a
                    sys.exit(1)
                if config['maxWorkers'] < 1:
                    print >>sys.stderr, "Worker count must be positive!"
                    print >>sys.stderr, "Provided value:", a
                    sys.exit(1)
            if o in [ '-h', '--help' ]:
                usage(0)
    except getopt.GetoptError:
        usage(1)

    return (config, args)

def nodes_to_ping(args):

    # Identify all nodes, and down nodes
    gi = group.GroupInfo()
    allNodes = gi.getNodeList()
    if not config['suppressDown']:
        downNodes = gi.getDownNodeList()
        allNodes = allNodes + downNodes

    # If we have an explicit node list, use that; otherwise, build one
    # by excluding down nodes from the list of all nodes.
    if config['explicitNodeList']:
        if args:
            try:
                nodeList = [ int(x) for x in args ]
            except ValueError:
                print >>sys.stderr, "Non-integer provided as node ID."
                print >>sys.stderr, "Args: ", args
                sys.exit(1)
        else:
            print >>sys.stderr, "Empty node list!"
            sys.exit(1)
    else:
        if args:
            print >>sys.stderr, "WARNING: node list provided but -n " \
                                "not specified."
            print >>sys.stderr, "Provided node list: ", args
            sys.exit(1)
        nodeList = []
        for thisNode in allNodes:
            if config['useArrayID']:
                nodeList.append(thisNode)
            else:
                nodeList.append(gi.getLnn(thisNode))

    # Build key order list - if we're using an explicit node list, use the
    # given order instead
    if not config['explicitNodeList']: nodeList.sort()

    return nodeList


# Queue up workers, take their output, and collate the results
def do_remote_pings(pingList):
    workQueue = Queue()
    resultsQueue = Queue()
    unsortedResults = {}
    finalResults = []

    # Queue up remote pings to perform
    for each in pingList: workQueue.put((each, pingList))
    
    # Spin up threads, launch 'em, and wait for the pings to complete
    workers = [ Thread(target=remote_ping_worker,
        args=(workQueue,resultsQueue)) for i in range(config['maxWorkers']) ]
    for w in workers:
        try:
            w.start()
        # The except clause doesn't work well, but sometimes helps
        except (KeyboardInterrupt, SystemExit):
            print >>sys.stderr, "Interrupted at keyboard"
            sys.exit(1)
    for w in workers:
        w.join()

    # Take our queued results and order them appropriately
    while True:
        try:
            (node, results) = resultsQueue.get(block=False)
            unsortedResults[node] = results
        except Empty:
            break
    allKeys = unsortedResults.keys()
    allKeys.sort()
    for each in allKeys:
        finalResults.append([str(each)] + unsortedResults[each])

    return finalResults

# Multithreading wrapper for pings
def remote_ping_worker(workItems,results):

    while True:
        try:
            (node, toPing) = workItems.get(block=False)
            pingResults = remote_ping(node, toPing)
            results.put((node, pingResults))
        # The except clause doesn't work well, but sometimes helps
        except (KeyboardInterrupt, SystemExit):
            print >>sys.stderr, "Interrupted at keyboard"
            sys.exit(1)
        except Empty:
            break
    return

# Connect to a remote node and run isi_internal_ping from there, using
# isi_rdo and a provided list of nodes.
def remote_ping(node, toPing):

    ar = array.getArray()

    thisCmd = sys.argv[0]

    flags = []
    flags.append('-l')  # Remote node execution always local mode
    flags.append('-s')  # with suppressed headers and
    flags.append('-n')  # an explicit node list
    flags.append('-C')  # results in CSV
    flags.append('-c')
    flags.append(str(config['pingCount']))
    if not config['useAverage']: flags.append('-t')
    if config['useArrayID']: flags.append('-i')
    if not config['suppressDown']: flags.append('-d')

    addNodes = [ str(x) for x in toPing ]

    commandHeader = [ pythonCmd, thisCmd ] 

    if config['useArrayID']:
        nodeName = [ str(ar.nodes.findByDevid(int(node)).name()) ]
    else:
        nodeName = [ str(ar.nodes.findByLnn(int(node)).name()) ]

    args = [ rdoCmd ] + nodeName + commandHeader + flags + addNodes

    proc = subprocess.Popen(args, stdout=subprocess.PIPE)
    (out, err) = proc.communicate()
    cookedOut = out.strip().split(',')
    return cookedOut

# Ping each node and get the results
def line_of_pings(nodeList):

    pingResults = {}
    ar = array.getArray()
    gi = group.GroupInfo()

    for thisNode in nodeList:
        if config['useArrayID']:
            nodeNumber = thisNode
        else:
            nodeNumber = ar.nodes.findByLnn(thisNode).id()
        pingResults[thisNode] = node_ping_time(nodeNumber)

    # assemble and return results
    outputLine = [ ]
    for eachResult in nodeList:
        outputLine.append(str(pingResults[eachResult]))

    return outputLine


# Build regex based on variables, so we can compile and use once.
def build_regex(handlers, sizes):

    # Match any enabled handler or size
    handler_re_frag = '|'.join(handlers)
    size_re_frag = '|'.join(config['whichSizes'])
    lineHeader = handler_re_frag + ', ' + size_re_frag

    if config['useAverage']:
        resultPattern = '.*\(([0-9]+) usec per RTT'
    else:
        resultPattern = '.*RTTs in ([0-9]+) msec and ([0-9]+) usec'
    
    fullRegex = lineHeader + resultPattern
    return re.compile(fullRegex)

# Return a single node's ping time in usec, getting from the test thread
def node_ping_time(targetNode):

    resultRegex = build_regex(whichHandlers, config['whichSizes'])
    pingThreadBase='rbm_latency('
    pingDetails = str(targetNode) + ',' + str(config['pingCount']) + ')'
    pingThread = pingThreadBase + pingDetails
    pingResults =  misc.test_thread(pingThread)

    if pingResults[0] != 0: return('ERROR')

    # parse the annoying output - include only handler/size pairs we
    # want to include
    for each_element in pingResults[1].split('\n'):
        results = re.match(resultRegex, each_element)
        if results:
            if config['useAverage']:
                pingTime = results.group(1)
            else:
                msec = int(results.group(1))
                usec = int(results.group(2))
                pingTime = (1000 * msec) + usec
            return str(pingTime)

# Print the results output.
def print_results(nodes, results):

    # Define our header columns
    header = [ str(i) for i in nodes ]
    if not config['localMode']: header.insert(0, 'Source')

    # Print the output
    if config['outputCSV']:
        writer=csv.writer(sys.stdout)
        if config['printHeader']: writer.writerow(header)
        if config['localMode']:
            writer.writerow(results)
        else:
            for each in results: writer.writerow(each)
    else:               # not CSV
        if config['printHeader']: print ' \t'.join(header)
        if config['localMode']:
            print '\t'.join(results)
        else:
            for each in results: print '\t'.join(each)

    return

if __name__ == '__main__':
    main()
