def test_span_context_sampled_no_handlers(
zipkin_logger_mock,
generate_string_mock,
thread_local_mock,
):
zipkin_attrs = ZipkinAttrs(
trace_id='1111111111111111',
span_id='2222222222222222',
parent_span_id='3333333333333333',
flags='flags',
is_sampled=True,
)
thread_local_mock.zipkin_attrs = [zipkin_attrs]
zipkin_logger_mock.handlers = []
generate_string_mock.return_value = '1'
context = zipkin.zipkin_span(
service_name='my_service',
port=5,
transport_handler=mock.Mock(),
sample_rate=0.0,
)
with context:
# Assert that the new ZipkinAttrs were saved
new_zipkin_attrs = get_zipkin_attrs()
assert new_zipkin_attrs.span_id == '1'
# Outside of the context, things should be returned to normal
assert get_zipkin_attrs() == zipkin_attrs
def test_span_context(
zipkin_logger_mock,
generate_string_128bit_mock,
generate_string_mock,
thread_local_mock,
span_func,
expected_annotations,
):
zipkin_attrs = ZipkinAttrs(
trace_id='1111111111111111',
span_id='2222222222222222',
parent_span_id='3333333333333333',
flags='flags',
is_sampled=True,
)
thread_local_mock.zipkin_attrs = [zipkin_attrs]
logging_handler = ZipkinLoggerHandler(zipkin_attrs)
assert logging_handler.parent_span_id is None
assert logging_handler.client_spans == []
zipkin_logger_mock.handlers = [logging_handler]
generate_string_mock.return_value = '1'
context = span_func(
service_name='svc',
span_name='span',
annotations={'something': 1},
binary_annotations={'foo': 'bar'},
)
with context:
# Assert that the new ZipkinAttrs were saved
new_zipkin_attrs = get_zipkin_attrs()
assert new_zipkin_attrs.span_id == '1'
# And that the logging handler has a parent_span_id
assert logging_handler.parent_span_id == '1'
# Outside of the context, things should be returned to normal,
# except a new client span is saved in the handler
assert logging_handler.parent_span_id is None
assert get_zipkin_attrs() == zipkin_attrs
client_span = logging_handler.client_spans.pop()
assert logging_handler.client_spans == []
# These reserved annotations are based on timestamps so pop em.
# This also acts as a check that they exist.
for annotation in expected_annotations:
client_span['annotations'].pop(annotation)
expected_client_span = {
'span_name': 'span',
'service_name': 'svc',
'parent_span_id': None,
'span_id': '1',
'annotations': {'something': 1},
'binary_annotations': {'foo': 'bar'},
'sa_binary_annotations': [],
}
assert client_span == expected_client_span
assert generate_string_128bit_mock.call_count == 0
def enrich_with_zipkin_data(data):
zipkin_attr = get_zipkin_attrs()
data['trace_id'] = zipkin_attr.trace_id
data['parent_span_id'] = zipkin_attr.parent_span_id
data['is_sampled'] = True if zipkin_attr.is_sampled else False
return data
def enrich_with_zipkin_data(data):
zipkin_attr = get_zipkin_attrs()
data[0]['trace_id'] = zipkin_attr.trace_id
data[0]['parent_span_id'] = zipkin_attr.span_id
data[0]['span_id'] = generate_random_64bit_string()
data[0]['is_sampled'] = True if zipkin_attr.is_sampled else False
return data
def enrich_with_zipkin_data(data):
# enrich == add key - value pair for data dict
# the following data will be sent to kafka, the next span
# so the kafka is the parent span of the next span
zipkin_attr = get_zipkin_attrs()
data['trace_id'] = zipkin_attr.trace_id
data['parent_span_id'] = zipkin_attr.parent_span_id
data['is_sampled'] = True if zipkin_attr.is_sampled else False
return data
def create_http_headers_for_new_span():
"""
Generate the headers for a new zipkin span.
.. note::
If the method is not called from within a zipkin_trace conext,
empty dict will be returned back.
:returns: dict containing (X-B3-TraceId, X-B3-SpanId, X-B3-ParentSpanId,
X-B3-Flags and X-B3-Sampled) keys OR an empty dict.
"""
zipkin_attrs = get_zipkin_attrs()
if not zipkin_attrs:
return {}
return {
'X-B3-TraceId': zipkin_attrs.trace_id,
'X-B3-SpanId': generate_random_64bit_string(),
'X-B3-ParentSpanId': zipkin_attrs.span_id,
'X-B3-Flags': '0',
'X-B3-Sampled': '1' if zipkin_attrs.is_sampled else '0',
}
def test_push_zipkin_attrs_adds_new_zipkin_attrs_to_list():
tracer = storage.get_default_tracer()
with mock.patch.object(tracer._context_stack, '_storage', ['foo']):
assert 'foo' == thread_local.get_zipkin_attrs()
thread_local.push_zipkin_attrs('bar')
assert 'bar' == thread_local.get_zipkin_attrs()
def start(self):
"""Enter the new span context. All annotations logged inside this
context will be attributed to this span. All new spans generated
inside this context will have this span as their parent.
In the unsampled case, this context still generates new span IDs and
pushes them onto the threadlocal stack, so downstream services calls
made will pass the correct headers. However, the logging handler is
never attached in the unsampled case, so the spans are never logged.
"""
self.do_pop_attrs = False
# If zipkin_attrs are passed in or this span is doing its own sampling,
# it will need to actually log spans at __exit__.
self.perform_logging = self.zipkin_attrs or self.sample_rate is not None
if self.sample_rate is not None:
# This clause allows for sampling this service independently
# of the passed-in zipkin_attrs.
if self.zipkin_attrs and not self.zipkin_attrs.is_sampled:
self.zipkin_attrs = create_attrs_for_span(
sample_rate=self.sample_rate,
trace_id=self.zipkin_attrs.trace_id,
)
elif not self.zipkin_attrs:
self.zipkin_attrs = create_attrs_for_span(
sample_rate=self.sample_rate, )
if not self.zipkin_attrs:
# This span is inside the context of an existing trace
existing_zipkin_attrs = get_zipkin_attrs()
if existing_zipkin_attrs:
self.zipkin_attrs = ZipkinAttrs(
trace_id=existing_zipkin_attrs.trace_id,
span_id=generate_random_64bit_string(),
parent_span_id=existing_zipkin_attrs.span_id,
flags=existing_zipkin_attrs.flags,
is_sampled=existing_zipkin_attrs.is_sampled,
)
# If zipkin_attrs are not set up by now, that means this span is not
# configured to perform logging itself, and it's not in an existing
# Zipkin trace. That means there's nothing else to do and it can exit
# early.
if not self.zipkin_attrs:
return self
push_zipkin_attrs(self.zipkin_attrs)
self.do_pop_attrs = True
self.start_timestamp = time.time()
# Set up logging if this is the root span
if self.perform_logging:
# Don't set up any logging if we're not sampling
if not self.zipkin_attrs.is_sampled:
return self
endpoint = create_endpoint(self.port, self.service_name)
self.log_handler = ZipkinLoggerHandler(self.zipkin_attrs)
self.logging_context = ZipkinLoggingContext(
self.zipkin_attrs,
endpoint,
self.log_handler,
self.span_name,
self.transport_handler,
self.binary_annotations,
add_logging_annotation=self.add_logging_annotation,
)
self.logging_context.start()
return self
else:
# In the sampled case, patch the ZipkinLoggerHandler.
if self.zipkin_attrs.is_sampled:
# Be defensive about logging setup. Since ZipkinAttrs are local to
# the thread, multithreaded frameworks can get in strange states.
# The logging is not going to be correct in these cases, so we set
# a flag that turns off logging on __exit__.
if len(zipkin_logger.handlers) > 0:
# Put span ID on logging handler. Assume there's only a single
# handler, since all logging should be set up in this package.
self.log_handler = zipkin_logger.handlers[0]
# Store the old parent_span_id, probably None, in case we have
# nested zipkin_spans
self.old_parent_span_id = self.log_handler.parent_span_id
self.log_handler.parent_span_id = self.zipkin_attrs.span_id
return self
def test_push_zipkin_attrs_adds_new_zipkin_attrs_to_list():
assert 'foo' == thread_local.get_zipkin_attrs()
thread_local.push_zipkin_attrs('bar')
assert 'bar' == thread_local.get_zipkin_attrs()
def test_pop_zipkin_attrs_removes_the_last_zipkin_attrs():
assert 'bar' == thread_local.pop_zipkin_attrs()
assert 'foo' == thread_local.get_zipkin_attrs()
def test_get_zipkin_attrs_returns_the_last_of_the_list():
assert 'foo' == thread_local.get_zipkin_attrs()
def test_pop_zipkin_attrs_removes_the_last_zipkin_attrs():
tracer = storage.get_default_tracer()
with mock.patch.object(tracer._context_stack, '_storage', ['foo', 'bar']):
assert 'bar' == thread_local.pop_zipkin_attrs()
assert 'foo' == thread_local.get_zipkin_attrs()
def test_span_context(
generate_string_128bit_mock,
generate_string_mock,
thread_local_mock,
span_func,
expected_annotations,
):
span_storage = SpanStorage()
generate_string_mock.return_value = '1'
with zipkin.zipkin_span(
service_name='root_span',
span_name='root_span',
sample_rate=100.0,
transport_handler=MockTransportHandler(),
span_storage=span_storage,
):
zipkin_attrs = ZipkinAttrs(
trace_id='1111111111111111',
span_id='2222222222222222',
parent_span_id='3333333333333333',
flags='flags',
is_sampled=True,
)
thread_local_mock.zipkin_attrs = [zipkin_attrs]
ts = time.time()
with mock.patch('time.time', return_value=ts):
with span_func(
service_name='svc',
span_name='span',
binary_annotations={'foo': 'bar'},
span_storage=span_storage,
):
# Assert that the new ZipkinAttrs were saved
new_zipkin_attrs = get_zipkin_attrs()
assert new_zipkin_attrs.span_id == '1'
# Outside of the context, things should be returned to normal
assert get_zipkin_attrs() == zipkin_attrs
client_span = span_storage.pop().build_v1_span()
# These reserved annotations are based on timestamps so pop em.
# This also acts as a check that they exist.
for annotation in expected_annotations:
client_span.annotations.pop(annotation)
expected_client_span = _V1Span(
trace_id='1111111111111111',
name='span',
parent_id='2222222222222222',
id='1',
timestamp=ts,
duration=0.0,
endpoint=None,
annotations={},
binary_annotations={'foo': 'bar'},
sa_endpoint=None,
)
assert client_span == expected_client_span
assert generate_string_128bit_mock.call_count == 0
def test_push_zipkin_attrs_adds_new_zipkin_attrs_to_list():
tracer = storage.get_default_tracer()
with mock.patch.object(tracer._context_stack, "_storage", ["foo"]):
assert "foo" == thread_local.get_zipkin_attrs()
thread_local.push_zipkin_attrs("bar")
assert "bar" == thread_local.get_zipkin_attrs()
def test_pop_zipkin_attrs_removes_the_last_zipkin_attrs():
tracer = storage.get_default_tracer()
with mock.patch.object(tracer._context_stack, "_storage", ["foo", "bar"]):
assert "bar" == thread_local.pop_zipkin_attrs()
assert "foo" == thread_local.get_zipkin_attrs()
def test_get_zipkin_attrs_returns_the_last_of_the_list():
tracer = storage.get_default_tracer()
with mock.patch.object(tracer._context_stack, "_storage", ["foo"]):
assert "foo" == thread_local.get_zipkin_attrs()
def test_get_zipkin_attrs_returns_none_if_no_zipkin_attrs():
assert not thread_local.get_zipkin_attrs()
