#! /usr/bin/python
#
# tune-vcpus          Pin and change number of Dom0 VCPUs on recommended PCPUs
#
# chkconfig: 2345 24 76
# description: Pin and change number of Dom0 VCPUs on recommended PCPUs
#
# Three pinning policies:
# P.A. nopin: no pinning at all
#   - legacy behaviour mode
# P.B. pin: pin all the time
#   - use recommended pcpus in /var/lib/xen/dom0-pcpus
#   - change xen cmdline to pin dom0 before dom0 starts
# P.C. detect: pin if necessary (default policy)
#   - start daemon that monitors dom0 activity, pinning when necessary according to activity
#   - the activity is based on the number n of VMs started in the last s seconds
#
# Three dynamic vcpu-plugging policies:
# D.A. min: constant minimum number of dom0 vcpus
#   - legacy behaviour mode, minimum number is determined by /etc/sysconfig/unplug-vcpus
# D.B. max: constant maximum number of dom0 vcpus
#   - maximum number of vcpus is determined by advise-pcpus daemon
# D.C. dynamic: dynamic number of dom0 vcpus according to cpu demand
#   - maximum vcpus in the xen cmd line as determined by advise-pcpus daemon
#   - starting with existing number of vcpus in dom0

import sys,os,signal,fcntl,errno
import getopt,re
from xcp.logger import LOG
from xcp.logger import logToSyslog
import xcp.cmd as cmd
import logging

dom0_pcpus_input = "/var/lib/xen/dom0-pcpus"
policy_file = "/var/lib/xen/tune-vcpus-policy"
pid_file = "/var/run/tune-vcpus.pid"
config_file = "/etc/sysconfig/tune-vcpus"
xl = "/usr/sbin/xl"

def logerr(msg):
    LOG.error(msg)

def logdbg(msg):
    LOG.debug(msg)

def loginfo(msg):
    LOG.info(msg)

def logexc(e):
    logerr(e)

def stdout(msg):
    sys.stdout.write(msg)

