uncloud/ucloud/host/main.py

import argparse
import multiprocessing as mp
import time

from ucloud.common.request import RequestEntry, RequestType
from ucloud.shared import shared
from ucloud.settings import settings
from ucloud.common.vm import VMStatus
from ucloud.vmm import VMM
from os.path import join as join_path

from . import virtualmachine, logger


def update_heartbeat(hostname):
    """Update Last HeartBeat Time for :param hostname: in etcd"""
    host_pool = shared.host_pool
    this_host = next(
        filter(lambda h: h.hostname == hostname, host_pool.hosts), None
    )
    while True:
        this_host.update_heartbeat()
        host_pool.put(this_host)
        time.sleep(10)


def maintenance(host):
    vmm = VMM()
    running_vms = vmm.discover()
    for vm_uuid in running_vms:
        if (
            vmm.is_running(vm_uuid)
            and vmm.get_status(vm_uuid) == "running"
        ):
            vm = shared.vm_pool.get(
                join_path(settings["etcd"]["vm_prefix"], vm_uuid)
            )
            vm.status = VMStatus.running
            vm.vnc_socket = vmm.get_vnc(vm_uuid)
            vm.hostname = host
            shared.vm_pool.put(vm)


def main(hostname):
    host_pool = shared.host_pool
    host = next(
        filter(lambda h: h.hostname == hostname, host_pool.hosts), None
    )
    assert host is not None, "No such host with name = {}".format(
        hostname
    )

    try:
        heartbeat_updating_process = mp.Process(
            target=update_heartbeat, args=(hostname,)
        )
        heartbeat_updating_process.start()
    except Exception as e:
        raise Exception(
            "ucloud-host heartbeat updating mechanism is not working"
        ) from e

    for events_iterator in [
        shared.etcd_client.get_prefix(
            settings["etcd"]["request_prefix"], value_in_json=True
        ),
        shared.etcd_client.watch_prefix(
            settings["etcd"]["request_prefix"],
            timeout=10,
            value_in_json=True,
        ),
    ]:
        for request_event in events_iterator:
            request_event = RequestEntry(request_event)

            if request_event.type == "TIMEOUT":
                maintenance(host.key)

            if request_event.hostname == host.key:
                logger.debug("VM Request: %s", request_event)

                shared.request_pool.client.client.delete(
                    request_event.key
                )
                vm_entry = shared.etcd_client.get(
                    join_path(
                        settings["etcd"]["vm_prefix"],
                        request_event.uuid,
                    )
                )

                if vm_entry:
                    vm = virtualmachine.VM(vm_entry)
                    if request_event.type == RequestType.StartVM:
                        vm.start()

                    elif request_event.type == RequestType.StopVM:
                        vm.stop()

                    elif request_event.type == RequestType.DeleteVM:
                        vm.delete()

                    elif (
                        request_event.type
                        == RequestType.InitVMMigration
                    ):
                        vm.start(destination_host_key=host.key)

                    elif request_event.type == RequestType.TransferVM:
                        host = host_pool.get(
                            request_event.destination_host_key
                        )
                        if host:
                            vm.migrate(
                                destination_host=host.hostname,
                                destination_sock_path=request_event.destination_sock_path,
                            )
                        else:
                            logger.error(
                                "Host %s not found!",
                                request_event.destination_host_key,
                            )
                else:
                    logger.info("VM Entry missing")


if __name__ == "__main__":
    argparser = argparse.ArgumentParser()
    argparser.add_argument(
        "hostname", help="Name of this host. e.g uncloud1.ungleich.ch"
    )
    args = argparser.parse_args()
    mp.set_start_method("spawn")
    main(args.hostname)