[core] Add import hook module (#769)

* [core] Add hook module with tests

* [core] linting

* [core] add more tests

* [core] linting

* fix bad merge

* [core] address comments
- fix test
- add dedup check and log

* [core] improve documentation for matcher

* [core] use module_name helper

* [core] deregister takes hook directly

* [core] enhance docstrings

Co-Authored-By: Kyle-Verhoog <[email protected]>

* [core] docstring improvement

Co-Authored-By: Kyle-Verhoog <[email protected]>

* [core] deregister returns outcome; docstring updates

* [core] remove stale comment

Author: Kyle-Verhoog

ddtrace/compat.py

@@ -18,6 +18,7 @@
 
 PYTHON_VERSION_INFO = sys.version_info
 PY2 = sys.version_info[0] == 2
+PY3 = sys.version_info[0] == 3
 
 # Infos about python passed to the trace agent through the header
 PYTHON_VERSION = platform.python_version()
@@ -34,6 +35,7 @@
 Queue = six.moves.queue.Queue
 iteritems = six.iteritems
 reraise = six.reraise
+reload_module = six.moves.reload_module
 
 stringify = six.text_type
 string_type = six.string_types[0]

ddtrace/utils/__init__.py

@@ -4,3 +4,12 @@ def merge_dicts(x, y):
     z = x.copy()  # start with x's keys and values
     z.update(y)  # modifies z with y's keys and values & returns None
     return z
+
+
+def get_module_name(module):
+    """Returns a module's name or None if one cannot be found.
+    Relevant PEP: https://www.python.org/dev/peps/pep-0451/
+    """
+    if hasattr(module, '__spec__'):
+        return module.__spec__.name
+    return getattr(module, '__name__', None)

ddtrace/utils/hook.py

@@ -0,0 +1,196 @@
+"""
+This module is based off of wrapt.importer (wrapt==1.11.0)
+https://github.com/GrahamDumpleton/wrapt/blob/4bcd190457c89e993ffcfec6dad9e9969c033e9e/src/wrapt/importer.py#L127-L136
+
+The reasoning for this is that wrapt.importer does not provide a mechanism to
+remove the import hooks and that wrapt removes the hooks after they are fired.
+
+So this module differs from wrapt.importer in that:
+    - removes unnecessary functionality (like allowing hooks to be import paths)
+    - deregister_post_import_hook is introduced to remove hooks
+    - the values of _post_import_hooks can only be lists (instead of allowing None)
+    - notify_module_loaded is modified to not remove the hooks when they are
+      fired.
+"""
+import logging
+import sys
+import threading
+
+from wrapt.decorators import synchronized
+
+from ddtrace.compat import PY3
+from ddtrace.utils import get_module_name
+
+
+log = logging.getLogger(__name__)
+
+
+_post_import_hooks = {}
+_post_import_hooks_init = False
+_post_import_hooks_lock = threading.RLock()
+
+
+@synchronized(_post_import_hooks_lock)
+def register_post_import_hook(name, hook):
+    """
+    Registers a module import hook, ``hook`` for a module with name ``name``.
+
+    If the module is already imported the hook is called immediately and a
+    debug message is logged since this should not be expected in our use-case.
+
+    :param name: Name of the module (full dotted path)
+    :type name: str
+    :param hook: Callable to be invoked with the module when it is imported.
+    :type hook: Callable
+    :return:
+    """
+    # Automatically install the import hook finder if it has not already
+    # been installed.
+    global _post_import_hooks_init
+
+    if not _post_import_hooks_init:
+        _post_import_hooks_init = True
+        sys.meta_path.insert(0, ImportHookFinder())
+
+    hooks = _post_import_hooks.get(name, [])
+
+    if hook in hooks:
+        log.debug('hook "{}" already exists on module "{}"'.format(hook, name))
+        return
+
+    module = sys.modules.get(name, None)
+
+    # If the module has been imported already fire the hook and log a debug msg.
+    if module:
+        log.debug('module "{}" already imported, firing hook'.format(name))
+        hook(module)
+
+    hooks.append(hook)
+    _post_import_hooks[name] = hooks
+
+
+@synchronized(_post_import_hooks_lock)
+def notify_module_loaded(module):
+    """
+    Indicate that a module has been loaded. Any post import hooks which were
+    registered for the target module will be invoked.
+
+    Any raised exceptions will be caught and an error message indicating that
+    the hook failed.
+
+    :param module: The module being loaded
+    :type module: ``types.ModuleType``
+    """
+    name = get_module_name(module)
+    hooks = _post_import_hooks.get(name, [])
+
+    for hook in hooks:
+        try:
+            hook(module)
+        except Exception as err:
+            log.warn('hook "{}" for module "{}" failed: {}'.format(hook, name, err))
+
+
+class _ImportHookLoader(object):
+    """
+    A custom module import finder. This intercepts attempts to import
+    modules and watches out for attempts to import target modules of
+    interest. When a module of interest is imported, then any post import
+    hooks which are registered will be invoked.
+    """
+    def load_module(self, fullname):
+        module = sys.modules[fullname]
+        notify_module_loaded(module)
+        return module
+
+
+class _ImportHookChainedLoader(object):
+    def __init__(self, loader):
+        self.loader = loader
+
+    def load_module(self, fullname):
+        module = self.loader.load_module(fullname)
+        notify_module_loaded(module)
+        return module
+
+
+class ImportHookFinder:
+    def __init__(self):
+        self.in_progress = {}
+
+    @synchronized(_post_import_hooks_lock)
+    def find_module(self, fullname, path=None):
+        # If the module being imported is not one we have registered
+        # post import hooks for, we can return immediately. We will
+        # take no further part in the importing of this module.
+
+        if fullname not in _post_import_hooks:
+            return None
+
+        # When we are interested in a specific module, we will call back
+        # into the import system a second time to defer to the import
+        # finder that is supposed to handle the importing of the module.
+        # We set an in progress flag for the target module so that on
+        # the second time through we don't trigger another call back
+        # into the import system and cause a infinite loop.
+
+        if fullname in self.in_progress:
+            return None
+
+        self.in_progress[fullname] = True
+
+        # Now call back into the import system again.
+
+        try:
+            if PY3:
+                # For Python 3 we need to use find_spec().loader
+                # from the importlib.util module. It doesn't actually
+                # import the target module and only finds the
+                # loader. If a loader is found, we need to return
+                # our own loader which will then in turn call the
+                # real loader to import the module and invoke the
+                # post import hooks.
+                try:
+                    import importlib.util
+                    loader = importlib.util.find_spec(fullname).loader
+                except (ImportError, AttributeError):
+                    loader = importlib.find_loader(fullname, path)
+                if loader:
+                    return _ImportHookChainedLoader(loader)
+            else:
+                # For Python 2 we don't have much choice but to
+                # call back in to __import__(). This will
+                # actually cause the module to be imported. If no
+                # module could be found then ImportError will be
+                # raised. Otherwise we return a loader which
+                # returns the already loaded module and invokes
+                # the post import hooks.
+                __import__(fullname)
+                return _ImportHookLoader()
+
+        finally:
+            del self.in_progress[fullname]
+
+
+@synchronized(_post_import_hooks_lock)
+def deregister_post_import_hook(modulename, hook):
+    """
+    Deregisters post import hooks for a module given the module name and a hook
+    that was previously installed.
+
+    :param modulename: Name of the module the hook is installed on.
+    :type: str
+    :param hook: The hook to remove (the function itself)
+    :type hook: Callable
+    :return: whether a hook was removed or not
+    """
+    if modulename not in _post_import_hooks:
+        return False
+
+    hooks = _post_import_hooks[modulename]
+
+    try:
+        hooks.remove(hook)
+        return True
+    except ValueError:
+        return False

tests/subprocesstest.py

@@ -84,15 +84,19 @@ def _run_test_in_subprocess(self, result):
             stderr=subprocess.PIPE,
             env=sp_test_env,
         )
-        _, stderr = sp.communicate()
+        stdout, stderr = sp.communicate()
 
         if sp.returncode:
             try:
                 cmdf = ' '.join(sp_test_cmd)
                 raise Exception('Subprocess Test "{}" Failed'.format(cmdf))
             except Exception:
                 exc_info = sys.exc_info()
-            sys.stderr.write(stderr)
+
+            # DEV: stderr, stdout are byte sequences so to print them nicely
+            #      back out they should be decoded.
+            sys.stderr.write(stderr.decode())
+            sys.stdout.write(stdout.decode())
             result.addFailure(self, exc_info)
         else:
             result.addSuccess(self)