From 8290efaac1cafd7bf8c3b4f13fb19c602487f5c6 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Fri, 28 Mar 2025 10:56:54 +0100 Subject: [PATCH 1/9] Always instrument methods --- .../datadog_checks/base/utils/tracing.py | 88 ++++++++++++------- 1 file changed, 58 insertions(+), 30 deletions(-) diff --git a/datadog_checks_base/datadog_checks/base/utils/tracing.py b/datadog_checks_base/datadog_checks/base/utils/tracing.py index 5899b0443cb51..4f82b8eb50abd 100644 --- a/datadog_checks_base/datadog_checks/base/utils/tracing.py +++ b/datadog_checks_base/datadog_checks/base/utils/tracing.py @@ -69,7 +69,7 @@ def wrapper(*args, **kwargs): def traced_warning(f, tracer): """ - Traces the AgentCheck.warning method + Traces the AgentCheck.warning method. The span is always an error span, including the current stack trace. The error message is set to the warning message. """ @@ -102,6 +102,31 @@ def wrapper(self, warning_message, *args, **kwargs): return f +def generate_tracer_context(f, tracer): + """ + Generate a tracer context for the given function with configurable sampling rate. + If not set or invalid, defaults to 0 (no sampling). + The tracer context is only set at entry point functions so we can attach a trace root to the span. + """ + apm_tracing_disabled = False + try: + integration_tracing, integration_tracing_exhaustive = tracing_enabled() + if integration_tracing or integration_tracing_exhaustive: + apm_tracing_disabled = True + except (ValueError, TypeError): + pass + + try: + # Update the tracer configuration to make sure we trace only if we really need to + tracer.configure( + context_provider=None, # TODO: add a context provider to the tracer, so we can add a trace root to the span + appsec_enabled=False, + apm_tracing_disabled=apm_tracing_disabled, + ) + except Exception: + pass + + def tracing_enabled(): """ :return: (integration_tracing, integration_tracing_exhaustive) @@ -116,44 +141,47 @@ def tracing_enabled(): return integration_tracing, integration_tracing_exhaustive - def traced_class(cls): - integration_tracing, integration_tracing_exhaustive = tracing_enabled() - if integration_tracing: - try: - integration_tracing_exhaustive = is_affirmative(datadog_agent.get_config('integration_tracing_exhaustive')) + """ + Decorator that adds tracing to all methods of a class. + Only traces specific methods by default, unless exhaustive tracing is enabled. + """ + _, integration_tracing_exhaustive = tracing_enabled() - from ddtrace import patch_all, tracer + try: + from ddtrace import patch_all, tracer - patch_all() + patch_all() - def decorate(cls): - for attr in cls.__dict__: - attribute = getattr(cls, attr) + def decorate(cls): + for attr in cls.__dict__: + attribute = getattr(cls, attr) - if not callable(attribute) or inspect.isclass(attribute): - continue + if not callable(attribute) or inspect.isclass(attribute): + continue - # Ignoring staticmethod and classmethod because they don't need cls in args - # also ignore nested classes - if isinstance(cls.__dict__[attr], staticmethod) or isinstance(cls.__dict__[attr], classmethod): - continue + # Ignoring staticmethod and classmethod because they don't need cls in args + # also ignore nested classes + if isinstance(cls.__dict__[attr], staticmethod) or isinstance(cls.__dict__[attr], classmethod): + continue - # Get rid of SnmpCheck._thread_factory and related - if getattr(attribute, '__module__', 'threading') in EXCLUDED_MODULES: - continue + # Get rid of SnmpCheck._thread_factory and related + if getattr(attribute, '__module__', 'threading') in EXCLUDED_MODULES: + continue - if not integration_tracing_exhaustive and attr not in AGENT_CHECK_DEFAULT_TRACED_METHODS: - continue + if not integration_tracing_exhaustive and attr not in AGENT_CHECK_DEFAULT_TRACED_METHODS: + continue - if attr == 'warning': - setattr(cls, attr, traced_warning(attribute, tracer)) - else: - setattr(cls, attr, tracing_method(attribute, tracer)) - return cls + if attr == 'warning': + setattr(cls, attr, traced_warning(attribute, tracer)) + else: + if attr == 'run' or attr == 'check': + generate_tracer_context(attribute, tracer) + setattr(cls, attr, tracing_method(attribute, tracer)) + return cls - return decorate(cls) - except Exception: - pass + return decorate(cls) + except Exception: + pass return cls From 98f5f7ce2f593c46c3a6152cd1774e64cf01dff4 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Fri, 28 Mar 2025 11:08:56 +0100 Subject: [PATCH 2/9] Rename --- datadog_checks_base/datadog_checks/base/utils/tracing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datadog_checks_base/datadog_checks/base/utils/tracing.py b/datadog_checks_base/datadog_checks/base/utils/tracing.py index 4f82b8eb50abd..21c70e3c14a35 100644 --- a/datadog_checks_base/datadog_checks/base/utils/tracing.py +++ b/datadog_checks_base/datadog_checks/base/utils/tracing.py @@ -102,7 +102,7 @@ def wrapper(self, warning_message, *args, **kwargs): return f -def generate_tracer_context(f, tracer): +def configure_tracer(tracer): """ Generate a tracer context for the given function with configurable sampling rate. If not set or invalid, defaults to 0 (no sampling). @@ -176,7 +176,7 @@ def decorate(cls): setattr(cls, attr, traced_warning(attribute, tracer)) else: if attr == 'run' or attr == 'check': - generate_tracer_context(attribute, tracer) + configure_tracer(tracer) setattr(cls, attr, tracing_method(attribute, tracer)) return cls From cb49d091d61ba092851a72db268616d7fc663b10 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Fri, 28 Mar 2025 11:14:05 +0100 Subject: [PATCH 3/9] Move inside --- .../datadog_checks/base/utils/tracing.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/datadog_checks_base/datadog_checks/base/utils/tracing.py b/datadog_checks_base/datadog_checks/base/utils/tracing.py index 21c70e3c14a35..aa393c4044d6b 100644 --- a/datadog_checks_base/datadog_checks/base/utils/tracing.py +++ b/datadog_checks_base/datadog_checks/base/utils/tracing.py @@ -45,11 +45,14 @@ def _get_integration_name(function_name, self, *args, **kwargs): return integration_name if integration_name else "UNKNOWN_INTEGRATION" -def tracing_method(f, tracer): +def tracing_method(f, tracer, is_entry_point): if inspect.signature(f).parameters.get('self'): @functools.wraps(f) def wrapper(self, *args, **kwargs): + if is_entry_point: + configure_tracer(tracer) + integration_name = _get_integration_name(f.__name__, self, *args, **kwargs) with tracer.trace(f.__name__, service=INTEGRATION_TRACING_SERVICE_NAME, resource=integration_name) as span: span.set_tag('_dd.origin', INTEGRATION_TRACING_SERVICE_NAME) @@ -59,6 +62,9 @@ def wrapper(self, *args, **kwargs): @functools.wraps(f) def wrapper(*args, **kwargs): + if is_entry_point: + configure_tracer(tracer) + integration_name = _get_integration_name(f.__name__, None, *args, **kwargs) with tracer.trace(f.__name__, service=INTEGRATION_TRACING_SERVICE_NAME, resource=integration_name) as span: span.set_tag('_dd.origin', INTEGRATION_TRACING_SERVICE_NAME) @@ -172,12 +178,12 @@ def decorate(cls): if not integration_tracing_exhaustive and attr not in AGENT_CHECK_DEFAULT_TRACED_METHODS: continue + is_entry_point = attr == 'run' or attr == 'check' + if attr == 'warning': setattr(cls, attr, traced_warning(attribute, tracer)) else: - if attr == 'run' or attr == 'check': - configure_tracer(tracer) - setattr(cls, attr, tracing_method(attribute, tracer)) + setattr(cls, attr, tracing_method(attribute, tracer, is_entry_point)) return cls return decorate(cls) From bf84bf1987aed78bdbe784dc70fb14392b597831 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Fri, 28 Mar 2025 11:38:00 +0100 Subject: [PATCH 4/9] Lint --- datadog_checks_base/datadog_checks/base/utils/tracing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/datadog_checks_base/datadog_checks/base/utils/tracing.py b/datadog_checks_base/datadog_checks/base/utils/tracing.py index aa393c4044d6b..31ff4173fa10e 100644 --- a/datadog_checks_base/datadog_checks/base/utils/tracing.py +++ b/datadog_checks_base/datadog_checks/base/utils/tracing.py @@ -147,6 +147,7 @@ def tracing_enabled(): return integration_tracing, integration_tracing_exhaustive + def traced_class(cls): """ Decorator that adds tracing to all methods of a class. From 49f035fbc4e680575989e7f90eabc0933d3ae051 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Thu, 3 Apr 2025 10:45:04 +0200 Subject: [PATCH 5/9] Configure tracing --- .../datadog_checks/base/utils/tracing.py | 32 ++++++++++++------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/datadog_checks_base/datadog_checks/base/utils/tracing.py b/datadog_checks_base/datadog_checks/base/utils/tracing.py index 31ff4173fa10e..ea09b7ede9f92 100644 --- a/datadog_checks_base/datadog_checks/base/utils/tracing.py +++ b/datadog_checks_base/datadog_checks/base/utils/tracing.py @@ -50,10 +50,10 @@ def tracing_method(f, tracer, is_entry_point): @functools.wraps(f) def wrapper(self, *args, **kwargs): + integration_name = _get_integration_name(f.__name__, self, *args, **kwargs) if is_entry_point: - configure_tracer(tracer) + configure_tracer(tracer, self) - integration_name = _get_integration_name(f.__name__, self, *args, **kwargs) with tracer.trace(f.__name__, service=INTEGRATION_TRACING_SERVICE_NAME, resource=integration_name) as span: span.set_tag('_dd.origin', INTEGRATION_TRACING_SERVICE_NAME) return f(self, *args, **kwargs) @@ -62,9 +62,6 @@ def wrapper(self, *args, **kwargs): @functools.wraps(f) def wrapper(*args, **kwargs): - if is_entry_point: - configure_tracer(tracer) - integration_name = _get_integration_name(f.__name__, None, *args, **kwargs) with tracer.trace(f.__name__, service=INTEGRATION_TRACING_SERVICE_NAME, resource=integration_name) as span: span.set_tag('_dd.origin', INTEGRATION_TRACING_SERVICE_NAME) @@ -108,27 +105,40 @@ def wrapper(self, warning_message, *args, **kwargs): return f -def configure_tracer(tracer): +def configure_tracer(tracer, self_check): """ Generate a tracer context for the given function with configurable sampling rate. If not set or invalid, defaults to 0 (no sampling). The tracer context is only set at entry point functions so we can attach a trace root to the span. """ - apm_tracing_disabled = False + apm_tracing_enabled = False + context_provider = None try: integration_tracing, integration_tracing_exhaustive = tracing_enabled() if integration_tracing or integration_tracing_exhaustive: - apm_tracing_disabled = True - except (ValueError, TypeError): + apm_tracing_enabled = True + + # If the check has a dd_trace_id and dd_parent_id, we can use it to create a trace root + dd_trace_id = self_check.instance.get("dd_trace_id") + dd_parent_id = self_check.instance.get("dd_parent_span_id") + if dd_trace_id and dd_parent_id: + from ddtrace.tracer.context import Context + apm_tracing_enabled = True + context_provider = Context( + trace_id=dd_trace_id, + span_id=dd_parent_id, + ) + except (ValueError, TypeError, AttributeError, ImportError): pass try: # Update the tracer configuration to make sure we trace only if we really need to tracer.configure( - context_provider=None, # TODO: add a context provider to the tracer, so we can add a trace root to the span appsec_enabled=False, - apm_tracing_disabled=apm_tracing_disabled, + enabled=apm_tracing_enabled, ) + if context_provider: + tracer.context_provider.activate(context_provider) except Exception: pass From e1a7c54d8ddda9ff99173d7f7a7ccfe6c083352e Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Thu, 3 Apr 2025 11:04:18 +0200 Subject: [PATCH 6/9] Lint --- datadog_checks_base/datadog_checks/base/utils/tracing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/datadog_checks_base/datadog_checks/base/utils/tracing.py b/datadog_checks_base/datadog_checks/base/utils/tracing.py index ea09b7ede9f92..b15a15664651d 100644 --- a/datadog_checks_base/datadog_checks/base/utils/tracing.py +++ b/datadog_checks_base/datadog_checks/base/utils/tracing.py @@ -123,6 +123,7 @@ def configure_tracer(tracer, self_check): dd_parent_id = self_check.instance.get("dd_parent_span_id") if dd_trace_id and dd_parent_id: from ddtrace.tracer.context import Context + apm_tracing_enabled = True context_provider = Context( trace_id=dd_trace_id, From 6d5d01b28aa1a49e3a44a7a5c409b70a99d4f89e Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Fri, 4 Apr 2025 15:47:05 +0200 Subject: [PATCH 7/9] Fix tests --- .../datadog_checks/base/utils/tracing.py | 20 +++-- .../tests/base/utils/test_tracing.py | 73 +++++++++++++------ 2 files changed, 66 insertions(+), 27 deletions(-) diff --git a/datadog_checks_base/datadog_checks/base/utils/tracing.py b/datadog_checks_base/datadog_checks/base/utils/tracing.py index b15a15664651d..4079449270672 100644 --- a/datadog_checks_base/datadog_checks/base/utils/tracing.py +++ b/datadog_checks_base/datadog_checks/base/utils/tracing.py @@ -119,10 +119,17 @@ def configure_tracer(tracer, self_check): apm_tracing_enabled = True # If the check has a dd_trace_id and dd_parent_id, we can use it to create a trace root - dd_trace_id = self_check.instance.get("dd_trace_id") - dd_parent_id = self_check.instance.get("dd_parent_span_id") + dd_parent_id = None + dd_trace_id = None + if hasattr(self_check, "instance") and self_check.instance: + dd_trace_id = self_check.instance.get("dd_trace_id", None) + dd_parent_id = self_check.instance.get("dd_parent_span_id", None) + elif hasattr(self_check, "instances") and self_check.instances and len(self_check.instances) > 0: + dd_trace_id = self_check.instances[0].get("dd_trace_id", None) + dd_parent_id = self_check.instances[0].get("dd_parent_span_id", None) + if dd_trace_id and dd_parent_id: - from ddtrace.tracer.context import Context + from ddtrace.context import Context apm_tracing_enabled = True context_provider = Context( @@ -130,7 +137,7 @@ def configure_tracer(tracer, self_check): span_id=dd_parent_id, ) except (ValueError, TypeError, AttributeError, ImportError): - pass + raise try: # Update the tracer configuration to make sure we trace only if we really need to @@ -138,7 +145,10 @@ def configure_tracer(tracer, self_check): appsec_enabled=False, enabled=apm_tracing_enabled, ) - if context_provider: + + # If the current trace context is not set or is set to an empty trace_id, activate the context provider + current_context = tracer.current_trace_context() + if (current_context is None or (current_context is not None and len(current_context.trace_id) == 0)) and context_provider: tracer.context_provider.activate(context_provider) except Exception: pass diff --git a/datadog_checks_base/tests/base/utils/test_tracing.py b/datadog_checks_base/tests/base/utils/test_tracing.py index 2ee9b88490972..e3f47665ec4b4 100644 --- a/datadog_checks_base/tests/base/utils/test_tracing.py +++ b/datadog_checks_base/tests/base/utils/test_tracing.py @@ -79,37 +79,66 @@ def traced_mock_classes(): 'integration_tracing_exhaustive', [pytest.param(False, id="exhaustive_false"), pytest.param(True, id="exhaustive_true")], ) -def test_traced_class(integration_tracing, integration_tracing_exhaustive, datadog_agent): +@pytest.mark.parametrize( + 'dd_trace_id', [pytest.param(None, id="no_trace_id"), pytest.param("123456789", id="with_trace_id")] +) +@pytest.mark.parametrize( + 'dd_parent_id', [pytest.param(None, id="no_parent_id"), pytest.param("987654321", id="with_parent_id")] +) +def test_traced_class(integration_tracing, integration_tracing_exhaustive, dd_trace_id, dd_parent_id, datadog_agent): def _get_config(key): return { 'integration_tracing': str(integration_tracing).lower(), 'integration_tracing_exhaustive': str(integration_tracing_exhaustive).lower(), }.get(key, None) + instance = {} + if dd_trace_id is not None: + instance['dd_trace_id'] = dd_trace_id + if dd_parent_id is not None: + instance['dd_parent_span_id'] = dd_parent_id + with mock.patch.object(datadog_agent, 'get_config', _get_config), mock.patch('ddtrace.tracer') as tracer: + # Track the last activated context + def mock_activate(context): + def mock_current_trace_context(): + return context + tracer.current_trace_context.side_effect = mock_current_trace_context + + tracer.context_provider.activate.side_effect = mock_activate + with traced_mock_classes(): - check = DummyCheck('dummy', {}, [{}]) + check = DummyCheck('dummy', {}, [instance]) check.run() - if integration_tracing: - called_services = {c.kwargs['service'] for c in tracer.trace.mock_calls if 'service' in c.kwargs} - called_methods = {c.args[0] for c in tracer.trace.mock_calls if c.args} - - assert called_services == {INTEGRATION_TRACING_SERVICE_NAME} - for m in AGENT_CHECK_DEFAULT_TRACED_METHODS: + called_services = {c.kwargs['service'] for c in tracer.trace.mock_calls if 'service' in c.kwargs} + called_methods = {c.args[0] for c in tracer.trace.mock_calls if c.args} + + assert called_services == {INTEGRATION_TRACING_SERVICE_NAME} + for m in AGENT_CHECK_DEFAULT_TRACED_METHODS: + assert m in called_methods + + warning_span_tag_calls = tracer.trace().__enter__().set_tag.call_args_list + assert mock.call('_dd.origin', INTEGRATION_TRACING_SERVICE_NAME) in warning_span_tag_calls + assert mock.call(ERROR_MSG, 'whoops oh no') in warning_span_tag_calls + assert mock.call(ERROR_TYPE, 'AgentCheck.warning') in warning_span_tag_calls + + # If dd_trace_id and dd_parent_id are set, verify context provider is activated + if dd_trace_id is not None and dd_parent_id is not None: + # Assert called once + tracer.context_provider.activate.assert_called_once() + context = tracer.context_provider.activate.call_args[0][0] + assert context.trace_id == dd_trace_id + assert context.span_id == dd_parent_id + + # Check that the tracer is configured with the correct enabled value + tracing = integration_tracing or integration_tracing_exhaustive or (dd_trace_id is not None and dd_parent_id is not None) + assert tracer.configure.call_args[1]['enabled'] is tracing + + exhaustive_only_methods = {'__init__', 'dummy_method'} + if integration_tracing_exhaustive: + for m in exhaustive_only_methods: assert m in called_methods - - warning_span_tag_calls = tracer.trace().__enter__().set_tag.call_args_list - assert mock.call('_dd.origin', INTEGRATION_TRACING_SERVICE_NAME) in warning_span_tag_calls - assert mock.call(ERROR_MSG, 'whoops oh no') in warning_span_tag_calls - assert mock.call(ERROR_TYPE, 'AgentCheck.warning') in warning_span_tag_calls - - exhaustive_only_methods = {'__init__', 'dummy_method'} - if integration_tracing_exhaustive: - for m in exhaustive_only_methods: - assert m in called_methods - else: - for m in exhaustive_only_methods: - assert m not in called_methods else: - tracer.trace.assert_not_called() + for m in exhaustive_only_methods: + assert m not in called_methods From 95905ab96c96135357f9bea68a6a325d8f9334e5 Mon Sep 17 00:00:00 2001 From: Paul Coignet Date: Fri, 4 Apr 2025 15:55:37 +0200 Subject: [PATCH 8/9] Lint --- .../datadog_checks/base/utils/tracing.py | 4 +++- datadog_checks_base/tests/base/utils/test_tracing.py | 11 ++++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/datadog_checks_base/datadog_checks/base/utils/tracing.py b/datadog_checks_base/datadog_checks/base/utils/tracing.py index 4079449270672..336576fd99a92 100644 --- a/datadog_checks_base/datadog_checks/base/utils/tracing.py +++ b/datadog_checks_base/datadog_checks/base/utils/tracing.py @@ -148,7 +148,9 @@ def configure_tracer(tracer, self_check): # If the current trace context is not set or is set to an empty trace_id, activate the context provider current_context = tracer.current_trace_context() - if (current_context is None or (current_context is not None and len(current_context.trace_id) == 0)) and context_provider: + if ( + current_context is None or (current_context is not None and len(current_context.trace_id) == 0) + ) and context_provider: tracer.context_provider.activate(context_provider) except Exception: pass diff --git a/datadog_checks_base/tests/base/utils/test_tracing.py b/datadog_checks_base/tests/base/utils/test_tracing.py index e3f47665ec4b4..c4925506e3a71 100644 --- a/datadog_checks_base/tests/base/utils/test_tracing.py +++ b/datadog_checks_base/tests/base/utils/test_tracing.py @@ -103,10 +103,11 @@ def _get_config(key): def mock_activate(context): def mock_current_trace_context(): return context + tracer.current_trace_context.side_effect = mock_current_trace_context - + tracer.context_provider.activate.side_effect = mock_activate - + with traced_mock_classes(): check = DummyCheck('dummy', {}, [instance]) check.run() @@ -132,7 +133,11 @@ def mock_current_trace_context(): assert context.span_id == dd_parent_id # Check that the tracer is configured with the correct enabled value - tracing = integration_tracing or integration_tracing_exhaustive or (dd_trace_id is not None and dd_parent_id is not None) + tracing = ( + integration_tracing + or integration_tracing_exhaustive + or (dd_trace_id is not None and dd_parent_id is not None) + ) assert tracer.configure.call_args[1]['enabled'] is tracing exhaustive_only_methods = {'__init__', 'dummy_method'} From 1753580e7fc053ec9a744ff7553abcc7fc20f730 Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 8 Apr 2025 14:17:13 +0200 Subject: [PATCH 9/9] Handle exceptions without raising errors --- datadog_checks_base/datadog_checks/base/utils/tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog_checks_base/datadog_checks/base/utils/tracing.py b/datadog_checks_base/datadog_checks/base/utils/tracing.py index 336576fd99a92..71bf6b097fe11 100644 --- a/datadog_checks_base/datadog_checks/base/utils/tracing.py +++ b/datadog_checks_base/datadog_checks/base/utils/tracing.py @@ -137,7 +137,7 @@ def configure_tracer(tracer, self_check): span_id=dd_parent_id, ) except (ValueError, TypeError, AttributeError, ImportError): - raise + pass try: # Update the tracer configuration to make sure we trace only if we really need to