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, Iterator, Optional, Tuple, List, Generator
  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):
  27         # type: (str) -> None
  28         self._directory_path = tempfile.mkdtemp(prefix=prefix)
  29
  30     def __del__(self):
  31         shutil.rmtree(self._directory_path, ignore_errors=True)
  32
  33     @property
  34     def path(self):
  35         # type: () -> pathlib.Path
  36         return pathlib.Path(self._directory_path)
  37
  38
  39 class _SignalWaitQueue:
  40     """
  41     Utility class useful to wait for a signal before proceeding.
  42
  43     Simply register the `signal` method as the handler for the signal you are
  44     interested in and call `wait_for_signal` to wait for its reception.
  45
  46     Registering a signal:
  47         signal.signal(signal.SIGWHATEVER, queue.signal)
  48
  49     Waiting for the signal:
  50         queue.wait_for_signal()
  51     """
  52
  53     def __init__(self):
  54         self._queue = queue.Queue()  # type: queue.Queue
  55
  56     def signal(
  57         self,
  58         signal_number,
  59         frame,  # type: Optional[FrameType]
  60     ):
  61         self._queue.put_nowait(signal_number)
  62
  63     def wait_for_signal(self):
  64         self._queue.get(block=True)
  65
  66     @contextlib.contextmanager
  67     def intercept_signal(self, signal_number):
  68         # type: (int) -> Generator[None, None, None]
  69         original_handler = signal.getsignal(signal_number)
  70         signal.signal(signal_number, self.signal)
  71         try:
  72             yield
  73         except:
  74             # Restore the original signal handler and forward the exception.
  75             raise
  76         finally:
  77             signal.signal(signal_number, original_handler)
  78
  79
  80 class WaitTraceTestApplication:
  81     """
  82     Create an application that waits before tracing. This allows a test to
  83     launch an application, get its PID, and get it to start tracing when it
  84     has completed its setup.
  85     """
  86
  87     def __init__(
  88         self,
  89         binary_path,  # type: pathlib.Path
  90         event_count,  # type: int
  91         environment,  # type: Environment
  92         wait_time_between_events_us=0,  # type: int
  93     ):
  94         self._environment = environment  # type: Environment
  95         if event_count % 5:
  96             # The test application currently produces 5 different events per iteration.
  97             raise ValueError("event count must be a multiple of 5")
  98         self._iteration_count = int(event_count / 5)  # type: int
  99         # File that the application will wait to see before tracing its events.
 100         self._app_start_tracing_file_path = pathlib.Path(
 101             tempfile.mktemp(
 102                 prefix="app_",
 103                 suffix="_start_tracing",
 104                 dir=self._compat_open_path(environment.lttng_home_location),
 105             )
 106         )
 107         self._has_returned = False
 108
 109         test_app_env = os.environ.copy()
 110         test_app_env["LTTNG_HOME"] = str(environment.lttng_home_location)
 111         # Make sure the app is blocked until it is properly registered to
 112         # the session daemon.
 113         test_app_env["LTTNG_UST_REGISTER_TIMEOUT"] = "-1"
 114
 115         # File that the application will create to indicate it has completed its initialization.
 116         app_ready_file_path = tempfile.mktemp(
 117             prefix="app_",
 118             suffix="_ready",
 119             dir=self._compat_open_path(environment.lttng_home_location),
 120         )  # type: str
 121
 122         test_app_args = [str(binary_path)]
 123         test_app_args.extend(
 124             shlex.split(
 125                 "--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(
 126                     iteration_count=self._iteration_count,
 127                     app_ready_file_path=app_ready_file_path,
 128                     app_start_tracing_file_path=self._app_start_tracing_file_path,
 129                     wait_time_between_events_us=wait_time_between_events_us,
 130                 )
 131             )
 132         )
 133
 134         self._process = subprocess.Popen(
 135             test_app_args,
 136             env=test_app_env,
 137         )  # type: subprocess.Popen
 138
 139         # Wait for the application to create the file indicating it has fully
 140         # initialized. Make sure the app hasn't crashed in order to not wait
 141         # forever.
 142         while True:
 143             if os.path.exists(app_ready_file_path):
 144                 break
 145
 146             if self._process.poll() is not None:
 147                 # Application has unexepectedly returned.
 148                 raise RuntimeError(
 149                     "Test application has unexepectedly returned during its initialization with return code `{return_code}`".format(
 150                         return_code=self._process.returncode
 151                     )
 152                 )
 153
 154             time.sleep(0.1)
 155
 156     def trace(self):
 157         # type: () -> None
 158         if self._process.poll() is not None:
 159             # Application has unexepectedly returned.
 160             raise RuntimeError(
 161                 "Test application has unexepectedly before tracing with return code `{return_code}`".format(
 162                     return_code=self._process.returncode
 163                 )
 164             )
 165         open(self._compat_open_path(self._app_start_tracing_file_path), mode="x")
 166
 167     def wait_for_exit(self):
 168         # type: () -> None
 169         if self._process.wait() != 0:
 170             raise RuntimeError(
 171                 "Test application has exit with return code `{return_code}`".format(
 172                     return_code=self._process.returncode
 173                 )
 174             )
 175         self._has_returned = True
 176
 177     @property
 178     def vpid(self):
 179         # type: () -> int
 180         return self._process.pid
 181
 182     @staticmethod
 183     def _compat_open_path(path):
 184         # type: (pathlib.Path) -> pathlib.Path | str
 185         """
 186         The builtin open() in python >= 3.6 expects a path-like object while
 187         prior versions expect a string or bytes object. Return the correct type
 188         based on the presence of the "__fspath__" attribute specified in PEP-519.
 189         """
 190         if hasattr(path, "__fspath__"):
 191             return path
 192         else:
 193             return str(path)
 194
 195     def __del__(self):
 196         if not self._has_returned:
 197             # This is potentially racy if the pid has been recycled. However,
 198             # we can't use pidfd_open since it is only available in python >= 3.9.
 199             self._process.kill()
 200             self._process.wait()
 201
 202
 203 class TraceTestApplication:
 204     """
 205     Create an application that emits events as soon as it is launched. In most
 206     scenarios, it is preferable to use a WaitTraceTestApplication.
 207     """
 208
 209     def __init__(self, binary_path, environment):
 210         # type: (pathlib.Path, Environment)
 211         self._environment = environment  # type: Environment
 212         self._has_returned = False
 213
 214         test_app_env = os.environ.copy()
 215         test_app_env["LTTNG_HOME"] = str(environment.lttng_home_location)
 216         # Make sure the app is blocked until it is properly registered to
 217         # the session daemon.
 218         test_app_env["LTTNG_UST_REGISTER_TIMEOUT"] = "-1"
 219
 220         test_app_args = [str(binary_path)]
 221
 222         self._process = subprocess.Popen(
 223             test_app_args, env=test_app_env
 224         )  # type: subprocess.Popen
 225
 226     def wait_for_exit(self):
 227         # type: () -> None
 228         if self._process.wait() != 0:
 229             raise RuntimeError(
 230                 "Test application has exit with return code `{return_code}`".format(
 231                     return_code=self._process.returncode
 232                 )
 233             )
 234         self._has_returned = True
 235
 236     def __del__(self):
 237         if not self._has_returned:
 238             # This is potentially racy if the pid has been recycled. However,
 239             # we can't use pidfd_open since it is only available in python >= 3.9.
 240             self._process.kill()
 241             self._process.wait()
 242
 243
 244 class ProcessOutputConsumer(threading.Thread, logger._Logger):
 245     def __init__(
 246         self,
 247         process,  # type: subprocess.Popen
 248         name,  # type: str
 249         log,  # type: Callable[[str], None]
 250     ):
 251         threading.Thread.__init__(self)
 252         self._prefix = name
 253         logger._Logger.__init__(self, log)
 254         self._process = process
 255
 256     def run(self):
 257         # type: () -> None
 258         while self._process.poll() is None:
 259             assert self._process.stdout
 260             line = self._process.stdout.readline().decode("utf-8").replace("\n", "")
 261             if len(line) != 0:
 262                 self._log("{prefix}: {line}".format(prefix=self._prefix, line=line))
 263
 264
 265 # Generate a temporary environment in which to execute a test.
 266 class _Environment(logger._Logger):
 267     def __init__(
 268         self,
 269         with_sessiond,  # type: bool
 270         log=None,  # type: Optional[Callable[[str], None]]
 271     ):
 272         super().__init__(log)
 273         signal.signal(signal.SIGTERM, self._handle_termination_signal)
 274         signal.signal(signal.SIGINT, self._handle_termination_signal)
 275
 276         # Assumes the project's hierarchy to this file is:
 277         # tests/utils/python/this_file
 278         self._project_root = (
 279             pathlib.Path(__file__).absolute().parents[3]
 280         )  # type: pathlib.Path
 281         self._lttng_home = TemporaryDirectory(
 282             "lttng_test_env_home"
 283         )  # type: Optional[TemporaryDirectory]
 284
 285         self._sessiond = (
 286             self._launch_lttng_sessiond() if with_sessiond else None
 287         )  # type: Optional[subprocess.Popen[bytes]]
 288
 289     @property
 290     def lttng_home_location(self):
 291         # type: () -> pathlib.Path
 292         if self._lttng_home is None:
 293             raise RuntimeError("Attempt to access LTTng home after clean-up")
 294         return self._lttng_home.path
 295
 296     @property
 297     def lttng_client_path(self):
 298         # type: () -> pathlib.Path
 299         return self._project_root / "src" / "bin" / "lttng" / "lttng"
 300
 301     def create_temporary_directory(self, prefix=None):
 302         # type: (Optional[str]) -> pathlib.Path
 303         # Simply return a path that is contained within LTTNG_HOME; it will
 304         # be destroyed when the temporary home goes out of scope.
 305         assert self._lttng_home
 306         return pathlib.Path(
 307             tempfile.mkdtemp(
 308                 prefix="tmp" if prefix is None else prefix,
 309                 dir=str(self._lttng_home.path),
 310             )
 311         )
 312
 313     # Unpack a list of environment variables from a string
 314     # such as "HELLO=is_it ME='/you/are/looking/for'"
 315     @staticmethod
 316     def _unpack_env_vars(env_vars_string):
 317         # type: (str) -> List[Tuple[str, str]]
 318         unpacked_vars = []
 319         for var in shlex.split(env_vars_string):
 320             equal_position = var.find("=")
 321             # Must have an equal sign and not end with an equal sign
 322             if equal_position == -1 or equal_position == len(var) - 1:
 323                 raise ValueError(
 324                     "Invalid sessiond environment variable: `{}`".format(var)
 325                 )
 326
 327             var_name = var[0:equal_position]
 328             var_value = var[equal_position + 1 :]
 329             # Unquote any paths
 330             var_value = var_value.replace("'", "")
 331             var_value = var_value.replace('"', "")
 332             unpacked_vars.append((var_name, var_value))
 333
 334         return unpacked_vars
 335
 336     def _launch_lttng_sessiond(self):
 337         # type: () -> Optional[subprocess.Popen]
 338         is_64bits_host = sys.maxsize > 2**32
 339
 340         sessiond_path = (
 341             self._project_root / "src" / "bin" / "lttng-sessiond" / "lttng-sessiond"
 342         )
 343         consumerd_path_option_name = "--consumerd{bitness}-path".format(
 344             bitness="64" if is_64bits_host else "32"
 345         )
 346         consumerd_path = (
 347             self._project_root / "src" / "bin" / "lttng-consumerd" / "lttng-consumerd"
 348         )
 349
 350         no_sessiond_var = os.environ.get("TEST_NO_SESSIOND")
 351         if no_sessiond_var and no_sessiond_var == "1":
 352             # Run test without a session daemon; the user probably
 353             # intends to run one under gdb for example.
 354             return None
 355
 356         # Setup the session daemon's environment
 357         sessiond_env_vars = os.environ.get("LTTNG_SESSIOND_ENV_VARS")
 358         sessiond_env = os.environ.copy()
 359         if sessiond_env_vars:
 360             self._log("Additional lttng-sessiond environment variables:")
 361             additional_vars = self._unpack_env_vars(sessiond_env_vars)
 362             for var_name, var_value in additional_vars:
 363                 self._log("  {name}={value}".format(name=var_name, value=var_value))
 364                 sessiond_env[var_name] = var_value
 365
 366         sessiond_env["LTTNG_SESSION_CONFIG_XSD_PATH"] = str(
 367             self._project_root / "src" / "common"
 368         )
 369
 370         assert self._lttng_home is not None
 371         sessiond_env["LTTNG_HOME"] = str(self._lttng_home.path)
 372
 373         wait_queue = _SignalWaitQueue()
 374         with wait_queue.intercept_signal(signal.SIGUSR1):
 375             self._log(
 376                 "Launching session daemon with LTTNG_HOME=`{home_dir}`".format(
 377                     home_dir=str(self._lttng_home.path)
 378                 )
 379             )
 380             process = subprocess.Popen(
 381                 [
 382                     str(sessiond_path),
 383                     consumerd_path_option_name,
 384                     str(consumerd_path),
 385                     "--sig-parent",
 386                 ],
 387                 stdout=subprocess.PIPE,
 388                 stderr=subprocess.STDOUT,
 389                 env=sessiond_env,
 390             )
 391
 392             if self._logging_function:
 393                 self._sessiond_output_consumer = ProcessOutputConsumer(
 394                     process, "lttng-sessiond", self._logging_function
 395                 )  # type: Optional[ProcessOutputConsumer]
 396                 self._sessiond_output_consumer.daemon = True
 397                 self._sessiond_output_consumer.start()
 398
 399             # Wait for SIGUSR1, indicating the sessiond is ready to proceed
 400             wait_queue.wait_for_signal()
 401
 402         return process
 403
 404     def _handle_termination_signal(self, signal_number, frame):
 405         # type: (int, Optional[FrameType]) -> None
 406         self._log(
 407             "Killed by {signal_name} signal, cleaning-up".format(
 408                 signal_name=signal.strsignal(signal_number)
 409             )
 410         )
 411         self._cleanup()
 412
 413     def launch_wait_trace_test_application(self, event_count):
 414         # type: (int) -> WaitTraceTestApplication
 415         """
 416         Launch an application that will wait before tracing `event_count` events.
 417         """
 418         return WaitTraceTestApplication(
 419             self._project_root
 420             / "tests"
 421             / "utils"
 422             / "testapp"
 423             / "gen-ust-nevents"
 424             / "gen-ust-nevents",
 425             event_count,
 426             self,
 427         )
 428
 429     def launch_trace_test_constructor_application(self):
 430         # type () -> TraceTestApplication
 431         """
 432         Launch an application that will trace from within constructors.
 433         """
 434         return TraceTestApplication(
 435             self._project_root
 436             / "tests"
 437             / "utils"
 438             / "testapp"
 439             / "gen-ust-events-constructor"
 440             / "gen-ust-events-constructor",
 441             self,
 442         )
 443
 444     # Clean-up managed processes
 445     def _cleanup(self):
 446         # type: () -> None
 447         if self._sessiond and self._sessiond.poll() is None:
 448             # The session daemon is alive; kill it.
 449             self._log(
 450                 "Killing session daemon (pid = {sessiond_pid})".format(
 451                     sessiond_pid=self._sessiond.pid
 452                 )
 453             )
 454
 455             self._sessiond.terminate()
 456             self._sessiond.wait()
 457             if self._sessiond_output_consumer:
 458                 self._sessiond_output_consumer.join()
 459                 self._sessiond_output_consumer = None
 460
 461             self._log("Session daemon killed")
 462             self._sessiond = None
 463
 464         self._lttng_home = None
 465
 466     def __del__(self):
 467         self._cleanup()
 468
 469
 470 @contextlib.contextmanager
 471 def test_environment(with_sessiond, log=None):
 472     # type: (bool, Optional[Callable[[str], None]]) -> Iterator[_Environment]
 473     env = _Environment(with_sessiond, log)
 474     try:
 475         yield env
 476     finally:
 477         env._cleanup()