def call(c,dbg=True):
    if dbg: loginfo("call: %s" % (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)
        logerr(msg)
        raise RuntimeError(msg)
    return out.splitlines()

def success():
    stdout("[  OK  ]\n")

def failure():
    stdout("[FAILED]\n")

def get_dom0_vcpus():
    lines = call([xl,"vcpu-list","0"])[1:]
    dom0_recs = [ x.split() for x in lines ]
    dom0_vcpus = sorted(set( int(x[2]) for x in dom0_recs ))
    return dom0_vcpus

def unpin_vcpus(verbose = True):
    if verbose: stdout("...")
    #unpin all dom0 vcpus
    try:
        call([xl,"vcpu-pin","0","all","all"])
    except Exception, e:
        logdbg("unpin_vcpus: %s" % (e,))

def pin_vcpus(verbose = True):
    if verbose: stdout("...")
    #to avoid pinning more than one vcpu to a pcpu, unpin before pinning
    unpin_vcpus(False)

    #pin each dom0 vcpu to a pcpu in /var/lib/xen/dom0-pcpus
    f = None
    try:
        try:
            f=open(dom0_pcpus_input)
            #make sure to use only unique pcpu ids, because pinning
            #more than one vcpu to a pcpu can lock up the dom0 kernel
            dom0_pcpus = sorted(set( int(x) for x in f.readlines() ))

            #create a vcpu-pcpu pair list:
            #1.use pcpus in the order they appear in DOM0_PCPUS
            #2.use vcpus in the order the appear in xl vcpu-list
            #3.only return a pcpu if it can be paired to an existing vcpu
            #4.also pin vcpus that are not enabled (but can be in the future)
            dom0_vcpus = get_dom0_vcpus()
            for vp, pp in zip(dom0_vcpus, dom0_pcpus):
                call([xl, "vcpu-pin", "0", str(vp), str(pp)])
        except Exception, e:
            loginfo("opening file %s: error %s" % (dom0_pcpus_input, e))
            logexc(e)
    finally:
        if f:
            f.close()

def unpin_xencmdline(verbose = True):
    if verbose: stdout("(xen boot) ")
    # unpin in xencmd line
    call(["/opt/xensource/libexec/xen-cmdline","--delete-xen","dom0_vcpus_pin"])

def pin_xencmdline(verbose = True):
    if verbose: stdout("(xen boot) ")
    # pin in xen cmdline
    call(["/opt/xensource/libexec/xen-cmdline","--set-xen","dom0_vcpus_pin"])

def save(filename,content):
    f = None
    try:
        try:
            f=open(filename, "w")
            f.write(content)
            loginfo("write '%s' to %s" % (" ".join(content.splitlines()), (filename,)))
        except Exception, e:
            logerr("writing file %s: error %s" % (filename, e))
    finally:
        if f:
            f.close()

def save_policy(policy):
    save(policy_file, policy)

def load(filename):
    r = None
    f = None
    try:
        try:
            f=open(filename, "r")
            r = f.readlines()
        except Exception, e:
            logerr("reading file %s: error %s" % (filename, e))
    finally:
        if f:
            f.close()
    return r

def rm(filename):
    try:
        os.unlink(filename)
    except Exception, e:
        logexc(e)

pid_file_lock_fd = None

def _lock_op(flags):
    global pid_file_lock_fd
    is_locked = False
    if not pid_file_lock_fd:
        try:
            # needs to be 'a' to avoid overwriting the contents
            pid_file_lock_fd = open(pid_file, "a") 
        except Exception, e:
            if e.errno == errno.ENOENT: 
                # file_not_found, nothing to lock
                pass
            else:
                logexc(e)
    if pid_file_lock_fd:
        try:
            fcntl.lockf(pid_file_lock_fd, flags)
        except IOError, e:
            logexc(e)
            if e.errno == errno.EACCES or e.errno == errno.EAGAIN:
                is_locked = True
    return is_locked

def unlock_pid():
    return _lock_op(fcntl.LOCK_UN)

def is_pid_locked():
    is_locked = _lock_op(fcntl.LOCK_EX | fcntl.LOCK_NB)
    # in this test, we do not want the lock to continue for this process
    unlock_pid()
    return is_locked

def lock_pid():
    if _lock_op(fcntl.LOCK_EX | fcntl.LOCK_NB):
        raise RuntimeError("%s is locked" % (pid_file,))

def killproc(pid):
    os.kill(pid,signal.SIGTERM)

def stop():
    try:
        pid = load(pid_file)
        stdout("Stopping Dom0 VCPU pinning daemon...")
        if pid and len(pid)>0:
            killproc(int(pid[0]))
            unlock_pid()
            rm(pid_file)
        success()
    except Exception, e:
        logexc(e)
        failure()

def status():
    stdout("Dom0 VCPU tuning daemon is")
    if is_pid_locked():
        stdout(" running.\n")
    else:
        stdout(" stopped.\n")

import XenAPI
import thread
import traceback
import time

class XapiEvent(object):
    def __init__(self):
        self.events = {}
        self.x = XenAPI.xapi_local()
        self.xapi_session = None
        self.is_listening = False

    def loginfo(self, msg):
        loginfo("XapiEvent: %s" % (msg,))
    def logdbg(self, msg):
        logdbg("XapiEvent: %s" % (msg,))
    def logerr(self, msg):
        logerr("XapiEvent: %s" % (msg,))

    def get_session(self):
        self.loginfo("obtaining session...")
        if self.xapi_session:
            # try to reuse any existing session
            try:
                r = self.x.session.get_last_active(self.xapi_session,self.xapi_session)
                if r["Status"] == "Failure":
                    raise XenAPI.Failure( r["ErrorDescription"] )
                self.loginfo("reusing existing valid session=%s" % (self.xapi_session,))
                return self.xapi_session
            except Exception, e:
                self.loginfo("invalid existing session %s: %s" % (self.xapi_session,e))            
                self.xapi_session = None
        self.loginfo("logging in...")
        s=self.x.session.login_with_password("__dom0__tune_vcpus","")['Value']
        self.loginfo("session=%s" % (s,))
        self.xapi_session = s
        return s

    def listen(self):
        def listen_xapi():
            try:
                try:
                    session = self.get_session()
                    try:
                        try:
                            self.process_events(session)
                        except Exception, e:
                            self.logerr("listen except: %s" % (e,))
                    finally:
                        self.x.session.logout(session)
                        self.loginfo("logout session=%s" % session)
                except Exception, e:
                    self.loginfo("listen session except: %s" % (e,))
            finally:
                self.is_listening = False

        self.loginfo("starting listen")
        self.is_listening = True
        thread.start_new_thread(listen_xapi,())

    def process_vm(self, session, ref, ref_type, data):
        if ref_type=="event":
            snapshot=data['snapshot']
        elif ref_type=="vm":
            snapshot=data
        else:
            self.logdbg("unknown ref_type=%s" % (ref_type,))
            return
        vm_uuid=snapshot['uuid']
        self.logdbg("vm_uuid=%s(%s),power_state=%s,current_operations=%s" % (snapshot['uuid'],snapshot['name_label'],snapshot['power_state'],snapshot["current_operations"].values()))
        update_active_vms(vm_uuid,snapshot)
        self.events[vm_uuid]=snapshot

    def process_events(self, session):
        def register():
            self.loginfo("registering for events")
            self.x.event.register(session,["VM"])
            # look at current state
            r_all_vms = self.x.VM.get_all_records(session)
            if r_all_vms["Status"] == "Failure":
                raise XenAPI.Failure( r_all_vms["ErrorDescription"] )
            all_vms = r_all_vms["Value"]
            for vm_ref in all_vms.keys():
                self.process_vm(session, vm_ref, "vm", all_vms[vm_ref])

        try:
            self.loginfo("process_events.start")
            register()
            self.complete = False
            while not self.complete:
                # Event loop
                try:
                    self.logdbg("calling event.next()")
                    r_evs = self.x.event.next(session)
                    if r_evs["Status"] == "Failure":
                        raise XenAPI.Failure( r_evs["ErrorDescription"] )
                    evs = r_evs["Value"]
                    for event in evs:
                        self.logdbg("received event op='%s' class='%s' ref='%s'" % (event['operation'], event['class'], event['ref']))
                        if event['class'] == 'vm' and event['operation'] == 'mod':
                            self.process_vm(session, event['ref'], "event", event)
                            continue

                except XenAPI.Failure, e:
                    self.logerr("process_events: e = [%s], e.details = [%s]" % (e, e.details) )
                    if len(e.details) > 0 and e.details[0] == 'EVENTS_LOST':
                        self.logerr("Caught EVENTS_LOST")
                        self.x.event.unregister(session,["VM"])
                        register()
                    else:
                        self.logerr("Non-EVENTS_LOST 'failure' exception: %s" % (traceback.format_exc(),))
                        self.logerr("trying re-registering anyway")
                        self.x.event.unregister(session,["VM"])
                        register()
                except:
                    self.logdbg("fatal exception: %s" % (traceback.format_exc(),))
                    self.complete = True
                    self.error = True
        finally:
            self.loginfo("process_events.end")

vm_starttimes = {}
active_vms = set()
SAMPLING_PERIOD       = 20 # in seconds
VM_START_TAIL         = 60 # expected tail of activity, in seconds
VM_START_NR_THRESHOLD = 5  # minimum number of active VMs

def update_active_vms(vm_uuid,snapshot):
    now = time.time()
    power_state=snapshot['power_state']
    current_operations=snapshot['current_operations'].values()
    is_vm_starting = (power_state=="Halted" and "start" in current_operations)
    is_vm_clean_shutdown = ("clean_shutdown" in current_operations)
    if is_vm_starting or is_vm_clean_shutdown:
        vm_starttimes[vm_uuid]=now

    if vm_uuid in vm_starttimes:
        vm_starttime = vm_starttimes[vm_uuid]

        #the more active VMs, the longer the VMs are considered active
        vm_start_tail = float(VM_START_TAIL) * float(1 + len(active_vms))

        is_active_vm = is_vm_starting or (power_state=="Running" and (now - vm_starttime < vm_start_tail))
        if is_active_vm:
            if vm_uuid not in active_vms:
                active_vms.add(vm_uuid)
                loginfo("detected new active vm: %s, active_vms=%s" % (vm_uuid, len(active_vms)))
        else:
            if vm_uuid in active_vms:
                active_vms.remove(vm_uuid)
                loginfo("removed active vm: %s, active_vms=%s" % (vm_uuid, len(active_vms)))

def pin_dom0_vcpus():
    pin_vcpus(verbose=False)

def unpin_dom0_vcpus():
    unpin_vcpus(verbose=False)

def apply_vcpu_pin_policy(is_pinning_vcpus):
    # policy
    nr_active_vms = len(active_vms)
    if nr_active_vms > float(VM_START_NR_THRESHOLD):
        # pin (at state transition)
        if not is_pinning_vcpus:
            pin_dom0_vcpus()
            return True
    else:
        # unpin (at state transition)
        if is_pinning_vcpus:
            unpin_dom0_vcpus()
            return False
    return is_pinning_vcpus

def read_config_file(config_file_name):
    try:
        config_f = open(config_file_name, "r")
        loginfo("reading configuration from %s" % (config_file_name,))
        for line in config_f:
            line = line.strip()
            if len(line) and line[0] is not "#":
                parts = line.split("=",1)
                if len(parts) == 2:
                    k = parts[0]
                    v = parts[1]
                    loginfo("setting %s=%s" % (k,v))
                    globals()[k]=v
    except Exception, e:
        logdbg("read_config_file: %s" % (e,))

# overridden by key-value pair in config_file
NR_DOMAIN0_VCPUS = 4
# overridden by key-value pair in config_file
MAX_NR_DOMAIN0_VCPUS = 8

def get_max_nr_dom0_pcpus():
    try:
        if POLICY_PIN == PP_OFF:
            #get number of pcpus in the host
            pcpus = len(call(["/usr/sbin/xenpm","get-cpu-topology"])[1:])
        else:
            #get number of pcpus in the same socket and node in the host
            pcpus = len([ x.strip() for x in load(dom0_pcpus_input) if len(x.strip()) ])
    except Exception, e:
        logexc(e)
        pcpus = int(NR_DOMAIN0_VCPUS)

    if (pcpus > int(MAX_NR_DOMAIN0_VCPUS)):
        pcpus = int(MAX_NR_DOMAIN0_VCPUS)
    return pcpus

def get_max_nr_dom0_current_vcpus():
    return int((load("/sys/devices/system/cpu/possible")[0]).split("-")[1]) + 1

def get_max_nr_dom0_vcpus():
    vcpus = get_max_nr_dom0_pcpus()
    max_current_vcpus = get_max_nr_dom0_current_vcpus()
    if (vcpus > max_current_vcpus):
        vcpus = max_current_vcpus
    return vcpus

def get_current_dom0_vcpus():
    return [x for x in load("/proc/cpuinfo") if "processor" in x]

def get_current_nr_dom0_vcpus():
    return len(get_current_dom0_vcpus())

def get_min_nr_dom0_vcpus():
    return int(NR_DOMAIN0_VCPUS)

def dyn_set_number_vcpus(n):
    if n < get_min_nr_dom0_vcpus(): n = get_min_nr_dom0_vcpus()
    if n > get_max_nr_dom0_vcpus(): n = get_max_nr_dom0_vcpus()
    online = "/sys/devices/system/cpu/cpu%s/online"
    # nb: vcpu0 online state cannot be changed
    rg_enabling = range(1, get_max_nr_dom0_vcpus() )
    for v in rg_enabling:
        if v < n:
            # enable if necessary
            if int(load(online % (v,))[0]) == 0:
                try:
                    save(online % (v,), "1")
                except Exception, e:
                    logexc(e)
    rg_disabling = range(1, get_max_nr_dom0_current_vcpus() )
    for v in reversed(rg_disabling):
        if v >= n:
            # disable if necessary
            if int(load(online % (v,))[0]) == 1:
                try:
                    save(online % (v,), "0")
                except Exception, e:
                    logexc(e)

def dyn_set_number_vcpus_xencmdline(n, verbose=True):
    if verbose: stdout("(xen boot) ")
    call(["/opt/xensource/libexec/xen-cmdline","--set-xen","dom0_max_vcpus=1-%s" % (n,)])

def dyn_set_min_vcpus(verbose=True):
    if verbose: stdout("...")
    dyn_set_number_vcpus( get_min_nr_dom0_vcpus() )
    dyn_set_number_vcpus_xencmdline( get_min_nr_dom0_vcpus() )

def dyn_set_max_vcpus(verbose=True):
    if verbose: stdout("...")
    dyn_set_number_vcpus( get_max_nr_dom0_vcpus() )
    dyn_set_number_vcpus_xencmdline( get_max_nr_dom0_pcpus() )

def dyn_set_dyn_vcpus():
    dyn_set_number_vcpus_xencmdline( get_max_nr_dom0_pcpus(), verbose=False )

# overriden by key-value pair in config_file
DOM0_VCPUS_MIN_IDLE = 20
DOM0_VCPUS_MAX_IDLE = 40
DOM0_VCPUS_STEP = 1
# possible pin vcpu policies
PP_ON = "pin"
PP_OFF = "nopin"
PP_DETECT = "detectpin"
# possible dynamic vcpu policies
PD_MIN = "min"
PD_MAX = "max"
PD_DYNAMIC = "dynamic"

def daemon_dyn_main(xapi_event, policy_dyn):
    sp = str(SAMPLING_PERIOD)
    #sample cpu for sp seconds before returning
    dom0_vcpus_idle = float(call(["/usr/bin/mpstat",sp,"1"],dbg=False)[-1:][0].split()[9])
    current_nr_dom0_vcpus = get_current_nr_dom0_vcpus()
    if dom0_vcpus_idle < float(DOM0_VCPUS_MIN_IDLE) or current_nr_dom0_vcpus < get_min_nr_dom0_vcpus():
        # >=80% vcpu usage in dom0 or lack of vcpus => increase dom0 vcpus
        if current_nr_dom0_vcpus < get_max_nr_dom0_vcpus():
            loginfo("vcpus=%s, avg idle=%s" % (current_nr_dom0_vcpus, dom0_vcpus_idle))
            dyn_set_number_vcpus(current_nr_dom0_vcpus + int(DOM0_VCPUS_STEP))
    elif dom0_vcpus_idle > float(DOM0_VCPUS_MAX_IDLE) or current_nr_dom0_vcpus > get_max_nr_dom0_vcpus():
        #<=60% vcpu usage in dom0 or excess vcpus => decrease dom0 vcpus
        if current_nr_dom0_vcpus > get_min_nr_dom0_vcpus():
            loginfo("vcpus=%s, avg idle=%s" % (current_nr_dom0_vcpus, dom0_vcpus_idle))
            dyn_set_number_vcpus(current_nr_dom0_vcpus - int(DOM0_VCPUS_STEP))

def daemon_pin_main(xapi_event, policy_pin, is_pinning_vcpus):
    try:
        for vm_uuid in xapi_event.events:
            snapshot=xapi_event.events[vm_uuid]
            update_active_vms(vm_uuid,snapshot)
        is_pinning_vcpus = apply_vcpu_pin_policy(is_pinning_vcpus)
    except Exception, e:
        logerr("daemon_pin_main: %s" % (e,))
    return is_pinning_vcpus

def daemon_main(policy_pin, policy_dyn):
    xapi_event = XapiEvent()
    while True:
        try:
            if not xapi_event.is_listening:
                xapi_event.listen()
            is_pinning_vcpus = False
            while xapi_event.is_listening:
                try:
                    #pin
                    if policy_pin == PP_DETECT:
                        is_pinning_vcpus = daemon_pin_main(xapi_event, policy_pin, is_pinning_vcpus)
                    #dyn
                    if policy_dyn == PD_DYNAMIC:
                        daemon_dyn_main(xapi_event, policy_dyn)
                    else:
                        try:
                            time.sleep(float(SAMPLING_PERIOD))
                        except Exception, e: # eg. due to invalid config
                            logerr("waiting next sample: %s" % (e,))
                            time.sleep(10)
                finally:
                    time.sleep(1)
        except Exception, e:
            logerr("daemon_main: %s" % (e,))
        time.sleep(1)

def set_policy(policy_pin, policy_dyn):
    # pin
    stdout("%s" % (policy_pin,))
    if policy_pin == PP_ON:
        pin_vcpus()
        pin_xencmdline()
        save_policy_pin_dyn(policy_pin, policy_dyn)
    elif policy_pin == PP_OFF:
        unpin_vcpus()
        unpin_xencmdline()
        save_policy_pin_dyn(policy_pin, policy_dyn)
    elif policy_pin == PP_DETECT:
        unpin_xencmdline()
    else:
        usage()
        sys.exit(1)
    # dyn
    stdout("%s" % (policy_dyn,))
    if policy_dyn == PD_MIN:
        dyn_set_min_vcpus()
        save_policy_pin_dyn(policy_pin, policy_dyn)
    elif policy_dyn == PD_MAX:
        dyn_set_max_vcpus()
        save_policy_pin_dyn(policy_pin, policy_dyn)
    elif policy_dyn == PD_DYNAMIC:
        dyn_set_dyn_vcpus()
    else:
        usage()
        sys.exit(1)

def daemon(policy_pin, policy_dyn):
    try:
        sys.stdout.flush()
        sys.stderr.flush()
        pid = os.fork()
        if pid == 0:
            # child: start daemon
            pid = str(os.getpid())
            save(pid_file, "%s\n" % (pid,))
            lock_pid()
            read_config_file(config_file)
            set_policy(policy_pin, policy_dyn)
            #detach std{in,out,err}
            devnull_fd = os.open(os.devnull, os.O_RDWR)
            os.close(0) #stdin
            os.close(1) #stdout
            os.close(2) #stderr
            os.dup2(devnull_fd,0)
            os.dup2(devnull_fd,1)
            os.dup2(devnull_fd,2)
            daemon_main(policy_pin, policy_dyn)
        else:
            # parent
           pass
    except Exception, e:
        logerr("daemon: %s" % (e,))


# default values when no policy has been defined yet
POLICY_PIN=PP_OFF
POLICY_DYN=PD_MIN
def save_policy_pin_dyn(policy_pin, policy_dyn):
    save_policy("POLICY_PIN=%s\nPOLICY_DYN=%s\n" % (policy_pin, policy_dyn))
def load_policy_pin_dyn():
    read_config_file(policy_file)

def usage():
    stdout("Usage: %s {start {pin|nopin|detectpin} {min|max|dynamic} | stop | status}\n" % (sys.argv[0],))

pin_policies = [PP_ON, PP_OFF, PP_DETECT]
dyn_policies = [PD_MIN, PD_MAX, PD_DYNAMIC]

def start(argv):
    global POLICY_PIN
    global POLICY_DYN
    if is_pid_locked():
        stdout("Found %s: already running..." % (pid_file,))
        failure()
    else:
        # load current or default policy if none defined in policy_file
        load_policy_pin_dyn()
        policy_pin = POLICY_PIN
        policy_dyn = POLICY_DYN
        for p in argv:
            if p not in pin_policies and p not in dyn_policies:
                usage()
                sys.exit(1)
            if p in pin_policies: policy_pin = p
            if p in dyn_policies: policy_dyn = p
        POLICY_PIN = policy_pin
        POLICY_DYN = policy_dyn

        # detect if daemon is necessary
        is_daemon = (policy_pin == PP_DETECT or policy_dyn == PD_DYNAMIC)
        
        try:
            if is_daemon:
                daemon(policy_pin, policy_dyn)
                save_policy_pin_dyn(policy_pin, policy_dyn)
            else: 
                stdout("Tuning: ")
                set_policy(policy_pin, policy_dyn)
            success()
        except Exception, e:
            logexc(e)
            failure()

def main():
    logToSyslog(level=logging.INFO)
    if len(sys.argv) > 1:
        if sys.argv[1] == "start":
            start(sys.argv[2:])
        elif sys.argv[1] == "stop":
            stop()
        elif sys.argv[1] == "status":
            status()
        else:
            usage()
    else:
        usage()

if __name__ == "__main__":
    main()
