#!/usr/bin/env python
#
# Read in all the packages installed in dom0 and sum up the
# "memory-requirement-mb" attributes of the <repository> nodes.  (Default
# is 0.)  Then work out what memory-target *should* be and set it.
#
# To calculate memory-target we add these to (memory-static-min + 100MiB) and
# truncate this value to between memory-static-min and memory-static-max.


import XenAPI
import sys
import os
import time
import traceback
import xml.dom.minidom
from xml.parsers.expat import ExpatError

INSTALLED_REPOS_DIR = '/etc/xensource/installed-repos'

def main():
    session = XenAPI.xapi_local()
    logged_in = False
    # read dom0 uuid out of inventory file
    uuid = filter(lambda x: x.startswith('CONTROL_DOMAIN_UUID'), open('/etc/xensource-inventory').readlines())[0].strip().split('=')[1].strip("'")
    oref = rec = None
    for tries in range(180):
        try:
            session.xenapi.login_with_password("", "")
            logged_in = True

            # get dom0's DB record
            oref = session.xenapi.VM.get_by_uuid(uuid)
            rec = session.xenapi.VM.get_record(oref)
            break
        except:
            # repeat after a delay
            time.sleep(1)
    if oref == None:
        if logged_in:
            session.xenapi.session.logout()
        print >> sys.stderr, "Xapi is not ready. Aborting."
        return 1

    static_min = int(rec['memory_static_min'])
    static_max = int(rec['memory_static_max'])
    current_target = int(rec['memory_target'])

    additional_mb = 0

    packs = os.listdir(INSTALLED_REPOS_DIR)
    for pack in packs:
        try:
            xmldoc = xml.dom.minidom.parse(os.path.join(INSTALLED_REPOS_DIR, pack, 'XS-REPOSITORY'))
            additional_mb_str = xmldoc.documentElement.getAttribute('memory-requirement-mb')
            if additional_mb_str:
                additional_mb += int(additional_mb_str)
        except:
            pass

    # new memory target
    memory_target = static_min + (100<<20) + (additional_mb<<20)

    print >> sys.stdout, "Memory required by all installed packages: %d" % memory_target

    if memory_target <= current_target:
        print >> sys.stdout, "Current target %d greater, skipping" % current_target
        return 0
    
    if memory_target > static_max:
        print >> sys.stdout, "Truncating to static-max: %d" % static_max
        memory_target = static_max
        
    print >> sys.stdout, "Setting VM.memory_target: %d" % memory_target

    # Next line commented out because this doesn't seem to work - maybe a bug in the Python bindings
    #session.xenapi.VM.set_memory_target_live(oref, memory_target)
    
    # Do this instead
    import subprocess
    subprocess.call(['xe','vm-memory-target-set', 'uuid=%s' % uuid, 'target=%d' % memory_target])

    session.xenapi.session.logout()
    return 0

if __name__ == '__main__':
    rc = 1
    try:
        rc = main()
    except:
        ex = sys.exc_info()
        err = traceback.format_exception(*ex)
        for exline in err:
            print >> sys.stderr, exline

    sys.exit(rc)
