Using Decorators to Instrument Python Code with OpenTelemetry Traces

Last week, I wrote a post on the topic of using the Decorator design pattern to help remove some of the boilerplate required to set up tracing. My example code was using .NET, which required some heavy lifting, using a DispatchProxy class to intercept method calls and inject the tracing logic. Python, however, makes our lives much easier with built-in support for function decorators.

Why Use a Tracing Decorator?

OpenTelemetry tracing is awesome! It allows defining specific segments within the code called ‘Spans’ and tracking their execution and dependencies in runtime. It does, however, require a lot of discipline in instrumenting the code in order for that information to be available. Basically, for each class/module/function, you would need to add code like this:

tracer = trace.get_tracer(__name__)

class UserService:

    def all(self, id):
        with tracer.start_as_current_span("retrieve users") as span:
            users = self.user_store.get_users(id)
            span.set_attribute('some attribute', 'some_value')
            
        ...

    async def validate(self, user_name):

        with tracer.start_as_current_span("validate users") as span:
            span.set_attribute('some attribute', 'some_value')
            try:
                await self.user_store.get_user(user_name)
            ...
    
    
     async def test(self, user_name):

        with tracer.start_as_current_span("test user") as span:
            span.set_attribute('some attribute', 'some_value')
            try:
                await self.user_store.get_user(user_name)
            ...

Beyond the repetition and the need to enforce adding the tracer.start_as_current_span calls everywhere, coming up with the right naming convention for spans, and ensuring they have a unique name and level of granularity can be quite challenging for large code bases.

It would be great if instead of the above code, we’d be able to write something like this:

@instrument({"some attribute", "some_value"})
class UserService:

    def all(self, id):
        with tracer.start_as_current_span("retrieve users"):
            users = self.user_store.get_users(id)
        ...

    @instrument(span_name="custom_name")
    async def validate(self, user_name):
      
        try:
            await self.user_store.get_user(user_name)
        ...

    
    async def test(self, user_name):

        try:
            await self.user_store.get_user(user_name)
        ...

In this manner, an entire class could automatically be instrumented and the added Span attributes already injected. Naming can be assigned by convention (for example, the function name) with an easy way to change that convention if needed.

The entire source code for this post is available on this repository. It is also available as a pypi package which includes the decorator implementation if you just want to be able to use this right away.

Implementing a Basic Tracing Decorator Using Python Decorators

Python decorators are really neat. If you’re looking for more comprehensive documentation on how to use them, I recommend Geir Arne Hjelle’s post Primer on Python Decorators, which does an excellent job covering the topic.

To get started, we’ll create a simple decorator which can be added to a function in order to automatically instrument it. This will remove the need to add the tracing boilerplate code, as well as take care of the default naming convention for the span (which we’ll generalize later). The implementation in Python is pretty straightforward and consists of returning the wrapper function that will be used to execute the decorated function. Here’s the code:

def instrument(_func=None, *,  span_name: str = "", record_exception: bool = True,
               attributes: Dict[str, str] = None, existing_tracer: Tracer = None):

    def span_decorator(func):
        tracer = existing_tracer or trace.get_tracer(func.__module__)

        def _set_attributes(span, attributes_dict):
            if attributes_dict:
                for att in attributes_dict:
                    span.set_attribute(att, attributes_dict[att])
                    
        @wraps(func)
        def wrap_with_span(*args, **kwargs):
            name = span_name or TracingDecoratorOptions.naming_scheme(func)
            with tracer.start_as_current_span(name, record_exception=record_exception) as span:
                _set_attributes(span, attributes)
                return func(*args, **kwargs)

        return wrap_with_span

    if _func is None:
        return span_decorator
    else:
        return span_decorator(_func)

Basically, at this stage, all we’re doing is returning the wrapper function wrap_with_span which will be applied as the function decorator by the interpreter. We are using the functools.wraps() decorator (line 12) which makes sure the wrapper function will have the same name and metadata as the original function, making the decorator transparent to the function caller.

The wrap_with_span function will automatically create and name a span as well as set its attributes. By default, we name the span based on the function name, but we can create extensibility around that. For example, this static class can let developers replace the default naming convention with other implementations:

