#! /usr/bin/python
#
# advise-pcpus          Advise PCPUs to Dom0
#
# chkconfig: 2345 20 80
# description: Advise PCPUs to Dom0

# Eligible pcpus for dom0:
# 1. same socket and memory node
# 2. use node with more pcpus or the one with lowest node if more than one

import sys,getopt,re
import xcp.logger as logger
import xcp.cmd as cmd
from logging.handlers import SysLogHandler

dom0_pcpus_output="/var/lib/xen/dom0-pcpus"

def call(c):
    rv, out, err = cmd.runCmd(c, with_stdout=True, with_stderr=True)
    if rv!=0 or len(err)>0:
        msg="cmd=%s: rv=%s, stderr=%s" % (c,rv,err)
        logger.error(msg)
        raise RuntimeError(msg)
    return out.splitlines()

def get_pcpus():
    # xenpm provides the host pcpu topology by pcpu, core, socket, node
    # obtain list of host pcpus, skipping xenpm header line
    ps=call(["/usr/sbin/xenpm","get-cpu-topology"])[1:]
    host_pcpus = [ x.split() for x in ps ]
    #eg.: [['CPU0', '0', '0', '0'], ['CPU1', '1', '0', '0'], ['CPU2', '2', '0', '0'], ['CPU3', '3', '0', '0'], ['CPU4', '4', '0', '0'],...

    # unique socket and node groups, sorted by decreasing number of pcpus in each group
    soc_nodes = [ (int(x[2]), int(x[3])) for x in host_pcpus ]
    #eg.: [ (0, 0), (0, 0), ...
    pcpus_soc_node = [ (soc_nodes.count(sn), sn[0], sn[1]) for sn in set(soc_nodes) ]
    #eg.:[(16, 0, 0), (8, 1, 2), (8, 1, 0), (8, 2, 0), (8, 2, 4), (8, 3, 0), (8, 3, 6)]

    # sort by number of pcpus, smaller sockets and nodes first
    pcpus_soc_node.sort(key=lambda x: x[1])
    pcpus_soc_node.sort(key=lambda x: x[2])
    pcpus_soc_node.sort(key=lambda x: x[0], reverse=True)
    #eg.:[(16, 0, 0), ...

    # choose the (first) socket-node group, i.e.  with largest number of pcpus
    soc_node=pcpus_soc_node[0][1:3]
    #eg. '0 0'

    # pcpus within the chosen socket-node group
    host_pcpus_in_node = filter( lambda x: int(x[2]) == soc_node[0] and int(x[3]) == soc_node[1], host_pcpus )
    #eg.: [['CPU0', '0', '0', '0'], ['CPU1', '1', '0', '0'], ...
    pcpus = [ x[0][3:] for x in host_pcpus_in_node ]

    # make sure that no pcpu duplication is present
    pcpus_uq=sorted(set(pcpus), key=lambda x:int(x))

    return pcpus_uq

def success():
    print "[  OK  ]"

def failure():
    print "[FAILED]"

def start():
    logger.openLog(SysLogHandler(facility=SysLogHandler.LOG_DAEMON))
    print "Advising PCPUs to Dom0: " ,
    f = None
    try:
        try:
            pcpus = get_pcpus()
            # write the output persistently for later use
            f = open(dom0_pcpus_output,"w")
            f.write("\n".join(pcpus))
            f.close()

            # print some information for the console user
            print " ".join(pcpus) ,
            success()
        except Exception, e:
            failure()
            logger.logException(e)
    finally:
        if f:
            f.close()

def usage():
    print "Usage: %s {start|stop}" % sys.argv[0]

def main():
    if len(sys.argv)>1:
        if sys.argv[1]=="start":
            start()
        elif sys.argv[1]=="stop":
            pass
        elif sys.argv[1]=="status":
            pass
        else:
            usage()
    else:
        usage()

if __name__ == "__main__":
    main()

