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 __exit__(self, _exc_type, _exc_value, _exc_traceback):
"""Exit the span context. The new zipkin attrs are pushed onto the
threadlocal stack regardless of sampling, so they always need to be
popped off. The actual logging of spans depends on sampling and that
the logging was correctly set up.
"""
if self.do_pop_attrs:
pop_zipkin_attrs()
if not (self.is_sampled and self.logging_initialized):
return
end_timestamp = time.time()
# Put the old parent_span_id back on the handler
self.handler.parent_span_id = self.old_parent_span_id
# To get a full span we just set cs=sr and ss=cr.
self.annotations.update({
'cs': self.start_timestamp,
'sr': self.start_timestamp,
'ss': end_timestamp,
'cr': end_timestamp,
})
# Store this span on the logging handler object.
self.handler.store_client_span(
span_name=self.span_name,
service_name=self.service_name,
annotations=self.annotations,
binary_annotations=self.binary_annotations,
span_id=self.span_id,
)
def __exit__(self, _type, _value, _traceback):
"""Actions to be taken post request handling.
1) Record the end timestamp.
2) Pop zipkin attributes from thread_local stack
3) Detach `zipkin_logger` handler.
4) And finally, if sampled, log the service annotations to scribe
"""
self.log_spans()
pop_zipkin_attrs()
zipkin_logger.removeHandler(self.handler)
def __exit__(self, _type, _value, _traceback):
"""Actions to be taken post request handling.
1) Record the end timestamp.
2) Pop zipkin attributes from thread_local stack
3) Detach `zipkin_logger` handler.
4) And finally, if sampled, log the service annotations to scribe
"""
self.log_spans()
pop_zipkin_attrs()
zipkin_logger.removeHandler(self.handler)
def __exit__(self, _exc_type, _exc_value, _exc_traceback):
if not self.is_sampled:
return
end_timestamp = time.time()
# Put the old parent_span_id back on the handler
self.handler.parent_span_id = self.old_parent_span_id
# Pop off the new zipkin attrs
pop_zipkin_attrs()
self.annotations["cs"] = self.start_timestamp
self.annotations["cr"] = end_timestamp
# Store this client span on the logging handler object
self.handler.store_client_span(
span_name=self.span_name,
service_name=self.service_name,
annotations=self.annotations,
binary_annotations=self.binary_annotations,
span_id=self.span_id,
)
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_pop_zipkin_attrs_does_nothing_if_no_requests():
assert not thread_local.pop_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_pop_zipkin_attrs_does_nothing_if_no_requests():
assert not thread_local.pop_zipkin_attrs()