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=self._compat_open_path(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_",
  99             suffix="_ready",
 100             dir=self._compat_open_path(environment.lttng_home_location),
 101         )
 102
 103         test_app_args = [str(binary_path)]
 104         test_app_args.extend(
 105             shlex.split(
 106                 "--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(
 107                     iteration_count=self._iteration_count,
 108                     app_ready_file_path=app_ready_file_path,
 109                     app_start_tracing_file_path=self._app_start_tracing_file_path,
 110                     wait_time_between_events_us=wait_time_between_events_us,
 111                 )
 112             )
 113         )
 114
 115         self._process: subprocess.Popen = subprocess.Popen(
 116             test_app_args,
 117             env=test_app_env,
 118         )
 119
 120         # Wait for the application to create the file indicating it has fully
 121         # initialized. Make sure the app hasn't crashed in order to not wait
 122         # forever.
 123         while True:
 124             if os.path.exists(app_ready_file_path):
 125                 break
 126
 127             if self._process.poll() is not None:
 128                 # Application has unexepectedly returned.
 129                 raise RuntimeError(
 130                     "Test application has unexepectedly returned during its initialization with return code `{return_code}`".format(
 131                         return_code=self._process.returncode
 132                     )
 133                 )
 134
 135             time.sleep(0.1)
 136
 137     def trace(self) -> None:
 138         if self._process.poll() is not None:
 139             # Application has unexepectedly returned.
 140             raise RuntimeError(
 141                 "Test application has unexepectedly before tracing with return code `{return_code}`".format(
 142                     return_code=self._process.returncode
 143                 )
 144             )
 145         open(self._compat_open_path(self._app_start_tracing_file_path), mode="x")
 146
 147     def wait_for_exit(self) -> None:
 148         if self._process.wait() != 0:
 149             raise RuntimeError(
 150                 "Test application has exit with return code `{return_code}`".format(
 151                     return_code=self._process.returncode
 152                 )
 153             )
 154         self._has_returned = True
 155
 156     @property
 157     def vpid(self) -> int:
 158         return self._process.pid
 159
 160     @staticmethod
 161     def _compat_open_path(path):
 162         # type: (pathlib.Path)
 163         """
 164         The builtin open() in python >= 3.6 expects a path-like object while
 165         prior versions expect a string or bytes object. Return the correct type
 166         based on the presence of the "__fspath__" attribute specified in PEP-519.
 167         """
 168         if hasattr(path, "__fspath__"):
 169             return path
 170         else:
 171             return str(path)
 172
 173     def __del__(self):
 174         if not self._has_returned:
 175             # This is potentially racy if the pid has been recycled. However,
 176             # we can't use pidfd_open since it is only available in python >= 3.9.
 177             self._process.kill()
 178             self._process.wait()
 179
 180
 181 class ProcessOutputConsumer(threading.Thread, logger._Logger):
 182     def __init__(
 183         self, process: subprocess.Popen, name: str, log: Callable[[str], None]
 184     ):
 185         threading.Thread.__init__(self)
 186         self._prefix = name
 187         logger._Logger.__init__(self, log)
 188         self._process = process
 189
 190     def run(self) -> None:
 191         while self._process.poll() is None:
 192             assert self._process.stdout
 193             line = self._process.stdout.readline().decode("utf-8").replace("\n", "")
 194             if len(line) != 0:
 195                 self._log("{prefix}: {line}".format(prefix=self._prefix, line=line))
 196
 197
 198 # Generate a temporary environment in which to execute a test.
 199 class _Environment(logger._Logger):
 200     def __init__(
 201         self, with_sessiond: bool, log: Optional[Callable[[str], None]] = None
 202     ):
 203         super().__init__(log)
 204         signal.signal(signal.SIGTERM, self._handle_termination_signal)
 205         signal.signal(signal.SIGINT, self._handle_termination_signal)
 206
 207         # Assumes the project's hierarchy to this file is:
 208         # tests/utils/python/this_file
 209         self._project_root: pathlib.Path = pathlib.Path(__file__).absolute().parents[3]
 210         self._lttng_home: Optional[TemporaryDirectory] = TemporaryDirectory(
 211             "lttng_test_env_home"
 212         )
 213
 214         self._sessiond: Optional[subprocess.Popen[bytes]] = (
 215             self._launch_lttng_sessiond() if with_sessiond else None
 216         )
 217
 218     @property
 219     def lttng_home_location(self) -> pathlib.Path:
 220         if self._lttng_home is None:
 221             raise RuntimeError("Attempt to access LTTng home after clean-up")
 222         return self._lttng_home.path
 223
 224     @property
 225     def lttng_client_path(self) -> pathlib.Path:
 226         return self._project_root / "src" / "bin" / "lttng" / "lttng"
 227
 228     def create_temporary_directory(self, prefix: Optional[str] = None) -> pathlib.Path:
 229         # Simply return a path that is contained within LTTNG_HOME; it will
 230         # be destroyed when the temporary home goes out of scope.
 231         assert self._lttng_home
 232         return pathlib.Path(
 233             tempfile.mkdtemp(
 234                 prefix="tmp" if prefix is None else prefix,
 235                 dir=str(self._lttng_home.path),
 236             )
 237         )
 238
 239     # Unpack a list of environment variables from a string
 240     # such as "HELLO=is_it ME='/you/are/looking/for'"
 241     @staticmethod
 242     def _unpack_env_vars(env_vars_string: str) -> List[Tuple[str, str]]:
 243         unpacked_vars = []
 244         for var in shlex.split(env_vars_string):
 245             equal_position = var.find("=")
 246             # Must have an equal sign and not end with an equal sign
 247             if equal_position == -1 or equal_position == len(var) - 1:
 248                 raise ValueError(
 249                     "Invalid sessiond environment variable: `{}`".format(var)
 250                 )
 251
 252             var_name = var[0:equal_position]
 253             var_value = var[equal_position + 1 :]
 254             # Unquote any paths
 255             var_value = var_value.replace("'", "")
 256             var_value = var_value.replace('"', "")
 257             unpacked_vars.append((var_name, var_value))
 258
 259         return unpacked_vars
 260
 261     def _launch_lttng_sessiond(self) -> Optional[subprocess.Popen]:
 262         is_64bits_host = sys.maxsize > 2**32
 263
 264         sessiond_path = (
 265             self._project_root / "src" / "bin" / "lttng-sessiond" / "lttng-sessiond"
 266         )
 267         consumerd_path_option_name = "--consumerd{bitness}-path".format(
 268             bitness="64" if is_64bits_host else "32"
 269         )
 270         consumerd_path = (
 271             self._project_root / "src" / "bin" / "lttng-consumerd" / "lttng-consumerd"
 272         )
 273
 274         no_sessiond_var = os.environ.get("TEST_NO_SESSIOND")
 275         if no_sessiond_var and no_sessiond_var == "1":
 276             # Run test without a session daemon; the user probably
 277             # intends to run one under gdb for example.
 278             return None
 279
 280         # Setup the session daemon's environment
 281         sessiond_env_vars = os.environ.get("LTTNG_SESSIOND_ENV_VARS")
 282         sessiond_env = os.environ.copy()
 283         if sessiond_env_vars:
 284             self._log("Additional lttng-sessiond environment variables:")
 285             additional_vars = self._unpack_env_vars(sessiond_env_vars)
 286             for var_name, var_value in additional_vars:
 287                 self._log("  {name}={value}".format(name=var_name, value=var_value))
 288                 sessiond_env[var_name] = var_value
 289
 290         sessiond_env["LTTNG_SESSION_CONFIG_XSD_PATH"] = str(
 291             self._project_root / "src" / "common"
 292         )
 293
 294         assert self._lttng_home is not None
 295         sessiond_env["LTTNG_HOME"] = str(self._lttng_home.path)
 296
 297         wait_queue = _SignalWaitQueue()
 298         signal.signal(signal.SIGUSR1, wait_queue.signal)
 299
 300         self._log(
 301             "Launching session daemon with LTTNG_HOME=`{home_dir}`".format(
 302                 home_dir=str(self._lttng_home.path)
 303             )
 304         )
 305         process = subprocess.Popen(
 306             [
 307                 str(sessiond_path),
 308                 consumerd_path_option_name,
 309                 str(consumerd_path),
 310                 "--sig-parent",
 311             ],
 312             stdout=subprocess.PIPE,
 313             stderr=subprocess.STDOUT,
 314             env=sessiond_env,
 315         )
 316
 317         if self._logging_function:
 318             self._sessiond_output_consumer: Optional[
 319                 ProcessOutputConsumer
 320             ] = ProcessOutputConsumer(process, "lttng-sessiond", self._logging_function)
 321             self._sessiond_output_consumer.daemon = True
 322             self._sessiond_output_consumer.start()
 323
 324         # Wait for SIGUSR1, indicating the sessiond is ready to proceed
 325         wait_queue.wait_for_signal()
 326         signal.signal(signal.SIGUSR1, wait_queue.signal)
 327
 328         return process
 329
 330     def _handle_termination_signal(
 331         self, signal_number: int, frame: Optional[FrameType]
 332     ) -> None:
 333         self._log(
 334             "Killed by {signal_name} signal, cleaning-up".format(
 335                 signal_name=signal.strsignal(signal_number)
 336             )
 337         )
 338         self._cleanup()
 339
 340     def launch_wait_trace_test_application(
 341         self, event_count: int
 342     ) -> WaitTraceTestApplication:
 343         """
 344         Launch an application that will wait before tracing `event_count` events.
 345         """
 346         return WaitTraceTestApplication(
 347             self._project_root
 348             / "tests"
 349             / "utils"
 350             / "testapp"
 351             / "gen-ust-nevents"
 352             / "gen-ust-nevents",
 353             event_count,
 354             self,
 355         )
 356
 357     # Clean-up managed processes
 358     def _cleanup(self) -> None:
 359         if self._sessiond and self._sessiond.poll() is None:
 360             # The session daemon is alive; kill it.
 361             self._log(
 362                 "Killing session daemon (pid = {sessiond_pid})".format(
 363                     sessiond_pid=self._sessiond.pid
 364                 )
 365             )
 366
 367             self._sessiond.terminate()
 368             self._sessiond.wait()
 369             if self._sessiond_output_consumer:
 370                 self._sessiond_output_consumer.join()
 371                 self._sessiond_output_consumer = None
 372
 373             self._log("Session daemon killed")
 374             self._sessiond = None
 375
 376         self._lttng_home = None
 377
 378     def __del__(self):
 379         self._cleanup()
 380
 381
 382 @contextlib.contextmanager
 383 def test_environment(with_sessiond: bool, log: Optional[Callable[[str], None]] = None):
 384     env = _Environment(with_sessiond, log)
 385     try:
 386         yield env
 387     finally:
 388         env._cleanup()