class TracingDecoratorOptions:
    class NamingSchemes:
        @staticmethod
        def function_qualified_name(func: Callable):
            return func.__qualname__

        default_scheme = function_qualified_name

    naming_scheme: Callable[[Callable], str] = NamingSchemes.default_scheme
    default_attributes: Dict[str, str] = {}

    @staticmethod
    def set_naming_scheme(naming_scheme: Callable[[Callable], str]):
        TracingDecoratorOptions.naming_scheme = naming_scheme

We also provide a span_name parameter developers can use to set a custom span name if needed.

Testing the Code

Let’s add a test to check our new decorator is working:

@classmethod
  def setup_class(cls):
      resource = Resource.create(attributes={SERVICE_NAME: "test"})
      provider = TracerProvider(resource=resource)
      trace.set_tracer_provider(provider)

  @instrument
  def test_decorated_function_gets_instrumented_automatically_with_span():
      assert trace.get_current_span().is_recording() is True

We set up OTEL so that the trace operation will have an effect, then validate that inside the test method, which has the new decorator applied, we have an active tracing span.

Success! The first iteration of the TracingDecorator is complete.

Instrumenting an Entire Class

Adding a decorator on a function-to-function basis might also become a bit tedious and repetitive. To help with that, we can modify the decorator to iterate over each function in the class and decorate it, ignoring private functions for now:

def instrument(_func_or_class=None, *, span_name: str = "", record_exception: bool = True,
               attributes: Dict[str, str] = None, existing_tracer: Tracer = None, ignore=False):
               
    def decorate_class(cls):
        for name, method in inspect.getmembers(cls, inspect.isfunction):
            # Ignore private functions, TODO: maybe make this a setting?
            if not name.startswith('_'):
                setattr(cls, name, instrument(record_exception=record_exception,
                                              attributes=attributes,
                                              existing_tracer=existing_tracer)(method))
        return cls

    # Check if this is a span or class decorator
    if inspect.isclass(_func_or_class):
        return decorate_class(_func_or_class)
    
    def span_decorator(func_or_class):

        if inspect.isclass(func_or_class):
            return decorate_class(func_or_class)
        ...
    
    if _func_or_class is None:
        return span_decorator
    else:
        return span_decorator(_func_or_class)

In the code above, we chose not to create a new separate decorator for classes but instead to overload the same decorator we use for functions. To accomplish that, we add a check testing whether the passed parameter is a function or a class (line 14,19), and apply the correct logic accordingly.

For classes, we iterate over the class functions and inject the decorator, returning the class object unchanged.

For functions, we follow the same logic as before to apply the wrapper.

With this code in place, we can now rewrite the original code to take advantage of the new decorator:

@instrument({"some attribute", "some_value"})
class UserService:

    def all(self, id):
        with tracer.start_as_current_span("retrieve users"):
            users = self.user_store.get_users(id)
        ...

    @instrument(span_name="custom_name")
    async def validate(self, user_name):
      
        try:
            await self.user_store.get_user(user_name)
        ...

    
    async def test(self, user_name):

        try:
            await self.user_store.get_user(user_name)

*Almost* Working

Notice we still have one small gotcha. The validate function in the example above will have both decorators applied to it. As a result, two spans will be created instead of one, which is not the behavior we were after.

To handle this situation, the decorator code must have some ‘memory’ of whether a function has already been decorated. There are several ways to go about it, in our implementation, we chose to save that information in the function meta-data. We, therefore, check before decorating the function whether it has already been decorated:

def span_decorator(func_or_class):

      if inspect.isclass(func_or_class):
          return decorate_class(func_or_class)

      # Check if already decorated (happens if both class and function
      # decorated). If so, we keep the function decorator settings only
      undecorated_func = getattr(func_or_class, '__tracing_unwrapped__', None)
      if undecorated_func:
          # We have already decorated this function, override
          return func_or_class

      setattr(func_or_class, '__tracing_unwrapped__', func_or_class)

If the function has already been decorated, we return it as is.

What Else Would You Add?

You can find the full source code in the Digma OpenTelemetry repository. Let me know if this is useful in your projects! Also, if there are any features you feel would be useful, feel free to reach out or open an issue or a PR on GitHub.

History

• 4^th July, 2023: Initial version