def test_derived_ctx_is_returned_for_failure(self):
"""Ensure returned context is derived from the given context."""
old_ctx = Context({"k2": "v2"})
new_ctx = self.get_propagator().extract({}, old_ctx)
self.assertNotIn("current-span", new_ctx)
for key, value in old_ctx.items(): # pylint:disable=no-member
self.assertIn(key, new_ctx)
# pylint:disable=unsubscriptable-object
self.assertEqual(new_ctx[key], value)
def test_extract_none_context(self):
"""Given no trace ID, do not modify context"""
old_ctx = None
carrier = {}
new_ctx = self.get_propagator().extract(carrier, old_ctx)
self.assertDictEqual(Context(), new_ctx)
def on_end(self, span: Span) -> None:
with Context.use(suppress_instrumentation=True):
try:
self.span_exporter.export((span, ))
# pylint: disable=broad-except
except Exception:
logger.exception("Exception while exporting Span.")
def extract(
self,
carrier: CarrierT,
context: typing.Optional[Context] = None,
getter: Getter = default_getter,
) -> Context:
return Context()
def __init__(self, name: str = "") -> None:
if name:
slot_name = "DistributedContext.{}".format(name)
else:
slot_name = "DistributedContext"
self._current_context = Context.register_slot(slot_name)
def extract(
self,
carrier: CarrierT,
context: typing.Optional[Context] = None,
getter: Getter = default_getter,
) -> Context:
if context is None:
context = Context()
header = getter.get(carrier, self.TRACE_ID_KEY)
if not header:
return context
context = self._extract_baggage(getter, carrier, context)
trace_id, span_id, flags = _parse_trace_id_header(header)
if (trace_id == trace.INVALID_TRACE_ID
or span_id == trace.INVALID_SPAN_ID):
return context
span = trace.NonRecordingSpan(
trace.SpanContext(
trace_id=trace_id,
span_id=span_id,
is_remote=True,
trace_flags=trace.TraceFlags(flags & trace.TraceFlags.SAMPLED),
))
return trace.set_span_in_context(span, context)
def test_extract_missing_parent_id_to_explicit_ctx(self):
"""If a parent id is missing, populate an invalid trace id."""
orig_ctx = Context({"k1": "v1"})
carrier = {FORMAT.TRACE_ID_KEY: self.serialized_trace_id}
ctx = FORMAT.extract(carrier, orig_ctx)
self.assertDictEqual(orig_ctx, ctx)
def test_derived_ctx_is_returned_for_success(self):
"""Ensure returned context is derived from the given context."""
old_ctx = Context({"k1": "v1"})
new_ctx = FORMAT.extract(
{
FORMAT.TRACE_ID_KEY: self.serialized_trace_id,
FORMAT.SPAN_ID_KEY: self.serialized_span_id,
FORMAT.FLAGS_KEY: "1",
},
old_ctx,
)
self.assertIn("current-span", new_ctx)
for key, value in old_ctx.items(): # pylint:disable=no-member
self.assertIn(key, new_ctx)
# pylint:disable=unsubscriptable-object
self.assertEqual(new_ctx[key], value)
def export(self) -> None:
"""Exports at most max_export_batch_size spans."""
idx = 0
notify_flush = False
# currently only a single thread acts as consumer, so queue.pop() will
# not raise an exception
while idx < self.max_export_batch_size and self.queue:
span = self.queue.pop()
if span is self._FLUSH_TOKEN_SPAN:
notify_flush = True
else:
self.spans_list[idx] = span
idx += 1
with Context.use(suppress_instrumentation=True):
try:
# Ignore type b/c the Optional[None]+slicing is too "clever"
# for mypy
self.span_exporter.export(
self.spans_list[:idx]) # type: ignore
# pylint: disable=broad-except
except Exception:
logger.exception("Exception while exporting Span batch.")
if notify_flush:
with self.flush_condition:
self.flush_condition.notify()
# clean up list
for index in range(idx):
self.spans_list[index] = None
def test_extract_missing_span_id_to_implicit_ctx(self):
carrier = {
FORMAT.TRACE_ID_KEY: self.serialized_trace_id,
FORMAT.FLAGS_KEY: "1",
}
new_ctx = FORMAT.extract(carrier)
self.assertDictEqual(Context(), new_ctx)
def test_extract_invalid_single_header_to_explicit_ctx(self):
"""Given unparsable header, do not modify context"""
old_ctx = Context({"k1": "v1"})
carrier = {FORMAT.SINGLE_HEADER_KEY: "0-1-2-3-4-5-6-7"}
new_ctx = FORMAT.extract(carrier, old_ctx)
self.assertDictEqual(new_ctx, old_ctx)
def with_current_context(cls, func):
# type: (Callable) -> Callable
"""Passes the current spans to the new context the function will be run in.
:param func: The function that will be run in the new context
:return: The target the pass in instead of the function
"""
return Context.with_current_context(func)
def test_extract_empty_carrier_to_explicit_ctx(self):
"""Given no headers at all, do not modify context"""
old_ctx = Context({"k1": "v1"})
carrier = {}
new_ctx = self.get_propagator().extract(carrier, old_ctx)
self.assertDictEqual(new_ctx, old_ctx)
def extract(
self,
carrier: CarrierT,
context: Optional[Context] = None,
getter: Getter = default_getter,
) -> Context:
if context is None:
context = Context()
traceid = _extract_identifier(
getter.get(carrier, OT_TRACE_ID_HEADER),
_valid_extract_traceid,
INVALID_TRACE_ID,
)
spanid = _extract_identifier(
getter.get(carrier, OT_SPAN_ID_HEADER),
_valid_extract_spanid,
INVALID_SPAN_ID,
)
sampled = _extract_first_element(
getter.get(carrier, OT_SAMPLED_HEADER)
)
if sampled == "true":
traceflags = TraceFlags.SAMPLED
else:
traceflags = TraceFlags.DEFAULT
if traceid != INVALID_TRACE_ID and spanid != INVALID_SPAN_ID:
context = set_span_in_context(
NonRecordingSpan(
SpanContext(
trace_id=traceid,
span_id=spanid,
is_remote=True,
trace_flags=TraceFlags(traceflags),
)
),
context,
)
baggage = get_all(context) or {}
for key in getter.keys(carrier):
if not key.startswith(OT_BAGGAGE_PREFIX):
continue
baggage[
key[len(OT_BAGGAGE_PREFIX) :]
] = _extract_first_element(getter.get(carrier, key))
for key, value in baggage.items():
context = set_baggage(key, value, context)
return context
def test_extract_missing_span_id_to_implicit_ctx(self):
propagator = self.get_propagator()
carrier = {
propagator.TRACE_ID_KEY: self.serialized_trace_id,
propagator.FLAGS_KEY: "1",
}
new_ctx = propagator.extract(carrier)
self.assertDictEqual(Context(), new_ctx)
def test_extract_malformed_headers_to_implicit_ctx(self):
malformed_trace_id_key = FORMAT.TRACE_ID_KEY + "-x"
malformed_parent_id_key = FORMAT.PARENT_ID_KEY + "-x"
context = FORMAT.extract({
malformed_trace_id_key:
self.serialized_trace_id,
malformed_parent_id_key:
self.serialized_parent_id,
})
self.assertDictEqual(Context(), context)
def test_extract_missing_span_id_to_explicit_ctx(self):
"""Given no span ID, do not modify context"""
old_ctx = Context({"k1": "v1"})
carrier = {
FORMAT.TRACE_ID_KEY: self.serialized_trace_id,
FORMAT.FLAGS_KEY: "1",
}
new_ctx = FORMAT.extract(carrier, old_ctx)
self.assertDictEqual(new_ctx, old_ctx)
def __init__(
self,
name: str = "",
sampler: sampling.Sampler = trace_api.sampling.ALWAYS_ON,
) -> None:
slot_name = "current_span"
if name:
slot_name = "{}.current_span".format(name)
self._current_span_slot = Context.register_slot(slot_name)
self._active_span_processor = MultiSpanProcessor()
self.sampler = sampler
def test_extract_missing_trace_id_to_explicit_ctx(self):
"""Given no trace ID, do not modify context"""
old_ctx = Context({"k1": "v1"})
propagator = self.get_propagator()
carrier = {
propagator.SPAN_ID_KEY: self.serialized_span_id,
propagator.FLAGS_KEY: "1",
}
new_ctx = propagator.extract(carrier, old_ctx)
self.assertDictEqual(new_ctx, old_ctx)
def test_span_processor_accepts_parent_context(self):
span_processor = mock.Mock(
wraps=datadog.DatadogExportSpanProcessor(self.exporter))
tracer_provider = trace.TracerProvider()
tracer_provider.add_span_processor(span_processor)
tracer = tracer_provider.get_tracer(__name__)
context = Context()
span = tracer.start_span("foo", context=context)
span_processor.on_start.assert_called_once_with(span,
parent_context=context)
def test_start_accepts_context(self):
# pylint: disable=no-self-use
span_processor = mock.Mock(spec=trace.SpanProcessor)
span = trace._Span(
"name",
mock.Mock(spec=trace_api.SpanContext),
span_processor=span_processor,
)
context = Context()
span.start(parent_context=context)
span_processor.on_start.assert_called_once_with(span,
parent_context=context)
def __init__(
self,
sampler: sampling.Sampler = trace_api.sampling.ALWAYS_ON,
shutdown_on_exit: bool = True,
):
# TODO: How should multiple TracerSources behave? Should they get their own contexts?
# This could be done by adding `str(id(self))` to the slot name.
self._current_span_slot = Context.register_slot("current_span")
self._active_span_processor = MultiSpanProcessor()
self.sampler = sampler
self._atexit_handler = None
if shutdown_on_exit:
self._atexit_handler = atexit.register(self.shutdown)
def test_on_start_accepts_context(self):
# pylint: disable=no-self-use
tracer_provider = trace.TracerProvider()
tracer = tracer_provider.get_tracer(__name__)
exporter = MySpanExporter([])
span_processor = mock.Mock(wraps=export.SimpleSpanProcessor(exporter))
tracer_provider.add_span_processor(span_processor)
context = Context()
span = tracer.start_span("foo", context=context)
span_processor.on_start.assert_called_once_with(span,
parent_context=context)
def test_extract_invalid_uber_trace_id_header_to_implicit_ctx(self):
trace_id_headers = [
"000000000000000000000000deadbeef:00000000deadbef0:00",
"00000000000000000000000000000000:00000000deadbef0:00:00",
"000000000000000000000000deadbeef:0000000000000000:00:00",
"000000000000000000000000deadbeef:0000000000000000:00:xyz",
]
for trace_id_header in trace_id_headers:
with self.subTest(trace_id_header=trace_id_header):
carrier = {"uber-trace-id": trace_id_header}
ctx = FORMAT.extract(carrier)
self.assertDictEqual(Context(), ctx)
def test_extract_malformed_headers_to_explicit_ctx(self):
"""Test with no Datadog headers"""
orig_ctx = Context({"k1": "v1"})
malformed_trace_id_key = FORMAT.TRACE_ID_KEY + "-x"
malformed_parent_id_key = FORMAT.PARENT_ID_KEY + "-x"
context = FORMAT.extract(
{
malformed_trace_id_key: self.serialized_trace_id,
malformed_parent_id_key: self.serialized_parent_id,
},
orig_ctx,
)
self.assertDictEqual(orig_ctx, context)
def test_extract_invalid_uber_trace_id_header_to_explicit_ctx(self):
trace_id_headers = [
"000000000000000000000000deadbeef:00000000deadbef0:00",
"00000000000000000000000000000000:00000000deadbef0:00:00",
"000000000000000000000000deadbeef:0000000000000000:00:00",
"000000000000000000000000deadbeef:0000000000000000:00:xyz",
]
for trace_id_header in trace_id_headers:
with self.subTest(trace_id_header=trace_id_header):
carrier = {"uber-trace-id": trace_id_header}
orig_ctx = Context({"k1": "v1"})
ctx = FORMAT.extract(carrier, orig_ctx)
self.assertDictEqual(orig_ctx, ctx)
def extract(
self,
carrier: CarrierT,
context: typing.Optional[Context] = None,
getter: Getter = default_getter,
) -> Context:
if context is None:
context = Context()
trace_header_list = getter.get(carrier, TRACE_HEADER_KEY)
if not trace_header_list or len(trace_header_list) != 1:
return context
trace_header = trace_header_list[0]
if not trace_header:
return context
try:
(
trace_id,
span_id,
sampled,
) = AwsXRayPropagator._extract_span_properties(trace_header)
except AwsParseTraceHeaderError as err:
_logger.debug(err.message)
return context
options = 0
if sampled:
options |= trace.TraceFlags.SAMPLED
span_context = trace.SpanContext(
trace_id=trace_id,
span_id=span_id,
is_remote=True,
trace_flags=trace.TraceFlags(options),
trace_state=trace.TraceState(),
)
if not span_context.is_valid:
_logger.debug(
"Invalid Span Extracted. Inserting INVALID span into provided context."
)
return context
return trace.set_span_in_context(
trace.NonRecordingSpan(span_context), context=context
)
def __init__(
self,
name: str = "",
sampler: sampling.Sampler = trace_api.sampling.ALWAYS_ON,
shutdown_on_exit: bool = True,
) -> None:
slot_name = "current_span"
if name:
slot_name = "{}.current_span".format(name)
self._current_span_slot = Context.register_slot(slot_name)
self._active_span_processor = MultiSpanProcessor()
self.sampler = sampler
self._atexit_handler = None
if shutdown_on_exit:
self._atexit_handler = atexit.register(self.shutdown)
def test_on_start(self):
multi_processor = self.create_multi_span_processor()
mocks = [mock.Mock(spec=trace.SpanProcessor) for _ in range(0, 5)]
for mock_processor in mocks:
multi_processor.add_span_processor(mock_processor)
span = self.create_default_span()
context = Context()
multi_processor.on_start(span, parent_context=context)
for mock_processor in mocks:
mock_processor.on_start.assert_called_once_with(
span, parent_context=context)
multi_processor.shutdown()
def test_extract_invalid_to_implicit_ctx(self):
trace_headers = [
"Root=1-12345678-abcdefghijklmnopqrstuvwx;Parent=53995c3f42cd8ad8;Sampled=0", # invalid trace id
"Root=1-8a3c60f7-d188f8fa79d48a391a778fa600;Parent=53995c3f42cd8ad8;Sampled=0", # invalid size trace id
"Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=abcdefghijklmnop;Sampled=0", # invalid span id
"Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad800;Sampled=0" # invalid size span id
"Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=", # no sampled flag
"Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=011", # invalid size sampled
"Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=a", # non numeric sampled flag
]
for trace_header in trace_headers:
with self.subTest(trace_header=trace_header):
ctx = AwsXRayPropagatorTest.XRAY_PROPAGATOR.extract(
CaseInsensitiveDict({TRACE_HEADER_KEY: trace_header}), )
self.assertDictEqual(Context(), ctx)