#!/usr/bin/python
# Copyright (C) 2006-2007 XenSource Ltd.
# Copyright (C) 2008-2009 Citrix Ltd.
#
# This program is free software; you can redistribute it and/or modify 
# it under the terms of the GNU Lesser General Public License as published 
# by the Free Software Foundation; version 2.1 only.
#
# This program is distributed in the hope that it will be useful, 
# but WITHOUT ANY WARRANTY; without even the implied warranty of 
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 
# GNU Lesser General Public License for more details.
#
# Script to coalesce the leaf VDIs for a given VM
#

import sys
sys.path.insert(0, "/opt/xensource/sm")
import getopt
import XenAPI
import cleanup

USAGE_STRING = \
"""Usage: %s -u/--uuid <UUID of VM whose VDIs should be leaf-coalesced>
This will coalesce each VDI attached to the given VM that consists of a pair
of VHD files into a single VHD file.
Only LVM SRs will be considered, and only VDIs whose VHD chain length equals 2.
Note that the VM will be suspended during the operation. DO NOT
start/resume/unpause the VM during the operation."""

def main():
    shortArgs  = "u:"
    longArgs   = ["uuid"]

    try:
        opts, args = getopt.getopt(sys.argv[1:], shortArgs, longArgs)
    except getopt.GetoptError:
        print USAGE_STRING % sys.argv[0]
        sys.exit(-1)

    uuid = None
    for o, a in opts:
        if o in ("-u", "--uuid"):
            uuid = a

    if not uuid:
        print USAGE_STRING % sys.argv[0]
        sys.exit(-1)

    session = XenAPI.xapi_local()
    session.xenapi.login_with_password('root', '')

    sr_uuid = None
    vdis = {}
    try:
        vm_ref = session.xenapi.VM.get_by_uuid(uuid)
    except XenAPI.Failure:
        print "Error: VM %s not found" % uuid
        sys.exit(2)
    vm_rec = session.xenapi.VM.get_record(vm_ref)
    vbd_recs = session.xenapi.VBD.get_all_records()
    for vbd_ref, vbd_rec in vbd_recs.iteritems():
        if vbd_rec["VM"] != vm_ref:
            continue
        vdi_ref = vbd_rec["VDI"]
        if not vdi_ref or vdi_ref == 'OpaqueRef:NULL':
            continue
        vdi_rec = session.xenapi.VDI.get_record(vdi_ref)
        sr_ref = vdi_rec["SR"]
        sr_rec = session.xenapi.SR.get_record(sr_ref)
        if sr_rec["type"].startswith("lv"):
            print "VDI to consider: %s" % vdi_rec["uuid"]
            if not vdis.get(sr_rec["uuid"]):
                vdis[sr_rec["uuid"]] = []
            vdis[sr_rec["uuid"]].append(vdi_rec["uuid"])
        else:
            print "Skipping non-LVM VDI: %s" % vdi_rec["uuid"]

    coalesceable_vdis = {}
    for sr_uuid, vdi_uuids in vdis.iteritems():
        try:
            vdi_list = cleanup.get_coalesceable_leaves(session, sr_uuid, vdi_uuids)
        except Exception, e:
            print "Error: %s" % e
            sys.exit(3)
        if len(vdi_list) > 0:
            coalesceable_vdis[sr_uuid] = vdi_list

    if len(coalesceable_vdis) == 0:
        print "The VM has no VDIs that could be leaf-coalesced"
        sys.exit(1)

    suspended = False
    if vm_rec["power_state"] == "Running":
        print "Suspending VM %s" % vm_rec["uuid"]
        try:
            session.xenapi.VM.suspend(vm_ref)
        except XenAPI.Failure:
            print "Failed to suspend the VM"
            sys.exit(2)
        suspended = True
    try:
        for sr_uuid, vdi_uuids in coalesceable_vdis.iteritems():
            print "Processing VDIs: %s" % vdi_uuids
            try:
                cleanup.coalesce_leaves(session, sr_uuid, vdi_uuids)
            except Exception, e:
                print "Failed: %s" % e
                cleanup.Util.logException("coalesce-leaf")
                sys.exit(3)
    finally:
        if suspended:
            print "Resuming VM %s" % vm_rec["uuid"]
            session.xenapi.VM.resume(vm_ref, False, False)
        cleanup.gc(session, sr_uuid, "lvhd", True)
    sys.exit(0)

if __name__ == '__main__':
    main()
