1import random
2
3from eventlet import wsgi
4from eventlet.zipkin import api
5from eventlet.zipkin._thrift.zipkinCore.constants import \
6    SERVER_RECV, SERVER_SEND
7from eventlet.zipkin.http import \
8    HDR_TRACE_ID, HDR_SPAN_ID, HDR_PARENT_SPAN_ID, HDR_SAMPLED
9
10
11_sampler = None
12__original_handle_one_response__ = wsgi.HttpProtocol.handle_one_response
13
14
15def _patched_handle_one_response(self):
16    api.init_trace_data()
17    trace_id = int_or_none(self.headers.getheader(HDR_TRACE_ID))
18    span_id = int_or_none(self.headers.getheader(HDR_SPAN_ID))
19    parent_id = int_or_none(self.headers.getheader(HDR_PARENT_SPAN_ID))
20    sampled = bool_or_none(self.headers.getheader(HDR_SAMPLED))
21    if trace_id is None:  # front-end server
22        trace_id = span_id = api.generate_trace_id()
23        parent_id = None
24        sampled = _sampler.sampling()
25    ip, port = self.request.getsockname()[:2]
26    ep = api.ZipkinDataBuilder.build_endpoint(ip, port)
27    trace_data = api.TraceData(name=self.command,
28                               trace_id=trace_id,
29                               span_id=span_id,
30                               parent_id=parent_id,
31                               sampled=sampled,
32                               endpoint=ep)
33    api.set_trace_data(trace_data)
34    api.put_annotation(SERVER_RECV)
35    api.put_key_value('http.uri', self.path)
36
37    __original_handle_one_response__(self)
38
39    if api.is_sample():
40        api.put_annotation(SERVER_SEND)
41
42
43class Sampler(object):
44    def __init__(self, sampling_rate):
45        self.sampling_rate = sampling_rate
46
47    def sampling(self):
48        # avoid generating unneeded random numbers
49        if self.sampling_rate == 1.0:
50            return True
51        r = random.random()
52        if r < self.sampling_rate:
53            return True
54        return False
55
56
57def int_or_none(val):
58    if val is None:
59        return None
60    return int(val, 16)
61
62
63def bool_or_none(val):
64    if val == '1':
65        return True
66    if val == '0':
67        return False
68    return None
69
70
71def patch(sampling_rate):
72    global _sampler
73    _sampler = Sampler(sampling_rate)
74    wsgi.HttpProtocol.handle_one_response = _patched_handle_one_response
75
76
77def unpatch():
78    wsgi.HttpProtocol.handle_one_response = __original_handle_one_response__
79