blob: 02144fa0759f4c708d6811c3fa9556ec32af49f0 [file] [log] [blame]
# -*- coding: utf-8 -*-
# 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."""
from __future__ import print_function
import contextlib
import errno
import functools
import json
import os
import random
import re
import sys
import google.protobuf.internal.well_known_types as types
from infra_libs import ts_mon
from chromite.lib import cros_logging as log
from chromite.lib import metrics
from chromite.lib import structured
assert sys.version_info >= (3, 6), 'This module requires Python 3.6+'
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.
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.
span: A Span instance to serialize.
with open(GetSpanLogFilePath(span), 'w') as fh:
# Catch various configuration errors.
# TODO(vapier): Drop IOError when we're Python 3-only.
# pylint: disable=overlapping-except
except (OSError, IOError) as error:
if error.errno == errno.EPERM:
'Received permissions error while trying to open the span log file.')
return None
elif error.errno == errno.ENOENT:
log.warning('/var/log/traces does not exist; skipping trace log.')
return None
def _RecordSpanMetrics(span):
"""Increments the count of spans logged.
span: The span to record.
m = metrics.Counter(
description='A count of spans logged by a client.',
# -- User-facing API -----------------------------------------------------------
class Span(structured.Structured):
"""An object corresponding to a cloud trace Span."""
'name', 'spanId', 'parentSpanId', 'labels',
'startTime', 'endTime', 'status')
def __init__(self, name, spanId=None, labels=None, parentSpanId=None,
"""Creates a Span object.
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 = 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()
def GenerateSpanId():
"""Returns a random 64-bit number as a string."""
return str(random.randint(0, 2**64))
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 = types.Timestamp()
self.startTime = start.ToJsonString()
return self
def __exit__(self, _type, _value, _traceback):
"""Exits the span context.
Record the end Timestamp.
end = types.Timestamp()
self.endTime = end.ToJsonString()
class SpanStack(object):
"""A stack of Span contexts."""
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.
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
def Span(self, name, **kwargs):
"""Enter a new Span context contained within the top Span of the stack.
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.
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
with span:
with self.EnvironmentContext():
yield span
self.last_span_id = old_span_id
# Log each span to a file for later processing.
if self.enabled:
# pylint: disable=docstring-misnamed-args
def Spanned(self, *span_args, **span_kwargs):
"""A decorator equivalent of 'with span_stack.Span(...)'
*span_args: *args to use with the .Span
**span_kwargs: **kwargs to use with the .Span
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."""
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 (
'X-Cloud-Trace-Context: TRACE_ID/SPAN_ID;o=TRACE_TRUE'
- 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 '' --header 'X-Cloud-Trace-Context:
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(
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)
os.environ[self.CLOUD_TRACE_CONTEXT_ENV] = (
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]
def _ParseCloudTraceContext(context):
"""Sets current_span_id and trace_id from the |context|.
See _GetCloudTraceContextHeader.
context: The context variable, either from X-Cloud-Trace-Context
or from the CLOUD_TRACE_CONTEXT environment variable.
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()
return {}