wally/start_vms.py

import re
import os
import stat
import time
import os.path
import logging
import subprocess
import collections

from concurrent.futures import ThreadPoolExecutor

from novaclient.exceptions import NotFound
from novaclient.client import Client as n_client
from cinderclient.v1.client import Client as c_client

import wally
from wally.discover import Node


logger = logging.getLogger("wally.vms")


STORED_OPENSTACK_CREDS = None
NOVA_CONNECTION = None
CINDER_CONNECTION = None


def is_connected():
    return NOVA_CONNECTION is not None


OSCreds = collections.namedtuple("OSCreds",
                                 ["name", "passwd",
                                  "tenant", "auth_url", "insecure"])


def ostack_get_creds():
    if STORED_OPENSTACK_CREDS is None:
        return OSCreds(os.environ.get('OS_USERNAME'),
                       os.environ.get('OS_PASSWORD'),
                       os.environ.get('OS_TENANT_NAME'),
                       os.environ.get('OS_AUTH_URL'),
                       os.environ.get('OS_INSECURE', False))
    else:
        return STORED_OPENSTACK_CREDS


def nova_connect(os_creds=None):
    global NOVA_CONNECTION
    global STORED_OPENSTACK_CREDS

    if NOVA_CONNECTION is None:
        if os_creds is None:
            os_creds = ostack_get_creds()
        else:
            STORED_OPENSTACK_CREDS = os_creds

        NOVA_CONNECTION = n_client('1.1',
                                   os_creds.name,
                                   os_creds.passwd,
                                   os_creds.tenant,
                                   os_creds.auth_url,
                                   insecure=os_creds.insecure)
    return NOVA_CONNECTION


def cinder_connect(os_creds=None):
    global CINDER_CONNECTION
    global STORED_OPENSTACK_CREDS

    if CINDER_CONNECTION is None:
        if os_creds is None:
            os_creds = ostack_get_creds()
        else:
            STORED_OPENSTACK_CREDS = os_creds
        CINDER_CONNECTION = c_client(os_creds.name,
                                     os_creds.passwd,
                                     os_creds.tenant,
                                     os_creds.auth_url,
                                     insecure=os_creds.insecure)
    return CINDER_CONNECTION


def prepare_os_subpr(nova, params, os_creds):
    if os_creds is None:
        os_creds = ostack_get_creds()

    MAX_VM_PER_NODE = 8
    serv_groups = " ".join(map(params['aa_group_name'].format,
                               range(MAX_VM_PER_NODE)))

    image_name = params['image']['name']
    env = os.environ.copy()

    env.update(dict(
        OS_USERNAME=os_creds.name,
        OS_PASSWORD=os_creds.passwd,
        OS_TENANT_NAME=os_creds.tenant,
        OS_AUTH_URL=os_creds.auth_url,
        OS_INSECURE="1" if os_creds.insecure else "0",

        FLAVOR_NAME=params['flavor']['name'],
        FLAVOR_RAM=str(params['flavor']['ram_size']),
        FLAVOR_HDD=str(params['flavor']['hdd_size']),
        FLAVOR_CPU_COUNT=str(params['flavor']['cpu_count']),

        SERV_GROUPS=serv_groups,

        SECGROUP=params['security_group'],

        IMAGE_NAME=image_name,
        IMAGE_URL=params['image']['url'],

        # KEYPAIR_NAME=params['keypair_name'],
        # KEY_FILE_NAME=params['keypair_file_private'],
    ))

    spath = os.path.dirname(os.path.dirname(wally.__file__))
    spath = os.path.join(spath, 'scripts/prepare.sh')

    cmd = "bash {spath} >/dev/null 2>&1 ".format(spath=spath)
    subprocess.check_call(cmd, shell=True, env=env)

    while True:
        status = nova.images.find(name=image_name).status
        if status == 'ACTIVE':
            break
        msg = "Image {0} is still in {1} state. Waiting 10 more seconds"
        logger.info(msg.format(image_name, status))
        time.sleep(10)

    create_keypair(nova,
                   params['keypair_name'],
                   params['keypair_file_public'],
                   params['keypair_file_private'])


