def test_tween_not_sampled_sets_zipkin_trace_id(is_tracing_mock):
# Tests that even in the unsampled case, the zipkin tween sets
# the zipkin_trace_id attr on the request.
is_tracing_mock.return_value = False
assert get_zipkin_attrs() is None
request = Request.blank('/', headers={'X-B3-TraceId': 'deadbeefdeadbeef'})
tween = zipkin.zipkin_tween(mock.Mock(), 'registry')
tween(request)
assert request.zipkin_trace_id == 'deadbeefdeadbeef'
# Make sure the tween doesn't leave zipkin attrs on threadlocal storage
assert get_zipkin_attrs() is None
def test_span_context(
zipkin_logger_mock,
generate_string_mock,
thread_local_mock,
):
zipkin_attrs = ZipkinAttrs(
'trace_id', 'span_id', 'parent_span_id', 'flags', True)
thread_local_mock.requests = [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 = zipkin.SpanContext(
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 ('cs', 'cr', 'ss', 'sr'):
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'},
}
assert client_span == expected_client_span
def create_headers_for_new_span():
"""
Generate the headers for a new zipkin span.
.. note::
If the method is not called from within a pyramid service call OR
pyramid_zipkin is not included as a pyramid tween, 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_span_id(),
'X-B3-ParentSpanId': zipkin_attrs.span_id,
'X-B3-Flags': '0',
'X-B3-Sampled': '1' if zipkin_attrs.is_sampled else '0',
}
def __enter__(self):
"""Enter the client context. All spans/annotations logged inside this
context will be attributed to this client span.
"""
zipkin_attrs = get_zipkin_attrs()
self.is_sampled = zipkin_attrs is not None and zipkin_attrs.is_sampled
if not self.is_sampled:
return self
self.start_timestamp = time.time()
self.span_id = generate_random_64bit_string()
# Put span ID on logging handler. Assume there's only a single handler
# on the logger, since all logging should be set up in this package.
self.handler = zipkin_logger.handlers[0]
# Store the old parent_span_id, probably None, in case we have
# nested ClientSpanContexts
self.old_parent_span_id = self.handler.parent_span_id
self.handler.parent_span_id = self.span_id
# Push new zipkin attributes onto the threadlocal stack, so that
# create_headers_for_new_span() performs as expected in this context.
# The only difference is that span_id is this new client span's ID
# and parent_span_id is the old span's ID.
new_zipkin_attrs = ZipkinAttrs(
trace_id=zipkin_attrs.trace_id,
span_id=self.span_id,
parent_span_id=zipkin_attrs.span_id,
flags=zipkin_attrs.flags,
is_sampled=zipkin_attrs.is_sampled,
)
push_zipkin_attrs(new_zipkin_attrs)
return self
def create_headers_for_new_span():
"""
Generate the headers for a new zipkin span.
.. note::
If the method is not called from within a pyramid service call OR
pyramid_zipkin is not included as a pyramid tween, 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_span_id(),
'X-B3-ParentSpanId': zipkin_attrs.span_id,
'X-B3-Flags': '0',
'X-B3-Sampled': '1' if zipkin_attrs.is_sampled else '0',
}
def test_span_context_sampled_no_handlers(
zipkin_logger_mock,
generate_string_mock,
thread_local_mock,
):
zipkin_attrs = ZipkinAttrs(
'trace_id', 'span_id', 'parent_span_id', 'flags', True)
thread_local_mock.requests = [zipkin_attrs]
zipkin_logger_mock.handlers = []
generate_string_mock.return_value = '1'
context = zipkin.SpanContext('svc', 'span')
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 __enter__(self):
"""Enter the new span context. All spans/annotations logged inside this
context will be attributed to this span.
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 it is left alone.
"""
zipkin_attrs = get_zipkin_attrs()
self.is_sampled = zipkin_attrs is not None and zipkin_attrs.is_sampled
self.span_id = generate_random_64bit_string()
self.start_timestamp = time.time()
# Push new zipkin attributes onto the threadlocal stack, so that
# create_headers_for_new_span() performs as expected in this context.
# The only difference is that span_id is this new span's ID
# and parent_span_id is the old span's ID. Checking for a None
# zipkin_attrs value is protecting against calling this outside of
# a zipkin logging context entirely (e.g. in a batch). If new attrs
# are stored, set a flag to pop them off at context exit.
self.do_pop_attrs = False
if zipkin_attrs is not None:
new_zipkin_attrs = ZipkinAttrs(
trace_id=zipkin_attrs.trace_id,
span_id=self.span_id,
parent_span_id=zipkin_attrs.span_id,
flags=zipkin_attrs.flags,
is_sampled=zipkin_attrs.is_sampled,
)
push_zipkin_attrs(new_zipkin_attrs)
self.do_pop_attrs = True
# In the sampled case, patch the ZipkinLoggerHandler.
if self.is_sampled:
# Be defensive about logging setup. Since ZipkinAttrs are local
# the a 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.handler = zipkin_logger.handlers[0]
# Store the old parent_span_id, probably None, in case we have
# nested SpanContexts
self.old_parent_span_id = self.handler.parent_span_id
self.handler.parent_span_id = self.span_id
self.logging_initialized = True
return self
def test_push_zipkin_attrs_adds_new_request_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_request():
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_get_zipkin_attrs_returns_none_if_no_requests():
assert not thread_local.get_zipkin_attrs()
def test_push_zipkin_attrs_adds_new_request_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_request():
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_get_zipkin_attrs_returns_none_if_no_requests():
assert not thread_local.get_zipkin_attrs()
