ClickHouse/tests/ci/tee_popen.py

Ignoring revisions in .git-blame-ignore-revs. Click here to bypass and see the normal blame view.

128 lines
3.9 KiB
Python
Raw Normal View History

2021-12-03 08:33:16 +00:00
#!/usr/bin/env python3
2024-02-26 17:46:15 +00:00
import logging
import os
import signal
2024-02-26 17:46:15 +00:00
import sys
2022-12-14 10:53:17 +00:00
from io import TextIOWrapper
from pathlib import Path
2024-02-26 17:46:15 +00:00
from subprocess import PIPE, STDOUT, Popen
from threading import Thread
from time import sleep
from typing import Optional, Union
2021-12-03 08:33:16 +00:00
2023-06-25 04:14:31 +00:00
# Very simple tee logic implementation. You can specify a shell command, output
2021-12-03 08:33:16 +00:00
# logfile and env variables. After TeePopen is created you can only wait until
# it finishes. stderr and stdout will be redirected both to specified file and
# stdout.
class TeePopen:
2022-12-14 10:53:17 +00:00
def __init__(
self,
command: str,
log_file: Union[str, Path],
2022-12-14 10:53:17 +00:00
env: Optional[dict] = None,
timeout: Optional[int] = None,
):
2021-12-03 08:33:16 +00:00
self.command = command
2022-12-14 10:53:17 +00:00
self._log_file_name = log_file
self._log_file = None # type: Optional[TextIOWrapper]
self.env = env or os.environ.copy()
2022-11-14 17:59:01 +00:00
self._process = None # type: Optional[Popen]
self.timeout = timeout
self.timeout_exceeded = False
self.terminated_by_sigterm = False
self.terminated_by_sigkill = False
2022-12-14 10:53:17 +00:00
def _check_timeout(self) -> None:
if self.timeout is None:
return
sleep(self.timeout)
logging.warning(
"Timeout exceeded. Send SIGTERM to process %s, timeout %s",
self.process.pid,
self.timeout,
)
self.send_signal(signal.SIGTERM)
time_wait = 0
self.terminated_by_sigterm = True
self.timeout_exceeded = True
while self.process.poll() is None and time_wait < 100:
print("wait...")
wait = 5
sleep(wait)
time_wait += wait
while self.process.poll() is None:
logging.error(
"Process is still running. Send SIGKILL",
)
self.send_signal(signal.SIGKILL)
self.terminated_by_sigkill = True
sleep(5)
2021-12-03 08:33:16 +00:00
2022-12-14 10:53:17 +00:00
def __enter__(self) -> "TeePopen":
self.process = Popen(
self.command,
shell=True,
universal_newlines=True,
env=self.env,
start_new_session=True, # signall will be sent to all children
stderr=STDOUT,
stdout=PIPE,
bufsize=1,
errors="backslashreplace",
)
sleep(1)
2024-08-10 08:01:16 +00:00
print(f"Subprocess started, pid [{self.process.pid}]")
if self.timeout is not None and self.timeout > 0:
t = Thread(target=self._check_timeout)
t.daemon = True # does not block the program from exit
t.start()
2021-12-03 08:33:16 +00:00
return self
2022-12-14 10:53:17 +00:00
def __exit__(self, exc_type, exc_value, traceback):
self.wait()
if self.timeout_exceeded:
exceeded_log = (
f"Command `{self.command}` has failed, "
f"timeout {self.timeout}s is exceeded"
)
if self.process.stdout is not None:
sys.stdout.write(exceeded_log)
self.log_file.write(exceeded_log)
2021-12-03 08:33:16 +00:00
self.log_file.close()
def wait(self) -> int:
2022-12-14 10:53:17 +00:00
if self.process.stdout is not None:
for line in self.process.stdout:
sys.stdout.write(line)
self.log_file.write(line)
2024-09-23 12:33:48 +00:00
self.log_file.flush()
2021-12-03 08:33:16 +00:00
2022-02-15 22:50:21 +00:00
return self.process.wait()
2022-11-14 17:59:01 +00:00
def poll(self):
return self.process.poll()
def send_signal(self, signal_num):
2024-08-10 08:01:16 +00:00
os.killpg(self.process.pid, signal_num)
2022-11-14 17:59:01 +00:00
@property
def process(self) -> Popen:
if self._process is not None:
return self._process
raise AttributeError("process is not created yet")
@process.setter
def process(self, process: Popen) -> None:
self._process = process
2022-12-14 10:53:17 +00:00
@property
def log_file(self) -> TextIOWrapper:
if self._log_file is None:
2024-02-26 17:46:15 +00:00
# pylint:disable-next=consider-using-with
2022-12-14 10:53:17 +00:00
self._log_file = open(self._log_file_name, "w", encoding="utf-8")
return self._log_file