def _start_client_span(self, client_call_details):
span = self.tracer.start_span(name=_get_span_name(client_call_details))
span.span_kind = span_module.SpanKind.CLIENT
# Add the component grpc to span attribute
self.tracer.add_attribute_to_current_span(
attribute_key=attributes_helper.COMMON_ATTRIBUTES.get(
ATTRIBUTE_COMPONENT),
attribute_value='grpc')
# Add the host:port info to span attribute
self.tracer.add_attribute_to_current_span(
attribute_key=attributes_helper.GRPC_ATTRIBUTES.get(
GRPC_HOST_PORT),
attribute_value=self.host_port)
# Add the method to span attribute
self.tracer.add_attribute_to_current_span(
attribute_key=attributes_helper.GRPC_ATTRIBUTES.get(GRPC_METHOD),
attribute_value=str(client_call_details.method))
execution_context.set_opencensus_tracer(self.tracer)
execution_context.set_current_span(span)
return span
def start_span(self, name='span'):
"""Start a span.
:type name: str
:param name: The name of the span.
:rtype: :class:`~opencensus.trace.span.Span`
:returns: The Span object.
"""
parent_span = self.current_span()
# If a span has remote parent span, then the parent_span.span_id
# should be the span_id from the request header.
if parent_span is None:
parent_span = base.NullContextManager(
span_id=self.span_context.span_id)
span = trace_span.Span(name,
parent_span=parent_span,
context_tracer=self)
with self._spans_list_condition:
self._spans_list.append(span)
self.span_context.span_id = span.span_id
execution_context.set_current_span(span)
span.start()
return span
def end_span(self, *args, **kwargs):
"""End a span. Update the span_id in SpanContext to the current span's
parent span id; Update the current span.
"""
cur_span = self.current_span()
if cur_span is None and self._spans_list:
cur_span = self._spans_list[-1]
if cur_span is None:
logging.warning('No active span, cannot do end_span.')
return
cur_span.finish()
self.span_context.span_id = cur_span.parent_span.span_id if \
cur_span.parent_span else None
if isinstance(cur_span.parent_span, trace_span.Span):
execution_context.set_current_span(cur_span.parent_span)
else:
execution_context.set_current_span(None)
with self._spans_list_condition:
if cur_span in self._spans_list:
span_datas = self.get_span_datas(cur_span)
self.exporter.export(span_datas)
self._spans_list.remove(cur_span)
return cur_span
def test_get_and_set_full_context(self):
mock_tracer_get = mock.Mock()
mock_span_get = mock.Mock()
execution_context.set_opencensus_tracer(mock_tracer_get)
execution_context.set_current_span(mock_span_get)
execution_context.set_opencensus_attr("test", "test_value")
tracer, span, attrs = execution_context.get_opencensus_full_context()
self.assertEqual(mock_tracer_get, tracer)
self.assertEqual(mock_span_get, span)
self.assertEqual({"test": "test_value"}, attrs)
mock_tracer_set = mock.Mock()
mock_span_set = mock.Mock()
execution_context.set_opencensus_full_context(mock_tracer_set,
mock_span_set, None)
self.assertEqual(mock_tracer_set,
execution_context.get_opencensus_tracer())
self.assertEqual(mock_span_set, execution_context.get_current_span())
self.assertEqual({}, execution_context.get_opencensus_attrs())
execution_context.set_opencensus_full_context(
mock_tracer_set, mock_span_set, {"test": "test_value"})
self.assertEqual("test_value",
execution_context.get_opencensus_attr("test"))
def _start_server_span(self, servicer_context):
metadata = servicer_context.invocation_metadata()
span_context = None
if metadata is not None:
propagator = binary_format.BinaryFormatPropagator()
metadata_dict = dict(metadata)
trace_header = metadata_dict.get(oc_grpc.GRPC_TRACE_KEY)
span_context = propagator.from_header(trace_header)
tracer = tracer_module.Tracer(span_context=span_context,
sampler=self.sampler,
exporter=self.exporter)
span = tracer.start_span(name=_get_span_name(servicer_context))
span.span_kind = span_module.SpanKind.SERVER
tracer.add_attribute_to_current_span(
attribute_key=attributes_helper.COMMON_ATTRIBUTES.get(
ATTRIBUTE_COMPONENT),
attribute_value='grpc')
execution_context.set_opencensus_tracer(tracer)
execution_context.set_current_span(span)
return span
def test_span_callbacks(self):
cb = mock.Mock()
execution_context.add_current_span_set_callback(cb)
mock_span = mock.Mock()
execution_context.set_current_span(mock_span)
cb.assert_called_once_with(mock_span)
def callback(future_response):
grpc_utils.add_message_event(
proto_message=future_response.result(),
span=current_span,
message_event_type=time_event.Type.RECEIVED,
)
execution_context.set_current_span(current_span)
self._trace_future_exception(future_response)
self.tracer.end_span()
def stop_trace():
"""Stop the current trace."""
if not execution_context:
return
execution_context.set_current_span(None)
tracer = execution_context.get_opencensus_tracer()
tracer.tracer = noop_tracer.NoopTracer()
tracer.span_context = SpanContext(trace_options=trace_options.TraceOptions(0))
tracer.sampler = always_off.AlwaysOffSampler()
def stop_trace():
"""Stop the current trace."""
if not execution_context:
return
execution_context.set_current_span(None)
tracer = execution_context.get_opencensus_tracer()
tracer.tracer = noop_tracer.NoopTracer()
tracer.span_context = SpanContext(
trace_options=trace_options.TraceOptions(0))
tracer.sampler = always_off.AlwaysOffSampler()
def change_context(cls, span):
# type: (Span) -> ContextManager
"""Change the context for the life of this context manager.
"""
original_span = cls.get_current_span()
try:
execution_context.set_current_span(span)
yield
finally:
execution_context.set_current_span(original_span)
def test_current_span_sampled(self):
from opencensus.trace import execution_context
sampler = mock.Mock()
sampler.should_sample.return_value = True
tracer = tracer_module.Tracer(sampler=sampler)
span = mock.Mock()
execution_context.set_current_span(span)
result = tracer.current_span()
self.assertEqual(result, span)
def test_clean_span(self):
mock_span = mock.Mock()
some_value = mock.Mock()
execution_context.set_current_span(mock_span)
thread_local = threading.local()
setattr(thread_local, 'random_non_oc_attr', some_value)
execution_context.clean()
self.assertNotEqual(mock_span, execution_context.get_current_span())
self.assertEqual(some_value, getattr(thread_local,
'random_non_oc_attr'))
def prerun_task_span(task_id=None, task=None, *args, **kwargs):
if settings.STACKDRIVER_TRACE_PROJECT_ID:
exporter = stackdriver_exporter.StackdriverExporter(
project_id=settings.STACKDRIVER_TRACE_PROJECT_ID,
transport=BackgroundThreadTransport)
sampler = probability.ProbabilitySampler(
rate=settings.CELERY_TRACE_SAMPLING_RATE)
tracer = tracer_module.Tracer(exporter=exporter, sampler=sampler)
span = tracer.start_span()
span.name = '[celery]{0}'.format(task.name)
execution_context.set_opencensus_tracer(tracer)
span.add_attribute('args', str(kwargs['args']))
span.add_attribute('kwargs', str(kwargs['kwargs']))
execution_context.set_current_span(span)
def test_add_attribute_to_current_span(self):
from opencensus.trace.span import Span
from opencensus.trace import execution_context
tracer = context_tracer.ContextTracer()
span1 = mock.Mock(spec=Span)
span1.attributes = {}
execution_context.set_current_span(span1)
attribute_key = 'key'
attribute_value = 'value'
tracer.add_attribute_to_current_span(attribute_key, attribute_value)
span1.add_attribute.assert_called_once_with(attribute_key,
attribute_value)
def test_end_span_sampled(self):
from opencensus.trace import execution_context
sampler = mock.Mock()
sampler.should_sample.return_value = True
tracer = tracer_module.Tracer(sampler=sampler)
span = mock.Mock()
span.attributes = {}
span.time_events = []
span.links = []
span.children = []
span.__iter__ = mock.Mock(return_value=iter([span]))
execution_context.set_current_span(span)
with mock.patch('opencensus.trace.span.utils.get_truncatable_str'):
tracer.end_span()
self.assertTrue(span.finish.called)
def set_current_span(cls, span):
# type: (Span) -> None
"""
Set the given span as the current span in the execution context.
:param span: The span to set the current span as
:type span: :class: opencensus.trace.Span
"""
return execution_context.set_current_span(span)
def set_current_span(cls, span):
# type: (Span) -> None
"""
Set the given span as the current span in the execution context.
:param span: The span to set the current span as
:type span: :class: opencensus.trace.Span
"""
warnings.warn("set_current_span is deprecated, use change_context instead", DeprecationWarning)
return execution_context.set_current_span(span)
def end_span(self, *args, **kwargs):
"""End a span. Update the span_id in SpanContext to the current span's
parent span id; Update the current span.
"""
cur_span = self.current_span()
if cur_span is None:
logging.warning('No active span, cannot do end_span.')
return
cur_span.finish()
self.span_context.span_id = cur_span.parent_span.span_id
if isinstance(cur_span.parent_span, trace_span.Span):
execution_context.set_current_span(cur_span.parent_span)
else:
execution_context.set_current_span(None)
span_json = self.get_trace_json(cur_span)
self.exporter.export(span_json)
return cur_span
def _end_span_between_context(self, current_span):
execution_context.set_current_span(current_span)
self.tracer.end_span()
def callback(future_response):
execution_context.set_current_span(current_span)
self._trace_future_exception(future_response)
self._tracer.end_span()
