1import warnings
2
3import six
4from eventlet.green import httplib
5from eventlet.zipkin import api
6
7
8# see https://twitter.github.io/zipkin/Instrumenting.html
9HDR_TRACE_ID = 'X-B3-TraceId'
10HDR_SPAN_ID = 'X-B3-SpanId'
11HDR_PARENT_SPAN_ID = 'X-B3-ParentSpanId'
12HDR_SAMPLED = 'X-B3-Sampled'
13
14
15if six.PY2:
16    __org_endheaders__ = httplib.HTTPConnection.endheaders
17    __org_begin__ = httplib.HTTPResponse.begin
18
19    def _patched_endheaders(self):
20        if api.is_tracing():
21            trace_data = api.get_trace_data()
22            new_span_id = api.generate_span_id()
23            self.putheader(HDR_TRACE_ID, hex_str(trace_data.trace_id))
24            self.putheader(HDR_SPAN_ID, hex_str(new_span_id))
25            self.putheader(HDR_PARENT_SPAN_ID, hex_str(trace_data.span_id))
26            self.putheader(HDR_SAMPLED, int(trace_data.sampled))
27            api.put_annotation('Client Send')
28
29        __org_endheaders__(self)
30
31    def _patched_begin(self):
32        __org_begin__(self)
33
34        if api.is_tracing():
35            api.put_annotation('Client Recv (%s)' % self.status)
36
37
38def patch():
39    if six.PY2:
40        httplib.HTTPConnection.endheaders = _patched_endheaders
41        httplib.HTTPResponse.begin = _patched_begin
42    if six.PY3:
43        warnings.warn("Since current Python thrift release \
44        doesn't support Python 3, eventlet.zipkin.http \
45        doesn't also support Python 3 (http.client)")
46
47
48def unpatch():
49    if six.PY2:
50        httplib.HTTPConnection.endheaders = __org_endheaders__
51        httplib.HTTPResponse.begin = __org_begin__
52    if six.PY3:
53        pass
54
55
56def hex_str(n):
57    """
58    Thrift uses a binary representation of trace and span ids
59    HTTP headers use a hexadecimal representation of the same
60    """
61    return '%0.16x' % (n,)
62