from __future__ import annotations import os import uuid from typing import Optional from tenacity import retry from tenacity.stop import stop_after_attempt from tenacity.wait import wait_fixed from frostfs_testlib import reporter from frostfs_testlib.shell import Shell from frostfs_testlib.shell.command_inspectors import SuInspector from frostfs_testlib.shell.interfaces import CommandInspector, CommandOptions class RemoteProcess: def __init__( self, cmd: str, process_dir: str, shell: Shell, cmd_inspector: Optional[CommandInspector], proc_id: str ): self.process_dir = process_dir self.cmd = cmd self.stdout_last_line_number = 0 self.stderr_last_line_number = 0 self.pid: Optional[str] = None self.proc_rc: Optional[int] = None self.proc_start_time: Optional[int] = None self.proc_end_time: Optional[int] = None self.saved_stdout: Optional[str] = None self.saved_stderr: Optional[str] = None self.shell = shell self.proc_id: str = proc_id self.cmd_inspectors: list[CommandInspector] = [cmd_inspector] if cmd_inspector else [] @classmethod @reporter.step("Create remote process") def create( cls, command: str, shell: Shell, working_dir: str = "/tmp", user: Optional[str] = None, proc_id: Optional[str] = None, ) -> RemoteProcess: """ Create a process on a remote host. Created dir for process with following files: command.sh: script to execute pid: contains process id rc: contains script return code stderr: contains script errors stdout: contains script output user: user on behalf whom command will be executed proc_id: process string identificator Args: shell: Shell instance command: command to be run on a remote host working_dir: working directory for the process Returns: RemoteProcess instance for further examination """ if proc_id is None: proc_id = f"{uuid.uuid4()}" cmd_inspector = SuInspector(user) if user else None remote_process = cls( cmd=command, process_dir=os.path.join(working_dir, f"proc_{proc_id}"), shell=shell, cmd_inspector=cmd_inspector, proc_id=proc_id, ) return remote_process @reporter.step("Start remote process") def start(self): """ Starts a process on a remote host. """ self._create_process_dir() self._generate_command_script() self._start_process() self.pid = self._get_pid() @reporter.step("Get process stdout") def stdout(self, full: bool = False) -> str: """ Method to get process stdout, either fresh info or full. Args: full: returns full stdout that we have to this moment Returns: Fresh stdout. By means of stdout_last_line_number only new stdout lines are returned. If process is finished (proc_rc is not None) saved stdout is returned """ if self.saved_stdout is not None: cur_stdout = self.saved_stdout else: terminal = self.shell.exec( f"cat {self.process_dir}/stdout", options=CommandOptions(no_log=True, extra_inspectors=self.cmd_inspectors), ) if self.proc_rc is not None: self.saved_stdout = terminal.stdout cur_stdout = terminal.stdout if full: return cur_stdout whole_stdout = cur_stdout.split("\n") if len(whole_stdout) > self.stdout_last_line_number: resulted_stdout = "\n".join(whole_stdout[self.stdout_last_line_number :]) self.stdout_last_line_number = len(whole_stdout) return resulted_stdout return "" @reporter.step("Get process stderr") def stderr(self, full: bool = False) -> str: """ Method to get process stderr, either fresh info or full. Args: full: returns full stderr that we have to this moment Returns: Fresh stderr. By means of stderr_last_line_number only new stderr lines are returned. If process is finished (proc_rc is not None) saved stderr is returned """ if self.saved_stderr is not None: cur_stderr = self.saved_stderr else: terminal = self.shell.exec( f"cat {self.process_dir}/stderr", options=CommandOptions(no_log=True, extra_inspectors=self.cmd_inspectors), ) if self.proc_rc is not None: self.saved_stderr = terminal.stdout cur_stderr = terminal.stdout if full: return cur_stderr whole_stderr = cur_stderr.split("\n") if len(whole_stderr) > self.stderr_last_line_number: resulted_stderr = "\n".join(whole_stderr[self.stderr_last_line_number :]) self.stderr_last_line_number = len(whole_stderr) return resulted_stderr return "" @reporter.step("Get process rc") def rc(self) -> Optional[int]: if self.proc_rc is not None: return self.proc_rc result = self._cat_proc_file("rc") if not result: return None self.proc_rc = int(result) return self.proc_rc @reporter.step("Get process start time") def start_time(self) -> Optional[int]: if self.proc_start_time is not None: return self.proc_start_time result = self._cat_proc_file("start_time") if not result: return None self.proc_start_time = int(result) return self.proc_start_time @reporter.step("Get process end time") def end_time(self) -> Optional[int]: if self.proc_end_time is not None: return self.proc_end_time result = self._cat_proc_file("end_time") if not result: return None self.proc_end_time = int(result) return self.proc_end_time def _cat_proc_file(self, file: str) -> Optional[str]: terminal = self.shell.exec( f"cat {self.process_dir}/{file}", CommandOptions(check=False, extra_inspectors=self.cmd_inspectors, no_log=True), ) if "No such file or directory" in terminal.stderr: return None elif terminal.stderr or terminal.return_code != 0: raise AssertionError(f"cat process {file} was not successful: {terminal.stderr}") return terminal.stdout @reporter.step("Check if process is running") def running(self) -> bool: return self.rc() is None @reporter.step("Send signal to process") def send_signal(self, signal: int) -> None: kill_res = self.shell.exec( f"kill -{signal} {self.pid}", CommandOptions(check=False, extra_inspectors=self.cmd_inspectors), ) if "No such process" in kill_res.stderr: return if kill_res.return_code: raise AssertionError(f"Signal {signal} not sent. Return code of kill: {kill_res.return_code}") @reporter.step("Stop process") def stop(self) -> None: self.send_signal(15) @reporter.step("Kill process") def kill(self) -> None: self.send_signal(9) @reporter.step("Clear process directory") def clear(self) -> None: if self.process_dir == "/": raise AssertionError(f"Invalid path to delete: {self.process_dir}") self.shell.exec(f"rm -rf {self.process_dir}", CommandOptions(extra_inspectors=self.cmd_inspectors)) @reporter.step("Start remote process") def _start_process(self) -> None: self.shell.exec( f"nohup {self.process_dir}/command.sh {self.process_dir}/stdout " f"2>{self.process_dir}/stderr &", CommandOptions(extra_inspectors=self.cmd_inspectors), ) @reporter.step("Create process directory") def _create_process_dir(self) -> None: self.shell.exec(f"mkdir -p {self.process_dir}", CommandOptions(extra_inspectors=self.cmd_inspectors)) self.shell.exec(f"chmod 777 {self.process_dir}", CommandOptions(extra_inspectors=self.cmd_inspectors)) terminal = self.shell.exec(f"realpath {self.process_dir}", CommandOptions(extra_inspectors=self.cmd_inspectors)) self.process_dir = terminal.stdout.strip() @reporter.step("Get pid") @retry(wait=wait_fixed(10), stop=stop_after_attempt(5), reraise=True) def _get_pid(self) -> str: terminal = self.shell.exec(f"cat {self.process_dir}/pid", CommandOptions(extra_inspectors=self.cmd_inspectors)) assert terminal.stdout, f"invalid pid: {terminal.stdout}" return terminal.stdout.strip() @reporter.step("Generate command script") def _generate_command_script(self) -> None: command = self.cmd.replace('"', '\\"').replace("\\", "\\\\") script = ( f"#!/bin/bash\n" f"cd {self.process_dir}\n" f"date +%s > {self.process_dir}/start_time\n" f"{command} &\n" f"pid=\$!\n" f"cd {self.process_dir}\n" f"echo \$pid > {self.process_dir}/pid\n" f"wait \$pid\n" f"echo $? > {self.process_dir}/rc\n" f"date +%s > {self.process_dir}/end_time\n" ) self.shell.exec( f'echo "{script}" > {self.process_dir}/command.sh', CommandOptions(extra_inspectors=self.cmd_inspectors), ) self.shell.exec( f"cat {self.process_dir}/command.sh", CommandOptions(extra_inspectors=self.cmd_inspectors), ) self.shell.exec( f"chmod +x {self.process_dir}/command.sh", CommandOptions(extra_inspectors=self.cmd_inspectors), )