uncloud-mravi/ucloud/common/vm.py

from contextlib import contextmanager
from datetime import datetime
from os.path import join

from .classes import SpecificEtcdEntryBase


class VMStatus:
    stopped = "STOPPED"  # After requested_shutdown
    killed = "KILLED"  # either host died or vm died itself
    running = "RUNNING"
    error = "ERROR"  # An error occurred that cannot be resolved automatically


class VMEntry(SpecificEtcdEntryBase):

    def __init__(self, e):
        self.owner = None  # type: str
        self.specs = None  # type: dict
        self.hostname = None  # type: str
        self.status = None  # type: str
        self.image_uuid = None  # type: str
        self.log = None  # type: list
        self.in_migration = None  # type: bool

        super().__init__(e)

    @property
    def uuid(self):
        return self.key.split("/")[-1]

    def declare_killed(self):
        self.hostname = ""
        self.in_migration = False
        if self.status == VMStatus.running:
            self.status = VMStatus.killed

    def declare_stopped(self):
        self.hostname = ""
        self.in_migration = False
        self.status = VMStatus.stopped

    def add_log(self, msg):
        self.log = self.log[:5]
        self.log.append("{} - {}".format(datetime.now().isoformat(), msg))


class VmPool:
    def __init__(self, etcd_client, vm_prefix):
        self.client = etcd_client
        self.prefix = vm_prefix

    @property
    def vms(self):
        _vms = self.client.get_prefix(self.prefix, value_in_json=True)
        return [VMEntry(vm) for vm in _vms]

    def by_host(self, host, _vms=None):
        if _vms is None:
            _vms = self.vms
        return list(filter(lambda x: x.hostname == host, _vms))

    def by_status(self, status, _vms=None):
        if _vms is None:
            _vms = self.vms
        return list(filter(lambda x: x.status == status, _vms))

    def by_owner(self, owner, _vms=None):
        if _vms is None:
            _vms = self.vms
        return list(filter(lambda x: x.owner == owner, _vms))

    def except_status(self, status, _vms=None):
        if _vms is None:
            _vms = self.vms
        return list(filter(lambda x: x.status != status, _vms))

    def get(self, key):
        if not key.startswith(self.prefix):
            key = join(self.prefix, key)
        v = self.client.get(key, value_in_json=True)
        if v:
            return VMEntry(v)
        return None

    def put(self, obj: VMEntry):
        self.client.put(obj.key, obj.value, value_in_json=True)

    @contextmanager
    def get_put(self, key) -> VMEntry:
        # Updates object at key on exit
        obj = self.get(key)
        yield obj
        if obj:
            self.put(obj)