diff --git a/ddtrace/tracer.py b/ddtrace/tracer.py index 10811186bcc..3f44033fa9f 100644 --- a/ddtrace/tracer.py +++ b/ddtrace/tracer.py @@ -1,4 +1,4 @@ - +import functools import logging import threading @@ -63,6 +63,53 @@ def configure(self, enabled=None, hostname=None, port=None, sampler=None): if sampler is not None: self.sampler = sampler + def wrap(self, name=None, service=None, resource=None, span_type=None): + """A decorator used to trace an entire function. + + :param str name: the name of the operation being traced. If not set, + defaults to the fully qualified function name. + :param str service: the name of the service being traced. If not set, + it will inherit the service from it's parent. + :param str resource: an optional name of the resource being tracked. + :param str span_type: an optional operation type. + + >>> @tracer.wrap('my.wrapped.function', service='my.service') + def run(): + return 'run' + >>> @tracer.wrap() # name will default to 'execute' if unset + def execute(): + return 'executed' + + You can access the parent span using `tracer.current_span()` to set + tags: + + >>> @tracer.wrap() + def execute(): + span = tracer.current_span() + span.set_tag('a', 'b') + + You can also create more spans within a traced function. These spans + will be children of the decorator's span: + + >>> @tracer.wrap('parent') + def parent_function(): + with tracer.trace('child'): + pass + """ + + def wrap_decorator(func): + if name is None: + span_name = '{}.{}'.format(func.__module__, func.__name__) + else: + span_name = name + + @functools.wraps(func) + def func_wrapper(*args, **kwargs): + with self.trace(span_name, service=service, resource=resource, span_type=span_type): + func(*args, **kwargs) + return func_wrapper + return wrap_decorator + def trace(self, name, service=None, resource=None, span_type=None): """Return a span that will trace an operation called `name`. diff --git a/tests/test_tracer.py b/tests/test_tracer.py index 0ae596fb6b7..5c0db2afadc 100644 --- a/tests/test_tracer.py +++ b/tests/test_tracer.py @@ -4,7 +4,7 @@ import time -from nose.tools import eq_ +from nose.tools import assert_raises, eq_ from ddtrace.tracer import Tracer @@ -77,6 +77,108 @@ def _make_cake(): for s in spans: assert s.trace_id != make.trace_id +def test_tracer_wrap(): + writer = DummyWriter() + tracer = Tracer() + tracer.writer = writer + + @tracer.wrap('decorated_function', service='s', resource='r', + span_type='t') + def f(tag_name, tag_value): + # make sure we can still set tags + span = tracer.current_span() + span.set_tag(tag_name, tag_value) + f('a', 'b') + + spans = writer.pop() + eq_(len(spans), 1) + s = spans[0] + eq_(s.name, 'decorated_function') + eq_(s.service, 's') + eq_(s.resource, 'r') + eq_(s.span_type, 't') + eq_(s.to_dict()['meta']['a'], 'b') + +def test_tracer_wrap_default_name(): + writer = DummyWriter() + tracer = Tracer() + tracer.writer = writer + + @tracer.wrap() + def f(): + pass + f() + + eq_(writer.spans[0].name, 'tests.test_tracer.f') + +def test_tracer_wrap_exception(): + writer = DummyWriter() + tracer = Tracer() + tracer.writer = writer + + @tracer.wrap() + def f(): + raise Exception('bim') + + assert_raises(Exception, f) + + eq_(len(writer.spans), 1) + eq_(writer.spans[0].error, 1) + +def test_tracer_wrap_multiple_calls(): + # Make sure that we create a new span each time the function is called + writer = DummyWriter() + tracer = Tracer() + tracer.writer = writer + + @tracer.wrap() + def f(): + pass + f() + f() + + spans = writer.pop() + eq_(len(spans), 2) + assert spans[0].span_id != spans[1].span_id + +def test_tracer_wrap_span_nesting(): + # Make sure that nested spans have the correct parents + writer = DummyWriter() + tracer = Tracer() + tracer.writer = writer + + @tracer.wrap('inner') + def inner(): + pass + @tracer.wrap('outer') + def outer(): + with tracer.trace('mid'): + inner() + outer() + + spans = writer.pop() + eq_(len(spans), 3) + + # sift through the list so we're not dependent on span ordering within the + # writer + for span in spans: + if span.name == 'outer': + outer_span = span + elif span.name == 'mid': + mid_span = span + elif span.name == 'inner': + inner_span = span + else: + assert False, 'unknown span found' # should never get here + + assert outer_span + assert mid_span + assert inner_span + + eq_(outer_span.parent_id, None) + eq_(mid_span.parent_id, outer_span.span_id) + eq_(inner_span.parent_id, mid_span.span_id) + def test_tracer_disabled(): # add some dummy tracing code. writer = DummyWriter()