lib/cloud_trace.py - third_party/chromite - Git at Google

 # Copyright 2017 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.

 """A library for emitting traces and spans to Google Cloud trace."""

 import collections
 import contextlib
 import errno
 import functools
 import json
 import logging
 import os
 import random
 import re

 from chromite.third_party.google.protobuf import timestamp_pb2
 from chromite.third_party.infra_libs import ts_mon

 from chromite.lib import metrics


 SPANS_LOG = '/var/log/trace/{pid}-{span_id}.json'
 _SPAN_COUNT_METRIC = 'chromeos/trace/client/logged_count'


 # --- Code for logging spans to a file for later processing. -------------------
 def GetSpanLogFilePath(span):
   """Gets the path to write a span to.

   Args:
     span: The span to write.
   """
   return SPANS_LOG.format(pid=os.getpid(), span_id=span.spanId)


 def LogSpan(span):
   """Serializes and logs a Span to a file.

   Args:
     span: A Span instance to serialize.
   """
   _RecordSpanMetrics(span)
   try:
     with open(GetSpanLogFilePath(span), 'w') as fh:
       fh.write(json.dumps(span, default=serialize_span))
   # Catch various configuration errors.
   except OSError as error:
     if error.errno == errno.EPERM:
       logging.warning(
           'Received permissions error while trying to open the span log file.')
       return None
     elif error.errno == errno.ENOENT:
       logging.warning('/var/log/traces does not exist; skipping trace log.')
       return None
     else:
       raise


 def serialize_span(obj):
   """Helper function to provide serialization of Spans for `json.dump`."""
   if isinstance(obj, Span):
     return dict(obj)
   else:
     raise TypeError(f'Unable to serialize an object of type {type(obj)}')


 def _RecordSpanMetrics(span):
   """Increments the count of spans logged.

   Args:
     span: The span to record.
   """
   m = metrics.Counter(
       _SPAN_COUNT_METRIC,
       description='A count of spans logged by a client.',
       field_spec=[ts_mon.StringField('name')])
   m.increment(fields={'name': span.name})


 # -- User-facing API -----------------------------------------------------------
 class Span(collections.abc.Mapping):
   """An object corresponding to a cloud trace Span."""

   __keys = ('name', 'spanId', 'parentSpanId', 'labels', 'startTime', 'endTime',
             'status')

   def __init__(self, name, spanId=None, labels=None, parentSpanId=None,
                traceId=None):
     """Creates a Span object.

     Args:
       name: The name of the span
       spanId: (optional) A 64-bit number as a string. If not provided, it will
           be generated randomly with .GenerateSpanId().
       labels: (optional) a dict<string, string> of key/values
       traceId: (optional) A 32 hex digit string referring to the trace
           containing this span. If not provided, a new trace will be created
           with a random id.
       parentSpanId: (optional) The spanId of the parent.
     """
     # Visible attributes
     self.name = name
     self.spanId = spanId or Span.GenerateSpanId()
     self.parentSpanId = parentSpanId
     self.labels = labels or {}
     self.startTime = None
     self.endTime = None
     # Non-visible attributes
     self.traceId = traceId or Span.GenerateTraceId()

   @staticmethod
   def GenerateSpanId():
     """Returns a random 64-bit number as a string."""
     return str(random.randint(0, 2**64))

   @staticmethod
   def GenerateTraceId():
     """Returns a random 128-bit number as a 32-byte hex string."""
     id_number = random.randint(0, 2**128)
     return '%0.32X' % id_number

   def __enter__(self):
     """Enters the span context.

     Side effect: Records the start time as a Timestamp.
     """
     start = timestamp_pb2.Timestamp()
     start.GetCurrentTime()
     self.startTime = start.ToJsonString()
     return self

   def __exit__(self, _type, _value, _traceback):
     """Exits the span context.

     Side-effect:
       Record the end Timestamp.
     """
     end = timestamp_pb2.Timestamp()
     end.GetCurrentTime()
     self.endTime = end.ToJsonString()

   def __getitem__(self, key: str):
     if not isinstance(key, str):
       raise TypeError('Span keys must be of type str')

     if key not in self.__keys:
       raise KeyError(key)

     val = getattr(self, key, None)
     if val is not None:
       return val
     else:
       raise KeyError(key)

   def keys(self):
     return (key for key in self.__keys if getattr(self, key, None) is not None)

   def __len__(self) -> int:
     return len(self.keys())

   def __iter__(self):
     return ((key, self[key]) for key in self.keys())


 class SpanStack(object):
   """A stack of Span contexts."""

   CLOUD_TRACE_CONTEXT_ENV = 'CLOUD_TRACE_CONTEXT'
   CLOUD_TRACE_CONTEXT_PATTERN = re.compile(
       r'(?P<traceId>[0-9A-Fa-f]+)'
       r'/'
       r'(?P<parentSpanId>\d+)'
       r';o=(?P<options>\d+)'
   )

   def __init__(self, global_context=None, traceId=None, parentSpanId=None,
                labels=None, enabled=True):
     """Initializes the Span.

       global_context: (optional) A global context str, perhaps read from the
           X-Cloud-Trace-Context header.
       traceId: (optional) A 32 hex digit string referring to the trace
           containing this span. If not provided, a new trace will be created
           with a random id.
       parentSpanId: (optional) The spanId of the parent.
       labels: (optional) a dict<string, string> of key/values to attach to
           each Span created, or None.
       enabled: (optional) a bool indicating whether we should log the spans
           to a file for later uploading by the cloud trace log consumer daemon.
     """
     self.traceId = traceId
     self.spans = []
     self.last_span_id = parentSpanId
     self.labels = labels
     self.enabled = enabled

     global_context = (global_context or
                       os.environ.get(self.CLOUD_TRACE_CONTEXT_ENV, ''))
     context = SpanStack._ParseCloudTraceContext(global_context)

     if traceId is None:
       self.traceId = context.get('traceId')
     if parentSpanId is None:
       self.last_span_id = context.get('parentSpanId')
     if context.get('options') == '0':
       self.enabled = False

   def _CreateSpan(self, name, **kwargs):
     """Creates a span instance, setting certain defaults.

     Args:
       name: The name of the span
       **kwargs: The keyword arguments to configure the span with.
     """
     kwargs.setdefault('traceId', self.traceId)
     kwargs.setdefault('labels', self.labels)
     kwargs.setdefault('parentSpanId', self.last_span_id)
     span = Span(name, **kwargs)
     if self.traceId is None:
       self.traceId = span.traceId
     return span

   @contextlib.contextmanager
   def Span(self, name, **kwargs):
     """Enter a new Span context contained within the top Span of the stack.

     Args:
       name: The name of the span to enter
       **kwargs: The kwargs to construct the span with.

     Side effect:
       Appends the new span object to |spans|, and yields span while in its
       context. Pops the span object when exiting the context.

     Returns:
       A contextmanager whose __enter__() returns the new Span.
     """
     span = self._CreateSpan(name, **kwargs)
     old_span_id, self.last_span_id = self.last_span_id, span.spanId
     self.spans.append(span)

     with span:
       with self.EnvironmentContext():
         yield span

     self.spans.pop()
     self.last_span_id = old_span_id

     # Log each span to a file for later processing.
     if self.enabled:
       LogSpan(span)

   # pylint: disable=docstring-misnamed-args
   def Spanned(self, *span_args, **span_kwargs):
     """A decorator equivalent of 'with span_stack.Span(...)'

     Args:
       *span_args: *args to use with the .Span
       **span_kwargs: **kwargs to use with the .Span

     Returns:
       A decorator to wrap the body of a function in a span block.
     """
     def SpannedDecorator(f):
       """Wraps the body of |f| with a .Span block."""
       @functools.wraps(f)
       def inner(*args, **kwargs):
         with self.Span(*span_args, **span_kwargs):
           f(*args, **kwargs)
       return inner
     return SpannedDecorator

   def _GetCloudTraceContextHeader(self):
     """Gets the Cloud Trace HTTP header context.

     From the cloud trace doc explaining this (
     https://cloud.google.com/trace/docs/support?hl=bg)

       'X-Cloud-Trace-Context: TRACE_ID/SPAN_ID;o=TRACE_TRUE'
       Where:
         - TRACE_ID is a 32-character hex value representing a 128-bit number.
         It should be unique between your requests, unless you intentionally
         want to bundle the requests together. You can use UUIDs.
         - SPAN_ID should be 0 for the first span in your trace. For
         subsequent requests, set SPAN_ID to the span ID of the parent
         request. See the description of TraceSpan (REST, RPC) for more
         information about nested traces.
         - TRACE_TRUE must be 1 to trace this request. Specify 0 to not trace
         the request. For example, to force a trace with cURL:
           curl 'http://www.example.com' --header 'X-Cloud-Trace-Context:
             105445aa7843bc8bf206b120001000/0;o=1'
     """
     if not self.traceId:
       return ''
     span_postfix = '/%s' % self.spans[-1].spanId if self.spans else ''
     enabled = '1' if self.enabled else '0'
     return '{trace_id}{span_postfix};o={enabled}'.format(
         trace_id=self.traceId,
         span_postfix=span_postfix,
         enabled=enabled)

   @contextlib.contextmanager
   def EnvironmentContext(self):
     """Sets CLOUD_TRACE_CONTEXT to the value of X-Cloud-Trace-Context.

     Cloud Trace uses an HTTP header to propagate trace context across RPC
     boundaries. This method does the same across process boundaries using an
     environment variable.
     """
     old_value = os.environ.get(self.CLOUD_TRACE_CONTEXT_ENV)
     try:
       os.environ[self.CLOUD_TRACE_CONTEXT_ENV] = (
           self._GetCloudTraceContextHeader())
       yield
     finally:
       if old_value is not None:
         os.environ[self.CLOUD_TRACE_CONTEXT_ENV] = old_value
       elif self.CLOUD_TRACE_CONTEXT_ENV in os.environ:
         del os.environ[self.CLOUD_TRACE_CONTEXT_ENV]

   @staticmethod
   def _ParseCloudTraceContext(context):
     """Sets current_span_id and trace_id from the |context|.

     See _GetCloudTraceContextHeader.

     Args:
       context: The context variable, either from X-Cloud-Trace-Context
           or from the CLOUD_TRACE_CONTEXT environment variable.

     Returns:
       A dictionary, which if the context string matches
       CLOUD_TRACE_CONTEXT_PATTERN, contains the matched groups. If not matched,
       returns an empty dictionary.
     """
     m = SpanStack.CLOUD_TRACE_CONTEXT_PATTERN.match(context)
     if m:
       return m.groupdict()
     else:
       return {}
	# Copyright 2017 The Chromium OS Authors. All rights reserved.
	# Use of this source code is governed by a BSD-style license that can be
	# found in the LICENSE file.

	"""A library for emitting traces and spans to Google Cloud trace."""

	import collections
	import contextlib
	import errno
	import functools
	import json
	import logging
	import os
	import random
	import re

	from chromite.third_party.google.protobuf import timestamp_pb2
	from chromite.third_party.infra_libs import ts_mon

	from chromite.lib import metrics


	SPANS_LOG = '/var/log/trace/{pid}-{span_id}.json'
	_SPAN_COUNT_METRIC = 'chromeos/trace/client/logged_count'


	# --- Code for logging spans to a file for later processing. -------------------
	def GetSpanLogFilePath(span):
	"""Gets the path to write a span to.

	Args:
	span: The span to write.
	"""
	return SPANS_LOG.format(pid=os.getpid(), span_id=span.spanId)


	def LogSpan(span):
	"""Serializes and logs a Span to a file.

	Args:
	span: A Span instance to serialize.
	"""
	_RecordSpanMetrics(span)
	try:
	with open(GetSpanLogFilePath(span), 'w') as fh:
	fh.write(json.dumps(span, default=serialize_span))
	# Catch various configuration errors.
	except OSError as error:
	if error.errno == errno.EPERM:
	logging.warning(
	'Received permissions error while trying to open the span log file.')
	return None
	elif error.errno == errno.ENOENT:
	logging.warning('/var/log/traces does not exist; skipping trace log.')
	return None
	else:
	raise


	def serialize_span(obj):
	"""Helper function to provide serialization of Spans for `json.dump`."""
	if isinstance(obj, Span):
	return dict(obj)
	else:
	raise TypeError(f'Unable to serialize an object of type {type(obj)}')


	def _RecordSpanMetrics(span):
	"""Increments the count of spans logged.

	Args:
	span: The span to record.
	"""
	m = metrics.Counter(
	_SPAN_COUNT_METRIC,
	description='A count of spans logged by a client.',
	field_spec=[ts_mon.StringField('name')])
	m.increment(fields={'name': span.name})


	# -- User-facing API -----------------------------------------------------------
	class Span(collections.abc.Mapping):
	"""An object corresponding to a cloud trace Span."""

	__keys = ('name', 'spanId', 'parentSpanId', 'labels', 'startTime', 'endTime',
	'status')

	def __init__(self, name, spanId=None, labels=None, parentSpanId=None,
	traceId=None):
	"""Creates a Span object.

	Args:
	name: The name of the span
	spanId: (optional) A 64-bit number as a string. If not provided, it will
	be generated randomly with .GenerateSpanId().
	labels: (optional) a dict<string, string> of key/values
	traceId: (optional) A 32 hex digit string referring to the trace
	containing this span. If not provided, a new trace will be created
	with a random id.
	parentSpanId: (optional) The spanId of the parent.
	"""
	# Visible attributes
	self.name = name
	self.spanId = spanId or Span.GenerateSpanId()
	self.parentSpanId = parentSpanId
	self.labels = labels or {}
	self.startTime = None
	self.endTime = None
	# Non-visible attributes
	self.traceId = traceId or Span.GenerateTraceId()

	@staticmethod
	def GenerateSpanId():
	"""Returns a random 64-bit number as a string."""
	return str(random.randint(0, 2**64))

	@staticmethod
	def GenerateTraceId():
	"""Returns a random 128-bit number as a 32-byte hex string."""
	id_number = random.randint(0, 2**128)
	return '%0.32X' % id_number

	def __enter__(self):
	"""Enters the span context.

	Side effect: Records the start time as a Timestamp.
	"""
	start = timestamp_pb2.Timestamp()
	start.GetCurrentTime()
	self.startTime = start.ToJsonString()
	return self

	def __exit__(self, _type, _value, _traceback):
	"""Exits the span context.

	Side-effect:
	Record the end Timestamp.
	"""
	end = timestamp_pb2.Timestamp()
	end.GetCurrentTime()
	self.endTime = end.ToJsonString()

	def __getitem__(self, key: str):
	if not isinstance(key, str):
	raise TypeError('Span keys must be of type str')

	if key not in self.__keys:
	raise KeyError(key)

	val = getattr(self, key, None)
	if val is not None:
	return val
	else:
	raise KeyError(key)

	def keys(self):
	return (key for key in self.__keys if getattr(self, key, None) is not None)

	def __len__(self) -> int:
	return len(self.keys())

	def __iter__(self):
	return ((key, self[key]) for key in self.keys())


	class SpanStack(object):
	"""A stack of Span contexts."""

	CLOUD_TRACE_CONTEXT_ENV = 'CLOUD_TRACE_CONTEXT'
	CLOUD_TRACE_CONTEXT_PATTERN = re.compile(
	r'(?P<traceId>[0-9A-Fa-f]+)'
	r'/'
	r'(?P<parentSpanId>\d+)'
	r';o=(?P<options>\d+)'
	)

	def __init__(self, global_context=None, traceId=None, parentSpanId=None,
	labels=None, enabled=True):
	"""Initializes the Span.

	global_context: (optional) A global context str, perhaps read from the
	X-Cloud-Trace-Context header.
	traceId: (optional) A 32 hex digit string referring to the trace
	containing this span. If not provided, a new trace will be created
	with a random id.
	parentSpanId: (optional) The spanId of the parent.
	labels: (optional) a dict<string, string> of key/values to attach to
	each Span created, or None.
	enabled: (optional) a bool indicating whether we should log the spans
	to a file for later uploading by the cloud trace log consumer daemon.
	"""
	self.traceId = traceId
	self.spans = []
	self.last_span_id = parentSpanId
	self.labels = labels
	self.enabled = enabled

	global_context = (global_context or
	os.environ.get(self.CLOUD_TRACE_CONTEXT_ENV, ''))
	context = SpanStack._ParseCloudTraceContext(global_context)

	if traceId is None:
	self.traceId = context.get('traceId')
	if parentSpanId is None:
	self.last_span_id = context.get('parentSpanId')
	if context.get('options') == '0':
	self.enabled = False

	def _CreateSpan(self, name, **kwargs):
	"""Creates a span instance, setting certain defaults.

	Args:
	name: The name of the span
	**kwargs: The keyword arguments to configure the span with.
	"""
	kwargs.setdefault('traceId', self.traceId)
	kwargs.setdefault('labels', self.labels)
	kwargs.setdefault('parentSpanId', self.last_span_id)
	span = Span(name, **kwargs)
	if self.traceId is None:
	self.traceId = span.traceId
	return span

	@contextlib.contextmanager
	def Span(self, name, **kwargs):
	"""Enter a new Span context contained within the top Span of the stack.

	Args:
	name: The name of the span to enter
	**kwargs: The kwargs to construct the span with.

	Side effect:
	Appends the new span object to \|spans\|, and yields span while in its
	context. Pops the span object when exiting the context.

	Returns:
	A contextmanager whose __enter__() returns the new Span.
	"""
	span = self._CreateSpan(name, **kwargs)
	old_span_id, self.last_span_id = self.last_span_id, span.spanId
	self.spans.append(span)

	with span:
	with self.EnvironmentContext():
	yield span

	self.spans.pop()
	self.last_span_id = old_span_id

	# Log each span to a file for later processing.
	if self.enabled:
	LogSpan(span)

	# pylint: disable=docstring-misnamed-args
	def Spanned(self, span_args, *span_kwargs):
	"""A decorator equivalent of 'with span_stack.Span(...)'

	Args:
	span_args: args to use with the .Span
	span_kwargs: kwargs to use with the .Span

	Returns:
	A decorator to wrap the body of a function in a span block.
	"""
	def SpannedDecorator(f):
	"""Wraps the body of \|f\| with a .Span block."""
	@functools.wraps(f)
	def inner(args, *kwargs):
	with self.Span(span_args, *span_kwargs):
	f(args, *kwargs)
	return inner
	return SpannedDecorator

	def _GetCloudTraceContextHeader(self):
	"""Gets the Cloud Trace HTTP header context.

	From the cloud trace doc explaining this (
	https://cloud.google.com/trace/docs/support?hl=bg)

	'X-Cloud-Trace-Context: TRACE_ID/SPAN_ID;o=TRACE_TRUE'
	Where:
	- TRACE_ID is a 32-character hex value representing a 128-bit number.
	It should be unique between your requests, unless you intentionally
	want to bundle the requests together. You can use UUIDs.
	- SPAN_ID should be 0 for the first span in your trace. For
	subsequent requests, set SPAN_ID to the span ID of the parent
	request. See the description of TraceSpan (REST, RPC) for more
	information about nested traces.
	- TRACE_TRUE must be 1 to trace this request. Specify 0 to not trace
	the request. For example, to force a trace with cURL:
	curl 'http://www.example.com' --header 'X-Cloud-Trace-Context:
	105445aa7843bc8bf206b120001000/0;o=1'
	"""
	if not self.traceId:
	return ''
	span_postfix = '/%s' % self.spans[-1].spanId if self.spans else ''
	enabled = '1' if self.enabled else '0'
	return '{trace_id}{span_postfix};o={enabled}'.format(
	trace_id=self.traceId,
	span_postfix=span_postfix,
	enabled=enabled)

	@contextlib.contextmanager
	def EnvironmentContext(self):
	"""Sets CLOUD_TRACE_CONTEXT to the value of X-Cloud-Trace-Context.

	Cloud Trace uses an HTTP header to propagate trace context across RPC
	boundaries. This method does the same across process boundaries using an
	environment variable.
	"""
	old_value = os.environ.get(self.CLOUD_TRACE_CONTEXT_ENV)
	try:
	os.environ[self.CLOUD_TRACE_CONTEXT_ENV] = (
	self._GetCloudTraceContextHeader())
	yield
	finally:
	if old_value is not None:
	os.environ[self.CLOUD_TRACE_CONTEXT_ENV] = old_value
	elif self.CLOUD_TRACE_CONTEXT_ENV in os.environ:
	del os.environ[self.CLOUD_TRACE_CONTEXT_ENV]

	@staticmethod
	def _ParseCloudTraceContext(context):
	"""Sets current_span_id and trace_id from the \|context\|.

	See _GetCloudTraceContextHeader.

	Args:
	context: The context variable, either from X-Cloud-Trace-Context
	or from the CLOUD_TRACE_CONTEXT environment variable.

	Returns:
	A dictionary, which if the context string matches
	CLOUD_TRACE_CONTEXT_PATTERN, contains the matched groups. If not matched,
	returns an empty dictionary.
	"""
	m = SpanStack.CLOUD_TRACE_CONTEXT_PATTERN.match(context)
	if m:
	return m.groupdict()
	else:
	return {}