Extract LdPreload action from Trace action to support preloading any lib

tracetools_launch/tracetools_launch/action.py

@@ -15,20 +15,18 @@
 """Module for the Trace action."""
 
 import re
-import subprocess
 from typing import List
 from typing import Optional
-from typing import Union
 
 from launch.action import Action
-from launch.actions import SetEnvironmentVariable
 from launch.event import Event
 from launch.event_handlers import OnShutdown
 from launch.launch_context import LaunchContext
 from tracetools_trace.tools import lttng
 from tracetools_trace.tools import names
 from tracetools_trace.tools import path
-from tracetools_trace.tools import tracing_supported
+
+from .actions.ld_preload import LdPreload
 
 
 class Trace(Action):
@@ -38,9 +36,12 @@ class Trace(Action):
     Sets up and enables tracing through a launch file description.
     """
 
-    PROFILE_LIB_NORMAL = 'liblttng-ust-cyg-profile.so'
-    PROFILE_LIB_FAST = 'liblttng-ust-cyg-profile-fast.so'
+    LIB_PROFILE_NORMAL = 'liblttng-ust-cyg-profile.so'
+    LIB_PROFILE_FAST = 'liblttng-ust-cyg-profile-fast.so'
+    LIB_MEMORY_UST = 'liblttng-ust-libc-wrapper.so'
+
     PROFILE_EVENT_PATTERN = '^lttng_ust_cyg_profile.*:func_.*'
+    MEMORY_UST_EVENT_PATTERN = '^lttng_ust_libc:.*'
 
     def __init__(
         self,
@@ -71,50 +72,53 @@ class Trace(Action):
         self.__events_ust = events_ust
         self.__events_kernel = events_kernel
         self.__profile_fast = profile_fast
-        self.__ld_preload_action = None
-        if self.has_profiling_events(events_ust):
-            profile_lib_name = self.PROFILE_LIB_FAST if profile_fast else self.PROFILE_LIB_NORMAL
-            profile_lib_path = self.get_shared_lib_path(profile_lib_name)
-            if profile_lib_path is not None:
-                self.__ld_preload_action = SetEnvironmentVariable(
-                    'LD_PRELOAD',
-                    profile_lib_path,
+        self.__ld_preload_actions = []
+        # Add LD_PRELOAD actions if corresponding events are enabled
+        if self.has_profiling_events(self.__events_ust):
+            self.__ld_preload_actions.append(
+                LdPreload(self.LIB_PROFILE_FAST if profile_fast else self.LIB_PROFILE_NORMAL)
+            )
+        if self.has_ust_memory_events(self.__events_ust):
+            self.__ld_preload_actions.append(
+                LdPreload(self.LIB_MEMORY_UST)
             )
 
-    @classmethod
-    def has_profiling_events(cls, events_ust: List[str]) -> bool:
-        """Check if the UST events list contains at least one profiling event."""
-        return any(re.match(cls.PROFILE_EVENT_PATTERN, event_name) for event_name in events_ust)
-
     @staticmethod
-    def get_shared_lib_path(lib_name: str) -> Union[str, None]:
+    def any_events_match(
+        name_pattern: str,
+        events: List[str],
+    ) -> bool:
         """
-        Get the full path to a given shared lib, if possible.
+        Check if any event name in the list matches the given pattern.
 
