uncloud-mravi/ucloud/host/main.py

import argparse
import multiprocessing as mp
import time
import sys

from ucloud.common.request import RequestEntry, RequestType
from ucloud.common.host import HostPool
from ucloud.shared import shared
from ucloud.settings import settings

from . import virtualmachine, logger

vmm = virtualmachine.VMM()


def update_heartbeat(hostname):
    """Update Last HeartBeat Time for :param hostname: in etcd"""
    
    client = shared.etcd_client
    host_pool = HostPool(client)
    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 main(hostname):
    host_pool = HostPool(shared.etcd_client)
    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:
        logger.exception(e)
        sys.exit("No Need To Go Further. ucloud-host heartbeat updating mechanism is not working")

    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":
                vmm.maintenance(host)
                continue

            # If the event is directed toward me OR I am destination of a InitVMMigration
            if request_event.hostname == host.key or request_event.destination == host.key:
                logger.debug("VM Request: %s", request_event)

                shared.request_pool.client.client.delete(request_event.key)
                vm_entry = shared.vm_pool.get(request_event.uuid)

                if vm_entry:
                    if request_event.type == RequestType.StartVM:
                        vmm.start(vm_entry)

                    elif request_event.type == RequestType.StopVM:
                        vmm.stop(vm_entry)

                    elif request_event.type == RequestType.DeleteVM:
                        vmm.delete(vm_entry)

                    elif request_event.type == RequestType.InitVMMigration:
                        vmm.start(vm_entry, host.key)

                    elif request_event.type == RequestType.TransferVM:
                        vmm.transfer(request_event)
                else:
                    logger.info("VM Entry missing")

                logger.info("Running VMs %s", vmm.running_vms)


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