def find_vms(nova, name_prefix):
    for srv in nova.servers.list():
        if srv.name.startswith(name_prefix):
            for ips in srv.addresses.values():
                for ip in ips:
                    if ip.get("OS-EXT-IPS:type", None) == 'floating':
                        yield ip['addr'], srv.id
                        break


def pause(ids):
    def pause_vm(conn, vm_id):
        vm = conn.servers.get(vm_id)
        if vm.status == 'ACTIVE':
            vm.pause()

    conn = nova_connect()
    with ThreadPoolExecutor(max_workers=16) as executor:
        futures = [executor.submit(pause_vm, conn, vm_id)
                   for vm_id in ids]
        for future in futures:
            future.result()


def unpause(ids, max_resume_time=10):
    def unpause(conn, vm_id):
        vm = conn.servers.get(vm_id)
        if vm.status == 'PAUSED':
            vm.unpause()

        for i in range(max_resume_time * 10):
            vm = conn.servers.get(vm_id)
            if vm.status != 'PAUSED':
                return
            time.sleep(0.1)
        raise RuntimeError("Can't unpause vm {0}".format(vm_id))

    conn = nova_connect()
    with ThreadPoolExecutor(max_workers=16) as executor:
        futures = [executor.submit(unpause, conn, vm_id)
                   for vm_id in ids]

        for future in futures:
            future.result()


def prepare_os(nova, params):
    allow_ssh(nova, params['security_group'])

    MAX_VM_PER_NODE = 8
    serv_groups = " ".join(map(params['aa_group_name'].format,
                               range(MAX_VM_PER_NODE)))

    shed_ids = []
    for shed_group in serv_groups:
        shed_ids.append(get_or_create_aa_group(nova, shed_group))

    create_keypair(nova,
                   params['keypair_name'],
                   params['keypair_name'] + ".pub",
                   params['keypair_name'] + ".pem")

    create_image(nova, params['image']['name'],
                 params['image']['url'])

    create_flavor(nova, **params['flavor'])


def create_keypair(nova, name, pub_key_path, priv_key_path):
    pub_key_exists = os.path.exists(pub_key_path)
    priv_key_exists = os.path.exists(priv_key_path)

    try:
        kpair = nova.keypairs.find(name=name)
        # if file not found- delete and recreate
    except NotFound:
        kpair = None

    if pub_key_exists and not priv_key_exists:
        raise EnvironmentError("Private key file doesn't exists")

    if not pub_key_exists and priv_key_exists:
        raise EnvironmentError("Public key file doesn't exists")

    if kpair is None:
        if pub_key_exists:
            with open(pub_key_path) as pub_key_fd:
                return nova.keypairs.create(name, pub_key_fd.read())
        else:
            key = nova.keypairs.create(name)

            with open(priv_key_path, "w") as priv_key_fd:
                priv_key_fd.write(key.private_key)
            os.chmod(priv_key_path, stat.S_IREAD | stat.S_IWRITE)

            with open(pub_key_path, "w") as pub_key_fd:
                pub_key_fd.write(key.public_key)
    elif not priv_key_exists:
        raise EnvironmentError("Private key file doesn't exists," +
                               " but key uploaded openstack." +
                               " Either set correct path to private key" +
                               " or remove key from openstack")


def get_or_create_aa_group(nova, name):
    try:
        group = nova.server_groups.find(name=name)
    except NotFound:
        group = nova.server_groups.create({'name': name,
                                           'policies': ['anti-affinity']})

    return group.id


def allow_ssh(nova, group_name):
    try:
        secgroup = nova.security_groups.find(name=group_name)
    except NotFound:
        secgroup = nova.security_groups.create(group_name,
                                               "allow ssh/ping to node")

    nova.security_group_rules.create(secgroup.id,
                                     ip_protocol="tcp",
                                     from_port="22",
                                     to_port="22",
                                     cidr="0.0.0.0/0")

    nova.security_group_rules.create(secgroup.id,
                                     ip_protocol="icmp",
                                     from_port=-1,
                                     cidr="0.0.0.0/0",
                                     to_port=-1)
    return secgroup.id


def create_image(nova, name, url):
    pass


def create_flavor(nova, name, ram_size, hdd_size, cpu_count):
    pass