-        :param lib_name: the name of the shared library
-        :return: the full path if found, `None` otherwise
+        :param name_pattern: the pattern to use for event names
+        :param events: the list of event names
+        :return true if there is a match, false otherwise
         """
-        if not tracing_supported():
-            return None
-        (exit_code, output) = subprocess.getstatusoutput(f'whereis -b {lib_name}')
-        if exit_code != 0:
-            return None
-        # Output of whereis is
-        # <input_lib_name>: <full path, if found>
-        # Filter out empty strings, in case lib is not found
-        output_split = [split_part for split_part in output.split(':') if len(split_part) > 0]
-        if len(output_split) != 2:
-            return None
-        return output_split[1].strip()
+        return any(re.match(name_pattern, event_name) for event_name in events)
+
+    @classmethod
+    def has_profiling_events(
+        cls,
+        events_ust: List[str],
+    ) -> bool:
+        """Check if the UST events list contains at least one profiling event."""
+        return cls.any_events_match(cls.PROFILE_EVENT_PATTERN, events_ust)
+
+    @classmethod
+    def has_ust_memory_events(
+        cls,
+        events_ust: List[str],
+    ) -> bool:
+        """Check if the UST events list contains at least one userspace memory event."""
+        return cls.any_events_match(cls.MEMORY_UST_EVENT_PATTERN, events_ust)
 
     def execute(self, context: LaunchContext) -> Optional[List[Action]]:
         # TODO make sure this is done as late as possible
         context.register_event_handler(OnShutdown(on_shutdown=self._destroy))
         # TODO make sure this is done as early as possible
         self._setup()
-        if self.__ld_preload_action is not None:
-            return [self.__ld_preload_action]
-        return None
+        return self.__ld_preload_actions
 
     def _setup(self) -> None:
         lttng.lttng_init(
@@ -132,7 +136,7 @@ class Trace(Action):
             f'session_name={self.__session_name}, '
             f'base_path={self.__base_path}, '
             f'num_events_ust={len(self.__events_ust)}, '
-            f'num_events_kernel={len(self.__events_kernel)}), '
-            f'profiling={self.__ld_preload_action is not None}, '
-            f'profile_fast={self.__profile_fast}'
+            f'num_events_kernel={len(self.__events_kernel)}, '
+            f'profile_fast={self.__profile_fast}, '
+            f'ld_preload_actions={self.__ld_preload_actions})'
         )

tracetools_launch/tracetools_launch/actions/ld_preload.py

@@ -0,0 +1,90 @@
+# Copyright 2019 Apex.AI, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module for the LdPreload action."""
+
+import subprocess
+from typing import List
+from typing import Optional
+from typing import Union
+
+from launch.action import Action
+from launch.actions import SetEnvironmentVariable
+from launch.launch_context import LaunchContext
+from tracetools_trace.tools import tracing_supported
+
+
+class LdPreload(Action):
+    """Action that adds a SetEnvironmentVariable action to preload a library."""
+
+    ENV_VAR_LD_PRELOAD = 'LD_PRELOAD'
+
+    def __init__(
+        self,
+        lib_name: str,
+        **kwargs,
+    ) -> None:
+        """
+        Create an LdPreload action.
+
+        :param lib_name: the name of the library
+        """
+        super().__init__(**kwargs)
+        self.__lib_name = lib_name
+        self.__set_env_action = None
+        # Try to find lib
+        self.__lib_path = self.get_shared_lib_path(self.__lib_name)
+        # And create action if found
+        if self.__lib_path is not None:
+            self.__set_env_action = SetEnvironmentVariable(
+                self.ENV_VAR_LD_PRELOAD,
+                self.__lib_path,
+            )
+
+    def lib_found(self) -> bool:
+        return self.__set_env_action is not None
+
+    def execute(self, context: LaunchContext) -> Optional[List[Action]]:
+        if self.lib_found():
+            return [self.__set_env_action]
+        return None
+
+    @staticmethod
+    def get_shared_lib_path(lib_name: str) -> Union[str, None]:
+        """
+        Get the full path to a given shared lib, if possible.
+
+        :param lib_name: the name of the shared library
+        :return: the full path if found, `None` otherwise
+        """
+        if not tracing_supported():
+            return None
+        (exit_code, output) = subprocess.getstatusoutput(f'whereis -b {lib_name}')
+        if exit_code != 0:
+            return None
+        # Output of whereis is
+        # <input_lib_name>: <full path, if found>
+        # Filter out empty strings, in case lib is not found
+        output_split = [split_part for split_part in output.split(':') if len(split_part) > 0]
+        if len(output_split) != 2:
+            return None
+        return output_split[1].strip()
+
+    def __repr__(self):
+        return (
+            'LdPreload('
+            f'lib_name={self.__lib_name}, '
+            f'lib found={self.lib_found()}, '
+            f'lib_path={self.__lib_path})'
+        )