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 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()