def create_volume(size, name):
    cinder = cinder_connect()
    vol = cinder.volumes.create(size=size, display_name=name)
    err_count = 0

    while vol.status != 'available':
        if vol.status == 'error':
            if err_count == 3:
                logger.critical("Fail to create volume")
                raise RuntimeError("Fail to create volume")
            else:
                err_count += 1
                cinder.volumes.delete(vol)
                time.sleep(1)
                vol = cinder.volumes.create(size=size, display_name=name)
                continue
        time.sleep(1)
        vol = cinder.volumes.get(vol.id)
    return vol


def wait_for_server_active(nova, server, timeout=300):
    t = time.time()
    while True:
        time.sleep(1)
        sstate = getattr(server, 'OS-EXT-STS:vm_state').lower()

        if sstate == 'active':
            return True

        if sstate == 'error':
            return False

        if time.time() - t > timeout:
            return False

        server = nova.servers.get(server)


class Allocate(object):
    pass


def get_floating_ips(nova, pool, amount):
    ip_list = nova.floating_ips.list()

    if pool is not None:
        ip_list = [ip for ip in ip_list if ip.pool == pool]

    return [ip for ip in ip_list if ip.instance_id is None][:amount]


def launch_vms(params, already_has_count=0):
    logger.debug("Calculating new vm count")
    count = params['count']
    nova = nova_connect()
    lst = nova.services.list(binary='nova-compute')
    srv_count = len([srv for srv in lst if srv.status == 'enabled'])

    if isinstance(count, basestring):
        if count.startswith("x"):
            count = srv_count * int(count[1:])
        else:
            assert count.startswith('=')
            count = int(count[1:]) - already_has_count

    if count <= 0:
        logger.debug("Not need new vms")
        return

    logger.debug("Starting new nodes on openstack")

    assert isinstance(count, (int, long))

    srv_params = "img: {image[name]}, flavor: {flavor[name]}".format(**params)
    msg_templ = "Will start {0} servers with next params: {1}"
    logger.info(msg_templ.format(count, srv_params))

    vm_params = dict(
        img_name=params['image']['name'],
        flavor_name=params['flavor']['name'],
        group_name=params['group_name'],
        keypair_name=params['keypair_name'],
        vol_sz=params.get('vol_sz'),
        network_zone_name=params.get("network_zone_name"),
        flt_ip_pool=params.get('flt_ip_pool'),
        name_templ=params.get('name_templ'),
        scheduler_hints={"group": params['aa_group_name']},
        security_group=params['security_group'],
        sec_group_size=srv_count
    )

    # precache all errors before start creating vms
    private_key_path = params['keypair_file_private']
    creds = params['image']['creds']
    creds.format(ip="1.1.1.1", private_key_path="/some_path/xx")

    for ip, os_node in create_vms_mt(NOVA_CONNECTION, count, **vm_params):

        conn_uri = creds.format(ip=ip, private_key_path=private_key_path)
        yield Node(conn_uri, []), os_node.id


def get_free_server_grpoups(nova, template=None):
    for g in nova.server_groups.list():
        if g.members == []:
            if re.match(template, g.name):
                yield str(g.name)


