diff --git a/krkn/rollback/__init__.py b/krkn/rollback/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/krkn/rollback/command.py b/krkn/rollback/command.py new file mode 100644 index 00000000..b1b0cdb8 --- /dev/null +++ b/krkn/rollback/command.py @@ -0,0 +1,121 @@ +import os +import logging +from typing import Optional, TYPE_CHECKING + +from krkn.rollback.config import RollbackConfig +from krkn.rollback.handler import execute_rollback_version_files, cleanup_rollback_version_files + + + +if TYPE_CHECKING: + from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift + + +def list_rollback(run_uuid: Optional[str]=None, scenario_type: Optional[str]=None): + """ + List rollback version files in a tree-like format. + + :param cfg: Configuration file path + :param run_uuid: Optional run UUID to filter by + :param scenario_type: Optional scenario type to filter by + :return: Exit code (0 for success, 1 for error) + """ + logging.info("Listing rollback version files") + + versions_directory = RollbackConfig().versions_directory + + logging.info(f"Rollback versions directory: {versions_directory}") + + # Check if the directory exists first + if not os.path.exists(versions_directory): + logging.info(f"Rollback versions directory does not exist: {versions_directory}") + return 0 + + # List all directories and files + try: + # Get all run directories + run_dirs = [] + for item in os.listdir(versions_directory): + item_path = os.path.join(versions_directory, item) + if os.path.isdir(item_path): + # Apply run_uuid filter if specified + if run_uuid is None or run_uuid in item: + run_dirs.append(item) + + if not run_dirs: + if run_uuid: + logging.info(f"No rollback directories found for run_uuid: {run_uuid}") + else: + logging.info("No rollback directories found") + return 0 + + # Sort directories for consistent output + run_dirs.sort() + + print(f"\n{versions_directory}/") + for i, run_dir in enumerate(run_dirs): + is_last_dir = (i == len(run_dirs) - 1) + dir_prefix = "└── " if is_last_dir else "├── " + print(f"{dir_prefix}{run_dir}/") + + # List files in this directory + run_dir_path = os.path.join(versions_directory, run_dir) + try: + files = [] + for file in os.listdir(run_dir_path): + file_path = os.path.join(run_dir_path, file) + if os.path.isfile(file_path): + # Apply scenario_type filter if specified + if scenario_type is None or file.startswith(scenario_type): + files.append(file) + + files.sort() + for j, file in enumerate(files): + is_last_file = (j == len(files) - 1) + file_prefix = " └── " if is_last_dir else "│ └── " if is_last_file else ("│ ├── " if not is_last_dir else " ├── ") + print(f"{file_prefix}{file}") + + except PermissionError: + file_prefix = " └── " if is_last_dir else "│ └── " + print(f"{file_prefix}[Permission Denied]") + + except Exception as e: + logging.error(f"Error listing rollback directory: {e}") + return 1 + + return 0 + + +def execute_rollback(telemetry_ocp: "KrknTelemetryOpenshift", run_uuid: Optional[str]=None, scenario_type: Optional[str]=None): + """ + Execute rollback version files and cleanup if successful. + + :param telemetry_ocp: Instance of KrknTelemetryOpenshift + :param run_uuid: Optional run UUID to filter by + :param scenario_type: Optional scenario type to filter by + :return: Exit code (0 for success, 1 for error) + """ + logging.info("Executing rollback version files") + + if not run_uuid: + logging.error("run_uuid is required for execute-rollback command") + return 1 + + if not scenario_type: + logging.warning("scenario_type is not specified, executing all scenarios in rollback directory") + + try: + # Execute rollback version files + logging.info(f"Executing rollback for run_uuid={run_uuid}, scenario_type={scenario_type or '*'}") + execute_rollback_version_files(telemetry_ocp, run_uuid, scenario_type) + + # If execution was successful, cleanup the version files + logging.info("Rollback execution completed successfully, cleaning up version files") + cleanup_rollback_version_files(run_uuid, scenario_type) + + logging.info("Rollback execution and cleanup completed successfully") + return 0 + + except Exception as e: + logging.error(f"Error during rollback execution: {e}") + return 1 diff --git a/krkn/rollback/config.py b/krkn/rollback/config.py new file mode 100644 index 00000000..76277f16 --- /dev/null +++ b/krkn/rollback/config.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, TYPE_CHECKING, Optional +from typing_extensions import TypeAlias +import time +import os +import logging + +from krkn_lib.utils import get_random_string + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift + +RollbackCallable: TypeAlias = Callable[ + ["RollbackContent", "KrknTelemetryOpenshift"], None +] + + +if TYPE_CHECKING: + from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift + +RollbackCallable: TypeAlias = Callable[ + ["RollbackContent", "KrknTelemetryOpenshift"], None +] + + +class SingletonMeta(type): + _instances = {} + + def __call__(cls, *args, **kwargs): + if cls not in cls._instances: + cls._instances[cls] = super().__call__(*args, **kwargs) + return cls._instances[cls] + + +@dataclass(frozen=True) +class RollbackContent: + """ + RollbackContent is a dataclass that defines the necessary fields for rollback operations. + """ + + resource_identifier: str + namespace: Optional[str] = None + + def __str__(self): + namespace = f'"{self.namespace}"' if self.namespace else "None" + resource_identifier = f'"{self.resource_identifier}"' + return f"RollbackContent(namespace={namespace}, resource_identifier={resource_identifier})" + + +class RollbackContext(str): + """ + RollbackContext is a string formatted as '-'. + It represents the context for rollback operations, uniquely identifying a run. + """ + + def __new__(cls, run_uuid: str): + return super().__new__(cls, f"{time.time_ns()}-{run_uuid}") + + +class RollbackConfig(metaclass=SingletonMeta): + """Configuration for the rollback scenarios.""" + + def __init__(self): + self._auto = False + self._versions_directory = "" + self._registered = False + + @property + def auto(self): + return self._auto + + @auto.setter + def auto(self, value): + if self._registered: + raise AttributeError("Can't modify 'auto' after registration") + self._auto = value + + @property + def versions_directory(self): + return self._versions_directory + + @versions_directory.setter + def versions_directory(self, value): + if self._registered: + raise AttributeError("Can't modify 'versions_directory' after registration") + self._versions_directory = value + @classmethod + def register(cls, auto=False, versions_directory=""): + """Initialize and return the singleton instance with given configuration.""" + instance = cls() + instance.auto = auto + instance.versions_directory = versions_directory + instance._registered = True + return instance + + @classmethod + def get_rollback_versions_directory(cls, rollback_context: RollbackContext) -> str: + """ + Get the rollback context directory for a given rollback context. + + :param rollback_context: The rollback context string. + :return: The path to the rollback context directory. + """ + return f"{cls().versions_directory}/{rollback_context}" + + @classmethod + def search_rollback_version_files(cls, run_uuid: str, scenario_type: str | None = None) -> list[str]: + """ + Search for rollback version files based on run_uuid and scenario_type. + + 1. Search directories with "run_uuid" in name under "cls.versions_directory". + 2. Search files in those directories that start with "scenario_type" in matched directories in step 1. + + :param run_uuid: Unique identifier for the run. + :param scenario_type: Type of the scenario. + :return: List of version file paths. + """ + rollback_context_directories = [ + dirname for dirname in os.listdir(cls().versions_directory) if run_uuid in dirname + ] + if len(rollback_context_directories) != 1: + raise ValueError( + f"Expected one directory for run UUID {run_uuid}, found: {rollback_context_directories}" + ) + rollback_context_directory = rollback_context_directories[0] + + version_files = [] + scenario_rollback_versions_directory = os.path.join( + cls().versions_directory, rollback_context_directory + ) + for file in os.listdir(scenario_rollback_versions_directory): + # assert all files start with scenario_type and end with .py + if file.endswith(".py") and (scenario_type is None or file.startswith(scenario_type)): + version_files.append( + os.path.join(scenario_rollback_versions_directory, file) + ) + else: + logger.warning( + f"File {file} does not match expected pattern for scenario type {scenario_type}" + ) + return version_files + +@dataclass(frozen=True) +class Version: + scenario_type: str + rollback_context: RollbackContext + timestamp: int = time.time_ns() # Get current timestamp in nanoseconds + hash_suffix: str = get_random_string(8) # Generate a random string of 8 characters + + @property + def version_file_name(self) -> str: + """ + Generate a version file name based on the timestamp and hash suffix. + :return: The generated version file name. + """ + return f"{self.scenario_type}_{self.timestamp}_{self.hash_suffix}.py" + + @property + def version_file_full_path(self) -> str: + """ + Get the full path for the version file based on the version object and current context. + + :return: The generated version file full path. + """ + return f"{RollbackConfig.get_rollback_versions_directory(self.rollback_context)}/{self.version_file_name}" + + @staticmethod + def new_version(scenario_type: str, rollback_context: RollbackContext) -> "Version": + """ + Get the current version of the rollback configuration. + :return: An instance of Version with the current timestamp and hash suffix. + """ + return Version( + scenario_type=scenario_type, + rollback_context=rollback_context, + ) diff --git a/krkn/rollback/handler.py b/krkn/rollback/handler.py new file mode 100644 index 00000000..01dd1998 --- /dev/null +++ b/krkn/rollback/handler.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +import logging +from typing import cast, TYPE_CHECKING +import os +import importlib.util +import inspect + +from krkn.rollback.config import RollbackConfig, RollbackContext, Version + + +logger = logging.getLogger(__name__) + + +if TYPE_CHECKING: + from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift + + from krkn.scenario_plugins.abstract_scenario_plugin import AbstractScenarioPlugin + from krkn.rollback.config import RollbackContent, RollbackCallable + from krkn.rollback.serialization import Serializer + + +def set_rollback_context_decorator(func): + """ + Decorator to automatically set and clear rollback context. + It extracts run_uuid from the function arguments and sets the context in rollback_handler + before executing the function, and clears it after execution. + + Usage: + + .. code-block:: python + from krkn.rollback.handler import set_rollback_context_decorator + # for any scenario plugin that inherits from AbstractScenarioPlugin + @set_rollback_context_decorator + def run( + self, + run_uuid: str, + scenario: str, + krkn_config: dict[str, any], + lib_telemetry: KrknTelemetryOpenshift, + scenario_telemetry: ScenarioTelemetry, + ): + # Your scenario logic here + pass + """ + + def wrapper(self, *args, **kwargs): + self = cast("AbstractScenarioPlugin", self) + # Since `AbstractScenarioPlugin.run_scenarios` will call `self.run` and pass all parameters as `kwargs` + logger.debug(f"kwargs of ScenarioPlugin.run: {kwargs}") + run_uuid = kwargs.get("run_uuid", None) + # so we can safely assume that `run_uuid` will be present in `kwargs` + assert run_uuid is not None, "run_uuid must be provided in kwargs" + + # Set context if run_uuid is available and rollback_handler exists + if run_uuid and hasattr(self, "rollback_handler"): + self.rollback_handler = cast("RollbackHandler", self.rollback_handler) + self.rollback_handler.set_context(run_uuid) + + try: + # Execute the `run` method with the original arguments + result = func(self, *args, **kwargs) + return result + finally: + # Clear context after function execution, regardless of success or failure + if hasattr(self, "rollback_handler"): + self.rollback_handler = cast("RollbackHandler", self.rollback_handler) + self.rollback_handler.clear_context() + + return wrapper + +def _parse_rollback_module(version_file_path: str) -> tuple[RollbackCallable, RollbackContent]: + """ + Parse a rollback module to extract the rollback function and RollbackContent. + + :param version_file_path: Path to the rollback version file + :return: Tuple of (rollback_callable, rollback_content) + """ + + # Create a unique module name based on the file path + module_name = f"rollback_module_{os.path.basename(version_file_path).replace('.py', '').replace('-', '_')}" + + # Load the module using importlib + spec = importlib.util.spec_from_file_location(module_name, version_file_path) + if spec is None or spec.loader is None: + raise ImportError(f"Could not load module from {version_file_path}") + + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + # Find the rollback function + rollback_callable = None + for name, obj in inspect.getmembers(module): + if inspect.isfunction(obj) and name.startswith('rollback_'): + # Check function signature + sig = inspect.signature(obj) + params = list(sig.parameters.values()) + if (len(params) == 2 and + 'RollbackContent' in str(params[0].annotation) and + 'KrknTelemetryOpenshift' in str(params[1].annotation)): + rollback_callable = obj + logger.debug(f"Found rollback function: {name}") + break + + if rollback_callable is None: + raise ValueError(f"No valid rollback function found in {version_file_path}") + + # Find the rollback_content variable + if not hasattr(module, 'rollback_content'): + raise ValueError("Could not find variable named 'rollback_content' in the module") + + rollback_content = getattr(module, 'rollback_content', None) + if rollback_content is None: + raise ValueError("Variable 'rollback_content' is None") + + logger.debug(f"Found rollback_content variable in module: {rollback_content}") + return rollback_callable, rollback_content + + +def execute_rollback_version_files(telemetry_ocp: "KrknTelemetryOpenshift", run_uuid: str, scenario_type: str | None = None): + """ + Execute rollback version files for the given run_uuid and scenario_type. + This function is called when a signal is received to perform rollback operations. + + :param run_uuid: Unique identifier for the run. + :param scenario_type: Type of the scenario being rolled back. + """ + + # Get the rollback versions directory + version_files = RollbackConfig.search_rollback_version_files(run_uuid, scenario_type) + + # Execute all version files in the directory + logger.info(f"Executing rollback version files for run_uuid={run_uuid}, scenario_type={scenario_type or '*'}") + for version_file in version_files: + try: + logger.info(f"Executing rollback version file: {version_file}") + + # Parse the rollback module to get function and content + rollback_callable, rollback_content = _parse_rollback_module(version_file) + # Execute the rollback function + logger.info('Executing rollback callable...') + rollback_callable(rollback_content, telemetry_ocp) + logger.info('Rollback completed.') + + logger.info(f"Executed {version_file} successfully.") + except Exception as e: + logger.error(f"Failed to execute rollback version file {version_file}: {e}") + raise + + +def cleanup_rollback_version_files(run_uuid: str, scenario_type: str): + """ + Cleanup rollback version files for the given run_uuid and scenario_type. + This function is called to remove the rollback version files after execution. + + :param run_uuid: Unique identifier for the run. + :param scenario_type: Type of the scenario being rolled back. + """ + + # Get the rollback versions directory + version_files = RollbackConfig.search_rollback_version_files(run_uuid, scenario_type) + + # Remove all version files in the directory + logger.info(f"Cleaning up rollback version files for run_uuid={run_uuid}, scenario_type={scenario_type}") + for version_file in version_files: + try: + os.remove(version_file) + logger.info(f"Removed {version_file} successfully.") + except Exception as e: + logger.error(f"Failed to remove rollback version file {version_file}: {e}") + raise + + +class RollbackHandler: + def __init__( + self, + scenario_type: str, + serializer: "Serializer", + ): + self.scenario_type = scenario_type + self.serializer = serializer + self.rollback_context: RollbackContext | None = ( + None # will be set when `set_context` is called + ) + + def set_context(self, run_uuid: str): + """ + Set the context for the rollback handler. + :param run_uuid: Unique identifier for the run. + """ + self.rollback_context = RollbackContext(run_uuid) + logger.info( + f"Set rollback_context: {self.rollback_context} for scenario_type: {self.scenario_type} RollbackHandler" + ) + + def clear_context(self): + """ + Clear the run_uuid context for the rollback handler. + """ + logger.debug( + f"Clear rollback_context {self.rollback_context} for scenario type {self.scenario_type} RollbackHandler" + ) + self.rollback_context = None + + def set_rollback_callable( + self, + callable: "RollbackCallable", + rollback_content: "RollbackContent", + ): + """ + Set the rollback callable to be executed after the scenario is finished. + + :param callable: The rollback callable to be set. + :param rollback_content: The rollback content for the callable. + """ + logger.debug( + f"Rollback callable set to {callable.__name__} for version directory {RollbackConfig.get_rollback_versions_directory(self.rollback_context)}" + ) + + version: Version = Version.new_version( + scenario_type=self.scenario_type, + rollback_context=self.rollback_context, + ) + + # Serialize the callable to a file + try: + version_file = self.serializer.serialize_callable( + callable, rollback_content, version + ) + logger.info(f"Rollback callable serialized to {version_file}") + except Exception as e: + logger.error(f"Failed to serialize rollback callable: {e}") diff --git a/krkn/rollback/serialization.py b/krkn/rollback/serialization.py new file mode 100644 index 00000000..f59e3549 --- /dev/null +++ b/krkn/rollback/serialization.py @@ -0,0 +1,123 @@ +import inspect +import os +import logging +from typing import TYPE_CHECKING + +from jinja2 import Environment, FileSystemLoader + +if TYPE_CHECKING: + from krkn.rollback.config import RollbackCallable, RollbackContent, Version + +logger = logging.getLogger(__name__) + + +class Serializer: + def __init__(self, scenario_type: str): + self.scenario_type = scenario_type + # Set up Jinja2 environment to load templates from the rollback directory + template_dir = os.path.join(os.path.dirname(__file__)) + env = Environment(loader=FileSystemLoader(template_dir)) + self.template = env.get_template("version_template.j2") + + def _parse_rollback_callable_code( + self, rollback_callable: "RollbackCallable" + ) -> tuple[str, str]: + """ + Parse the rollback callable code to extract its implementation. + :param rollback_callable: The callable function to parse (can be staticmethod or regular function). + :return: A tuple containing (function_name, function_code). + """ + # Get the implementation code of the rollback_callable + rollback_callable_code = inspect.getsource(rollback_callable) + + # Split into lines for processing + code_lines = rollback_callable_code.split("\n") + cleaned_lines = [] + function_name = None + + # Find the function definition line and extract function name + def_line_index = None + for i, line in enumerate(code_lines): + # Skip decorators (including @staticmethod) + if line.strip().startswith("@"): + continue + + # Look for function definition + if line.strip().startswith("def "): + def_line_index = i + # Extract function name from the def line + def_line = line.strip() + if "(" in def_line: + function_name = def_line.split("def ")[1].split("(")[0].strip() + break + + if def_line_index is None or function_name is None: + raise ValueError( + "Could not find function definition in callable source code" + ) + + # Get the base indentation level from the def line + def_line = code_lines[def_line_index] + base_indent_level = len(def_line) - len(def_line.lstrip()) + + # Process all lines starting from the def line + for i in range(def_line_index, len(code_lines)): + line = code_lines[i] + + # Handle empty lines + if not line.strip(): + cleaned_lines.append("") + continue + + # Calculate current line's indentation + current_indent = len(line) - len(line.lstrip()) + + # Remove the base indentation to normalize to function level + if current_indent >= base_indent_level: + # Remove base indentation + normalized_line = line[base_indent_level:] + cleaned_lines.append(normalized_line) + else: + # This shouldn't happen in well-formed code, but handle it gracefully + cleaned_lines.append(line.lstrip()) + + # Reconstruct the code and clean up trailing whitespace + function_code = "\n".join(cleaned_lines).rstrip() + + return function_name, function_code + + def serialize_callable( + self, + rollback_callable: "RollbackCallable", + rollback_content: "RollbackContent", + version: "Version", + ) -> str: + """ + Serialize a callable function to a file with its arguments and keyword arguments. + :param rollback_callable: The callable to serialize. + :param rollback_content: The rollback content for the callable. + :param version: The version representing the rollback context and file path for the rollback. + :return: Path to the serialized callable file. + """ + + rollback_callable_name, rollback_callable_code = ( + self._parse_rollback_callable_code(rollback_callable) + ) + + # Render the template with the required variables + file_content = self.template.render( + rollback_callable_name=rollback_callable_name, + rollback_callable_code=rollback_callable_code, + rollback_content=str(rollback_content), + ) + + # Write the file to the version directory + os.makedirs(os.path.dirname(version.version_file_full_path), exist_ok=True) + + logger.debug("Creating version file at %s", version.version_file_full_path) + logger.debug("Version file content:\n%s", file_content) + with open(version.version_file_full_path, "w") as f: + f.write(file_content) + logger.info(f"Serialized callable written to {version.version_file_full_path}") + + return version.version_file_full_path diff --git a/krkn/rollback/signal.py b/krkn/rollback/signal.py new file mode 100644 index 00000000..54ac13e1 --- /dev/null +++ b/krkn/rollback/signal.py @@ -0,0 +1,106 @@ +from typing import Dict, Any, Optional +import threading +import signal +import sys +import logging +from contextlib import contextmanager + +from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift + +from krkn.rollback.handler import execute_rollback_version_files + +logger = logging.getLogger(__name__) + +class SignalHandler: + # Class-level variables for signal handling (shared across all instances) + _signal_handlers_installed = False # No need for thread-safe variable due to _signal_lock + _original_handlers: Dict[int, Any] = {} + _signal_lock = threading.Lock() + + # Thread-local storage for context + _local = threading.local() + + @classmethod + def _set_context(cls, run_uuid: str, scenario_type: str, telemetry_ocp: KrknTelemetryOpenshift): + """Set the current execution context for this thread.""" + cls._local.run_uuid = run_uuid + cls._local.scenario_type = scenario_type + cls._local.telemetry_ocp = telemetry_ocp + logger.debug(f"Set signal context set for thread {threading.current_thread().name} - run_uuid={run_uuid}, scenario_type={scenario_type}") + + @classmethod + def _get_context(cls) -> tuple[Optional[str], Optional[str], Optional[KrknTelemetryOpenshift]]: + """Get the current execution context for this thread.""" + run_uuid = getattr(cls._local, 'run_uuid', None) + scenario_type = getattr(cls._local, 'scenario_type', None) + telemetry_ocp = getattr(cls._local, 'telemetry_ocp', None) + return run_uuid, scenario_type, telemetry_ocp + + @classmethod + def _signal_handler(cls, signum: int, frame): + """Handle signals with current thread context information.""" + signal_name = signal.Signals(signum).name + run_uuid, scenario_type, telemetry_ocp = cls._get_context() + if not run_uuid or not scenario_type or not telemetry_ocp: + logger.warning(f"Signal {signal_name} received without complete context, skipping rollback.") + return + + # Clear the context for the next signal, as another signal may arrive before the rollback completes. + # This ensures that the rollback is performed only once. + cls._set_context(None, None, telemetry_ocp) + + # Perform rollback + logger.info(f"Performing rollback for signal {signal_name} with run_uuid={run_uuid}, scenario_type={scenario_type}") + execute_rollback_version_files(telemetry_ocp, run_uuid, scenario_type) + + # Call original handler if it exists + if signum not in cls._original_handlers: + logger.info(f"Signal {signal_name} has no registered handler, exiting...") + return + + original_handler = cls._original_handlers[signum] + if callable(original_handler): + logger.info(f"Calling original handler for {signal_name}") + original_handler(signum, frame) + elif original_handler == signal.SIG_DFL: + # Restore default behavior + logger.info(f"Restoring default signal handler for {signal_name}") + signal.signal(signum, signal.SIG_DFL) + signal.raise_signal(signum) + + @classmethod + def _register_signal_handler(cls): + """Register signal handlers once (called by first instance).""" + with cls._signal_lock: # Lock protects _signal_handlers_installed from race conditions + if cls._signal_handlers_installed: + return + + signals_to_handle = [signal.SIGINT, signal.SIGTERM] + if hasattr(signal, 'SIGHUP'): + signals_to_handle.append(signal.SIGHUP) + + for sig in signals_to_handle: + try: + original_handler = signal.signal(sig, cls._signal_handler) + cls._original_handlers[sig] = original_handler + logger.debug(f"SignalHandler: Registered signal handler for {signal.Signals(sig).name}") + except (OSError, ValueError) as e: + logger.warning(f"AbstractScenarioPlugin: Could not register handler for signal {sig}: {e}") + + cls._signal_handlers_installed = True + logger.info("Signal handlers registered globally") + + @classmethod + @contextmanager + def signal_context(cls, run_uuid: str, scenario_type: str, telemetry_ocp: KrknTelemetryOpenshift): + """Context manager to set the signal context for the current thread.""" + cls._set_context(run_uuid, scenario_type, telemetry_ocp) + cls._register_signal_handler() + try: + yield + finally: + # Clear context after exiting the context manager + cls._set_context(None, None, telemetry_ocp) + + +signal_handler = SignalHandler() \ No newline at end of file diff --git a/krkn/rollback/version_template.j2 b/krkn/rollback/version_template.j2 new file mode 100644 index 00000000..105cdc76 --- /dev/null +++ b/krkn/rollback/version_template.j2 @@ -0,0 +1,55 @@ +# This file is auto-generated by krkn-lib. +# It contains the rollback callable and its arguments for the scenario plugin. + +from dataclasses import dataclass +import os +import logging +from typing import Optional + +from krkn_lib.utils import SafeLogger +from krkn_lib.ocp import KrknOpenshift +from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift + +@dataclass(frozen=True) +class RollbackContent: + resource_identifier: str + namespace: Optional[str] = None + +# Actual rollback callable +{{ rollback_callable_code }} + +# Create necessary variables for execution +lib_openshift = None +lib_telemetry = None +rollback_content = {{ rollback_content }} + + +# Main entry point for execution +if __name__ == '__main__': + # setup logging + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + handlers=[ + logging.StreamHandler(), + ] + ) + + # setup logging and get kubeconfig path + kubeconfig_path = os.getenv("KUBECONFIG", "~/.kube/config") + log_directory = os.path.dirname(os.path.abspath(__file__)) + os.makedirs(os.path.join(log_directory, 'logs'), exist_ok=True) + # setup SafeLogger for telemetry + telemetry_log_path = os.path.join(log_directory, 'logs', 'telemetry.log') + safe_logger = SafeLogger(telemetry_log_path) + # setup krkn-lib objects + lib_openshift = KrknOpenshift(kubeconfig_path=kubeconfig_path) + lib_telemetry = KrknTelemetryOpenshift(safe_logger=safe_logger, lib_openshift=lib_openshift) + + # execute + logging.info('Executing rollback callable...') + {{ rollback_callable_name }}( + rollback_content, + lib_telemetry + ) + logging.info('Rollback completed.') \ No newline at end of file diff --git a/krkn/scenario_plugins/abstract_scenario_plugin.py b/krkn/scenario_plugins/abstract_scenario_plugin.py index e17b15c2..eed363a3 100644 --- a/krkn/scenario_plugins/abstract_scenario_plugin.py +++ b/krkn/scenario_plugins/abstract_scenario_plugin.py @@ -5,9 +5,26 @@ from krkn_lib.models.telemetry import ScenarioTelemetry from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift from krkn import utils - +from krkn.rollback.handler import ( + RollbackHandler, + execute_rollback_version_files, + cleanup_rollback_version_files +) +from krkn.rollback.signal import signal_handler +from krkn.rollback.serialization import Serializer class AbstractScenarioPlugin(ABC): + + def __init__(self, scenario_type: str): + """Initializes the AbstractScenarioPlugin with the scenario type and rollback configuration. + + :param scenario_type: the scenario type defined in the config.yaml + """ + serializer = Serializer( + scenario_type=scenario_type, + ) + self.rollback_handler = RollbackHandler(scenario_type, serializer) + @abstractmethod def run( self, @@ -74,24 +91,38 @@ class AbstractScenarioPlugin(ABC): scenario_telemetry, scenario_config ) - try: - logging.info( - f"Running {self.__class__.__name__}: {self.get_scenario_types()} -> {scenario_config}" - ) - return_value = self.run( - run_uuid, - scenario_config, - krkn_config, - telemetry, - scenario_telemetry, - ) - except Exception as e: - logging.error( - f"uncaught exception on scenario `run()` method: {e} " - f"please report an issue on https://github.com/krkn-chaos/krkn" - ) - return_value = 1 + with signal_handler.signal_context( + run_uuid=run_uuid, + scenario_type=scenario_telemetry.scenario_type, + telemetry_ocp=telemetry + ): + try: + logging.info( + f"Running {self.__class__.__name__}: {self.get_scenario_types()} -> {scenario_config}" + ) + # pass all the parameters by kwargs to make `set_rollback_context_decorator` get the `run_uuid` and `scenario_type` + return_value = self.run( + run_uuid=run_uuid, + scenario=scenario_config, + krkn_config=krkn_config, + lib_telemetry=telemetry, + scenario_telemetry=scenario_telemetry, + ) + except Exception as e: + logging.error( + f"uncaught exception on scenario `run()` method: {e} " + f"please report an issue on https://github.com/krkn-chaos/krkn" + ) + return_value = 1 + # execute rollback files based on the return value + if return_value != 0: + execute_rollback_version_files( + telemetry, run_uuid, scenario_telemetry.scenario_type + ) + cleanup_rollback_version_files( + run_uuid, scenario_telemetry.scenario_type + ) scenario_telemetry.exit_status = return_value scenario_telemetry.end_timestamp = time.time() utils.collect_and_put_ocp_logs( diff --git a/krkn/scenario_plugins/application_outage/application_outage_scenario_plugin.py b/krkn/scenario_plugins/application_outage/application_outage_scenario_plugin.py index 92b3208a..1b96a906 100644 --- a/krkn/scenario_plugins/application_outage/application_outage_scenario_plugin.py +++ b/krkn/scenario_plugins/application_outage/application_outage_scenario_plugin.py @@ -7,9 +7,12 @@ from krkn_lib.utils import get_yaml_item_value, get_random_string from jinja2 import Template from krkn import cerberus from krkn.scenario_plugins.abstract_scenario_plugin import AbstractScenarioPlugin +from krkn.rollback.config import RollbackContent +from krkn.rollback.handler import set_rollback_context_decorator class ApplicationOutageScenarioPlugin(AbstractScenarioPlugin): + @set_rollback_context_decorator def run( self, run_uuid: str, @@ -57,6 +60,13 @@ class ApplicationOutageScenarioPlugin(AbstractScenarioPlugin): # Block the traffic by creating network policy logging.info("Creating the network policy") + self.rollback_handler.set_rollback_callable( + self.rollback_network_policy, + RollbackContent( + namespace=namespace, + resource_identifier=policy_name, + ), + ) lib_telemetry.get_lib_kubernetes().create_net_policy( yaml_spec, namespace ) @@ -89,5 +99,26 @@ class ApplicationOutageScenarioPlugin(AbstractScenarioPlugin): else: return 0 + @staticmethod + def rollback_network_policy( + rollback_content: RollbackContent, + lib_telemetry: KrknTelemetryOpenshift, + ): + """Rollback function to delete the network policy created during the scenario. + + :param rollback_content: Rollback content containing namespace and resource_identifier. + :param lib_telemetry: Instance of KrknTelemetryOpenshift for Kubernetes operations. + """ + try: + namespace = rollback_content.namespace + policy_name = rollback_content.resource_identifier + logging.info( + f"Rolling back network policy: {policy_name} in namespace: {namespace}" + ) + lib_telemetry.get_lib_kubernetes().delete_net_policy(policy_name, namespace) + logging.info("Network policy rollback completed successfully.") + except Exception as e: + logging.error(f"Failed to rollback network policy: {e}") + def get_scenario_types(self) -> list[str]: return ["application_outages_scenarios"] diff --git a/krkn/scenario_plugins/hogs/hogs_scenario_plugin.py b/krkn/scenario_plugins/hogs/hogs_scenario_plugin.py index b82e9acd..6b7ced73 100644 --- a/krkn/scenario_plugins/hogs/hogs_scenario_plugin.py +++ b/krkn/scenario_plugins/hogs/hogs_scenario_plugin.py @@ -16,9 +16,13 @@ from krkn_lib.k8s import KrknKubernetes from krkn_lib.utils import get_random_string from krkn.scenario_plugins.abstract_scenario_plugin import AbstractScenarioPlugin +from krkn.rollback.config import RollbackContent +from krkn.rollback.handler import set_rollback_context_decorator class HogsScenarioPlugin(AbstractScenarioPlugin): + + @set_rollback_context_decorator def run(self, run_uuid: str, scenario: str, krkn_config: dict[str, any], lib_telemetry: KrknTelemetryOpenshift, scenario_telemetry: ScenarioTelemetry) -> int: try: @@ -79,6 +83,13 @@ class HogsScenarioPlugin(AbstractScenarioPlugin): config.node_selector = f"kubernetes.io/hostname={node}" pod_name = f"{config.type.value}-hog-{get_random_string(5)}" node_resources_start = lib_k8s.get_node_resources_info(node) + self.rollback_handler.set_rollback_callable( + self.rollback_hog_pod, + RollbackContent( + namespace=config.namespace, + resource_identifier=pod_name, + ), + ) lib_k8s.deploy_hog(pod_name, config) start = time.time() # waiting 3 seconds before starting sample collection @@ -150,3 +161,22 @@ class HogsScenarioPlugin(AbstractScenarioPlugin): raise exception except queue.Empty: pass + + @staticmethod + def rollback_hog_pod(rollback_content: RollbackContent, lib_telemetry: KrknTelemetryOpenshift): + """ + Rollback function to delete hog pod. + + :param rollback_content: Rollback content containing namespace and resource_identifier. + :param lib_telemetry: Instance of KrknTelemetryOpenshift for Kubernetes operations + """ + try: + namespace = rollback_content.namespace + pod_name = rollback_content.resource_identifier + logging.info( + f"Rolling back hog pod: {pod_name} in namespace: {namespace}" + ) + lib_telemetry.get_lib_kubernetes().delete_pod(pod_name, namespace) + logging.info("Rollback of hog pod completed successfully.") + except Exception as e: + logging.error(f"Failed to rollback hog pod: {e}") diff --git a/krkn/scenario_plugins/kubevirt_vm_outage/kubevirt_vm_outage_scenario_plugin.py b/krkn/scenario_plugins/kubevirt_vm_outage/kubevirt_vm_outage_scenario_plugin.py index 1efec9b5..28790cf4 100644 --- a/krkn/scenario_plugins/kubevirt_vm_outage/kubevirt_vm_outage_scenario_plugin.py +++ b/krkn/scenario_plugins/kubevirt_vm_outage/kubevirt_vm_outage_scenario_plugin.py @@ -20,7 +20,8 @@ class KubevirtVmOutageScenarioPlugin(AbstractScenarioPlugin): This plugin simulates a VM crash or outage scenario and supports automated or manual recovery. """ - def __init__(self): + def __init__(self, scenario_type: str): + super().__init__(scenario_type) self.k8s_client = None self.original_vmi = None diff --git a/krkn/scenario_plugins/scenario_plugin_factory.py b/krkn/scenario_plugins/scenario_plugin_factory.py index bf945435..28dede5c 100644 --- a/krkn/scenario_plugins/scenario_plugin_factory.py +++ b/krkn/scenario_plugins/scenario_plugin_factory.py @@ -33,7 +33,7 @@ class ScenarioPluginFactory: inherits from the AbstractScenarioPlugin abstract class """ if scenario_type in self.loaded_plugins: - return self.loaded_plugins[scenario_type]() + return self.loaded_plugins[scenario_type](scenario_type) else: raise ScenarioPluginNotFound( f"Failed to load the {scenario_type} scenario plugin. " @@ -61,7 +61,10 @@ class ScenarioPluginFactory: continue cls = getattr(module, name) - instance = cls() + # The AbstractScenarioPlugin constructor requires a scenario_type. + # However, since we only need to call `get_scenario_types()` here, + # it is acceptable to use a placeholder value. + instance = cls("placeholder_scenario_type") get_scenario_type = getattr(instance, "get_scenario_types") scenario_types = get_scenario_type() has_duplicates = False diff --git a/run_kraken.py b/run_kraken.py index 3dedfe3b..341e77d8 100644 --- a/run_kraken.py +++ b/run_kraken.py @@ -11,6 +11,7 @@ import uuid import time import queue import threading +from typing import Optional from krkn_lib.elastic.krkn_elastic import KrknElastic from krkn_lib.models.elastic import ElasticChaosRunTelemetry @@ -32,6 +33,11 @@ from krkn.scenario_plugins.scenario_plugin_factory import ( ScenarioPluginFactory, ScenarioPluginNotFound, ) +from krkn.rollback.config import RollbackConfig +from krkn.rollback.command import ( + list_rollback as list_rollback_command, + execute_rollback as execute_rollback_command, +) # removes TripleDES warning import warnings @@ -39,13 +45,13 @@ warnings.filterwarnings(action='ignore', module='.*paramiko.*') report_file = "" - # Main function -def main(cfg) -> int: +def main(options, command: Optional[str]) -> int: # Start kraken print(pyfiglet.figlet_format("kraken")) logging.info("Starting kraken") + cfg = options.cfg # Parse and read the config if os.path.isfile(cfg): with open(cfg, "r") as f: @@ -61,6 +67,18 @@ def main(cfg) -> int: config["kraken"], "publish_kraken_status", False ) port = get_yaml_item_value(config["kraken"], "port", 8081) + RollbackConfig.register( + auto=get_yaml_item_value( + config["kraken"], + "auto_rollback", + False + ), + versions_directory=get_yaml_item_value( + config["kraken"], + "rollback_versions_directory", + "/tmp/kraken-rollback" + ), + ) signal_address = get_yaml_item_value( config["kraken"], "signal_address", "0.0.0.0" ) @@ -231,6 +249,19 @@ def main(cfg) -> int: logging.info("Server URL: %s" % kubecli.get_host()) + if command == "list-rollback": + sys.exit( + list_rollback_command( + options.run_uuid, options.scenario_type + ) + ) + elif command == "execute-rollback": + sys.exit( + execute_rollback_command( + telemetry_ocp, options.run_uuid, options.scenario_type + ) + ) + # Initialize the start iteration to 0 iteration = 0 @@ -519,7 +550,13 @@ def main(cfg) -> int: if __name__ == "__main__": # Initialize the parser to read the config - parser = optparse.OptionParser() + parser = optparse.OptionParser( + usage="%prog [options] [command]\n\n" + "Commands:\n" + " list-rollback List rollback version files in a tree-like format\n" + " execute-rollback Execute rollback version files and cleanup if successful\n\n" + "If no command is specified, kraken will run chaos scenarios.", + ) parser.add_option( "-c", "--config", @@ -556,7 +593,26 @@ if __name__ == "__main__": default=None, ) + # Add rollback command options + parser.add_option( + "-r", + "--run_uuid", + dest="run_uuid", + help="run UUID to filter rollback operations", + default=None, + ) + + parser.add_option( + "-s", + "--scenario_type", + dest="scenario_type", + help="scenario type to filter rollback operations", + default=None, + ) + (options, args) = parser.parse_args() + + # If no command or regular execution, continue with existing logic report_file = options.output tee_handler = TeeLogHandler() handlers = [ @@ -625,7 +681,9 @@ if __name__ == "__main__": if option_error: retval = 1 else: - retval = main(options.cfg) + # Check if command is provided as positional argument + command = args[0] if args else None + retval = main(options, command) junit_endtime = time.time() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/rollback_scenario_plugins/__init__.py b/tests/rollback_scenario_plugins/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/rollback_scenario_plugins/simple.py b/tests/rollback_scenario_plugins/simple.py new file mode 100644 index 00000000..f493ef04 --- /dev/null +++ b/tests/rollback_scenario_plugins/simple.py @@ -0,0 +1,63 @@ +import logging + +from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift +from krkn_lib.models.telemetry import ScenarioTelemetry + +from krkn.scenario_plugins.abstract_scenario_plugin import AbstractScenarioPlugin +from krkn.rollback.config import RollbackContent +from krkn.rollback.handler import set_rollback_context_decorator + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) +logger.addHandler(logging.StreamHandler()) + + +class SimpleRollbackScenarioPlugin(AbstractScenarioPlugin): + """ + Mock implementation of RollbackScenarioPlugin for testing purposes. + This plugin does not perform any actual rollback operations. + """ + + @set_rollback_context_decorator + def run( + self, + run_uuid: str, + scenario: str, + krkn_config: dict[str, any], + lib_telemetry: KrknTelemetryOpenshift, + scenario_telemetry: ScenarioTelemetry, + ) -> int: + logger.info( + f"Setting rollback callable for run {run_uuid} with scenario {scenario}." + ) + logger.debug(f"Krkn config: {krkn_config}") + self.rollback_handler.set_rollback_callable( + self.rollback_callable, + RollbackContent( + resource_identifier=run_uuid, + ), + ) + logger.info("Rollback callable set successfully.") + print("Rollback callable has been set for the scenario.") + return 0 + + def get_scenario_types(self) -> list[str]: + """ + Returns the scenario types that this plugin supports. + :return: a list of scenario types + """ + return ["simple_rollback_scenario"] + + @staticmethod + def rollback_callable( + rollback_context: RollbackContent, lib_telemetry: KrknTelemetryOpenshift + ): + """ + Simple rollback callable that simulates a rollback operation. + """ + run_uuid = rollback_context.resource_identifier + + print(f"Rollback called for run {run_uuid}.") + # Simulate a rollback operation + # In a real scenario, this would contain logic to revert changes made during the scenario execution. + print("Rollback operation completed successfully.") diff --git a/tests/test_rollback.py b/tests/test_rollback.py new file mode 100644 index 00000000..cb6657b1 --- /dev/null +++ b/tests/test_rollback.py @@ -0,0 +1,160 @@ +import pytest +import logging +import os +import sys +import uuid +import subprocess + +from krkn_lib.k8s import KrknKubernetes +from krkn_lib.ocp import KrknOpenshift +from krkn_lib.telemetry.ocp import KrknTelemetryOpenshift +from krkn_lib.models.telemetry import ScenarioTelemetry +from krkn_lib.utils import SafeLogger +from krkn.rollback.config import RollbackConfig + +sys.path.append( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +) # Adjust path to include krkn +TEST_LOGS_DIR = "/tmp/krkn_test_rollback_logs_directory" +TEST_VERSIONS_DIR = "/tmp/krkn_test_rollback_versions_directory" + + +class TestRollbackScenarioPlugin: + def validate_rollback_directory( + self, run_uuid: str, scenario: str, versions: int = 1 + ) -> list[str]: + """ + Validate that the rollback directory exists and contains version files. + + :param run_uuid: The UUID for current run, used to identify the rollback context directory. + :param scenario: The name of the scenario to validate. + :param versions: The expected number of version files. + :return: List of version files in full path. + """ + rollback_context_directories = [ + dirname for dirname in os.listdir(TEST_VERSIONS_DIR) if run_uuid in dirname + ] + assert len(rollback_context_directories) == 1, ( + f"Expected one directory for run UUID {run_uuid}, found: {rollback_context_directories}" + ) + + scenario_rollback_versions_directory = os.path.join( + TEST_VERSIONS_DIR, rollback_context_directories[0] + ) + version_files = os.listdir(scenario_rollback_versions_directory) + assert len(version_files) == versions, ( + f"Expected {versions} version files, found: {len(version_files)}" + ) + for version_file in version_files: + assert version_file.startswith(scenario), ( + f"Version file {version_file} does not start with '{scenario}'" + ) + assert version_file.endswith(".py"), ( + f"Version file {version_file} does not end with '.py'" + ) + + return [ + os.path.join(scenario_rollback_versions_directory, vf) + for vf in version_files + ] + + def execute_version_file(self, version_file: str): + """ + Execute a rollback version file using subprocess. + + :param version_file: The path to the version file to execute. + """ + print(f"Executing rollback version file: {version_file}") + result = subprocess.run( + [sys.executable, version_file], + capture_output=True, + text=True, + ) + assert result.returncode == 0, ( + f"Rollback version file {version_file} failed with return code {result.returncode}. " + f"Output: {result.stdout}, Error: {result.stderr}" + ) + print( + f"Rollback version file executed successfully: {version_file} with output: {result.stdout}" + ) + + @pytest.fixture(autouse=True) + def setup_logging(self): + os.makedirs(TEST_LOGS_DIR, exist_ok=True) + # setup logging + logging.basicConfig( + level=logging.DEBUG, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[ + logging.FileHandler(os.path.join(TEST_LOGS_DIR, "test_rollback.log")), + logging.StreamHandler(), + ], + ) + + @pytest.fixture(scope="module") + def kubeconfig_path(self): + # Provide the path to the kubeconfig file for testing + return os.getenv("KUBECONFIG", "~/.kube/config") + + @pytest.fixture(scope="module") + def safe_logger(self): + os.makedirs(TEST_LOGS_DIR, exist_ok=True) + with open(os.path.join(TEST_LOGS_DIR, "telemetry.log"), "w") as f: + pass # Create the file if it doesn't exist + yield SafeLogger(filename=os.path.join(TEST_LOGS_DIR, "telemetry.log")) + + @pytest.fixture(scope="module") + def kubecli(self, kubeconfig_path): + yield KrknKubernetes(kubeconfig_path=kubeconfig_path) + + @pytest.fixture(scope="module") + def lib_openshift(self, kubeconfig_path): + yield KrknOpenshift(kubeconfig_path=kubeconfig_path) + + @pytest.fixture(scope="module") + def lib_telemetry(self, lib_openshift, safe_logger): + yield KrknTelemetryOpenshift( + safe_logger=safe_logger, + lib_openshift=lib_openshift, + ) + + @pytest.fixture(scope="module") + def scenario_telemetry(self): + yield ScenarioTelemetry() + + @pytest.fixture(scope="module") + def setup_rollback_config(self): + RollbackConfig.register( + auto=False, + versions_directory=TEST_VERSIONS_DIR, + ) + + @pytest.mark.usefixtures("setup_rollback_config") + def test_simple_rollback_scenario_plugin(self, lib_telemetry, scenario_telemetry): + from tests.rollback_scenario_plugins.simple import SimpleRollbackScenarioPlugin + + scenario_type = "simple_rollback_scenario" + simple_rollback_scenario_plugin = SimpleRollbackScenarioPlugin( + scenario_type=scenario_type, + ) + run_uuid = str(uuid.uuid4()) + simple_rollback_scenario_plugin.run( + run_uuid=run_uuid, + scenario="test_scenario", + krkn_config={ + "key1": "value", + "key2": False, + "key3": 123, + "key4": ["value1", "value2", "value3"], + }, + lib_telemetry=lib_telemetry, + scenario_telemetry=scenario_telemetry, + ) + # Validate the rollback directory and version files do exist + version_files = self.validate_rollback_directory( + run_uuid, + scenario_type, + ) + # Execute the rollback version file + for version_file in version_files: + self.execute_version_file(version_file)