sr.py

from __future__ import annotations

import pytest

import enum
import functools
import logging
import re
import shlex
import time

import lib.commands as commands
from lib.common import (
    GiB,
    _param_add,
    _param_clear,
    _param_get,
    _param_remove,
    _param_set,
    prefix_object_name,
    randid,
    safe_split,
    strtobool,
    wait_for,
    wait_for_not,
)
from lib.vdi import VDI, ImageFormat

from typing import TYPE_CHECKING, Literal, overload

if TYPE_CHECKING:
    from lib.host import Host
    from lib.pool import Pool

QUICKTEST_BIN = "/opt/xensource/debug/quicktest"

QUICKTEST_SR_SUITES = (
    "cbt,copy,SR tests,Quicktest_vdi,Quicktest_async_calls,"
    "Quicktest_vm_import_export,Quicktest_vm_lifecycle,Quicktest_vm_snapshot,"
    "Quicktest_vdi_ops_data_integrity,Quicktest_max_vdi_size,Quicktest_static_vdis"
)

QUICKTEST_COMMON_SUITES = (
    "Quicktest_example,Quicktest_message,xenstore,event,import_raw_vdi,"
    "Quicktest_date,Quicktest_crypt_r,http,unixext,Timer"
)

class QuicktestScoping(enum.Enum):
    WITH_TAG_PARAM = enum.auto()
    RUN_ONLY_PARAM = enum.auto()
    NO_PARAM = enum.auto()

@functools.lru_cache(maxsize=None)
def _quicktest_scoping(hostname_or_ip: str) -> QuicktestScoping:
    tags_output = commands.ssh_with_result(hostname_or_ip, f"{QUICKTEST_BIN} -list-tags")
    if tags_output.returncode == 0 and re.search(r"^sr:", tags_output.stdout, re.MULTILINE):
        return QuicktestScoping.WITH_TAG_PARAM

    help_output = commands.ssh(hostname_or_ip, f"{QUICKTEST_BIN} -help", check=False)
    if "-run-only" in help_output:
        return QuicktestScoping.RUN_ONLY_PARAM

    return QuicktestScoping.NO_PARAM

