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
class zipkin_span(object):
"""Context manager/decorator for all of your zipkin tracing needs.
Usage #1: Start a trace with a given sampling rate
This begins the zipkin trace and also records the root span. The required
params are service_name, transport_handler, and sample_rate.
# Start a trace with do_stuff() as the root span
def some_batch_job(a, b):
with zipkin_span(
service_name='my_service',
span_name='my_span_name',
transport_handler=some_handler,
port=22,
sample_rate=0.05,
):
do_stuff()
Usage #2: Trace a service call.
The typical use case is instrumenting a framework like Pyramid or Django. Only
ss and sr times are recorded for the root span. Required params are
service_name, zipkin_attrs, transport_handler, and port.
# Used in a pyramid tween
def tween(request):
zipkin_attrs = some_zipkin_attr_creator(request)
with zipkin_span(
service_name='my_service,'
span_name='my_span_name',
zipkin_attrs=zipkin_attrs,
transport_handler=some_handler,
port=22,
) as zipkin_context:
response = handler(request)
zipkin_context.update_binary_annotations_for_root_span(
some_binary_annotations)
return response
Usage #3: Log a span within the context of a zipkin trace
If you're already in a zipkin trace, you can use this to log a span inside. The
only required param is service_name. If you're not in a zipkin trace, this
won't do anything.
# As a decorator
@zipkin_span(service_name='my_service', span_name='my_function')
def my_function():
do_stuff()
# As a context manager
def my_function():
with zipkin_span(service_name='my_service', span_name='do_stuff'):
do_stuff()
"""
def __init__(
self,
service_name,
span_name='span',
zipkin_attrs=None,
transport_handler=None,
annotations=None,
binary_annotations=None,
port=0,
sample_rate=None,
include=('client', 'server'),
add_logging_annotation=False,
):
"""Logs a zipkin span. If this is the root span, then a zipkin
trace is started as well.
:param service_name: The name of the called service
:type service_name: string
:param span_name: Optional name of span, defaults to 'span'
:type span_name: string
:param zipkin_attrs: Optional set of zipkin attributes to be used
:type zipkin_attrs: ZipkinAttrs
:param transport_handler: Callback function that takes a message parameter
and handles logging it
:type transport_handler: function
:param annotations: Optional dict of str -> timestamp annotations
:type annotations: dict of str -> int
:param binary_annotations: Optional dict of str -> str span attrs
:type binary_annotations: dict of str -> str
:param port: The port number of the service. Defaults to 0.
:type port: int
:param sample_rate: Rate at which to sample; 0.0 - 100.0. If passed-in
zipkin_attrs have is_sampled=False and the sample_rate param is > 0,
a new span will be generated at this rate. This means that if you
propagate sampling decisions to downstream services, but still have
sample_rate > 0 in those services, the actual rate of generated
spans for those services will be > sampling_rate.
:type sample_rate: float
:param include: which annotations to include
can be one of {'client', 'server'}
corresponding to ('cs', 'cr') and ('ss', 'sr') respectively
:type include: iterable
:param add_logging_annotation: Whether to add a 'start_logging'
annotation when py_zipkin starts logging spans
:type add_logging_annotation: boolean
"""
self.service_name = service_name
self.span_name = span_name
self.zipkin_attrs = zipkin_attrs
self.transport_handler = transport_handler
self.annotations = annotations or {}
self.binary_annotations = binary_annotations or {}
self.port = port
self.logging_context = None
self.sample_rate = sample_rate
self.include = include
self.add_logging_annotation = add_logging_annotation
# Validation checks
if self.zipkin_attrs or self.sample_rate is not None:
if self.transport_handler is None:
raise ZipkinError(
'Root spans require a transport handler to be given')
if self.sample_rate is not None and not (0.0 <= self.sample_rate <=
100.0):
raise ZipkinError('Sample rate must be between 0.0 and 100.0')
if not set(include).issubset(STANDARD_ANNOTATIONS_KEYS):
raise ZipkinError('Only %s are supported as annotations' %
STANDARD_ANNOTATIONS_KEYS)
else:
# get a list of all of the mapped annotations
self.annotation_filter = set()
for include_name in include:
self.annotation_filter.update(
STANDARD_ANNOTATIONS[include_name])
def __call__(self, f):
@functools.wraps(f)
def decorated(*args, **kwargs):
with zipkin_span(
service_name=self.service_name,
span_name=self.span_name,
zipkin_attrs=self.zipkin_attrs,
transport_handler=self.transport_handler,
annotations=self.annotations,
binary_annotations=self.binary_annotations,
port=self.port,
sample_rate=self.sample_rate,
include=self.include,
):
return f(*args, **kwargs)
return decorated
def __enter__(self):
return self.start()
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 __exit__(self, _exc_type, _exc_value, _exc_traceback):
self.stop(_exc_type, _exc_value, _exc_traceback)
def stop(self, _exc_type=None, _exc_value=None, _exc_traceback=None):
"""Exit the span context. 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.
"""
# Always remove the stored zipkin_attrs
if self.do_pop_attrs:
pop_zipkin_attrs()
# Exit early if this request is not being sampled
if not self.zipkin_attrs or not self.zipkin_attrs.is_sampled:
return
# If this is the root span, exit the context (which will handle logging)
if self.logging_context:
self.logging_context.stop()
self.logging_context = None
return
end_timestamp = time.time()
# Put the old parent_span_id back on the handler
self.log_handler.parent_span_id = self.old_parent_span_id
# To get a full span we just set cs=sr and ss=cr.
full_annotations = {
'cs': self.start_timestamp,
'sr': self.start_timestamp,
'ss': end_timestamp,
'cr': end_timestamp,
}
# But we filter down if we only want to emit some of the annotations
filtered_annotations = {
k: v
for k, v in full_annotations.items() if k in self.annotation_filter
}
self.annotations.update(filtered_annotations)
# Store this span on the logging handler object.
self.log_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.zipkin_attrs.span_id,
)
def update_binary_annotations_for_root_span(self, extra_annotations):
"""Updates the binary annotations for the root span of the trace.
If this trace is not being sampled then this is a no-op.
"""
if not self.zipkin_attrs.is_sampled:
return
if not self.logging_context:
raise ZipkinError('No logging context available')
self.logging_context.binary_annotations_dict.update(extra_annotations)
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
class zipkin_span(object):
"""Context manager/decorator for all of your zipkin tracing needs.
Usage #1: Start a trace with a given sampling rate
This begins the zipkin trace and also records the root span. The required
params are service_name, transport_handler, and sample_rate.
# Start a trace with do_stuff() as the root span
def some_batch_job(a, b):
with zipkin_span(
service_name='my_service',
span_name='my_span_name',
transport_handler=some_handler,
port=22,
sample_rate=0.05,
):
do_stuff()
Usage #2: Trace a service call.
The typical use case is instrumenting a framework like Pyramid or Django. Only
ss and sr times are recorded for the root span. Required params are
service_name, zipkin_attrs, transport_handler, and port.
# Used in a pyramid tween
def tween(request):
zipkin_attrs = some_zipkin_attr_creator(request)
with zipkin_span(
service_name='my_service,'
span_name='my_span_name',
zipkin_attrs=zipkin_attrs,
transport_handler=some_handler,
port=22,
) as zipkin_context:
response = handler(request)
zipkin_context.update_binary_annotations(
some_binary_annotations)
return response
Usage #3: Log a span within the context of a zipkin trace
If you're already in a zipkin trace, you can use this to log a span inside. The
only required param is service_name. If you're not in a zipkin trace, this
won't do anything.
# As a decorator
@zipkin_span(service_name='my_service', span_name='my_function')
def my_function():
do_stuff()
# As a context manager
def my_function():
with zipkin_span(service_name='my_service', span_name='do_stuff'):
do_stuff()
"""
def __init__(self,
service_name,
span_name='span',
zipkin_attrs=None,
transport_handler=None,
annotations=None,
binary_annotations=None,
port=0,
sample_rate=None,
include=('client', 'server'),
add_logging_annotation=False,
report_root_timestamp=False,
use_128bit_trace_id=False,
host=None):
"""Logs a zipkin span. If this is the root span, then a zipkin
trace is started as well.
:param service_name: The name of the called service
:type service_name: string
:param span_name: Optional name of span, defaults to 'span'
:type span_name: string
:param zipkin_attrs: Optional set of zipkin attributes to be used
:type zipkin_attrs: ZipkinAttrs
:param transport_handler: Callback function that takes a message parameter
and handles logging it
:type transport_handler: function
:param annotations: Optional dict of str -> timestamp annotations
:type annotations: dict of str -> int
:param binary_annotations: Optional dict of str -> str span attrs
:type binary_annotations: dict of str -> str
:param port: The port number of the service. Defaults to 0.
:type port: int
:param sample_rate: Rate at which to sample; 0.0 - 100.0. If passed-in
zipkin_attrs have is_sampled=False and the sample_rate param is > 0,
a new span will be generated at this rate. This means that if you
propagate sampling decisions to downstream services, but still have
sample_rate > 0 in those services, the actual rate of generated
spans for those services will be > sampling_rate.
:type sample_rate: float
:param include: which annotations to include
can be one of {'client', 'server'}
corresponding to ('cs', 'cr') and ('ss', 'sr') respectively
:type include: iterable
:param add_logging_annotation: Whether to add a 'logging_end'
annotation when py_zipkin finishes logging spans
:type add_logging_annotation: boolean
:param report_root_timestamp: Whether the span should report timestamp
and duration. Only applies to "root" spans in this local context,
so spans created inside other span contexts will always log
timestamp/duration. Note that this is only an override for spans
that have zipkin_attrs passed in. Spans that make their own
sampling decisions (i.e. are the root spans of entire traces) will
always report timestamp/duration.
:type report_root_timestamp: boolean
:param use_128bit_trace_id: If true, generate 128-bit trace_ids
:type use_128bit_trace_id: boolean
:param host: Contains the ipv4 value of the host. The ipv4 value isn't
automatically determined in a docker environment
:type host: string
"""
self.service_name = service_name
self.span_name = span_name
self.zipkin_attrs = zipkin_attrs
self.transport_handler = transport_handler
self.annotations = annotations or {}
self.binary_annotations = binary_annotations or {}
self.port = port
self.logging_context = None
self.sample_rate = sample_rate
self.include = include
self.add_logging_annotation = add_logging_annotation
self.report_root_timestamp_override = report_root_timestamp
self.use_128bit_trace_id = use_128bit_trace_id
self.host = host
# Spans that log a 'cs' timestamp can additionally record
# 'sa' binary annotations that show where the request is going.
# This holds a list of 'sa' binary annotations.
self.sa_binary_annotations = []
# Validation checks
if self.zipkin_attrs or self.sample_rate is not None:
if self.transport_handler is None:
raise ZipkinError(
'Root spans require a transport handler to be given')
if self.sample_rate is not None and not (0.0 <= self.sample_rate <=
100.0):
raise ZipkinError('Sample rate must be between 0.0 and 100.0')
if not set(include).issubset(STANDARD_ANNOTATIONS_KEYS):
raise ZipkinError('Only %s are supported as annotations' %
STANDARD_ANNOTATIONS_KEYS)
else:
# get a list of all of the mapped annotations
self.annotation_filter = set()
for include_name in include:
self.annotation_filter.update(
STANDARD_ANNOTATIONS[include_name])
def __call__(self, f):
@functools.wraps(f)
def decorated(*args, **kwargs):
with zipkin_span(service_name=self.service_name,
span_name=self.span_name,
zipkin_attrs=self.zipkin_attrs,
transport_handler=self.transport_handler,
annotations=self.annotations,
binary_annotations=self.binary_annotations,
port=self.port,
sample_rate=self.sample_rate,
include=self.include,
host=self.host):
return f(*args, **kwargs)
return decorated
def __enter__(self):
return self.start()
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
report_root_timestamp = False
if self.sample_rate is not None:
if self.zipkin_attrs and not self.zipkin_attrs.is_sampled:
report_root_timestamp = True
self.zipkin_attrs = create_attrs_for_span(
sample_rate=self.sample_rate,
trace_id=self.zipkin_attrs.trace_id,
use_128bit_trace_id=self.use_128bit_trace_id,
)
elif not self.zipkin_attrs:
report_root_timestamp = True
self.zipkin_attrs = create_attrs_for_span(
sample_rate=self.sample_rate,
use_128bit_trace_id=self.use_128bit_trace_id,
)
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()
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.host)
client_context = set(self.include) == {'client'}
self.log_handler = ZipkinLoggerHandler(self.zipkin_attrs)
self.logging_context = ZipkinLoggingContext(
self.zipkin_attrs,
endpoint,
self.log_handler,
self.span_name,
self.transport_handler,
report_root_timestamp or self.report_root_timestamp_override,
binary_annotations=self.binary_annotations,
add_logging_annotation=self.add_logging_annotation,
client_context=client_context)
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 __exit__(self, _exc_type, _exc_value, _exc_traceback):
self.stop(_exc_type, _exc_value, _exc_traceback)
def stop(self, _exc_type=None, _exc_value=None, _exc_traceback=None):
"""Exit the span context. 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.zipkin_attrs or not self.zipkin_attrs.is_sampled:
return
# Logging context is only initialized for "root" spans of the local
# process (i.e. this zipkin_span not inside of any other local
# zipkin_spans)
if self.logging_context:
self.logging_context.stop()
self.logging_context = None
return
# If we've gotten here, that means that this span is a child span of
# this context's root span (i.e. it's a zipkin_span inside another
# zipkin_span).
end_timestamp = time.time()
self.log_handler.parent_span_id = self.old_parent_span_id
# We are simulating a full two-part span locally, so set cs=sr and ss=cr
full_annotations = {
'cs': self.start_timestamp,
'sr': self.start_timestamp,
'ss': end_timestamp,
'cr': end_timestamp,
}
# But we filter down if we only want to emit some of the annotations
filtered_annotations = {
k: v
for k, v in full_annotations.items() if k in self.annotation_filter
}
self.annotations.update(filtered_annotations)
self.log_handler.store_local_span(
span_name=self.span_name,
service_name=self.service_name,
annotations=self.annotations,
binary_annotations=self.binary_annotations,
sa_binary_annotations=self.sa_binary_annotations,
span_id=self.zipkin_attrs.span_id,
)
def update_binary_annotations(self, extra_annotations):
"""Updates the binary annotations for the current span.
If this trace is not being sampled then this is a no-op.
"""
if not self.zipkin_attrs:
return
if not self.zipkin_attrs.is_sampled:
return
if not self.logging_context:
# This is not the root span, so binary annotations will be added
# to the log handler when this span context exits.
self.binary_annotations.update(extra_annotations)
else:
# Otherwise, we're in the context of the root span, so just update
# the binary annotations for the logging context directly.
self.logging_context.binary_annotations_dict.update(
extra_annotations)
def add_sa_binary_annotation(
self,
port=0,
service_name='unknown',
host='127.0.0.1',
):
"""Adds a 'sa' binary annotation to the current span.
'sa' binary annotations are useful for situations where you need to log
where a request is going but the destination doesn't support zipkin.
Note that the span must have 'cs'/'cr' annotations.
:param port: The port number of the destination
:type port: int
:param service_name: The name of the destination service
:type service_name: str
:param host: Host address of the destination
:type host: str
"""
if not self.zipkin_attrs or not self.zipkin_attrs.is_sampled:
return
if 'client' not in self.include:
# TODO: trying to set a sa binary annotation for a non-client span
# should result in a logged error
return
sa_endpoint = create_endpoint(
port=port,
service_name=service_name,
host=host,
)
sa_binary_annotation = create_binary_annotation(
key=zipkin_core.SERVER_ADDR,
value=SERVER_ADDR_VAL,
annotation_type=zipkin_core.AnnotationType.BOOL,
host=sa_endpoint,
)
if not self.logging_context:
self.sa_binary_annotations.append(sa_binary_annotation)
else:
self.logging_context.sa_binary_annotations.append(
sa_binary_annotation)
class zipkin_span(object):
"""Context manager/decorator for all of your zipkin tracing needs.
Usage #1: Start a trace with a given sampling rate
This begins the zipkin trace and also records the root span. The required
params are service_name, transport_handler, and sample_rate.
# Start a trace with do_stuff() as the root span
def some_batch_job(a, b):
with zipkin_span(
service_name='my_service',
span_name='my_span_name',
transport_handler=some_handler,
port=22,
sample_rate=0.05,
):
do_stuff()
Usage #2: Trace a service call.
The typical use case is instrumenting a framework like Pyramid or Django. Only
ss and sr times are recorded for the root span. Required params are
service_name, zipkin_attrs, transport_handler, and port.
# Used in a pyramid tween
def tween(request):
zipkin_attrs = some_zipkin_attr_creator(request)
with zipkin_span(
service_name='my_service,'
span_name='my_span_name',
zipkin_attrs=zipkin_attrs,
transport_handler=some_handler,
port=22,
) as zipkin_context:
response = handler(request)
zipkin_context.update_binary_annotations_for_root_span(
some_binary_annotations)
return response
Usage #3: Log a span within the context of a zipkin trace
If you're already in a zipkin trace, you can use this to log a span inside. The
only required param is service_name. If you're not in a zipkin trace, this
won't do anything.
# As a decorator
@zipkin_span(service_name='my_service', span_name='my_function')
def my_function():
do_stuff()
# As a context manager
def my_function():
with zipkin_span(service_name='my_service', span_name='do_stuff'):
do_stuff()
"""
def __init__(
self,
service_name,
span_name='span',
zipkin_attrs=None,
transport_handler=None,
annotations=None,
binary_annotations=None,
port=0,
sample_rate=None,
):
"""Logs a zipkin span. If this is the root span, then a zipkin
trace is started as well.
:param service_name: The name of the called service
:type service_name: string
:param span_name: Optional name of span, defaults to 'span'
:type span_name: string
:param zipkin_attrs: Optional set of zipkin attributes to be used
:type zipkin_attrs: ZipkinAttrs
:param transport_handler: Callback function that takes a message parameter
and handles logging it
:type transport_handler: function
:param annotations: Optional dict of str -> timestamp annotations
:type annotations: dict of str -> int
:param binary_annotations: Optional dict of str -> str span attrs
:type binary_annotations: dict of str -> str
:param port: The port number of the service. Defaults to 0.
:type port: int
:param sample_rate: Custom sampling rate (between 100.0 and 0.0) if
this is the root of the trace
:type sample_rate: float
"""
self.service_name = service_name
self.span_name = span_name
self.zipkin_attrs = zipkin_attrs
self.transport_handler = transport_handler
self.annotations = annotations or {}
self.binary_annotations = binary_annotations or {}
self.port = port
self.logging_context = None
self.sample_rate = sample_rate
# Validation checks
if self.zipkin_attrs or self.sample_rate is not None:
if self.transport_handler is None:
raise ZipkinError(
'Root spans require a transport handler to be given')
if self.sample_rate is not None and not (0.0 <= self.sample_rate <= 100.0):
raise ZipkinError('Sample rate must be between 0.0 and 100.0')
def __call__(self, f):
@functools.wraps(f)
def decorated(*args, **kwargs):
with self:
return f(*args, **kwargs)
return decorated
def __enter__(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 this span is the first span to be recorded for a service, then
# logging will need to be set up.
self.is_root = False
if self.zipkin_attrs:
self.is_root = True
if self.sample_rate is not None:
self.is_root = True
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,
)
else:
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,
)
# Don't do anything if zipkin attributes are not set up
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.is_root:
# 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,
)
self.logging_context.__enter__()
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 __exit__(self, _exc_type, _exc_value, _exc_traceback):
"""Exit the span context. 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.
"""
# Always remove the stored zipkin_attrs
if self.do_pop_attrs:
pop_zipkin_attrs()
# Exit early if this request is not being sampled
if not self.zipkin_attrs or not self.zipkin_attrs.is_sampled:
return
# If this is the root span, exit the context (which will handle logging)
if self.logging_context:
self.logging_context.__exit__(_exc_type, _exc_value, _exc_traceback)
self.logging_context = None
return
end_timestamp = time.time()
# Put the old parent_span_id back on the handler
self.log_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.log_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.zipkin_attrs.span_id,
)
def update_binary_annotations_for_root_span(self, extra_annotations):
"""Updates the binary annotations for the root span of the trace.
If this trace is not being sampled then this is a no-op.
"""
if not self.zipkin_attrs.is_sampled:
return
if not self.logging_context:
raise ZipkinError('No logging context available')
self.logging_context.binary_annotations_dict.update(extra_annotations)
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 = bool(self.zipkin_attrs
or self.sample_rate is not None
or self.firehose_handler is not None)
report_root_timestamp = False
if self.sample_rate is not None:
if self.zipkin_attrs and not self.zipkin_attrs.is_sampled:
report_root_timestamp = True
self.zipkin_attrs = create_attrs_for_span(
sample_rate=self.sample_rate,
trace_id=self.zipkin_attrs.trace_id,
use_128bit_trace_id=self.use_128bit_trace_id,
)
elif not self.zipkin_attrs:
report_root_timestamp = True
self.zipkin_attrs = create_attrs_for_span(
sample_rate=self.sample_rate,
use_128bit_trace_id=self.use_128bit_trace_id,
)
if not self.zipkin_attrs:
# This span is inside the context of an existing trace
existing_zipkin_attrs = self._context_stack.get()
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,
)
elif self.firehose_handler is not None:
# If it has gotten here, the only thing that is
# causing a trace is the firehose. So we force a trace
# with sample rate of 0
report_root_timestamp = True
self.zipkin_attrs = create_attrs_for_span(
sample_rate=0.0,
use_128bit_trace_id=self.use_128bit_trace_id,
)
# 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
self._context_stack.push(self.zipkin_attrs)
self.do_pop_attrs = True
self.start_timestamp = time.time()
if self.perform_logging:
# Don't set up any logging if we're not sampling
if not self.zipkin_attrs.is_sampled and not self.firehose_handler:
return self
endpoint = create_endpoint(self.port, self.service_name, self.host)
client_context = set(self.include) == {'client'}
self.log_handler = ZipkinLoggerHandler(self.zipkin_attrs)
self.logging_context = ZipkinLoggingContext(
self.zipkin_attrs,
endpoint,
self.log_handler,
self.span_name,
self.transport_handler,
report_root_timestamp or self.report_root_timestamp_override,
binary_annotations=self.binary_annotations,
add_logging_annotation=self.add_logging_annotation,
client_context=client_context,
max_span_batch_size=self.max_span_batch_size,
firehose_handler=self.firehose_handler,
)
self.logging_context.start()
self.logging_configured = True
return self
else:
# Patch the ZipkinLoggerHandler.
# 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__.
try:
# Assume there's only a single handler, since all logging
# should be set up in this package.
log_handler = zipkin_logger.handlers[0]
except IndexError:
return self
# Make sure it's not a NullHandler or something
if not isinstance(log_handler, ZipkinLoggerHandler):
return self
# Put span ID on logging handler.
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
self.logging_configured = True
return self