import copy from typing import Optional import frostfs_testlib.resources.optionals as optionals from frostfs_testlib import reporter from frostfs_testlib.load.interfaces.scenario_runner import ScenarioRunner from frostfs_testlib.load.load_config import EndpointSelectionStrategy, LoadParams, LoadScenario, LoadType from frostfs_testlib.load.load_report import LoadReport from frostfs_testlib.load.load_verifiers import LoadVerifier from frostfs_testlib.storage.cluster import ClusterNode from frostfs_testlib.storage.dataclasses.frostfs_services import S3Gate, StorageNode from frostfs_testlib.storage.dataclasses.wallet import WalletInfo from frostfs_testlib.testing.test_control import run_optionally class BackgroundLoadController: k6_dir: str load_params: LoadParams original_load_params: LoadParams verification_params: LoadParams cluster_nodes: list[ClusterNode] nodes_under_load: list[ClusterNode] load_counter: int loaders_wallet: WalletInfo load_summaries: dict endpoints: list[str] runner: ScenarioRunner started: bool def __init__( self, k6_dir: str, load_params: LoadParams, loaders_wallet: WalletInfo, cluster_nodes: list[ClusterNode], nodes_under_load: list[ClusterNode], runner: ScenarioRunner, ) -> None: self.k6_dir = k6_dir self.original_load_params = load_params self.load_params = copy.deepcopy(self.original_load_params) self.cluster_nodes = cluster_nodes self.nodes_under_load = nodes_under_load self.load_counter = 1 self.loaders_wallet = loaders_wallet self.runner = runner self.started = False if load_params.endpoint_selection_strategy is None: raise RuntimeError("endpoint_selection_strategy should not be None") @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED, []) def _get_endpoints(self, load_type: LoadType, endpoint_selection_strategy: EndpointSelectionStrategy): all_endpoints = { LoadType.gRPC: { EndpointSelectionStrategy.ALL: list( set( endpoint for node_under_load in self.nodes_under_load for endpoint in node_under_load.service(StorageNode).get_all_rpc_endpoint() ) ), EndpointSelectionStrategy.FIRST: list( set( node_under_load.service(StorageNode).get_rpc_endpoint() for node_under_load in self.nodes_under_load ) ), }, # for some reason xk6 appends http protocol on its own LoadType.S3: { EndpointSelectionStrategy.ALL: list( set( endpoint for node_under_load in self.nodes_under_load for endpoint in node_under_load.service(S3Gate).get_all_endpoints() ) ), EndpointSelectionStrategy.FIRST: list( set(node_under_load.service(S3Gate).get_endpoint() for node_under_load in self.nodes_under_load) ), }, } return all_endpoints[load_type][endpoint_selection_strategy] @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) @reporter.step("Prepare load instances") def prepare(self): self.endpoints = self._get_endpoints(self.load_params.load_type, self.load_params.endpoint_selection_strategy) self.runner.prepare(self.load_params, self.cluster_nodes, self.nodes_under_load, self.k6_dir) self.runner.init_k6_instances(self.load_params, self.endpoints, self.k6_dir) @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) def start(self): with reporter.step(f"Start load on nodes {self.nodes_under_load}"): self.runner.start() self.started = True @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) @reporter.step("Stop load") def stop(self): self.runner.stop() @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED, True) def is_running(self) -> bool: return self.runner.is_running @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) @reporter.step("Reset load") def _reset_for_consequent_load(self): """This method is required if we want to run multiple loads during test run. Raise load counter by 1 and append it to load_id """ self.load_counter += 1 self.load_params = copy.deepcopy(self.original_load_params) self.load_params.set_id(f"{self.load_params.load_id}_{self.load_counter}") @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) @reporter.step("Startup load") def startup(self): self.prepare() self.preset() self.start() @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) def preset(self): self.runner.preset() @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) @reporter.step("Stop and get results of load") def teardown(self, load_report: Optional[LoadReport] = None): if not self.started: return self.stop() self.load_summaries = self._get_results() self.started = False if load_report: load_report.add_summaries(self.load_summaries) @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) @reporter.step("Run post-load verification") def verify(self): try: load_issues = self._collect_load_issues() if self.load_params.verify: load_issues.extend(self._run_verify_scenario()) assert not load_issues, "\n".join(load_issues) finally: self._reset_for_consequent_load() @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) @reporter.step("Collect load issues") def _collect_load_issues(self): verifier = LoadVerifier(self.load_params) return verifier.collect_load_issues(self.load_summaries) @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) def wait_until_finish(self, soft_timeout: int = 0): self.runner.wait_until_finish(soft_timeout) @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) @reporter.step("Verify loaded objects") def _run_verify_scenario(self) -> list[str]: self.verification_params = LoadParams( verify_clients=self.load_params.verify_clients, scenario=LoadScenario.VERIFY, read_from=self.load_params.read_from, registry_file=self.load_params.registry_file, verify_time=self.load_params.verify_time, load_type=self.load_params.load_type, load_id=self.load_params.load_id, vu_init_time=0, working_dir=self.load_params.working_dir, endpoint_selection_strategy=self.load_params.endpoint_selection_strategy, k6_process_allocation_strategy=self.load_params.k6_process_allocation_strategy, setup_timeout="1s", ) if self.verification_params.verify_time is None: raise RuntimeError("verify_time should not be none") self.runner.init_k6_instances(self.verification_params, self.endpoints, self.k6_dir) with reporter.step("Run verify scenario"): self.runner.start() self.runner.wait_until_finish() with reporter.step("Collect verify issues"): verification_summaries = self._get_results() verifier = LoadVerifier(self.load_params) return verifier.collect_verify_issues(self.load_summaries, verification_summaries) @run_optionally(optionals.OPTIONAL_BACKGROUND_LOAD_ENABLED) def _get_results(self) -> dict: with reporter.step(f"Get {self.load_params.scenario.value} scenario results"): return self.runner.get_results() def __str__(self) -> str: return self.load_params.__str__() def __repr__(self) -> str: return repr(self.load_params)