|
- import os
- import sys
- import time
- import struct
- import socket
- import random
-
- from eventlet.green import threading
- from eventlet.zipkin._thrift.zipkinCore import ttypes
- from eventlet.zipkin._thrift.zipkinCore.constants import SERVER_SEND
-
-
- client = None
- _tls = threading.local() # thread local storage
-
-
- def put_annotation(msg, endpoint=None):
- """ This is annotation API.
- You can add your own annotation from in your code.
- Annotation is recorded with timestamp automatically.
- e.g.) put_annotation('cache hit for %s' % request)
-
- :param msg: String message
- :param endpoint: host info
- """
- if is_sample():
- a = ZipkinDataBuilder.build_annotation(msg, endpoint)
- trace_data = get_trace_data()
- trace_data.add_annotation(a)
-
-
- def put_key_value(key, value, endpoint=None):
- """ This is binary annotation API.
- You can add your own key-value extra information from in your code.
- Key-value doesn't have a time component.
- e.g.) put_key_value('http.uri', '/hoge/index.html')
-
- :param key: String
- :param value: String
- :param endpoint: host info
- """
- if is_sample():
- b = ZipkinDataBuilder.build_binary_annotation(key, value, endpoint)
- trace_data = get_trace_data()
- trace_data.add_binary_annotation(b)
-
-
- def is_tracing():
- """ Return whether the current thread is tracking or not """
- return hasattr(_tls, 'trace_data')
-
-
- def is_sample():
- """ Return whether it should record trace information
- for the request or not
- """
- return is_tracing() and _tls.trace_data.sampled
-
-
- def get_trace_data():
- if is_tracing():
- return _tls.trace_data
-
-
- def set_trace_data(trace_data):
- _tls.trace_data = trace_data
-
-
- def init_trace_data():
- if is_tracing():
- del _tls.trace_data
-
-
- def _uniq_id():
- """
- Create a random 64-bit signed integer appropriate
- for use as trace and span IDs.
- XXX: By experimentation zipkin has trouble recording traces with ids
- larger than (2 ** 56) - 1
- """
- return random.randint(0, (2 ** 56) - 1)
-
-
- def generate_trace_id():
- return _uniq_id()
-
-
- def generate_span_id():
- return _uniq_id()
-
-
- class TraceData(object):
-
- END_ANNOTATION = SERVER_SEND
-
- def __init__(self, name, trace_id, span_id, parent_id, sampled, endpoint):
- """
- :param name: RPC name (String)
- :param trace_id: int
- :param span_id: int
- :param parent_id: int or None
- :param sampled: lets the downstream servers know
- if I should record trace data for the request (bool)
- :param endpoint: zipkin._thrift.zipkinCore.ttypes.EndPoint
- """
- self.name = name
- self.trace_id = trace_id
- self.span_id = span_id
- self.parent_id = parent_id
- self.sampled = sampled
- self.endpoint = endpoint
- self.annotations = []
- self.bannotations = []
- self._done = False
-
- def add_annotation(self, annotation):
- if annotation.host is None:
- annotation.host = self.endpoint
- if not self._done:
- self.annotations.append(annotation)
- if annotation.value == self.END_ANNOTATION:
- self.flush()
-
- def add_binary_annotation(self, bannotation):
- if bannotation.host is None:
- bannotation.host = self.endpoint
- if not self._done:
- self.bannotations.append(bannotation)
-
- def flush(self):
- span = ZipkinDataBuilder.build_span(name=self.name,
- trace_id=self.trace_id,
- span_id=self.span_id,
- parent_id=self.parent_id,
- annotations=self.annotations,
- bannotations=self.bannotations)
- client.send_to_collector(span)
- self.annotations = []
- self.bannotations = []
- self._done = True
-
-
- class ZipkinDataBuilder:
- @staticmethod
- def build_span(name, trace_id, span_id, parent_id,
- annotations, bannotations):
- return ttypes.Span(
- name=name,
- trace_id=trace_id,
- id=span_id,
- parent_id=parent_id,
- annotations=annotations,
- binary_annotations=bannotations
- )
-
- @staticmethod
- def build_annotation(value, endpoint=None):
- if isinstance(value, unicode):
- value = value.encode('utf-8')
- return ttypes.Annotation(time.time() * 1000 * 1000,
- str(value), endpoint)
-
- @staticmethod
- def build_binary_annotation(key, value, endpoint=None):
- annotation_type = ttypes.AnnotationType.STRING
- return ttypes.BinaryAnnotation(key, value, annotation_type, endpoint)
-
- @staticmethod
- def build_endpoint(ipv4=None, port=None, service_name=None):
- if ipv4 is not None:
- ipv4 = ZipkinDataBuilder._ipv4_to_int(ipv4)
- if service_name is None:
- service_name = ZipkinDataBuilder._get_script_name()
- return ttypes.Endpoint(
- ipv4=ipv4,
- port=port,
- service_name=service_name
- )
-
- @staticmethod
- def _ipv4_to_int(ipv4):
- return struct.unpack('!i', socket.inet_aton(ipv4))[0]
-
- @staticmethod
- def _get_script_name():
- return os.path.basename(sys.argv[0])
|