Added new functions_to_trace option for celtral way of performance instrumentation (#1960)

Have a list of functions that can be passed to "sentry_sdk.init()". When the SDK starts it goes through the list and instruments all the functions in the list.

functions_to_trace = [
    {"qualified_name": "tests.test_basics._hello_world_counter"},
    {"qualified_name": "time.sleep"},
    {"qualified_name": "collections.Counter.most_common"},
]

sentry_sdk.init(
    dsn="...",
    traces_sample_rate=1.0,
    functions_to_trace=functions_to_trace,
)

Author: antonpirker

sentry_sdk/client.py

@@ -1,3 +1,4 @@
+from importlib import import_module
 import os
 import uuid
 import random
@@ -17,6 +18,7 @@
     logger,
 )
 from sentry_sdk.serializer import serialize
+from sentry_sdk.tracing import trace
 from sentry_sdk.transport import make_transport
 from sentry_sdk.consts import (
     DEFAULT_OPTIONS,
@@ -38,6 +40,7 @@
     from typing import Callable
     from typing import Dict
     from typing import Optional
+    from typing import Sequence
 
     from sentry_sdk.scope import Scope
     from sentry_sdk._types import Event, Hint
@@ -118,6 +121,14 @@ def _get_options(*args, **kwargs):
     return rv
 
 
+try:
+    # Python 3.6+
+    module_not_found_error = ModuleNotFoundError
+except Exception:
+    # Older Python versions
+    module_not_found_error = ImportError  # type: ignore
+
+
 class _Client(object):
     """The client is internally responsible for capturing the events and
     forwarding them to sentry through the configured transport.  It takes
@@ -140,6 +151,52 @@ def __setstate__(self, state):
         self.options = state["options"]
         self._init_impl()
 
+    def _setup_instrumentation(self, functions_to_trace):
+        # type: (Sequence[Dict[str, str]]) -> None
+        """
+        Instruments the functions given in the list `functions_to_trace` with the `@sentry_sdk.tracing.trace` decorator.
+        """
+        for function in functions_to_trace:
+            class_name = None
+            function_qualname = function["qualified_name"]
+            module_name, function_name = function_qualname.rsplit(".", 1)
+
+            try:
+                # Try to import module and function
+                # ex: "mymodule.submodule.funcname"
+
+                module_obj = import_module(module_name)
+                function_obj = getattr(module_obj, function_name)
+                setattr(module_obj, function_name, trace(function_obj))
+                logger.debug("Enabled tracing for %s", function_qualname)
+
+            except module_not_found_error:
+                try:
+                    # Try to import a class
+                    # ex: "mymodule.submodule.MyClassName.member_function"
+
+                    module_name, class_name = module_name.rsplit(".", 1)
+                    module_obj = import_module(module_name)
+                    class_obj = getattr(module_obj, class_name)
+                    function_obj = getattr(class_obj, function_name)
+                    setattr(class_obj, function_name, trace(function_obj))
+                    setattr(module_obj, class_name, class_obj)
+                    logger.debug("Enabled tracing for %s", function_qualname)
+
+                except Exception as e:
+                    logger.warning(
+                        "Can not enable tracing for '%s'. (%s) Please check your `functions_to_trace` parameter.",
+                        function_qualname,
+                        e,
+                    )
+
+            except Exception as e:
+                logger.warning(
+                    "Can not enable tracing for '%s'. (%s) Please check your `functions_to_trace` parameter.",
+                    function_qualname,
+                    e,
+                )
+
     def _init_impl(self):
         # type: () -> None
         old_debug = _client_init_debug.get(False)
@@ -184,6 +241,8 @@ def _capture_envelope(envelope):
             except ValueError as e:
                 logger.debug(str(e))
 
+        self._setup_instrumentation(self.options.get("functions_to_trace", []))
+
     @property
     def dsn(self):
         # type: () -> Optional[str]

sentry_sdk/consts.py

@@ -133,6 +133,7 @@ def __init__(
         trace_propagation_targets=[  # noqa: B006
             MATCH_ALL
         ],  # type: Optional[Sequence[str]]
+        functions_to_trace=[],  # type: Sequence[str]  # noqa: B006
         event_scrubber=None,  # type: Optional[sentry_sdk.scrubber.EventScrubber]
     ):
         # type: (...) -> None

tests/test_basics.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import sys
+import time
 
 import pytest
 
@@ -618,3 +619,70 @@ def foo(event, hint):
 )
 def test_get_sdk_name(installed_integrations, expected_name):
     assert get_sdk_name(installed_integrations) == expected_name
+
+
+def _hello_world(word):
+    return "Hello, {}".format(word)
+
+
+def test_functions_to_trace(sentry_init, capture_events):
+    functions_to_trace = [
+        {"qualified_name": "tests.test_basics._hello_world"},
+        {"qualified_name": "time.sleep"},
+    ]
+
+    sentry_init(
+        traces_sample_rate=1.0,
+        functions_to_trace=functions_to_trace,
+    )
+
+    events = capture_events()
+
+    with start_transaction(name="something"):
+        time.sleep(0)
+
+        for word in ["World", "You"]:
+            _hello_world(word)
+
+    assert len(events) == 1
+
+    (event,) = events
+
+    assert len(event["spans"]) == 3
+    assert event["spans"][0]["description"] == "time.sleep"
+    assert event["spans"][1]["description"] == "tests.test_basics._hello_world"
+    assert event["spans"][2]["description"] == "tests.test_basics._hello_world"
+
+
+class WorldGreeter:
+    def __init__(self, word):
+        self.word = word
+
+    def greet(self, new_word=None):
+        return "Hello, {}".format(new_word if new_word else self.word)
+
+
+def test_functions_to_trace_with_class(sentry_init, capture_events):
+    functions_to_trace = [
+        {"qualified_name": "tests.test_basics.WorldGreeter.greet"},
+    ]
+
+    sentry_init(
+        traces_sample_rate=1.0,
+        functions_to_trace=functions_to_trace,
+    )
+
+    events = capture_events()
+
+    with start_transaction(name="something"):
+        wg = WorldGreeter("World")
+        wg.greet()
+        wg.greet("You")
+
+    assert len(events) == 1
+
+    (event,) = events
+
+    assert len(event["spans"]) == 2
+    assert event["spans"][0]["description"] == "tests.test_basics.WorldGreeter.greet"
+    assert event["spans"][1]["description"] == "tests.test_basics.WorldGreeter.greet"

tox.ini

@@ -177,6 +177,7 @@ deps =
     arq: arq>=0.23.0
     arq: fakeredis>=2.2.0,<2.8
     arq: pytest-asyncio
+    arq: async-timeout
 
     # Asgi
     asgi: pytest-asyncio