frostfs-testlib/src/frostfs_testlib/storage/controllers/background_load_controller.py
Andrey Berezin dc6b0e407f [#133] Change reporter usage
Signed-off-by: Andrey Berezin <a.berezin@yadro.com>
2023-11-29 15:27:17 +03:00

203 lines
8 KiB
Python

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)