def tween(request):
# Creates zipkin attributes, attaches a zipkin_trace_id attr to the
# request, and pushes the attrs onto threadlocal stack.
zipkin_attrs = create_zipkin_attr(request)
push_zipkin_attrs(zipkin_attrs)
try:
# If this request isn't sampled, don't go through the work
# of initializing the rest of the zipkin attributes
if not zipkin_attrs.is_sampled:
return handler(request)
# If the request IS sampled, we create thrift objects and
# enter zipkin logging context
thrift_endpoint = create_endpoint(request)
log_handler = ZipkinLoggerHandler(zipkin_attrs)
with ZipkinLoggingContext(zipkin_attrs, thrift_endpoint,
log_handler, request) as context:
response = handler(request)
context.response_status_code = response.status_code
context.binary_annotations_dict = get_binary_annotations(
request, response)
return response
finally:
# Regardless of what happens in the request we want to pop attrs
pop_zipkin_attrs()
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 __enter__(self):
"""Actions to be taken before request is handled.
1) Attach `zipkin_logger` to :class:`ZipkingLoggerHandler` object.
2) Push zipkin attributes to thread_local stack.
3) Record the start timestamp.
"""
zipkin_logger.addHandler(self.handler)
push_zipkin_attrs(self.zipkin_attrs)
self.start_timestamp = time.time()
return self
def __enter__(self):
"""Actions to be taken before request is handled.
1) Attach `zipkin_logger` to :class:`ZipkingLoggerHandler` object.
2) Push zipkin attributes to thread_local stack.
3) Record the start timestamp.
"""
zipkin_logger.addHandler(self.handler)
push_zipkin_attrs(self.zipkin_attrs)
self.start_timestamp = time.time()
return self
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_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()
