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