ucloud-vm/main.py

# TODO

# For QEMU Monitor Protocol Commands Information, See
# https://qemu.weilnetz.de/doc/qemu-doc.html#pcsys_005fmonitor

import json
import argparse
import qmp
import logging
import os
import shutil

from etcd3_wrapper import Etcd3Wrapper
from dataclasses import dataclass
from typing import Union
from functools import wraps
from decouple import config

running_vms = []
client = Etcd3Wrapper()


@dataclass
class VM:
    key: str
    vm: qmp.QEMUMachine


logging.basicConfig(
    level=logging.DEBUG,
    filename="log.txt",
    filemode="a",
    format="%(asctime)s: %(levelname)s - %(message)s",
    datefmt="%d-%b-%y %H:%M:%S",
)


def need_running_vm(func):
    @wraps(func)
    def wrapper(e):
        vm = get_vm(running_vms, e.key)
        if vm:
            try:
                status = vm.vm.command("query-status")
                logging.debug(f"VM Status Check - {status}")
            except OSError:
                logging.info(f"{func.__name__} failed - VM {e.key} - Unknown Error")

            return func(e)
        else:
            logging.info(f"{func.__name__} failed because VM {e.key} is not running")
            return
    return wrapper


def create_vm(owner_dir, vm_uuid, e):
    image = client.get(
        f"/v1/image/{e.value['image_uuid']}", value_in_json=True
    )
    if image:
        logging.debug(image)
        image_uuid = e.value["image_uuid"]

        logging.info("Creating New VM...")

        os.makedirs(f"{owner_dir}/.vm", exist_ok=True)

        if not os.path.isfile(f"{owner_dir}/.vm/{vm_uuid}.raw"):
            shutil.copy(
                f"/var/vm/{image_uuid}.raw", f"{owner_dir}/.vm/{vm_uuid}.raw"
            )

        e.value["status"] = "REQUESTED_START"
    client.put(e.key, json.dumps(e.value))


def start_vm(vm_path, e):
    if not os.path.isfile(vm_path):
        logging.info(f"Image file of vm {e.key} does not exists")
        logging.info(f"Setting vm {e.key} status to DELETED")
        e.value["status"] = "DELETED"
        client.put(e.key, json.dumps(e.value))
        return

    _vm = get_vm(running_vms, e.key)
    if _vm:
        logging.info(f"{e.key} already running")
        e.value["status"] = "RUNNING"
        client.put(e.key, json.dumps(e.value))
        return

    vm = qmp.QEMUMachine(
        "/usr/bin/qemu-system-x86_64",
        test_dir="vm_socklog",
        args=[
            vm_path,
            "-boot",
            "c",
            "-net",
            "nic",
            "-net",
            "user",
            "-m",
            "1024",
            "-vnc",
            f":{len(running_vms)}",
        ],
    )
    try:
        logging.info(f"Starting {e.key}")
        vm.launch()
        if vm.is_running():
            running_vms.append(VM(e.key, vm))
            e.value["status"] = "RUNNING"
            client.put(e.key, e.value, value_in_json=True)
        else:
            e.value["status"] = "KILLED"
            client.put(e.key, e.value, value_in_json=True)
            return
    except (qmp.QEMUMachineError, TypeError):
        logging.info(f"Machine Error Occurred on {e.key}")
        e.value["status"] = "KILLED"
        client.put(e.key, e.value, value_in_json=True)
    else:
        logging.info(f"Started Successfully {e.key}")


@need_running_vm
def suspend_vm(e):
    vm = get_vm(running_vms, e.key)
    vm.vm.command("stop")
    if vm.vm.command("query-status")["status"] == "paused":
        e.value["status"] = "SUSPENDED"
        client.put(e.key, json.dumps(e.value))
        logging.info(f"Successfully suspended VM {e.key}")
    else:
        logging.info(f"Suspending VM {e.key} failed")


@need_running_vm
def resume_vm(e):
    vm = get_vm(running_vms, e.key)
    vm.vm.command("cont")
    if vm.vm.command("query-status")["status"] == "running":
        e.value["status"] = "RUNNING"
        client.put(e.key, json.dumps(e.value))
        logging.info(f"Successfully resumed VM {e.key}")
    else:
        logging.info(f"Resuming VM {e.key} failed")


@need_running_vm
def shutdown_vm(e):
    vm = get_vm(running_vms, e.key)
    vm.vm.shutdown()
    if not vm.vm.is_running():
        logging.info(f"VM {e.key} shutdown successfully")
        e.value["status"] = "STOPPED"
        client.put(e.key, json.dumps(e.value))
        running_vms.remove(vm)


def delete_vm(e):
    logging.info(f"Deleting VM {e.key}")
    shutdown_vm(e)
    vm = client.get(e.key, value_in_json=True)
    if vm:
        vm_id = e.key.split('/')[-1]
        vm_owner = e.value['owner']
        vm_path = f"{config('BASE_DIR')}/{vm_owner}/.vm/{vm_id}"

        if os.path.exists(vm_path):
            os.remove(vm_path)
        client.client.delete(e.key)
        logging.info(f"VM {vm.key} deleted")
    else:
        logging.info(f"Cannot delete key {e.key} because it doesn't exists")


def get_vm(vm_list: list, vm_key) -> Union[VM, None]:
    return next((vm for vm in vm_list if vm.key == vm_key), None)


def main(hostname):
    events = client.watch_prefix("/v1/vm/", timeout=10)

    # events = client.get_prefix("/v1/vm/")
    for e in events:
        try:
            e.value = json.loads(e.value)
        except json.JSONDecodeError:
            logging.error(f"Invalid JSON {e.value}")
            continue

        e_status = e.value["status"]

        if e_status == "TIMEOUT":
            logging.info("Timeout")
            _vms = filter(lambda v: v.value["hostname"] == hostname, client.get_prefix("/v1/vm", value_in_json=True))
            alleged_running_vms = filter(lambda v: v.value["status"] == "RUNNING", _vms)
            for vm in alleged_running_vms:
                _vm = get_vm(running_vms, vm.key)
                if (_vm and not _vm.vm.is_running()) or _vm is None:
                    logging.debug(f"{_vm} {vm.key}")
                    logging.info(f"{vm.key} is not running but is said to be running")
                    logging.info(f"Updating {vm.key} status to KILLED")
                    vm.value["status"] = "KILLED"
                    client.put(vm.key, json.dumps(vm.value))
            continue

        e_hostname = e.value["hostname"]
        vm_uuid = e.key.split("/")[-1]
        owner_dir = f"{config('BASE_DIR')}/{e.value['owner']}"

        # If it is not for me then skip it
        if e_hostname != hostname:
            continue

        logging.debug(f"EVENT: {e}")

        if e_status == "SCHEDULED_DEPLOY":
            create_vm(owner_dir, vm_uuid, e)

        elif e_status == "REQUESTED_SUSPEND":
            suspend_vm(e)

        elif e_status == "REQUESTED_RESUME":
            resume_vm(e)

        elif e_status == "REQUESTED_START":
            vm_path = f"{owner_dir}/.vm/{vm_uuid}.raw"
            start_vm(vm_path, e)

        elif e_status == "REQUESTED_SHUTDOWN":
            shutdown_vm(e)

        elif e_status == "DELETED":
            delete_vm(e)

        logging.info(f"Running VMs {running_vms}")


argparser = argparse.ArgumentParser()
argparser.add_argument("hostname", help="Name of this host. e.g /v1/host/1")
args = argparser.parse_args()

main(args.hostname)