tests/utils/lttngtest/environment.py

   1 #!/usr/bin/env python3
   2 #
   3 # Copyright (C) 2022 Jérémie Galarneau <jeremie.galarneau@efficios.com>
   4 #
   5 # SPDX-License-Identifier: GPL-2.0-only
   6 #
   7
   8 from types import FrameType
   9 from typing import Callable, Optional, Tuple, List
  10 import sys
  11 import pathlib
  12 import signal
  13 import subprocess
  14 import shlex
  15 import shutil
  16 import os
  17 import queue
  18 import tempfile
  19 from . import logger
  20 import time
  21 import threading
  22 import contextlib
  23
  24
  25 class TemporaryDirectory:
  26     def __init__(self, prefix: str):
  27         self._directory_path = tempfile.mkdtemp(prefix=prefix)
  28
  29     def __del__(self):
  30         shutil.rmtree(self._directory_path, ignore_errors=True)
  31
  32     @property
  33     def path(self) -> pathlib.Path:
  34         return pathlib.Path(self._directory_path)
  35
  36
  37 class _SignalWaitQueue:
  38     """
  39     Utility class useful to wait for a signal before proceeding.
  40
  41     Simply register the `signal` method as the handler for the signal you are
  42     interested in and call `wait_for_signal` to wait for its reception.
  43
  44     Registering a signal:
  45         signal.signal(signal.SIGWHATEVER, queue.signal)
  46
  47     Waiting for the signal:
  48         queue.wait_for_signal()
  49     """
  50
  51     def __init__(self):
  52         self._queue: queue.Queue = queue.Queue()
  53
  54     def signal(self, signal_number, frame: Optional[FrameType]):
  55         self._queue.put_nowait(signal_number)
  56
  57     def wait_for_signal(self):
  58         self._queue.get(block=True)
  59
  60
  61 class WaitTraceTestApplication:
  62     """
  63     Create an application that waits before tracing. This allows a test to
  64     launch an application, get its PID, and get it to start tracing when it
  65     has completed its setup.
  66     """
  67
  68     def __init__(
  69         self,
  70         binary_path: pathlib.Path,
  71         event_count: int,
  72         environment: "Environment",
  73         wait_time_between_events_us: int = 0,
  74     ):
  75         self._environment: Environment = environment
  76         if event_count % 5:
  77             # The test application currently produces 5 different events per iteration.
  78             raise ValueError("event count must be a multiple of 5")
  79         self._iteration_count: int = int(event_count / 5)
  80         # File that the application will wait to see before tracing its events.
  81         self._app_start_tracing_file_path: pathlib.Path = pathlib.Path(
  82             tempfile.mktemp(
  83                 prefix="app_",
  84                 suffix="_start_tracing",
  85                 dir=environment.lttng_home_location,
  86             )
  87         )
  88         self._has_returned = False
  89
  90         test_app_env = os.environ.copy()
  91         test_app_env["LTTNG_HOME"] = str(environment.lttng_home_location)
  92         # Make sure the app is blocked until it is properly registered to
  93         # the session daemon.
  94         test_app_env["LTTNG_UST_REGISTER_TIMEOUT"] = "-1"
  95
  96         # File that the application will create to indicate it has completed its initialization.
  97         app_ready_file_path: str = tempfile.mktemp(
  98             prefix="app_", suffix="_ready", dir=environment.lttng_home_location
  99         )
 100
 101         test_app_args = [str(binary_path)]
 102         test_app_args.extend(
 103             shlex.split(
 104                 "--iter {iteration_count} --create-in-main {app_ready_file_path} --wait-before-first-event {app_start_tracing_file_path} --wait {wait_time_between_events_us}".format(
 105                     iteration_count=self._iteration_count,
 106                     app_ready_file_path=app_ready_file_path,
 107                     app_start_tracing_file_path=self._app_start_tracing_file_path,
 108                     wait_time_between_events_us=wait_time_between_events_us,
 109                 )
 110             )
 111         )
 112
 113         self._process: subprocess.Popen = subprocess.Popen(
 114             test_app_args,
 115             env=test_app_env,
 116         )
 117
 118         # Wait for the application to create the file indicating it has fully
 119         # initialized. Make sure the app hasn't crashed in order to not wait
 120         # forever.
 121         while True:
 122             if os.path.exists(app_ready_file_path):
 123                 break
 124
 125             if self._process.poll() is not None:
 126                 # Application has unexepectedly returned.
 127                 raise RuntimeError(
 128                     "Test application has unexepectedly returned during its initialization with return code `{return_code}`".format(
 129                         return_code=self._process.returncode
 130                     )
 131                 )
 132
 133             time.sleep(0.1)
 134
 135     def trace(self) -> None:
 136         if self._process.poll() is not None:
 137             # Application has unexepectedly returned.
 138             raise RuntimeError(
 139                 "Test application has unexepectedly before tracing with return code `{return_code}`".format(
 140                     return_code=self._process.returncode
 141                 )
 142             )
 143         open(self._app_start_tracing_file_path, mode="x")
 144
 145     def wait_for_exit(self) -> None:
 146         if self._process.wait() != 0:
 147             raise RuntimeError(
 148                 "Test application has exit with return code `{return_code}`".format(
 149                     return_code=self._process.returncode
 150                 )
 151             )
 152         self._has_returned = True
 153
 154     @property
 155     def vpid(self) -> int:
 156         return self._process.pid
 157
 158     def __del__(self):
 159         if not self._has_returned:
 160             # This is potentially racy if the pid has been recycled. However,
 161             # we can't use pidfd_open since it is only available in python >= 3.9.
 162             self._process.kill()
 163             self._process.wait()
 164
 165
 166 class ProcessOutputConsumer(threading.Thread, logger._Logger):
 167     def __init__(
 168         self, process: subprocess.Popen, name: str, log: Callable[[str], None]
 169     ):
 170         threading.Thread.__init__(self)
 171         self._prefix = name
 172         logger._Logger.__init__(self, log)
 173         self._process = process
 174
 175     def run(self) -> None:
 176         while self._process.poll() is None:
 177             assert self._process.stdout
 178             line = self._process.stdout.readline().decode("utf-8").replace("\n", "")
 179             if len(line) != 0:
 180                 self._log("{prefix}: {line}".format(prefix=self._prefix, line=line))
 181
 182
 183 # Generate a temporary environment in which to execute a test.
 184 class _Environment(logger._Logger):
 185     def __init__(
 186         self, with_sessiond: bool, log: Optional[Callable[[str], None]] = None
 187     ):
 188         super().__init__(log)
 189         signal.signal(signal.SIGTERM, self._handle_termination_signal)
 190         signal.signal(signal.SIGINT, self._handle_termination_signal)
 191
 192         # Assumes the project's hierarchy to this file is:
 193         # tests/utils/python/this_file
 194         self._project_root: pathlib.Path = pathlib.Path(__file__).absolute().parents[3]
 195         self._lttng_home: Optional[TemporaryDirectory] = TemporaryDirectory(
 196             "lttng_test_env_home"
 197         )
 198
 199         self._sessiond: Optional[subprocess.Popen[bytes]] = (
 200             self._launch_lttng_sessiond() if with_sessiond else None
 201         )
 202
 203     @property
 204     def lttng_home_location(self) -> pathlib.Path:
 205         if self._lttng_home is None:
 206             raise RuntimeError("Attempt to access LTTng home after clean-up")
 207         return self._lttng_home.path
 208
 209     @property
 210     def lttng_client_path(self) -> pathlib.Path:
 211         return self._project_root / "src" / "bin" / "lttng" / "lttng"
 212
 213     def create_temporary_directory(self, prefix: Optional[str] = None) -> pathlib.Path:
 214         # Simply return a path that is contained within LTTNG_HOME; it will
 215         # be destroyed when the temporary home goes out of scope.
 216         assert self._lttng_home
 217         return pathlib.Path(
 218             tempfile.mkdtemp(
 219                 prefix="tmp" if prefix is None else prefix,
 220                 dir=str(self._lttng_home.path),
 221             )
 222         )
 223
 224     # Unpack a list of environment variables from a string
 225     # such as "HELLO=is_it ME='/you/are/looking/for'"
 226     @staticmethod
 227     def _unpack_env_vars(env_vars_string: str) -> List[Tuple[str, str]]:
 228         unpacked_vars = []
 229         for var in shlex.split(env_vars_string):
 230             equal_position = var.find("=")
 231             # Must have an equal sign and not end with an equal sign
 232             if equal_position == -1 or equal_position == len(var) - 1:
 233                 raise ValueError(
 234                     "Invalid sessiond environment variable: `{}`".format(var)
 235                 )
 236
 237             var_name = var[0:equal_position]
 238             var_value = var[equal_position + 1 :]
 239             # Unquote any paths
 240             var_value = var_value.replace("'", "")
 241             var_value = var_value.replace('"', "")
 242             unpacked_vars.append((var_name, var_value))
 243
 244         return unpacked_vars
 245
 246     def _launch_lttng_sessiond(self) -> Optional[subprocess.Popen]:
 247         is_64bits_host = sys.maxsize > 2**32
 248
 249         sessiond_path = (
 250             self._project_root / "src" / "bin" / "lttng-sessiond" / "lttng-sessiond"
 251         )
 252         consumerd_path_option_name = "--consumerd{bitness}-path".format(
 253             bitness="64" if is_64bits_host else "32"
 254         )
 255         consumerd_path = (
 256             self._project_root / "src" / "bin" / "lttng-consumerd" / "lttng-consumerd"
 257         )
 258
 259         no_sessiond_var = os.environ.get("TEST_NO_SESSIOND")
 260         if no_sessiond_var and no_sessiond_var == "1":
 261             # Run test without a session daemon; the user probably
 262             # intends to run one under gdb for example.
 263             return None
 264
 265         # Setup the session daemon's environment
 266         sessiond_env_vars = os.environ.get("LTTNG_SESSIOND_ENV_VARS")
 267         sessiond_env = os.environ.copy()
 268         if sessiond_env_vars:
 269             self._log("Additional lttng-sessiond environment variables:")
 270             additional_vars = self._unpack_env_vars(sessiond_env_vars)
 271             for var_name, var_value in additional_vars:
 272                 self._log("  {name}={value}".format(name=var_name, value=var_value))
 273                 sessiond_env[var_name] = var_value
 274
 275         sessiond_env["LTTNG_SESSION_CONFIG_XSD_PATH"] = str(
 276             self._project_root / "src" / "common"
 277         )
 278
 279         assert self._lttng_home is not None
 280         sessiond_env["LTTNG_HOME"] = str(self._lttng_home.path)
 281
 282         wait_queue = _SignalWaitQueue()
 283         signal.signal(signal.SIGUSR1, wait_queue.signal)
 284
 285         self._log(
 286             "Launching session daemon with LTTNG_HOME=`{home_dir}`".format(
 287                 home_dir=str(self._lttng_home.path)
 288             )
 289         )
 290         process = subprocess.Popen(
 291             [
 292                 str(sessiond_path),
 293                 consumerd_path_option_name,
 294                 str(consumerd_path),
 295                 "--sig-parent",
 296             ],
 297             stdout=subprocess.PIPE,
 298             stderr=subprocess.STDOUT,
 299             env=sessiond_env,
 300         )
 301
 302         if self._logging_function:
 303             self._sessiond_output_consumer: Optional[
 304                 ProcessOutputConsumer
 305             ] = ProcessOutputConsumer(process, "lttng-sessiond", self._logging_function)
 306             self._sessiond_output_consumer.daemon = True
 307             self._sessiond_output_consumer.start()
 308
 309         # Wait for SIGUSR1, indicating the sessiond is ready to proceed
 310         wait_queue.wait_for_signal()
 311         signal.signal(signal.SIGUSR1, wait_queue.signal)
 312
 313         return process
 314
 315     def _handle_termination_signal(
 316         self, signal_number: int, frame: Optional[FrameType]
 317     ) -> None:
 318         self._log(
 319             "Killed by {signal_name} signal, cleaning-up".format(
 320                 signal_name=signal.strsignal(signal_number)
 321             )
 322         )
 323         self._cleanup()
 324
 325     def launch_wait_trace_test_application(
 326         self, event_count: int
 327     ) -> WaitTraceTestApplication:
 328         """
 329         Launch an application that will wait before tracing `event_count` events.
 330         """
 331         return WaitTraceTestApplication(
 332             self._project_root
 333             / "tests"
 334             / "utils"
 335             / "testapp"
 336             / "gen-ust-nevents"
 337             / "gen-ust-nevents",
 338             event_count,
 339             self,
 340         )
 341
 342     # Clean-up managed processes
 343     def _cleanup(self) -> None:
 344         if self._sessiond and self._sessiond.poll() is None:
 345             # The session daemon is alive; kill it.
 346             self._log(
 347                 "Killing session daemon (pid = {sessiond_pid})".format(
 348                     sessiond_pid=self._sessiond.pid
 349                 )
 350             )
 351
 352             self._sessiond.terminate()
 353             self._sessiond.wait()
 354             if self._sessiond_output_consumer:
 355                 self._sessiond_output_consumer.join()
 356                 self._sessiond_output_consumer = None
 357
 358             self._log("Session daemon killed")
 359             self._sessiond = None
 360
 361         self._lttng_home = None
 362
 363     def __del__(self):
 364         self._cleanup()
 365
 366
 367 @contextlib.contextmanager
 368 def test_environment(with_sessiond: bool, log: Optional[Callable[[str], None]] = None):
 369     env = _Environment(with_sessiond, log)
 370     try:
 371         yield env
 372     finally:
 373         env._cleanup()