def create_vms_mt(nova, amount, group_name, keypair_name, img_name,
                  flavor_name, vol_sz=None, network_zone_name=None,
                  flt_ip_pool=None, name_templ='wally-{id}',
                  scheduler_hints=None, security_group=None,
                  sec_group_size=None):

    with ThreadPoolExecutor(max_workers=16) as executor:
        if network_zone_name is not None:
            network_future = executor.submit(nova.networks.find,
                                             label=network_zone_name)
        else:
            network_future = None

        fl_future = executor.submit(nova.flavors.find, name=flavor_name)
        img_future = executor.submit(nova.images.find, name=img_name)

        if flt_ip_pool is not None:
            ips_future = executor.submit(get_floating_ips,
                                         nova, flt_ip_pool, amount)
            logger.debug("Wait for floating ip")
            ips = ips_future.result()
            ips += [Allocate] * (amount - len(ips))
        else:
            ips = [None] * amount

        logger.debug("Getting flavor object")
        fl = fl_future.result()
        logger.debug("Getting image object")
        img = img_future.result()

        if network_future is not None:
            logger.debug("Waiting for network results")
            nics = [{'net-id': network_future.result().id}]
        else:
            nics = None

        names = []
        for i in range(amount):
            names.append(name_templ.format(group=group_name, id=i))

        futures = []
        logger.debug("Requesting new vm's")

        orig_scheduler_hints = scheduler_hints.copy()

        MAX_SHED_GROUPS = 32
        for start_idx in range(MAX_SHED_GROUPS):
            pass

        group_name_template = scheduler_hints['group'].format("\\d+")
        groups = list(get_free_server_grpoups(nova, group_name_template + "$"))
        groups.sort()

        for idx, (name, flt_ip) in enumerate(zip(names, ips), 2):

            scheduler_hints = None
            if orig_scheduler_hints is not None and sec_group_size is not None:
                if "group" in orig_scheduler_hints:
                    scheduler_hints = orig_scheduler_hints.copy()
                    scheduler_hints['group'] = groups[idx // sec_group_size]

            if scheduler_hints is None:
                scheduler_hints = orig_scheduler_hints.copy()

            params = (nova, name, keypair_name, img, fl,
                      nics, vol_sz, flt_ip, scheduler_hints,
                      flt_ip_pool, [security_group])

            futures.append(executor.submit(create_vm, *params))
        res = [future.result() for future in futures]
        logger.debug("Done spawning")
        return res


def create_vm(nova, name, keypair_name, img,
              fl, nics, vol_sz=None,
              flt_ip=False,
              scheduler_hints=None,
              pool=None,
              security_groups=None):
    for i in range(3):
        srv = nova.servers.create(name,
                                  flavor=fl,
                                  image=img,
                                  nics=nics,
                                  key_name=keypair_name,
                                  scheduler_hints=scheduler_hints,
                                  security_groups=security_groups)

        if not wait_for_server_active(nova, srv):
            msg = "Server {0} fails to start. Kill it and try again"
            logger.debug(msg.format(srv))
            nova.servers.delete(srv)

            try:
                for j in range(120):
                    srv = nova.servers.get(srv.id)
                    time.sleep(1)
                else:
                    msg = "Server {0} delete timeout".format(srv.id)
                    raise RuntimeError(msg)
            except NotFound:
                pass
        else:
            break
    else:
        raise RuntimeError("Failed to start server".format(srv.id))

    if vol_sz is not None:
        vol = create_volume(vol_sz, name)
        nova.volumes.create_server_volume(srv.id, vol.id, None)

    if flt_ip is Allocate:
        flt_ip = nova.floating_ips.create(pool)

    if flt_ip is not None:
        srv.add_floating_ip(flt_ip)

    return flt_ip.ip, nova.servers.get(srv.id)


def clear_nodes(nodes_ids):
    clear_all(NOVA_CONNECTION, nodes_ids, None)


def clear_all(nova, ids=None, name_templ=None):

    def need_delete(srv):
        if name_templ is not None:
            return re.match(name_templ.format("\\d+"), srv.name) is not None
        else:
            return srv.id in ids

    volumes_to_delete = []
    cinder = cinder_connect()
    for vol in cinder.volumes.list():
        for attachment in vol.attachments:
            if attachment['server_id'] in ids:
                volumes_to_delete.append(vol)
                break

    deleted_srvs = set()
    for srv in nova.servers.list():
        if need_delete(srv):
            logger.debug("Deleting server {0}".format(srv.name))
            nova.servers.delete(srv)
            deleted_srvs.add(srv.id)

    count = 0
    while True:
        if count % 60 == 0:
            logger.debug("Waiting till all servers are actually deleted")
        all_id = set(srv.id for srv in nova.servers.list())
        if len(all_id.intersection(deleted_srvs)) == 0:
            break
        count += 1
        time.sleep(1)
    logger.debug("Done, deleting volumes")

    # wait till vm actually deleted

    # logger.warning("Volume deletion commented out")
    for vol in volumes_to_delete:
        logger.debug("Deleting volume " + vol.display_name)
        cinder.volumes.delete(vol)

    logger.debug("Clearing done (yet some volumes may still deleting)")