class SR:
    xe_prefix = 'sr'

    def __init__(self, uuid: str, pool: Pool):
        self.uuid = uuid
        self.pool = pool
        self._is_shared: bool | None = None # cached value for is_shared()
        self._main_host: Host | None = None # cached value for main_host()
        self._type: str | None = None # cache value for get_type()

    def pbd_uuids(self) -> list[str]:
        return safe_split(self.pool.master.xe('pbd-list', {'sr-uuid': self.uuid}, minimal=True))

    def pbd_for_host(self, host: Host) -> str:
        return safe_split(self.pool.master.xe(
            'pbd-list',
            {'sr-uuid': self.uuid, 'host-uuid': host.uuid},
            minimal=True
        ))[0]

    def unplug_pbd(self, pbd_uuid: str, force: bool = False) -> None:
        try:
            self.pool.master.xe('pbd-unplug', {'uuid': pbd_uuid})
        except commands.SSHCommandFailed as e:
            # We must be sure to execute correctly "unplug" on unplugged VDIs without error
            # if force is set.
            if not force:
                raise
            logging.warning('Ignore exception during PBD unplug: {}'.format(e))

    def unplug_pbds(self, force: bool = False) -> None:
        logging.info(f"Unplug PBDs for SR {self.uuid}")
        for pbd_uuid in self.pbd_uuids():
            self.unplug_pbd(pbd_uuid, force=force)

    def all_pbds_attached(self) -> bool:
        all_attached = True
        for pbd_uuid in self.pbd_uuids():
            all_attached = all_attached and strtobool(self.pool.master.xe('pbd-param-get',
                                                                          {'uuid': pbd_uuid,
                                                                           'param-name': 'currently-attached',
                                                                           }))
        return all_attached

    def plug_pbd(self, pbd_uuid: str) -> None:
        self.pool.master.xe('pbd-plug', {'uuid': pbd_uuid})

    def plug_pbds(self, verify: bool = True) -> None:
        logging.info("Attach PBDs")
        for pbd_uuid in self.pbd_uuids():
            self.plug_pbd(pbd_uuid)
        if verify:
            wait_for(self.all_pbds_attached, "Wait for PBDs attached")

    def try_plug_pbds(self) -> bool:
        try:
            self.plug_pbds(verify=True)
            return True
        except commands.SSHCommandFailed:
            return False

    def vdi_uuids(self, managed: bool = False, name_label: str | None = None) -> list[str]:
        args: dict[str, str | bool | dict[str, str]] = {
            'sr-uuid': self.uuid,
            'managed': managed
        }
        if name_label is not None:
            args['name-label'] = name_label
        return safe_split(self.pool.master.xe('vdi-list', args, minimal=True))

    def destroy(self, verify: bool = False, force: bool = False) -> None:
        logging.info(f"Will attempt SR destroy on {self.uuid}...")
        # Rescan SR to improve the chances of the forced GC run triggered by sr-destroy
        # remove all VDIs in one pass and such have sr-destroy working on first try.
        self.scan()
        max_tries = 5
        for i in range(1, max_tries + 1): # [1, 2, ..., max_tries]
            self.unplug_pbds(force)
            logging.info(f"Destroy SR {self.uuid} (attempt {i})")
            try:
                # Note: sr-destroy triggers ONE forced GC run
                # This may not be enough in some cases
                # (when VDIs to GC are not all leafs and would require several runs)
                self.pool.master.xe('sr-destroy', {'uuid': self.uuid})
            except commands.SSHCommandFailed as e:
                if "the SR is not empty" not in e.stdout:
                    raise
                else:
                    logging.info(f"SR destroy failed with message: {e.stdout}")
                    try:
                        self.plug_pbds()
                        # rescan for an up to date list of VDIs
                        self.scan()
                    except commands.SSHCommandFailed:
                        raise Exception("SR destroy failed and then pbd-plug failed too. Can't continue further.")
                    output = self.vdi_uuids(managed=True)
                    if len(output) > 0:
                        raise Exception("SR destroy failed due to SR not empty, "
                                        "and there are indeed managed VDIs left on the SR.")
                    else:
                        logging.info("SR destroy failed due to SR not empty but there aren't any managed VDIs left.")
                        if i < max_tries:
                            if i == max_tries - 1:
                                # We tried already 4 times to destroy the SR, and there still are hidden VDIs that
                                # couldn't be force-GCed. In this case, we likely need to give time to the normal GC
                                # to run, which might also coalesce some VDIs if that's what it really needs.
                                # The GC should kick approximately 5 minutes after the last operation we did, so let's
                                # give it these 5 minutes plus extra time to complete.
                                gc_delay = 600
                                logging.warning(f"SR destroy failed {i} times in a row. "
                                                f"Wait for {gc_delay}s, hoping GC fully runs before next try")
                                time.sleep(gc_delay)
                            logging.info("Retrying sr-destroy in case it previously failed due to incomplete GC.")
                            continue
                        else:
                            raise Exception(f"Could not destroy the SR even after {i} attempts.")
            if verify:
                wait_for_not(self.exists, "Wait for SR destroyed")
            # Everything apparently went fine. Get out of the retry loop.
            break

    def forget(self, force: bool = False) -> None:
        self.unplug_pbds(force)
        logging.info("Forget SR " + self.uuid)
        self.pool.master.xe('sr-forget', {'uuid': self.uuid})

    def exists(self) -> bool:
        return self.pool.master.xe('sr-list', {'uuid': self.uuid}, minimal=True) == self.uuid

    def scan(self) -> None:
        logging.info("Scan SR " + self.uuid)
        self.pool.master.xe('sr-scan', {'uuid': self.uuid})

    def hosts_uuids(self) -> list[str]:
        return safe_split(self.pool.master.xe('pbd-list', {'sr-uuid': self.uuid, 'params': 'host-uuid'}, minimal=True))

    def attached_to_host(self, host: Host) -> bool:
        return host.uuid in self.hosts_uuids()

    def main_host(self) -> Host:
        """ Returns the host in case of a local SR, the master host in case of a shared SR. """
        if self._main_host is None:
            if self.is_shared():
                self._main_host = self.pool.master
            else:
                self._main_host = self.pool.get_host_by_uuid(self.hosts_uuids()[0])
        return self._main_host

    @overload
    def param_get(self, param_name: str, key: str | None = ..., accept_unknown_key: Literal[False] = ...) -> str:
        ...

    @overload
    def param_get(
        self, param_name: str, key: str | None = ..., accept_unknown_key: Literal[True] = ...
    ) -> str | None:
        ...

    def param_get(self, param_name: str, key: str | None = None, accept_unknown_key: bool = False) -> str | None:
        return _param_get(self.pool.master, self.xe_prefix, self.uuid, param_name, key, accept_unknown_key)

    def param_set(self, param_name: str, value: str | bool | dict[str, str], key: str | None = None) -> None:
        _param_set(self.pool.master, self.xe_prefix, self.uuid, param_name, value, key)

    def param_remove(self, param_name: str, key: str, accept_unknown_key: bool = False) -> None:
        _param_remove(self.pool.master, self.xe_prefix, self.uuid, param_name, key, accept_unknown_key)

    def param_add(self, param_name: str, value: str, key: str | None = None) -> None:
        _param_add(self.pool.master, self.xe_prefix, self.uuid, param_name, value, key)

    def param_clear(self, param_name: str) -> None:
        _param_clear(self.pool.master, self.xe_prefix, self.uuid, param_name)

    def content_type(self) -> str:
        return self.param_get('content-type')

    def is_shared(self) -> bool:
        if self._is_shared is None:
            self._is_shared = strtobool(self.param_get('shared'))
        return self._is_shared

    def get_type(self) -> str:
        if self._type is None:
            self._type = self.param_get('type')
        return self._type

    def get_name_label(self) -> str:
        return self.param_get('name-label')

    def create_vdi(
        self, name_label: str | None = None, virtual_size: int = 1 * GiB, image_format: ImageFormat | None = None
    ) -> VDI:
        name_label = name_label or f'test-vdi-{randid()}'
        logging.info("Create VDI %r on SR %s", name_label, self.uuid)
        args: dict[str, str | bool | dict[str, str]] = {
            'name-label': prefix_object_name(name_label),
            'virtual-size': str(virtual_size),
            'sr-uuid': self.uuid,
        }
        if image_format:
            args["sm-config:image-format"] = image_format
        vdi_uuid = self.pool.master.xe('vdi-create', args)
        return VDI(vdi_uuid, sr=self)

    def run_quicktest(self, sr_specific: bool = True) -> None:
        scoping = _quicktest_scoping(self.pool.master.hostname_or_ip)
        cmd = f"{QUICKTEST_BIN} -sr {self.uuid}"

        if scoping is QuicktestScoping.WITH_TAG_PARAM:
            cmd += " -with-tag sr" if sr_specific else " -without-tag sr"
        elif scoping is QuicktestScoping.RUN_ONLY_PARAM:
            suites = QUICKTEST_SR_SUITES if sr_specific else QUICKTEST_COMMON_SUITES
            cmd += f" -run-only {shlex.quote(suites)}"
        elif not sr_specific:
            # QuicktestScoping.NO_PARAM: no way to select just the common suites, and every
            # per-SR pass on this host is already unfiltered, so this run adds nothing.
            pytest.skip("quicktest has no scoping support on this host; "
                        "common suites are already covered by the per-SR runs.")

        logging.info(f"Run quicktest on SR {self.uuid}: {cmd}")

        # Always display the output of quicktest, failed or not.
        # This will duplicate the output in some cases, but it ensures we always have it for failure analysis,
        # even when quicktest leaves SRs in a state which makes teardown fail (in this case, pytest often doesn't
        # manage to display the details of the failed command, for a reason unknown - no usable reproducer found)
        try:
            output = self.pool.master.ssh(cmd)
            logging.info(f"Quicktest output: {output}")
        except commands.SSHCommandFailed as e:
            logging.error(f"Quicktest output: {e.stdout}")
            raise