Skip to content

Commit

Permalink
botocore: always use x-ray for http header injection (#1741)
Browse files Browse the repository at this point in the history
Co-authored-by: Diego Hurtado <ocelotl@users.noreply.github.com>
  • Loading branch information
tsloughter and ocelotl authored May 22, 2023
1 parent 530650d commit 42e8f8f
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 34 deletions.
5 changes: 5 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- `opentelemetry-instrumentation-logging` Add `otelTraceSampled` to instrumetation-logging
([#1773](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1773))

### Changed

- `opentelemetry-instrumentation-botocore` now uses the AWS X-Ray propagator by
default
([#1741](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1741))

### Fixed

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ dependencies = [
"opentelemetry-api ~= 1.12",
"opentelemetry-instrumentation == 0.40b0.dev",
"opentelemetry-semantic-conventions == 0.40b0.dev",
"opentelemetry-propagator-aws-xray == 1.0.1",
]

[project.optional-dependencies]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -101,22 +101,14 @@ def response_hook(span, service_name, operation_name, result):
_SUPPRESS_INSTRUMENTATION_KEY,
unwrap,
)
from opentelemetry.propagate import inject
from opentelemetry.propagators.aws.aws_xray_propagator import AwsXRayPropagator
from opentelemetry.semconv.trace import SpanAttributes
from opentelemetry.trace import get_tracer
from opentelemetry.trace.span import Span

logger = logging.getLogger(__name__)


# pylint: disable=unused-argument
def _patched_endpoint_prepare_request(wrapped, instance, args, kwargs):
request = args[0]
headers = request.headers
inject(headers)
return wrapped(*args, **kwargs)


class BotocoreInstrumentor(BaseInstrumentor):
"""An instrumentor for Botocore.
Expand All @@ -127,6 +119,7 @@ def __init__(self):
super().__init__()
self.request_hook = None
self.response_hook = None
self.propagator = AwsXRayPropagator()

def instrumentation_dependencies(self) -> Collection[str]:
return _instruments
Expand All @@ -140,6 +133,10 @@ def _instrument(self, **kwargs):
self.request_hook = kwargs.get("request_hook")
self.response_hook = kwargs.get("response_hook")

propagator = kwargs.get("propagator")
if propagator is not None:
self.propagator = propagator

wrap_function_wrapper(
"botocore.client",
"BaseClient._make_api_call",
Expand All @@ -149,13 +146,26 @@ def _instrument(self, **kwargs):
wrap_function_wrapper(
"botocore.endpoint",
"Endpoint.prepare_request",
_patched_endpoint_prepare_request,
self._patched_endpoint_prepare_request,
)

def _uninstrument(self, **kwargs):
unwrap(BaseClient, "_make_api_call")
unwrap(Endpoint, "prepare_request")

# pylint: disable=unused-argument
def _patched_endpoint_prepare_request(
self, wrapped, instance, args, kwargs
):
request = args[0]
headers = request.headers

# Only the x-ray header is propagated by AWS services. Using any
# other propagator will lose the trace context.
self.propagator.inject(headers)

return wrapped(*args, **kwargs)

# pylint: disable=too-many-branches
def _patched_api_call(self, original_func, instance, args, kwargs):
if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,11 @@
from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
from opentelemetry.propagate import get_global_textmap, set_global_textmap
from opentelemetry.propagators.aws.aws_xray_propagator import TRACE_HEADER_KEY
from opentelemetry.semconv.trace import SpanAttributes
from opentelemetry.test.mock_textmap import MockTextMapPropagator
from opentelemetry.test.test_base import TestBase
from opentelemetry.trace.span import format_span_id, format_trace_id

_REQUEST_ID_REGEX_MATCH = r"[A-Z0-9]{52}"

Expand Down Expand Up @@ -225,27 +227,21 @@ def test_unpatch(self):
@mock_ec2
def test_uninstrument_does_not_inject_headers(self):
headers = {}
previous_propagator = get_global_textmap()
try:
set_global_textmap(MockTextMapPropagator())

def intercept_headers(**kwargs):
headers.update(kwargs["request"].headers)
def intercept_headers(**kwargs):
headers.update(kwargs["request"].headers)

ec2 = self._make_client("ec2")
ec2 = self._make_client("ec2")

BotocoreInstrumentor().uninstrument()
BotocoreInstrumentor().uninstrument()

ec2.meta.events.register_first(
"before-send.ec2.DescribeInstances", intercept_headers
)
with self.tracer_provider.get_tracer("test").start_span("parent"):
ec2.describe_instances()
ec2.meta.events.register_first(
"before-send.ec2.DescribeInstances", intercept_headers
)
with self.tracer_provider.get_tracer("test").start_span("parent"):
ec2.describe_instances()

self.assertNotIn(MockTextMapPropagator.TRACE_ID_KEY, headers)
self.assertNotIn(MockTextMapPropagator.SPAN_ID_KEY, headers)
finally:
set_global_textmap(previous_propagator)
self.assertNotIn(TRACE_HEADER_KEY, headers)

@mock_sqs
def test_double_patch(self):
Expand Down Expand Up @@ -306,20 +302,42 @@ def check_headers(**kwargs):
"EC2", "DescribeInstances", request_id=request_id
)

self.assertIn(MockTextMapPropagator.TRACE_ID_KEY, headers)
self.assertEqual(
str(span.get_span_context().trace_id),
headers[MockTextMapPropagator.TRACE_ID_KEY],
# only x-ray propagation is used in HTTP requests
self.assertIn(TRACE_HEADER_KEY, headers)
xray_context = headers[TRACE_HEADER_KEY]
formated_trace_id = format_trace_id(
span.get_span_context().trace_id
)
self.assertIn(MockTextMapPropagator.SPAN_ID_KEY, headers)
self.assertEqual(
str(span.get_span_context().span_id),
headers[MockTextMapPropagator.SPAN_ID_KEY],
formated_trace_id = (
formated_trace_id[:8] + "-" + formated_trace_id[8:]
)

self.assertEqual(
xray_context.lower(),
f"root=1-{formated_trace_id};parent={format_span_id(span.get_span_context().span_id)};sampled=1".lower(),
)
finally:
set_global_textmap(previous_propagator)

@mock_ec2
def test_override_xray_propagator_injects_into_request(self):
headers = {}

def check_headers(**kwargs):
nonlocal headers
headers = kwargs["request"].headers

BotocoreInstrumentor().instrument()

ec2 = self._make_client("ec2")
ec2.meta.events.register_first(
"before-send.ec2.DescribeInstances", check_headers
)
ec2.describe_instances()

self.assertNotIn(MockTextMapPropagator.TRACE_ID_KEY, headers)
self.assertNotIn(MockTextMapPropagator.SPAN_ID_KEY, headers)

@mock_xray
def test_suppress_instrumentation_xray_client(self):
xray_client = self._make_client("xray")
Expand Down

0 comments on commit 42e8f8f

Please sign in to comment.